Update app.py

app.py

@@ -1,7 +1,6 @@
-from transformers import MarianTokenizer, MarianMTModel, Seq2SeqTrainingArguments, Seq2SeqTrainer
+from transformers import MarianTokenizer, MarianMTModel, Seq2SeqTrainingArguments, Seq2SeqTrainer, DataCollatorForSeq2Seq
 from datasets import Dataset, DatasetDict
 import pandas as pd
-import torch
 
 # Load the dataset
 file_path = "hindi_dataset.tsv"  # Update with your actual file path
@@ -24,23 +23,21 @@ model_name = "Helsinki-NLP/opus-mt-en-hi"  # Pre-trained English-to-Hindi model
 tokenizer = MarianTokenizer.from_pretrained(model_name)
 model = MarianMTModel.from_pretrained(model_name)
 
-# Tokenize source (English) text
+# Tokenize source and target text
 def tokenize_function(examples):
-    return tokenizer(examples['source'], truncation=True, padding='max_length', max_length=128)
-
-# Tokenize target (Hindi) text
-def tokenize_target_function(examples):
+    model_inputs = tokenizer(examples['english'], truncation=True, padding='max_length', max_length=128)
     with tokenizer.as_target_tokenizer():
-        return tokenizer(examples['target'], truncation=True, padding='max_length', max_length=128)
+        labels = tokenizer(examples['hindi'], truncation=True, padding='max_length', max_length=128)
+        model_inputs['labels'] = labels['input_ids']
+    return model_inputs
 
 # Apply tokenization to the dataset
 tokenized_datasets = dataset.map(tokenize_function, batched=True)
-tokenized_datasets = tokenized_datasets.map(tokenize_target_function, batched=True)
 
 # Define the training arguments
 training_args = Seq2SeqTrainingArguments(
     output_dir="./results",
-    eval_strategy="epoch",
+    evaluation_strategy="epoch",
     learning_rate=2e-5,
     per_device_train_batch_size=16,
     per_device_eval_batch_size=16,
@@ -53,16 +50,8 @@ training_args = Seq2SeqTrainingArguments(
     save_steps=500
 )
 
-# Data collator to pad sequences to the same length
-def data_collator(features):
-    keys = ["input_ids", "attention_mask", "labels"]
-    max_length = max(len(feature[key]) for feature in features for key in keys if key in feature)
-    for feature in features:
-        for key in keys:
-            if key in feature:
-                padding = [0] * (max_length - len(feature[key]))
-                feature[key].extend(padding)
-    return {key: torch.tensor([f[key] for f in features]) for key in keys}
+# Use the DataCollatorForSeq2Seq for padding
+data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)
 
 # Define the Trainer
 trainer = Seq2SeqTrainer(