chatbot updated

chatbot/chatbot.py

@@ -1,10 +1,11 @@
 # codingo/chatbot/chatbot.py
-"""Interactive chatbot using Flan-T5 for dynamic responses"""
+"""Interactive chatbot using Microsoft Phi-2 for efficient, quality responses"""
 
 import os
 import shutil
 from typing import List
 import torch
+import re
 
 os.environ.setdefault("HF_HOME", "/tmp/huggingface")
 os.environ.setdefault("TRANSFORMERS_CACHE", "/tmp/huggingface/transformers")
@@ -14,40 +15,45 @@ _model = None
 _tokenizer = None
 _chatbot_embedder = None
 _chatbot_collection = None
+_knowledge_chunks = []
 
 _current_dir = os.path.dirname(os.path.abspath(__file__))
 _knowledge_base_path = os.path.join(_current_dir, "chatbot.txt")
 _chroma_db_dir = "/tmp/chroma_db"
 
-# Using Flan-T5 - it's small, fast, and great for Q&A
-MODEL_NAME = "google/flan-t5-small"
+# Phi-2: 2.7B params, great performance, fits easily on T4
+MODEL_NAME = "microsoft/phi-2"
 
 def _init_model():
     global _model, _tokenizer
     if _model is not None and _tokenizer is not None:
         return
     
-    print("Loading Flan-T5 model...")
-    from transformers import T5ForConditionalGeneration, T5Tokenizer
+    print("Loading Phi-2 model...")
+    from transformers import AutoModelForCausalLM, AutoTokenizer
     
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     print(f"Using device: {device}")
     
-    tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME)
-    model = T5ForConditionalGeneration.from_pretrained(
+    # Load tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+    tokenizer.pad_token = tokenizer.eos_token
+    
+    # Load model
+    model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME,
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-        low_cpu_mem_usage=True
+        torch_dtype=torch.float16,
+        device_map="auto",
+        trust_remote_code=True
     )
-    model = model.to(device)
     model.eval()
     
     _model = model
     _tokenizer = tokenizer
-    print("Model loaded successfully!")
+    print("Phi-2 loaded successfully!")
 
 def _init_vector_store():
-    global _chatbot_embedder, _chatbot_collection
+    global _chatbot_embedder, _chatbot_collection, _knowledge_chunks
     if _chatbot_embedder is not None and _chatbot_collection is not None:
         return
 
@@ -57,10 +63,6 @@ def _init_vector_store():
     import chromadb
     from chromadb.config import Settings
 
-    # Clean and create directory
-    shutil.rmtree(_chroma_db_dir, ignore_errors=True)
-    os.makedirs(_chroma_db_dir, exist_ok=True)
-    
     # Load knowledge base
     try:
         with open(_knowledge_base_path, encoding="utf-8") as f:
@@ -73,13 +75,13 @@ def _init_vector_store():
     # Split into chunks
     splitter = RecursiveCharacterTextSplitter(chunk_size=400, chunk_overlap=50)
     docs = [doc.strip() for doc in splitter.split_text(raw_text) if doc.strip()]
-    print(f"Created {len(docs)} document chunks")
+    _knowledge_chunks = docs  # Store for reference
     
     # Create embeddings
     embedder = SentenceTransformer("all-MiniLM-L6-v2")
     embeddings = embedder.encode(docs, show_progress_bar=False)
     
-    # Create ChromaDB collection
+    # Create ChromaDB collection (in-memory)
     client = chromadb.Client(Settings(anonymized_telemetry=False, is_persistent=False))
     
     try:
@@ -93,12 +95,31 @@ def _init_vector_store():
     
     _chatbot_embedder = embedder
     _chatbot_collection = collection
-    print("Vector store ready!")
+    print(f"Vector store ready with {len(docs)} chunks!")
+
+def extract_faq_answer(query: str, docs: List[str]) -> str:
+    """Try to find direct FAQ answers"""
+    query_lower = query.lower()
+    
+    for doc in docs:
+        # Look for Q&A patterns
+        if "Q:" in doc and "A:" in doc:
+            lines = doc.split('\n')
+            for i, line in enumerate(lines):
+                if line.strip().startswith('Q:'):
+                    question = line[2:].strip().lower()
+                    # Check similarity
+                    if any(word in question for word in query_lower.split() if len(word) > 3):
+                        # Find the answer
+                        for j in range(i+1, min(i+5, len(lines))):
+                            if lines[j].strip().startswith('A:'):
+                                return lines[j][2:].strip()
+    return None
 
 def get_chatbot_response(query: str) -> str:
     try:
         if not query or not query.strip():
-            return "Hi! I'm LUNA AI. Ask me anything about Codingo!"
+            return "Hello! I'm LUNA AI, your Codingo assistant. I can help you with questions about our AI recruitment platform, job matching, CV tips, and more!"
         
         print(f"\nProcessing: '{query}'")
         
@@ -120,85 +141,110 @@ def get_chatbot_response(query: str) -> str:
         retrieved_docs = results.get("documents", [[]])[0] if results else []
         print(f"Found {len(retrieved_docs)} relevant chunks")
         
-        # Combine the most relevant information
-        context = " ".join(retrieved_docs[:2]) if retrieved_docs else "Codingo is an AI recruitment platform."
+        # Try to find FAQ answer first
+        faq_answer = extract_faq_answer(query, retrieved_docs)
+        if faq_answer:
+            print("Found FAQ match!")
+            return faq_answer
         
-        # Create a prompt for Flan-T5
-        prompt = f"""Answer the question based on the context about Codingo.
+        # Build context from retrieved docs
+        context = "\n".join(retrieved_docs[:2]) if retrieved_docs else ""
+        
+        # Create an instruction-following prompt for Phi-2
+        prompt = f"""Instruct: You are LUNA AI, a helpful assistant for Codingo recruitment platform. 
+Use the following information to answer the user's question:
 
-Context: {context}
+{context}
 
-Question: {query}
+User Question: {query}
 
-Answer:"""
+Output: Based on the information provided, """
         
-        # Tokenize
+        # Tokenize with appropriate length
         inputs = _tokenizer(
-            prompt,
-            max_length=512,
-            truncation=True,
-            return_tensors="pt"
-        ).to(_model.device)
+            prompt, 
+            return_tensors="pt", 
+            truncation=True, 
+            max_length=800,
+            padding=True
+        )
+        inputs = {k: v.to(_model.device) for k, v in inputs.items()}
         
         # Generate response
         with torch.no_grad():
             outputs = _model.generate(
                 **inputs,
-                max_new_tokens=150,
-                num_beams=4,
+                max_new_tokens=200,
                 temperature=0.7,
                 do_sample=True,
                 top_p=0.9,
-                repetition_penalty=1.2
+                repetition_penalty=1.15,
+                pad_token_id=_tokenizer.pad_token_id,
+                eos_token_id=_tokenizer.eos_token_id,
+                early_stopping=True
             )
         
         # Decode response
-        response = _tokenizer.decode(outputs[0], skip_special_tokens=True)
-        print(f"Generated: '{response}'")
-        
-        # Make sure we have a good response
-        if not response or len(response) < 5:
-            # Fallback: try a simpler prompt
-            simple_prompt = f"Question about Codingo: {query}\nAnswer:"
-            inputs = _tokenizer(simple_prompt, max_length=256, truncation=True, return_tensors="pt").to(_model.device)
-            
-            with torch.no_grad():
-                outputs = _model.generate(**inputs, max_new_tokens=100, temperature=0.8)
-            
-            response = _tokenizer.decode(outputs[0], skip_special_tokens=True)
+        full_response = _tokenizer.decode(outputs[0], skip_special_tokens=True)
         
-        # Clean up the response
-        response = response.strip()
+        # Extract only the generated part
+        response = full_response.split("Output:")[-1].strip()
         
-        # If still too short, provide a helpful response
-        if len(response) < 10:
-            if "hello" in query.lower() or "hi" in query.lower():
-                return "Hello! I'm LUNA AI, your Codingo assistant. I can help you with questions about our AI recruitment platform, job matching, CV tips, and more!"
+        # Clean up common artifacts
+        response = response.replace("Based on the information provided,", "").strip()
+        
+        # Remove the original prompt if it appears
+        if prompt in response:
+            response = response.replace(prompt, "").strip()
+        
+        # Ensure quality response
+        if len(response) < 20 or response.lower() == query.lower():
+            # Generate a contextual response
+            query_lower = query.lower()
+            
+            if "hello" in query_lower or "hi" in query_lower:
+                return "Hello! Welcome to Codingo! I'm LUNA AI, here to help you navigate our AI-powered recruitment platform. You can ask me about creating profiles, job matching, improving your CV, or any of our features!"
+            
+            elif "what" in query_lower and "codingo" in query_lower:
+                return "Codingo is an innovative AI-driven recruitment platform that transforms how companies hire and how candidates find jobs. We use advanced algorithms to match skills with opportunities, provide instant CV feedback, and streamline the entire hiring process."
+            
+            elif "how" in query_lower and ("work" in query_lower or "use" in query_lower):
+                return "Here's how Codingo works: As a candidate, you create a profile, upload your resume, and add portfolio links. Our AI then analyzes your skills and matches you with suitable jobs. You'll receive personalized recommendations and CV improvement tips. For employers, we offer smart candidate filtering and automated screening insights!"
+            
+            elif "feature" in query_lower or "special" in query_lower:
+                return "What makes Codingo special is our combination of AI-powered job matching, real-time CV analysis, bias-aware algorithms, and focus on tech professionals. We support various roles from developers to designers, making the hiring process smarter, faster, and fairer for everyone."
+            
             else:
-                return f"I can help you with that! Based on what I know about Codingo: {retrieved_docs[0][:200] if retrieved_docs else 'Codingo is an AI-powered recruitment platform that helps match candidates with jobs.'}"
+                # Use context to create a response
+                if retrieved_docs:
+                    return f"Let me help you with that! {retrieved_docs[0][:250]}..."
+                else:
+                    return "I'd be happy to help you learn more about Codingo! Could you please ask about specific features like job matching, CV tips, supported job types, or how our platform works?"
         
+        print(f"Generated: {response[:100]}...")
         return response
         
     except Exception as e:
         print(f"Error: {e}")
         import traceback
         traceback.print_exc()
-        return "I'm having a technical issue. Please try asking your question again!"
+        return "I apologize for the technical issue. Please feel free to ask me about Codingo's features, job matching process, or how to get started!"
 
-# Test function
+# Test the chatbot
 if __name__ == "__main__":
-    # Test the chatbot
+    print("Testing Codingo Chatbot...")
     test_queries = [
+        "Hello!",
         "What is Codingo?",
         "How does it work?",
-        "What makes Codingo special?",
-        "How can I improve my profile?",
-        "Is it free?"
+        "What job types do you support?",
+        "How can I improve my match score?",
+        "Is Codingo free?",
+        "Tell me about CV tips"
     ]
     
-    print("Testing chatbot...")
     for q in test_queries:
         response = get_chatbot_response(q)
-        print(f"\nQ: {q}")
-        print(f"A: {response}")
-        print("-" * 50)
+        print(f"\nUser: {q}")
+        print(f"LUNA: {response}")
+        print("-" * 80)