region:us
---
{}
---
import torch
from datasets import load_dataset
from transformers import (
    AutoTokenizer, 
    AutoModelForSeq2SeqLM,
    DataCollatorForSeq2Seq,
    Seq2SeqTrainingArguments,
    Seq2SeqTrainer
)
import numpy as np
from transformers.cache_utils import EncoderDecoderCache

# Load dataset
print("Loading dataset...")
dataset = load_dataset("SKNahin/bengali-transliteration-data")

# Create train-validation split
split_dataset = dataset["train"].train_test_split(test_size=0.1, seed=42)
train_dataset = split_dataset["train"]
val_dataset = split_dataset["test"]

# Initialize tokenizer and model
print("Initializing model and tokenizer...")
model_name = "google/mt5-small"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# Maximum length for input and output sequences
max_input_length = 128
max_target_length = 128

def preprocess_function(examples):
    inputs = examples['rm']
    targets = examples['bn']
    
    model_inputs = tokenizer(
        inputs, 
        max_length=max_input_length, 
        truncation=True, 
        padding="max_length"
    )
    
    labels = tokenizer(
        targets,
        max_length=max_target_length,
        truncation=True,
        padding="max_length"
    )
    
    model_inputs["labels"] = labels["input_ids"]
    return model_inputs

print("Preprocessing datasets...")
tokenized_train = train_dataset.map(
    preprocess_function,
    batched=True,
    remove_columns=train_dataset.column_names
)
tokenized_val = val_dataset.map(
    preprocess_function,
    batched=True,
    remove_columns=val_dataset.column_names
)

print("Setting up training arguments...")
training_args = Seq2SeqTrainingArguments(
    output_dir="./banglish-to-bengali-translator",
    evaluation_strategy="steps",
    eval_steps=100,
    logging_strategy="steps",
    logging_steps=50,
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    weight_decay=0.01,
    save_total_limit=3,
    num_train_epochs=3,
    predict_with_generate=True,
    fp16=True,
    push_to_hub=False,
    load_best_model_at_end=True,
    metric_for_best_model="eval_loss",
    report_to="tensorboard"
)

data_collator = DataCollatorForSeq2Seq(
    tokenizer=tokenizer,
    model=model,
    padding=True
)

def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
    
    accuracy = sum(1 for x, y in zip(decoded_preds, decoded_labels) if x.strip() == y.strip()) / len(decoded_labels)
    return {
        "accuracy": accuracy,
        "samples": decoded_preds[:5]  # Show first 5 predictions for monitoring
    }

print("Initializing trainer...")
trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train,
    eval_dataset=tokenized_val,
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

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

print("Saving model...")
model.save_pretrained("./final-model")
tokenizer.save_pretrained("./final-model")

def translate_text(text):
    inputs = tokenizer(text, return_tensors="pt", padding=True)
    outputs = model.generate(
        **inputs, 
        max_length=128,
        cache_class=EncoderDecoderCache
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Test the model
print("\nTesting the model:")
test_texts = [
    "ami tomake valobashi",
    "bangladesh amar desh",
    "kemon acho bondhu"
]

print("\nTranslation examples:")
for text in test_texts:
    translated = translate_text(text)
    print(f"Input: {text}")
    print(f"Output: {translated}\n")
