Update chat.py

chat.py

@@ -1,148 +1,73 @@
 import torch
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
-from typing import List, Tuple, Optional
 import logging
 import json
 import os
 import numpy as np
 import re
 
-# Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 class KiswahiliChatbot:
-    def __init__(self, model_name: str = "bert-base-multilingual-cased", device: str = None):
-        """
-        BERT-based Kiswahili chatbot with response selection
-        """
-        try:
-            self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
-            logger.info(f"Inatumia kifaa: {self.device}")
+    def __init__(self, model_path="./trained_bert_model", device=None, threshold=0.6):
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        logger.info(f"Using device: {self.device}")
 
-            logger.info(f"Inapakia modeli ya BERT '{model_name}'...")
-            
-            # Load model and tokenizer
-            model_path = "./trained_bert_model"
-            if os.path.exists(model_path):
-                self.tokenizer = AutoTokenizer.from_pretrained(model_path)
-                self.model = AutoModelForSequenceClassification.from_pretrained(model_path)
-                logger.info("✅ Modeli iliyofunzwa imepakika!")
-            else:
-                logger.info("ℹ️ Modeli ya msingi ya BERT inatumika")
-                self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-                self.model = AutoModelForSequenceClassification.from_pretrained(
-                    model_name,
-                    num_labels=2
-                )
-            
-            self.model.to(self.device)
-            self.model.eval()
-            
-            # Load response bank
-            self.responses = self._load_response_bank()
-            logger.info(f"📋 Benki ya majibu: {len(self.responses)} majibu")
-            
-        except Exception as e:
-            logger.error(f"❌ Hitilafu wakati wa kupakia modeli: {e}")
-            raise
+        # Load model
+        if os.path.exists(model_path):
+            self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+            self.model = AutoModelForSequenceClassification.from_pretrained(model_path)
+            logger.info("✅ Trained model loaded!")
+        else:
+            raise FileNotFoundError(f"{model_path} not found. Please train the model first.")
+
+        self.model.to(self.device)
+        self.model.eval()
+        self.threshold = threshold  # minimum probability to accept a response
+
+        # Load responses
+        self.responses = self._load_response_bank()
+        logger.info(f"📋 Loaded {len(self.responses)} responses")
 
     def _load_response_bank(self):
-        """Load response bank from file or use defaults"""
         response_file = "./trained_bert_model/responses.json"
-        responses = []
-        
         if os.path.exists(response_file):
-            try:
-                with open(response_file, 'r', encoding='utf-8') as f:
-                    data = json.load(f)
-                    responses = data.get('responses', [])
-            except Exception as e:
-                logger.error(f"❌ Hitilafu wakati wa kusoma faili ya majibu: {e}")
-        
-        # Add fallback responses if empty
-        if not responses:
-            responses = [
-                "Habari yako? Naitwa KiswahiliChetu, naweza kukusaidia na Kiswahili.",
-                "Asante kwa kuuliza! Ninafurahi kukusaidia na maswali yako ya Kiswahili.",
-                "Samahani, sielewi swali lako. Unaweza kuuliza kwa Kiswahili?",
-                "Ninajua Kiswahili vizuri. Nitaweza kukujibu maswali yako.",
-                "Tanzania ni nchi nzuri yenye utamaduni mwingi na lugha ya Kiswahili.",
-                "Hakuna matata inamaanisha 'hamna shida' kwa Kiswahili.",
-                "Unauliza kuhusu nini hasa? Ninaweza kukusaidia na Kiswahili.",
-                "Karibu katika masomo ya Kiswahili! Nianzie na swali lako."
-            ]
-        
-        return responses
+            with open(response_file, 'r', encoding='utf-8') as f:
+                data = json.load(f)
+                return data.get("responses", [])
+        return []
+
+    def _clean_text(self, text: str) -> str:
+        text = re.sub(r'[^\w\s?]', '', text)
+        return ' '.join(text.split()).lower()
+
+    def chat(self, user_input: str) -> str:
+        user_input_clean = self._clean_text(user_input)
+        if not user_input_clean:
+            return "Tafadhali andika ujumbe."
+
+        best_response = None
+        best_score = 0.0
 
-    def _select_best_response(self, user_input: str) -> str:
-        """Select the best response using BERT scoring"""
-        if not self.responses:
-            return "Samahani, sijafunzwa majibu bado. Tafadhali fanya mafunzo kwanza."
-        
-        # Score all responses
-        scores = []
         for response in self.responses:
-            # Format input for BERT
-            text = f"{user_input} [SEP] {response}"
-            inputs = self.tokenizer(
-                text,
-                return_tensors="pt",
-                truncation=True,
-                max_length=256,
-                padding=True
-            ).to(self.device)
-            
-            # Get prediction
+            combined_text = f"{user_input_clean} [SEP] {response}"
+            inputs = self.tokenizer(combined_text, return_tensors="pt", truncation=True, max_length=256, padding=True).to(self.device)
             with torch.no_grad():
                 outputs = self.model(**inputs)
-                prediction = torch.softmax(outputs.logits, dim=1)
-                score = prediction[0][1].item()  # Probability it's a good response
-            
-            scores.append((response, score))
-        
-        # Sort by score and return best response
-        scores.sort(key=lambda x: x[1], reverse=True)
-        
-        # Return the best response
-        return scores[0][0]
+                probs = torch.softmax(outputs.logits, dim=1)
+                score = probs[0][1].item()  # probability of being the correct response
 
-    def _clean_input(self, text: str) -> str:
-        """Clean user input"""
-        text = re.sub(r'[^\w\s?]', '', text)  # Remove special chars except spaces and ?
-        text = ' '.join(text.split())  # Remove extra spaces
-        return text.lower()
+            if score > best_score:
+                best_score = score
+                best_response = response
 
-    def _clean_response(self, response: str) -> str:
-        """Clean up the response"""
-        response = response.strip()
-        # Ensure proper punctuation
-        if response and not response.endswith(('.', '!', '?')):
-            response += '.'
-        # Capitalize first letter
-        if response:
-            response = response[0].upper() + response[1:]
-        return response
+        if best_score < self.threshold:
+            return "Samahani, sielewi. Unaweza kuuliza kwa njia nyingine?"
 
-    def chat(self, message: str) -> str:
-        """
-        Main chat method
-        """
-        try:
-            if not message.strip():
-                return "Tafadhali andika ujumbe..."
-            
-            # Clean and preprocess input
-            cleaned_input = self._clean_input(message)
-            
-            # Select best response
-            response = self._select_best_response(cleaned_input)
-            
-            # Final cleanup
-            response = self._clean_response(response)
-            
-            return response
-            
-        except Exception as e:
-            logger.error(f"❌ Hitilafu wakati wa kukokotoa jibu: {e}")
-            return "Samahani, kuna hitilafu ya kiufundi. Tafadhali jaribu tena."
+        # Capitalize first letter and ensure punctuation
+        best_response = best_response.strip()
+        if best_response and not best_response.endswith(('.', '!', '?')):
+            best_response += '.'
+        best_response = best_response[0].upper() + best_response[1:]
+        return best_response