from transformers import T5Tokenizer, T5ForConditionalGeneration, Trainer, TrainingArguments, DataCollatorForSeq2Seq
from datasets import load_dataset
import torch

# ✅ Load dataset
dataset = load_dataset("json", data_files="data/hrms_dataset.jsonl")["train"]  # Fixed path for typical execution

# ✅ Initialize tokenizer and model
model_name = "t5-small"
# model_name="models/t5_model"
tokenizer = T5Tokenizer.from_pretrained(model_name)
model = T5ForConditionalGeneration.from_pretrained(model_name)

# ✅ Preprocessing
def preprocess(example):
    input_enc = tokenizer(example["question"], padding="max_length", truncation=True, max_length=64)
    target_enc = tokenizer(example["sql"], padding="max_length", truncation=True, max_length=64)

    return {
        "input_ids": input_enc["input_ids"],
        "attention_mask": input_enc["attention_mask"],
        "labels": target_enc["input_ids"]
    }

encoded_dataset = dataset.map(preprocess, remove_columns=dataset.column_names)

# ✅ Training arguments
training_args = TrainingArguments(
    output_dir="models/t5_model",        # Saved model folder
    num_train_epochs=10,                 # Better results with more epochs
    per_device_train_batch_size=4,
    save_strategy="epoch",
    save_total_limit=2,
    logging_steps=50,
    remove_unused_columns=False,
    report_to="none",                    # Prevent wandb or hub errors
    fp16=False                           # Keep off if using CPU
)

# ✅ Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=encoded_dataset,
    tokenizer=tokenizer,
    data_collator=DataCollatorForSeq2Seq(tokenizer, model)
)

# ✅ Train and save
trainer.train()
# model.save_pretrained("models/t5_model_v2")
# tokenizer.save_pretrained("models/t5_model_v2")

trainer.model.save_pretrained("models/t5_model")
tokenizer.save_pretrained("models/t5_model")

print("✅ Model trained and saved in models/t5_model_v2/")