from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, Trainer, TrainingArguments
from datasets import load_dataset

# Загрузить модель и токенизатор
model_name = "HaveAI/FlareNew"  # Ваша модель
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# Загрузить свой набор данных (можно использовать Hugging Face Datasets или загрузить свои данные)
dataset = load_dataset("path/to/your_dataset")

# Преобразование данных для модели
def tokenize_function(examples):
    return tokenizer(examples['text'], padding="max_length", truncation=True)

tokenized_datasets = dataset.map(tokenize_function, batched=True)

# Настройка аргументов для обучения
training_args = TrainingArguments(
    output_dir="./results",
    evaluation_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=3,
    weight_decay=0.01,
)

# Использование Trainer для обучения
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
)

# Обучение модели
trainer.train()

# Сохранение обученной модели
model.save_pretrained("./flarenew_finetuned")
tokenizer.save_pretrained("./flarenew_finetuned")