Update app.py

app.py

@@ -1,103 +1,65 @@
-import os
-import pandas as pd
-from torch.utils.data import Dataset
-from transformers import (
-    MBartTokenizer,
-    MBartForConditionalGeneration,
-    Trainer,
-    TrainingArguments,
-    
-)
-from huggingface_hub import HfFolder
-# Save the Hugging Face token (if not already saved)
-token = os.getenv("HF_TOKEN")
-if token:
-    HfFolder.save_token(token)
-    print("Token saved successfully!")
-else:
-    print("HF_TOKEN environment variable not set. Ensure your token is saved for authentication.")
+from datasets import load_dataset
+from transformers import MarianMTModel, MarianTokenizer, TrainingArguments, Trainer, DataCollatorForSeq2Seq
 
-# Step 1: Define Dataset Class
-class HindiDataset(Dataset):
-    def __init__(self, data_path, tokenizer, max_length=512):
-        """
-        Dataset class for Hindi translation tasks.
+# Load dataset
+dataset = load_dataset('csv', data_files='dataset.tsv', delimiter='\t')
 
-        Args:
-            data_path (str): Path to the dataset file (e.g., TSV with source-target pairs).
-            tokenizer (MBartTokenizer): Tokenizer for mBART.
-            max_length (int): Maximum sequence length for tokenization.
-        """
-        self.data = pd.read_csv(data_path, sep="\t")
-        self.tokenizer = tokenizer
-        self.max_length = max_length
+# Load MarianMT tokenizer for translation task
+tokenizer = MarianTokenizer.from_pretrained('Helsinki-NLP/opus-mt-en-hi')
 
-    def __len__(self):
-        return len(self.data)
+# Tokenize the English text (source language)
+def tokenize_function(examples):
+    return tokenizer(examples['english'], truncation=True, padding='max_length', max_length=128)
 
-    def __getitem__(self, idx):
-        source = self.data.iloc[idx]["source"]
-        target = self.data.iloc[idx]["target"]
+# Tokenize both English and Hindi sentences
+tokenized_datasets = dataset.map(tokenize_function, batched=True)
 
-        source_encodings = self.tokenizer(
-            source, max_length=self.max_length, truncation=True, padding="max_length", return_tensors="pt"
-        )
-        target_encodings = self.tokenizer(
-            target, max_length=self.max_length, truncation=True, padding="max_length", return_tensors="pt"
-        )
+def tokenize_target_function(examples):
+    return tokenizer(examples['hindi'], truncation=True, padding='max_length', max_length=128)
 
-        return {
-            "input_ids": source_encodings["input_ids"].squeeze(),
-            "attention_mask": source_encodings["attention_mask"].squeeze(),
-            "labels": target_encodings["input_ids"].squeeze(),
-        }
+tokenized_datasets = tokenized_datasets.map(tokenize_target_function, batched=True)
 
-# Step 2: Load Tokenizer and Dataset
-data_path = "hindi_dataset.tsv"  # Path to your dataset file
-tokenizer = MBartTokenizer.from_pretrained("facebook/mbart-large-50")
-train_dataset = HindiDataset(data_path, tokenizer)
+# Data Collator for padding sequences
+data_collator = DataCollatorForSeq2Seq(tokenizer, model=None)
 
-# Step 3: Load Pre-trained mBART Model
-model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50")
+# Load MarianMT model for translation
+model = MarianMTModel.from_pretrained('Helsinki-NLP/opus-mt-en-hi')
 
-# Step 4: Define Training Arguments
+# Define training arguments
 training_args = TrainingArguments(
-    output_dir="./mbart-hindi",           # Output directory for model checkpoints
-    per_device_train_batch_size=4,       # Training batch size per GPU
-    per_device_eval_batch_size=4,        # Evaluation batch size per GPU
-    evaluation_strategy="steps",         # Evaluate every 'save_steps'
-    save_steps=500,                      # Save model every 500 steps
-    save_total_limit=2,                  # Keep only 2 checkpoints
-    logging_dir="./logs",                # Directory for training logs
-    num_train_epochs=3,                  # Number of training epochs
-    learning_rate=5e-5,                  # Learning rate
-    weight_decay=0.01,                   # Weight decay for optimizer
-    report_to="none"                     # Disable third-party logging
+    output_dir='./results',
+    evaluation_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,
+    predict_with_generate=True,
 )
 
-# Step 5: Initialize Trainer
+# Initialize Trainer
 trainer = Trainer(
     model=model,
     args=training_args,
-    train_dataset=train_dataset,
-    tokenizer=tokenizer
+    train_dataset=tokenized_datasets['train'],
+    eval_dataset=tokenized_datasets['test'],
+    tokenizer=tokenizer,
+    data_collator=data_collator,
 )
 
-# Step 6: Train the Model
-print("Starting training...")
+# Start training
 trainer.train()
 
-# Step 7: Save the Fine-Tuned Model
-output_dir = "./mbart-hindi-model"
-print(f"Saving fine-tuned model to {output_dir}...")
-trainer.save_model(output_dir)
+# Save the model
+trainer.save_model('./my_hindi_translation_model')
 
-# Step 8: Test the Fine-Tuned Model
-print("Testing the fine-tuned model...")
-model = MBartForConditionalGeneration.from_pretrained(output_dir)
-tokenizer = MBartTokenizer.from_pretrained(output_dir)
+# Evaluate the model
+results = trainer.evaluate()
+print(results)
 
-test_text = "Translate this to Hindi."
-inputs = tokenizer(test_text, return_tensors="pt")
-outputs = model.generate(**inputs)
-print("Generated Translation:", tokenizer.decode(outputs[0], skip_special_tokens=True))
+# 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))