# scripts/train_text_generation.py
from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
from datasets import load_dataset

# Replace with a public model available on Hugging Face
tokenizer = AutoTokenizer.from_pretrained("aubmindlab/bert-base-arabertv2")
model = AutoModelForCausalLM.from_pretrained("aubmindlab/bert-base-arabertv2")

# Load and tokenize dataset
dataset = load_dataset("text", data_files="../data/cleaned_data/iraqi_data_cleaned.txt")

def tokenize_function(examples):
    return tokenizer(examples["text"], return_tensors="pt", truncation=True)

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

# Training arguments
training_args = TrainingArguments(
    output_dir="../models/text_generation_model",
    evaluation_strategy="epoch",
    per_device_train_batch_size=2,
    num_train_epochs=3
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset['train']
)

# Train the model
trainer.train()

# Save the model after training
trainer.save_model("../models/text_generation_model")
tokenizer.save_pretrained("../models/text_generation_model")