import os
import time
import torch
from datasets import Dataset
from transformers import (
    AutoTokenizer, 
    AutoModelForCausalLM, 
    TrainingArguments, 
    Trainer, 
    DataCollatorForLanguageModeling
)
from huggingface_hub import login

login(os.getenv("HUGGINGFACE_HUB_TOKEN"))

start_time = time.time()

# Створюємо директорії з правами на запис для всіх користувачів
os.makedirs('./cache', exist_ok=True)
os.chmod('./cache', 0o777)
os.makedirs('./results', exist_ok=True)
os.chmod('./results', 0o777)
os.makedirs('./fine_tuned_model', exist_ok=True)
os.chmod('./fine_tuned_model', 0o777)

# Завантаження моделі та токенізатора
model_name = "meta-llama/Meta-Llama-3.1-8B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

tokenizer.pad_token = tokenizer.eos_token

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

# Завантаження данних з локального тестового файлу
with open("m_and_m.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=32,
    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()

print("LLM is finetunned");

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

end_time = time.time()

training_time = end_time - start_time
print(f"Донавчання зайняло {training_time:.2f} секунд")

print("Донавчання пройшло вдало, модель сбережена у './fine_tuned_model'")