Upload 7 files

src/evaluate_sinhala.py

@@ -0,0 +1,58 @@
+# src/evaluate_sinhala.py
+
+import torch
+import evaluate # The new, preferred Hugging Face library for metrics
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from tqdm import tqdm # A library to create smart progress bars
+
+def evaluate_model():
+    """
+    Loads a fine-tuned model and evaluates its performance on the test set using the BLEU score.
+    """
+    # --- 1. Configuration ---
+    MODEL_PATH = "thilina/mt5-sinhalese-english"
+    TEST_DIR = "data/test_sets"
+    SOURCE_LANG_FILE = f"{TEST_DIR}/test.si"
+    TARGET_LANG_FILE = f"{TEST_DIR}/test.en"
+    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    # --- 2. Load Model, Tokenizer, and Metric ---
+    print("Loading model, tokenizer, and evaluation metric...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+    model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_PATH).to(DEVICE)
+    bleu_metric = evaluate.load("sacrebleu")
+
+    # --- 3. Load Test Data ---
+    with open(SOURCE_LANG_FILE, "r", encoding="utf-8") as f:
+        source_sentences = [line.strip() for line in f.readlines()]
+    with open(TARGET_LANG_FILE, "r", encoding="utf-8") as f:
+        # The BLEU metric expects references to be a list of lists
+        reference_translations = [[line.strip()] for line in f.readlines()]
+
+    # --- 4. Generate Predictions ---
+    print(f"Generating translations for {len(source_sentences)} test sentences...")
+    predictions = []
+    for sentence in tqdm(source_sentences):
+        inputs = tokenizer(sentence, return_tensors="pt").to(DEVICE)
+        
+        generated_tokens = model.generate(
+            **inputs,
+            max_length=128
+        )
+        
+        translation = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+        predictions.append(translation)
+
+    # --- 5. Compute BLEU Score ---
+    print("Calculating BLEU score...")
+    results = bleu_metric.compute(predictions=predictions, references=reference_translations)
+    
+    # The result is a dictionary. The 'score' key holds the main BLEU score.
+    bleu_score = results["score"]
+
+    print("\n--- Evaluation Complete ---")
+    print(f"BLEU Score: {bleu_score:.2f}")
+    print("---------------------------")
+
+if __name__ == "__main__":
+    evaluate_model()

src/evaluation.py

@@ -0,0 +1,64 @@
+# src/evaluate.py
+
+import torch
+import evaluate # The new, preferred Hugging Face library for metrics
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from tqdm import tqdm # A library to create smart progress bars
+import argparse
+
+def evaluate_model():
+    """
+    Loads a fine-tuned model and evaluates its performance on the test set using the BLEU score.
+    """
+    parser = argparse.ArgumentParser(description="Evaluate a translation model.")
+    parser.add_argument("--model_path", type=str, required=True, help="Path to the fine-tuned model directory")
+    parser.add_argument("--source_lang_file", type=str, required=True, help="Path to the source language test file")
+    parser.add_argument("--target_lang_file", type=str, required=True, help="Path to the target language test file")
+    parser.add_argument("--source_lang_tokenizer", type=str, required=True, help="Source language code for tokenizer (e.g., 'nep_Npan')")
+    args = parser.parse_args()
+
+    # --- 1. Configuration ---
+    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    # --- 2. Load Model, Tokenizer, and Metric ---
+    print("Loading model, tokenizer, and evaluation metric...")
+    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
+    model = AutoModelForSeq2SeqLM.from_pretrained(args.model_path).to(DEVICE)
+    bleu_metric = evaluate.load("sacrebleu")
+
+    # --- 3. Load Test Data ---
+    with open(args.source_lang_file, "r", encoding="utf-8") as f:
+        source_sentences = [line.strip() for line in f.readlines()]
+    with open(args.target_lang_file, "r", encoding="utf-8") as f:
+        # The BLEU metric expects references to be a list of lists
+        reference_translations = [[line.strip()] for line in f.readlines()]
+
+    # --- 4. Generate Predictions ---
+    print(f"Generating translations for {len(source_sentences)} test sentences...")
+    predictions = []
+    for sentence in tqdm(source_sentences):
+        tokenizer.src_lang = args.source_lang_tokenizer
+        inputs = tokenizer(sentence, return_tensors="pt").to(DEVICE)
+        
+        generated_tokens = model.generate(
+            **inputs,
+            forced_bos_token_id=tokenizer.vocab["eng_Latn"],
+            max_length=128
+        )
+        
+        translation = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+        predictions.append(translation)
+
+    # --- 5. Compute BLEU Score ---
+    print("Calculating BLEU score...")
+    results = bleu_metric.compute(predictions=predictions, references=reference_translations)
+    
+    # The result is a dictionary. The 'score' key holds the main BLEU score.
+    bleu_score = results["score"]
+
+    print("\n--- Evaluation Complete ---")
+    print(f"BLEU Score: {bleu_score:.2f}")
+    print("---------------------------")
+
+if __name__ == "__main__":
+    evaluate_model()

src/train.py

@@ -0,0 +1,109 @@
+# src/train.py
+
+import os
+import argparse
+from datasets import Dataset
+from transformers import (
+    AutoModelForSeq2SeqLM,
+    AutoTokenizer,
+    DataCollatorForSeq2Seq,
+    Seq2SeqTrainingArguments,
+    Seq2SeqTrainer,
+)
+
+def train_model():
+    """
+    Fine-tunes a pre-trained NLLB model on a parallel dataset.
+    """
+    parser = argparse.ArgumentParser(description="Fine-tune a translation model.")
+    parser.add_argument("--model_checkpoint", type=str, default="facebook/nllb-200-distilled-600M")
+    parser.add_argument("--source_lang", type=str, required=True, help="Source language code (e.g., 'ne')")
+    parser.add_argument("--target_lang", type=str, default="en")
+    parser.add_argument("--source_lang_tokenizer", type=str, required=True, help="Source language code for tokenizer (e.g., 'nep_Npan')")
+    parser.add_argument("--train_file_source", type=str, required=True, help="Path to the source language training file")
+    parser.add_argument("--train_file_target", type=str, required=True, help="Path to the target language training file")
+    parser.add_argument("--output_dir", type=str, required=True, help="Directory to save the fine-tuned model")
+    parser.add_argument("--epochs", type=int, default=3)
+    parser.add_argument("--batch_size", type=int, default=8)
+
+    args = parser.parse_args()
+
+    # --- 1. Configuration ---
+    MODEL_CHECKPOINT = args.model_checkpoint
+    SOURCE_LANG = args.source_lang
+    TARGET_LANG = args.target_lang
+    MODEL_OUTPUT_DIR = args.output_dir
+
+    # --- 2. Load Tokenizer and Model ---
+    print("Loading tokenizer and model...")
+    tokenizer = AutoTokenizer.from_pretrained(
+        MODEL_CHECKPOINT, src_lang=args.source_lang_tokenizer, tgt_lang="eng_Latn"
+    )
+    model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_CHECKPOINT)
+
+    # --- 3. Load and Preprocess Data (Memory-Efficiently) ---
+    print("Loading and preprocessing data...")
+    
+    def generate_examples():
+        with open(args.train_file_source, "r", encoding="utf-8") as f_src, \
+             open(args.train_file_target, "r", encoding="utf-8") as f_tgt:
+            for src_line, tgt_line in zip(f_src, f_tgt):
+                yield {"translation": {SOURCE_LANG: src_line.strip(), TARGET_LANG: tgt_line.strip()}}
+
+    dataset = Dataset.from_generator(generate_examples)
+    
+    split_datasets = dataset.train_test_split(train_size=0.95, seed=42)
+    split_datasets["validation"] = split_datasets.pop("test")
+
+    def preprocess_function(examples):
+        inputs = [ex[SOURCE_LANG] for ex in examples["translation"]]
+        targets = [ex[TARGET_LANG] for ex in examples["translation"]]
+        
+        model_inputs = tokenizer(inputs, text_target=targets, max_length=128, truncation=True)
+        return model_inputs
+
+    tokenized_datasets = split_datasets.map(
+        preprocess_function,
+        batched=True,
+        remove_columns=split_datasets["train"].column_names,
+    )
+
+    # --- 4. Set Up Training Arguments ---
+    print("Setting up training arguments...")
+    training_args = Seq2SeqTrainingArguments(
+        output_dir=MODEL_OUTPUT_DIR,
+        eval_strategy="epoch",
+        learning_rate=2e-5,
+        per_device_train_batch_size=args.batch_size,
+        per_device_eval_batch_size=args.batch_size,
+        weight_decay=0.01,
+        save_total_limit=3,
+        num_train_epochs=args.epochs,
+        predict_with_generate=True,
+        fp16=False, # Set to True if you have a compatible GPU
+    )
+
+    # --- 5. Create the Trainer ---
+    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)
+    
+    trainer = Seq2SeqTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_datasets["train"],
+        eval_dataset=tokenized_datasets["validation"],
+        tokenizer=tokenizer,
+        data_collator=data_collator,
+    )
+
+    # --- 6. Start Training ---
+    print("\n--- Starting model fine-tuning ---")
+    trainer.train()
+    print("--- Training complete ---")
+
+    # --- 7. Save the Final Model ---
+    print(f"Saving final model to {MODEL_OUTPUT_DIR}")
+    trainer.save_model()
+    print("Model saved successfully!")
+
+if __name__ == "__main__":
+    train_model()

src/train_nepali.py

@@ -0,0 +1,95 @@
+# src/train_nepali.py
+
+import os
+from datasets import load_dataset, DatasetDict, concatenate_datasets
+from transformers import (
+    AutoModelForSeq2SeqLM,
+    AutoTokenizer,
+    DataCollatorForSeq2Seq,
+    Seq2SeqTrainingArguments,
+    Seq2SeqTrainer,
+)
+
+def train_nepali_model():
+    """
+    Fine-tunes a pre-trained NLLB model on the Nepali parallel dataset.
+    """
+    # --- 1. Configuration ---
+    MODEL_CHECKPOINT = "facebook/nllb-200-distilled-600M"
+    DATA_DIR = "data/processed"
+    MODEL_OUTPUT_DIR = "D:\\SIH\\models\\nllb-finetuned-nepali-en"
+
+    # --- 2. Load Tokenizer and Model ---
+    print("Loading tokenizer and model...")
+    tokenizer = AutoTokenizer.from_pretrained(
+        MODEL_CHECKPOINT, src_lang="nep_Npan", tgt_lang="eng_Latn"
+    )
+    model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_CHECKPOINT)
+
+    # --- 3. Load and Preprocess Data ---
+    print("Loading and preprocessing data...")
+    nepali_dataset = load_dataset("text", data_files=os.path.join(DATA_DIR, "nepali.ne"))["train"]
+    english_dataset = load_dataset("text", data_files=os.path.join(DATA_DIR, "nepali.en"))["train"]
+
+    # rename the 'text' column to 'ne' and 'en'
+    nepali_dataset = nepali_dataset.rename_column("text", "ne")
+    english_dataset = english_dataset.rename_column("text", "en")
+
+    # combine the datasets
+    raw_datasets = concatenate_datasets([nepali_dataset, english_dataset], axis=1)
+    
+    split_datasets = raw_datasets.train_test_split(train_size=0.95, seed=42)
+    split_datasets["validation"] = split_datasets.pop("test")
+
+    def preprocess_function(examples):
+        inputs = examples["ne"]
+        targets = examples["en"]
+        
+        model_inputs = tokenizer(inputs, text_target=targets, max_length=128, truncation=True)
+        return model_inputs
+
+    tokenized_datasets = split_datasets.map(
+        preprocess_function,
+        batched=True,
+        remove_columns=split_datasets["train"].column_names,
+    )
+
+    # --- 4. Set Up Training Arguments ---
+    print("Setting up training arguments...")
+    training_args = Seq2SeqTrainingArguments(
+        output_dir=MODEL_OUTPUT_DIR,
+        eval_strategy="epoch",
+        learning_rate=2e-5,
+        per_device_train_batch_size=8,
+        per_device_eval_batch_size=8,
+        weight_decay=0.01,
+        save_total_limit=3,
+        num_train_epochs=3, # Reduced for faster training, can be increased
+        predict_with_generate=True,
+        fp16=False, # Set to True if you have a compatible GPU
+    )
+
+    # --- 5. Create the Trainer ---
+    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)
+    
+    trainer = Seq2SeqTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_datasets["train"],
+        eval_dataset=tokenized_datasets["validation"],
+        tokenizer=tokenizer,
+        data_collator=data_collator,
+    )
+
+    # --- 6. Start Training ---
+    print(f"\n--- Starting model fine-tuning for Nepali-English ---")
+    trainer.train()
+    print("--- Training complete ---")
+
+    # --- 7. Save the Final Model ---
+    print(f"Saving final model to {MODEL_OUTPUT_DIR}")
+    trainer.save_model()
+    print("Model saved successfully!")
+
+if __name__ == "__main__":
+    train_nepali_model()

src/translate.py

@@ -0,0 +1,52 @@
+# src/translate.py
+
+# src/translate.py
+
+import torch
+from transformers import MBartForConditionalGeneration, NllbTokenizer
+import argparse
+
+# --- 1. Configuration ---
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# --- 2. Load Models and Tokenizers ---
+print(f"Loading models on {DEVICE.upper()}...")
+models = {
+    "nepali": MBartForConditionalGeneration.from_pretrained("models/nllb-finetuned-nepali-en").to(DEVICE)
+}
+tokenizers = {
+    "nepali": NllbTokenizer.from_pretrained("models/nllb-finetuned-nepali-en")
+}
+print("All models loaded successfully!")
+
+def translate_text(text_to_translate: str, source_language: str) -> str:
+    """
+    Translates a single string of text to English using our fine-tuned models.
+    """
+    model = models[source_language]
+    tokenizer = tokenizers[source_language]
+
+    tokenizer.src_lang = "nep_Npan"
+
+    inputs = tokenizer(text_to_translate, return_tensors="pt").to(DEVICE)
+
+    generated_tokens = model.generate(
+        **inputs,
+        forced_bos_token_id=tokenizer.convert_tokens_to_ids("eng_Latn"),
+        max_length=128
+    )
+
+    translation = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+    return translation
+
+# --- 3. Example Usage ---
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Translate text using a fine-tuned model.")
+    parser.add_argument("--text", type=str, required=True, help="Text to translate.")
+    parser.add_argument("--lang", type=str, required=True, choices=["nepali"], help="Source language: 'nepali'.")
+    args = parser.parse_args()
+
+    translated_sentence = translate_text(args.text, args.lang)
+    
+    print(f"\nOriginal ({args.lang}): {args.text}")
+    print(f"Translated (en): {translated_sentence}")