Update app.py

app.py

@@ -1,65 +1,94 @@
-from datasets import load_dataset
-from transformers import MarianMTModel, MarianTokenizer, TrainingArguments, Trainer, DataCollatorForSeq2Seq
+from transformers import MarianTokenizer, MarianMTModel, Seq2SeqTrainingArguments, Seq2SeqTrainer
+from datasets import Dataset, DatasetDict
+import pandas as pd
+import torch
 
-# Load dataset
-dataset = load_dataset('csv', data_files='hindi_dataset.tsv', delimiter='\t')
+# Load the dataset
+file_path = "hindi_dataset.tsv"  # Update with your actual file path
+data = pd.read_csv(file_path, delimiter="\t")
 
-# Load MarianMT tokenizer for translation task
-tokenizer = MarianTokenizer.from_pretrained('Helsinki-NLP/opus-mt-en-hi')
+# Convert the dataset to Hugging Face Dataset
+hf_dataset = Dataset.from_pandas(data)
 
-# Tokenize the English text (source language)
-def tokenize_function(examples):
-    return tokenizer(examples['english'], truncation=True, padding='max_length', max_length=128)
+# Split the dataset into train and test subsets
+split_dataset = hf_dataset.train_test_split(test_size=0.2)
 
-# Tokenize both English and Hindi sentences
-tokenized_datasets = dataset.map(tokenize_function, batched=True)
+# Create a DatasetDict with train and test splits
+dataset = DatasetDict({
+    "train": split_dataset["train"],
+    "test": split_dataset["test"]
+})
 
-def tokenize_target_function(examples):
-    return tokenizer(examples['hindi'], truncation=True, padding='max_length', max_length=128)
+# Load the tokenizer and model
+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)
 
-tokenized_datasets = tokenized_datasets.map(tokenize_target_function, batched=True)
+# Tokenize source (English) text
+def tokenize_function(examples):
+    return tokenizer(examples['source'], truncation=True, padding='max_length', max_length=128)
 
-# Data Collator for padding sequences
-data_collator = DataCollatorForSeq2Seq(tokenizer, model=None)
+# Tokenize target (Hindi) text
+def tokenize_target_function(examples):
+    with tokenizer.as_target_tokenizer():
+        return tokenizer(examples['target'], truncation=True, padding='max_length', max_length=128)
 
-# Load MarianMT model for translation
-model = MarianMTModel.from_pretrained('Helsinki-NLP/opus-mt-en-hi')
+# Apply tokenization to the dataset
+tokenized_datasets = dataset.map(tokenize_function, batched=True)
+tokenized_datasets = tokenized_datasets.map(tokenize_target_function, batched=True)
 
-# Define training arguments
-training_args = TrainingArguments(
-    output_dir='./results',
-    evaluation_strategy="epoch",
+# Define the training arguments
+training_args = Seq2SeqTrainingArguments(
+    output_dir="./results",
+    eval_strategy="epoch",
     learning_rate=2e-5,
     per_device_train_batch_size=16,
     per_device_eval_batch_size=16,
     num_train_epochs=3,
     weight_decay=0.01,
-    save_total_limit=2,
+    save_total_limit=3,
     predict_with_generate=True,
+    logging_dir="./logs",
+    logging_steps=10,
+    save_steps=500
 )
 
-# Initialize Trainer
-trainer = Trainer(
+# 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}
+
+# Define the Trainer
+trainer = Seq2SeqTrainer(
     model=model,
     args=training_args,
     train_dataset=tokenized_datasets['train'],
     eval_dataset=tokenized_datasets['test'],
     tokenizer=tokenizer,
-    data_collator=data_collator,
+    data_collator=data_collator
 )
 
-# Start training
+# Train the model
 trainer.train()
 
-# Save the model
-trainer.save_model('./my_hindi_translation_model')
-
 # Evaluate the model
-results = trainer.evaluate()
-print(results)
+eval_results = trainer.evaluate()
+print("Evaluation Results:", eval_results)
+
+# Test the model with sample inputs
+def translate_text(text):
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=128)
+    translated = model.generate(**inputs)
+    return [tokenizer.decode(t, skip_special_tokens=True) for t in translated]
 
-# Generate a prediction
-model.eval()
-inputs = tokenizer("How are you?", return_tensors="pt")
-outputs = model.generate(inputs["input_ids"], max_length=128)
-print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+# Test translation
+sample_text = "How are you?"
+hindi_translation = translate_text(sample_text)
+print(f"English: {sample_text}")
+print(f"Hindi: {hindi_translation[0]}")