chatbot/chatbot.py

# codingo/chatbot/chatbot.py
"""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")
os.environ.setdefault("HUGGINGFACE_HUB_CACHE", "/tmp/huggingface/hub")

_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"

# 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 Phi-2 model...")
    from transformers import AutoModelForCausalLM, AutoTokenizer
    
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")
    
    # 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,
        device_map="auto",
        trust_remote_code=True
    )
    model.eval()
    
    _model = model
    _tokenizer = tokenizer
    print("Phi-2 loaded successfully!")

def _init_vector_store():
    global _chatbot_embedder, _chatbot_collection, _knowledge_chunks
    if _chatbot_embedder is not None and _chatbot_collection is not None:
        return

    print("Initializing vector store...")
    from langchain.text_splitter import RecursiveCharacterTextSplitter
    from sentence_transformers import SentenceTransformer
    import chromadb
    from chromadb.config import Settings

    # Load knowledge base
    try:
        with open(_knowledge_base_path, encoding="utf-8") as f:
            raw_text = f.read()
            print(f"Loaded knowledge base: {len(raw_text)} characters")
    except FileNotFoundError:
        print("Knowledge base not found!")
        raw_text = "Codingo is an AI recruitment platform."

    # 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()]
    _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 (in-memory)
    client = chromadb.Client(Settings(anonymized_telemetry=False, is_persistent=False))
    
    try:
        client.delete_collection("chatbot")
    except:
        pass
    
    collection = client.create_collection("chatbot")
    ids = [f"doc_{i}" for i in range(len(docs))]
    collection.add(documents=docs, embeddings=embeddings.tolist(), ids=ids)
    
    _chatbot_embedder = embedder
    _chatbot_collection = collection
    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 "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}'")
        
        # Clear GPU cache
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        
        # Initialize
        _init_vector_store()
        _init_model()
        
        # Search for relevant context
        query_embedding = _chatbot_embedder.encode([query])[0]
        results = _chatbot_collection.query(
            query_embeddings=[query_embedding.tolist()], 
            n_results=3
        )
        
        retrieved_docs = results.get("documents", [[]])[0] if results else []
        print(f"Found {len(retrieved_docs)} relevant chunks")
        
        # Try to find FAQ answer first
        faq_answer = extract_faq_answer(query, retrieved_docs)
        if faq_answer:
            print("Found FAQ match!")
            return faq_answer
        
        # 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}

User Question: {query}

Output: Based on the information provided, """
        
        # Tokenize with appropriate length
        inputs = _tokenizer(
            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=200,
                temperature=0.7,
                do_sample=True,
                top_p=0.9,
                repetition_penalty=1.15,
                pad_token_id=_tokenizer.pad_token_id,
                eos_token_id=_tokenizer.eos_token_id,
                early_stopping=True
            )
        
        # Decode response
        full_response = _tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Extract only the generated part
        response = full_response.split("Output:")[-1].strip()
        
        # 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:
                # 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 apologize for the technical issue. Please feel free to ask me about Codingo's features, job matching process, or how to get started!"

# Test the chatbot
if __name__ == "__main__":
    print("Testing Codingo Chatbot...")
    test_queries = [
        "Hello!",
        "What is Codingo?",
        "How does it work?",
        "What job types do you support?",
        "How can I improve my match score?",
        "Is Codingo free?",
        "Tell me about CV tips"
    ]
    
    for q in test_queries:
        response = get_chatbot_response(q)
        print(f"\nUser: {q}")
        print(f"LUNA: {response}")
        print("-" * 80)