from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, TextDataset, DataCollatorForLanguageModeling
import torch

# Load the tokenizer and model
model_name = "aubmindlab/aragpt2-base"  # Arabic GPT-2 model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Add a padding token if not present
if tokenizer.pad_token is None:
    tokenizer.add_special_tokens({'pad_token': '[PAD]'})  # Adds a [PAD] token if it's not defined
    model.resize_token_embeddings(len(tokenizer))  # Resize model's embeddings to account for the new token

# Load and prepare the data for training
def load_dataset(file_path, tokenizer, block_size=128):
    # Read the text data from the file
    with open(file_path, "r", encoding="utf-8") as f:
        text = f.read()
    
    # Tokenize the text data
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=block_size)
    
    # Return as a PyTorch dataset
    dataset = TextDataset(tokenizer=tokenizer, file_path=file_path, block_size=block_size)
    return dataset

# Define the path to your text data file
data_file = "data/raw_data/iraqi_data.txt"  # Update this path if necessary
train_dataset = load_dataset(data_file, tokenizer)

# Set up training arguments
training_args = TrainingArguments(
    output_dir="./models",
    overwrite_output_dir=True,
    num_train_epochs=3,
    per_device_train_batch_size=2,
    save_steps=10_000,
    save_total_limit=2,
)

# Define the data collator for language modeling
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

# Initialize the Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    data_collator=data_collator,
)

# Save the trained model and tokenizer
trainer.save_model("models/iraqi_proverb_model")
tokenizer.save_pretrained("models/iraqi_proverb_model")
print("Training complete. Model and tokenizer saved to 'models/iraqi_proverb_model'.")