Update training/train_trenslation.py

training/train_trenslation.py

@@ -1,52 +1,55 @@
-import os
-import json
-from datasets import load_dataset
-from transformers import MarianTokenizer, MarianMTModel, Seq2SeqTrainer, Seq2SeqTrainingArguments, DataCollatorForSeq2Seq
-
-MODEL_NAME = "Helsinki-NLP/opus-mt-ha-en"   # Hausa-English base model
-OUTPUT_DIR = "./training/outputs/model"
-
-def train_from_jsonl(jsonl_path):
-    dataset = load_dataset("json", data_files={"train": jsonl_path}, split="train")
-
-    # Train/validation split
-    dataset = dataset.train_test_split(test_size=0.1)
-
-    tokenizer = MarianTokenizer.from_pretrained(MODEL_NAME)
-    model = MarianMTModel.from_pretrained(MODEL_NAME)
-
-    def preprocess(batch):
-        inputs = tokenizer(batch["src"], truncation=True, padding="max_length", max_length=128)
-        targets = tokenizer(batch["tgt"], truncation=True, padding="max_length", max_length=128)
-        inputs["labels"] = targets["input_ids"]
-        return inputs
-
-    tokenized = dataset.map(preprocess, batched=True, remove_columns=["src", "tgt"])
-    data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
-
-    training_args = Seq2SeqTrainingArguments(
-        output_dir=OUTPUT_DIR,
-        evaluation_strategy="epoch",
-        learning_rate=5e-5,
-        per_device_train_batch_size=8,
-        per_device_eval_batch_size=8,
-        num_train_epochs=3,
-        weight_decay=0.01,
-        save_total_limit=2,
-        predict_with_generate=True,
-        logging_dir="./training/logs",
-    )
-
-    trainer = Seq2SeqTrainer(
-        model=model,
-        args=training_args,
-        train_dataset=tokenized["train"],
-        eval_dataset=tokenized["test"],
-        tokenizer=tokenizer,
-        data_collator=data_collator,
-    )
-
-    trainer.train()
-    trainer.save_model(OUTPUT_DIR)
-    tokenizer.save_pretrained(OUTPUT_DIR)
-    print("✅ Training complete. Model saved at", OUTPUT_DIR)
+# voice_translator/training/train_translation.py
+
+import os
+from datasets import load_dataset, Dataset
+from transformers import (
+    MarianTokenizer,
+    MarianMTModel,
+    Seq2SeqTrainingArguments,
+    Seq2SeqTrainer,
+    DataCollatorForSeq2Seq,
+)
+
+MODEL_NAME = "Helsinki-NLP/opus-mt-mul-en"
+OUTPUT_DIR = "./training/outputs/model"
+
+def train_from_jsonl(file_path):
+    # Load dataset
+    dataset = load_dataset("json", data_files=file_path, split="train")
+
+    tokenizer = MarianTokenizer.from_pretrained(MODEL_NAME)
+    model = MarianMTModel.from_pretrained(MODEL_NAME)
+
+    def preprocess(batch):
+        inputs = tokenizer(batch["src"], truncation=True, padding="max_length", max_length=128)
+        targets = tokenizer(batch["tgt"], truncation=True, padding="max_length", max_length=128)
+        inputs["labels"] = targets["input_ids"]
+        return inputs
+
+    tokenized = dataset.map(preprocess, batched=True)
+
+    collator = DataCollatorForSeq2Seq(tokenizer, model=model)
+
+    args = Seq2SeqTrainingArguments(
+        output_dir=OUTPUT_DIR,
+        evaluation_strategy="no",
+        learning_rate=5e-5,
+        per_device_train_batch_size=8,
+        num_train_epochs=3,
+        save_total_limit=1,
+        predict_with_generate=True,
+    )
+
+    trainer = Seq2SeqTrainer(
+        model=model,
+        args=args,
+        train_dataset=tokenized,
+        tokenizer=tokenizer,
+        data_collator=collator,
+    )
+
+    trainer.train()
+    trainer.save_model(OUTPUT_DIR)
+    tokenizer.save_pretrained(OUTPUT_DIR)
+
+    return f"✅ Model trained and saved to {OUTPUT_DIR}"