Update train_model.py

train_model.py

@@ -1,203 +1,105 @@
 import torch
-from transformers import (
-    AutoModelForSequenceClassification,
-    AutoTokenizer, 
-    TrainingArguments, 
-    Trainer,
-    DataCollatorWithPadding
-)
-from datasets import Dataset
-import os
-import glob
+from torch.utils.data import Dataset, DataLoader
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, AdamW
+import torch.nn.functional as F
 import json
-from typing import List, Dict
-import numpy as np
-
-class KiswahiliBERTTrainer:
-    def __init__(self, model_name="bert-base-multilingual-cased"):
-        """
-        BERT-based chatbot trainer for response selection
-        """
-        self.model_name = model_name
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        
-        # Initialize model
-        self.model = AutoModelForSequenceClassification.from_pretrained(
-            model_name,
-            num_labels=2  # Binary classification: good response vs bad response
-        )
-        
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.model.to(self.device)
-        self.responses = []  # Store predefined responses
-        print("✅ BERT Model and tokenizer loaded successfully!")
-    
-    def load_training_data(self, data_dir="./training_data"):
-        """Load training data and create response selection pairs"""
-        conversation_pairs = []
-        text_files = glob.glob(os.path.join(data_dir, "*.txt"))
-        
-        if not text_files:
-            print(f"⚠️ Hakuna faili za .txt katika {data_dir}")
-            return []
-        
-        for file_path in text_files:
-            print(f"📖 Inapakia data kutoka: {file_path}")
-            try:
-                with open(file_path, 'r', encoding='utf-8') as f:
-                    content = f.read().strip()
-                    
-                    # Split by conversation blocks
-                    blocks = [b.strip() for b in content.split('\n\n') if b.strip()]
-                    
-                    for block in blocks:
-                        lines = block.split('\n')
-                        user_input = None
-                        assistant_response = None
-                        
-                        for line in lines:
-                            if line.startswith('User:'):
-                                user_input = line.replace('User:', '').strip()
-                            elif line.startswith('Assistant:'):
-                                assistant_response = line.replace('Assistant:', '').strip()
-                        
-                        if user_input and assistant_response:
-                            # Store as positive example
-                            conversation_pairs.append({
-                                'user_input': user_input,
-                                'response': assistant_response,
-                                'label': 1  # Positive example
-                            })
-                            # Also store the response for later use
-                            if assistant_response not in self.responses:
-                                self.responses.append(assistant_response)
-                            
-            except Exception as e:
-                print(f"❌ Hitilafu wakati wa kusoma {file_path}: {e}")
-        
-        print(f"📊 Imepakika jozi {len(conversation_pairs)} za mazungumzo")
-        print(f"📝 Imepatikana majibu {len(self.responses)} ya kipekee")
-        return conversation_pairs
+import os
+from sklearn.model_selection import train_test_split
+
+# --- Config ---
+MODEL_NAME = "bert-base-multilingual-cased"
+TRAINING_FILE = "./training_data/greetings.txt"
+SAVE_PATH = "./trained_bert_model"
+EPOCHS = 3
+BATCH_SIZE = 8
+MAX_LEN = 64
+LEARNING_RATE = 2e-5
+
+# --- Load training data ---
+def load_training_data(file_path):
+    inputs, responses = [], []
+    if not os.path.exists(file_path):
+        raise FileNotFoundError(f"{file_path} not found!")
+
+    with open(file_path, "r", encoding="utf-8") as f:
+        lines = [line.strip() for line in f if line.strip()]
     
-    def create_training_pairs(self, conversation_pairs):
-        """Create training data with positive and negative examples"""
-        training_examples = []
-        
-        for pair in conversation_pairs:
-            # Positive example
-            training_examples.append({
-                'text': f"{pair['user_input']} [SEP] {pair['response']}",
-                'label': 1
-            })
-            
-            # Create negative examples (random wrong responses)
-            for _ in range(2):  # 2 negative examples per positive
-                if len(self.responses) > 1:
-                    wrong_responses = [r for r in self.responses if r != pair['response']]
-                    if wrong_responses:
-                        wrong_response = np.random.choice(wrong_responses)
-                        training_examples.append({
-                            'text': f"{pair['user_input']} [SEP] {wrong_response}",
-                            'label': 0
-                        })
-        
-        return training_examples
+    for i in range(0, len(lines), 2):
+        user_input = lines[i].replace("User:", "").strip()
+        assistant_response = lines[i+1].replace("Assistant:", "").strip()
+        inputs.append(user_input)
+        responses.append(assistant_response)
     
-    def prepare_dataset(self, training_examples):
-        """Tokenize the dataset for BERT"""
-        texts = [example['text'] for example in training_examples]
-        labels = [example['label'] for example in training_examples]
-        
-        # Tokenize all texts
-        encodings = self.tokenizer(
-            texts,
+    return inputs, responses
+
+# --- Dataset ---
+class KiswahiliDataset(Dataset):
+    def __init__(self, inputs, responses, tokenizer):
+        self.inputs = inputs
+        self.responses = responses
+        self.tokenizer = tokenizer
+
+    def __len__(self):
+        return len(self.inputs)
+
+    def __getitem__(self, idx):
+        text = f"{self.inputs[idx]} [SEP] {self.responses[idx]}"
+        encoding = self.tokenizer(
+            text,
             truncation=True,
-            padding=True,
-            max_length=256,
-            return_tensors="pt"
-        )
-        
-        # Create dataset
-        class CustomDataset(torch.utils.data.Dataset):
-            def __init__(self, encodings, labels):
-                self.encodings = encodings
-                self.labels = labels
-            
-            def __getitem__(self, idx):
-                item = {key: val[idx] for key, val in self.encodings.items()}
-                item['labels'] = torch.tensor(self.labels[idx])
-                return item
-            
-            def __len__(self):
-                return len(self.labels)
-        
-        return CustomDataset(encodings, labels)
-    
-    def train(self, data_dir="./training_data", output_dir="./trained_bert_model"):
-        """Train BERT for response selection"""
-        conversation_pairs = self.load_training_data(data_dir=data_dir)
-        
-        if not conversation_pairs:
-            print("❌ Hakuna data ya mafunzo! Tafadhali weka faili za .txt katika training_data/")
-            return
-        
-        # Create training examples
-        training_examples = self.create_training_pairs(conversation_pairs)
-        dataset = self.prepare_dataset(training_examples)
-        
-        # BERT-specific training arguments
-        training_args = TrainingArguments(
-            output_dir=output_dir,
-            num_train_epochs=5,
-            per_device_train_batch_size=8,
-            per_device_eval_batch_size=8,
-            warmup_steps=100,
-            weight_decay=0.01,
-            logging_dir='./logs',
-            logging_steps=10,
-            evaluation_strategy="no",
-            save_strategy="epoch",
-            load_best_model_at_end=False,
-            fp16=torch.cuda.is_available(),
-        )
-        
-        data_collator = DataCollatorWithPadding(tokenizer=self.tokenizer)
-        
-        trainer = Trainer(
-            model=self.model,
-            args=training_args,
-            train_dataset=dataset,
-            tokenizer=self.tokenizer,
-            data_collator=data_collator,
+            padding='max_length',
+            max_length=MAX_LEN,
+            return_tensors='pt'
         )
-        
-        print("🚀 Mafunzo ya BERT yanaanza...")
-        trainer.train()
-        
-        print("💾 Inahifadhi modeli...")
-        trainer.save_model(output_dir)
-        self.tokenizer.save_pretrained(output_dir)
-        
-        # Save the response bank
-        response_data = {
-            'responses': self.responses,
-            'model_type': 'bert-response-selector'
+        # Label 1 = positive example
+        label = torch.tensor(1)
+        return {
+            'input_ids': encoding['input_ids'].squeeze(),
+            'attention_mask': encoding['attention_mask'].squeeze(),
+            'labels': label
         }
-        with open(os.path.join(output_dir, 'responses.json'), 'w', encoding='utf-8') as f:
-            json.dump(response_data, f, ensure_ascii=False, indent=2)
-        
-        print(f"✅ Mafunzo ya BERT yamekamilika! Modeli imehifadhiwa katika {output_dir}")
-        print(f"📋 Benki ya majibu imehifadhiwa na majibu {len(self.responses)}")
-
-# Alternative BERT models for Kiswahili
-BERT_MODEL_OPTIONS = {
-    "multilingual": "bert-base-multilingual-cased",
-    "afriberta": "castorini/afriberta_base",
-    "xlm-roberta": "xlm-roberta-base",
-    "swahili-bert": "ds4sd/ixbert-swahili-base",  # Swahili-specific BERT
-}
+
+# --- Main training ---
+def main():
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    inputs, responses = load_training_data(TRAINING_FILE)
+
+    dataset = KiswahiliDataset(inputs, responses, tokenizer)
+    train_loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
+
+    model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=2)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model.to(device)
+
+    optimizer = AdamW(model.parameters(), lr=LEARNING_RATE)
+
+    model.train()
+    for epoch in range(EPOCHS):
+        total_loss = 0
+        for batch in train_loader:
+            optimizer.zero_grad()
+            input_ids = batch['input_ids'].to(device)
+            attention_mask = batch['attention_mask'].to(device)
+            labels = batch['labels'].to(device)
+
+            outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
+            loss = outputs.loss
+            total_loss += loss.item()
+            loss.backward()
+            optimizer.step()
+        print(f"Epoch {epoch+1}/{EPOCHS} - Loss: {total_loss/len(train_loader):.4f}")
+
+    # Save model
+    if not os.path.exists(SAVE_PATH):
+        os.makedirs(SAVE_PATH)
+    model.save_pretrained(SAVE_PATH)
+    tokenizer.save_pretrained(SAVE_PATH)
+
+    # Save responses for chatbot
+    with open(os.path.join(SAVE_PATH, "responses.json"), "w", encoding="utf-8") as f:
+        json.dump({"responses": responses}, f, ensure_ascii=False, indent=4)
+
+    print(f"✅ Training complete. Model saved to {SAVE_PATH}")
 
 if __name__ == "__main__":
-    # Example usage with a Swahili-optimized BERT model
-    trainer = KiswahiliBERTTrainer(model_name=BERT_MODEL_OPTIONS["multilingual"])
-    trainer.train()
+    main()