# Fine-tuning on Hugging Face Notebook
# This runs in a HF notebook with GPU

# Install requirements
!pip install transformers datasets pandas scikit-learn accelerate

import pandas as pd
from transformers import (
    AutoTokenizer, 
    AutoModelForSeq2SeqLM,
    DataCollatorForSeq2Seq,
    Seq2SeqTrainingArguments,
    Seq2SeqTrainer
)
from datasets import Dataset
from sklearn.model_selection import train_test_split

# Load your dataset
df = pd.read_csv('Data_1.csv')
print(f"Dataset loaded: {len(df)} rows")

# Data preprocessing
df = df.dropna()
train_df, val_df = train_test_split(df, test_size=0.2, random_state=42)

print(f"Training samples: {len(train_df)}")
print(f"Validation samples: {len(val_df)}")

# Load model and tokenizer
model_name = "Helsinki-NLP/opus-mt-en-fr"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# Preprocessing function
def preprocess_function(examples):
    inputs = []
    for i in range(len(examples['year'])):
        input_text = f"Generate French title: Year: {examples['year'][i]} City: {examples['city'][i]} Event: {examples['original_title'][i]}"
        inputs.append(input_text)
    
    targets = examples['target_title']
    
    model_inputs = tokenizer(inputs, max_length=128, truncation=True, padding=True)
    with tokenizer.as_target_tokenizer():
        labels = tokenizer(targets, max_length=128, truncation=True, padding=True)
    
    model_inputs["labels"] = labels["input_ids"]
    return model_inputs

# Prepare datasets
train_dataset = Dataset.from_pandas(train_df)
val_dataset = Dataset.from_pandas(val_df)

tokenized_train = train_dataset.map(preprocess_function, batched=True)
tokenized_val = val_dataset.map(preprocess_function, batched=True)

# Training arguments
training_args = Seq2SeqTrainingArguments(
    output_dir="./french-title-model",
    evaluation_strategy="epoch",
    learning_rate=5e-5,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    weight_decay=0.01,
    save_total_limit=3,
    num_train_epochs=10,
    predict_with_generate=True,
    push_to_hub=True,
    hub_model_id="sufya12/french-title-model",  # Change to your username
)

# Data collator
data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)

# Trainer
trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train,
    eval_dataset=tokenized_val,
    tokenizer=tokenizer,
    data_collator=data_collator,
)

# Start training
print("Starting training...")
trainer.train()

# Save and push to Hub
trainer.push_to_hub()
print("Training completed and model pushed to Hugging Face Hub!")