finetunning.py

import torch
from datasets import Dataset # , load_dataset
from transformers import (
    AutoTokenizer, 
    AutoModelForCausalLM, 
    TrainingArguments, 
    Trainer, 
    DataCollatorForLanguageModeling
)

# Завантаження моделі та токенізатора
model_name = "facebook/opt-350m"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Завантаження датасета (приклад з використанням датасета wiki_text)
# dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")

# Завантаження данних з локального тестового файлу
with open("ilya_klimov_data.txt", "r", encoding="utf-8") as file:
    text_data = file.read().strip()

# Створення датасету
dataset = Dataset.from_dict({"text": [text_data]})

# Функція для токенізації данних
def tokenize_function(examples):
    return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=512)

# Токенізация датасету
tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=dataset.column_names)

# Створення data collator
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

# Налаштування параметрів навчання
training_args = TrainingArguments(
    output_dir="./results",
    overwrite_output_dir=True,
    num_train_epochs=3,
    per_device_train_batch_size=4,
    save_steps=10_000,
    save_total_limit=2,
    prediction_loss_only=True,
    learning_rate=2e-5,
)

# Ініціалізація Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    data_collator=data_collator,
)

# Запуск fine-tuning
trainer.train()

# Зберігання моделі
trainer.save_model("./fine_tuned_model")