update

chatbot/chatbot.py

@@ -1,121 +1,65 @@
-"""
-Chatbot module for Codingo
-==========================
+# codingo/chatbot/chatbot.py
+"""Chatbot module for Codingo …
 
-This module encapsulates all functionality required to serve answers to
-questions about the Codingo platform.  It loads a small conversational
-model from Hugging Face and a lightweight vector database populated from
-``chatbot.txt``.  When a user asks a question, the module retrieves
-relevant snippets from the knowledge base and feeds them into the
-language model to generate a friendly response.
-
-Key features:
-
-* Completely self‑contained: there are no OpenAI or external API
-  dependencies.  Only free, locally hosted Hugging Face models are used.
-* Lazy initialisation: the model and vector store are loaded on the
-  first call to ``get_chatbot_response``.  Subsequent calls reuse
-  existing objects, avoiding expensive reloads.
-* GPU support: if a CUDA device is available, the model is automatically
-  moved onto the GPU for faster inference.
-
-This file lives inside ``codingo/chatbot`` alongside ``chatbot.txt``.
-``chatbot.txt`` should contain a plain‑text knowledge base of
-Codingo‑specific information and FAQs.  Feel free to update the
-contents of that file without touching any code here.
-
-"""
-
-from __future__ import annotations
+Default model changed to blenderbot-400M-distill; generation uses max_new_tokens; fallback between causal and seq2seq models."""
 
 import os
 import shutil
 from typing import List
 
-# -----------------------------------------------------------------------------
-# Environment configuration
-#
-# We set a few environment variables to force Hugging Face to store model
-# weights and tokeniser files inside ``/tmp``.  Hugging Face Spaces
-# provisions a read‑only file system outside of ``/tmp``, so without these
-# settings the transformers library might attempt to write into
-# unwritable locations.  These variables have no effect if the same
-# variables are already set by the hosting environment.
-
 os.environ.setdefault("HF_HOME", "/tmp/huggingface")
 os.environ.setdefault("TRANSFORMERS_CACHE", "/tmp/huggingface/transformers")
 os.environ.setdefault("HUGGINGFACE_HUB_CACHE", "/tmp/huggingface/hub")
 
-# -----------------------------------------------------------------------------
-# Module‑level state
-_hf_model = None  # type: ignore[assignment]
-_hf_tokenizer = None  # type: ignore[assignment]
-_chatbot_embedder = None  # type: ignore[assignment]
-_chatbot_collection = None  # type: ignore[assignment]
+_hf_model = None
+_hf_tokenizer = None
+_chatbot_embedder = None
+_chatbot_collection = None
 
-# Paths
 _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"
 
-# Default Hugging Face model for FAQ‑style Q&A.  You can override this
-# behaviour at deployment time by setting the ``HF_CHATBOT_MODEL``
-# environment variable.  DialoGPT is a lightweight conversational model
-# suitable for generating coherent short answers.  If you need more
-# open‑domain capability, consider ``facebook/blenderbot-400M-distill``.
-DEFAULT_MODEL_NAME = "microsoft/DialoGPT-medium"
-
+DEFAULT_MODEL_NAME = "facebook/blenderbot-400M-distill"
 
 def _init_hf_model() -> None:
-    """Load the Hugging Face model and tokenizer if not already loaded."""
+    from transformers import (
+        AutoModelForCausalLM,
+        AutoModelForSeq2SeqLM,
+        AutoTokenizer,
+    )
+    import torch
+
     global _hf_model, _hf_tokenizer
     if _hf_model is not None and _hf_tokenizer is not None:
         return
 
-    from transformers import AutoModelForCausalLM, AutoTokenizer  # slow import
-    import torch
-
     model_name = os.getenv("HF_CHATBOT_MODEL", DEFAULT_MODEL_NAME)
-    # Choose GPU if available; otherwise CPU
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-    # Download and load tokenizer and model.  They will be cached under
-    # the directories specified above.  If running for the first time on
-    # Hugging Face Spaces, model download may take a while.
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
-
-    # Ensure the pad token is defined.  Many casual conversation models
-    # reuse the end‑of‑sentence token for padding.
+    try:
+        model = AutoModelForCausalLM.from_pretrained(model_name)
+    except Exception:
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+    model = model.to(device)
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 
     _hf_model = model
     _hf_tokenizer = tokenizer
 
-
 def _init_vector_store() -> None:
-    """Initialise the Chroma vector store from ``chatbot.txt`` if needed."""
     global _chatbot_embedder, _chatbot_collection
     if _chatbot_embedder is not None and _chatbot_collection is not None:
         return
 
-    # Import heavy dependencies lazily to reduce module import time
     from langchain.text_splitter import RecursiveCharacterTextSplitter
     from sentence_transformers import SentenceTransformer
     import chromadb
     from chromadb.config import Settings
 
-    # Clear out any legacy database path that might be unwritable.  Previous
-    # versions of this project wrote under ``/app/chatbot/chroma_db`` which
-    # fails on Hugging Face Spaces.  The ``ignore_errors=True`` flag
-    # suppresses FileNotFoundError.
     shutil.rmtree("/app/chatbot/chroma_db", ignore_errors=True)
-
     os.makedirs(_chroma_db_dir, exist_ok=True)
-
-    # Read the knowledge base file.  If the file is missing, fall back to a
-    # minimal description of Codingo so the chatbot can still respond.
     try:
         with open(_knowledge_base_path, encoding="utf-8") as f:
             raw_text = f.read()
@@ -127,27 +71,16 @@ def _init_vector_store() -> None:
             "and intelligent recommendations."
         )
 
-    # Split the knowledge base into overlapping chunks for semantic search.
     splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=100)
     docs: List[str] = [doc.strip() for doc in splitter.split_text(raw_text) if doc.strip()]
-
-    # Embed the chunks using a small sentence transformer.  This model is
-    # lightweight (~80 MB) and works well for semantic similarity tasks.
     embedder = SentenceTransformer("all-MiniLM-L6-v2")
     embeddings = embedder.encode(docs, show_progress_bar=False, batch_size=32)
-
-    # Initialise a persistent Chroma client.  We disable anonymous telemetry
-    # because the environment has no outbound internet access.
     client = chromadb.Client(Settings(
         persist_directory=_chroma_db_dir,
         anonymized_telemetry=False,
         is_persistent=True,
     ))
-
-    # Create or retrieve the "chatbot" collection within the database.
     collection = client.get_or_create_collection("chatbot")
-
-    # If no documents are present, populate the collection with our chunks.
     try:
         existing = collection.get(limit=1)
         if not existing.get("documents"):
@@ -159,77 +92,39 @@ def _init_vector_store() -> None:
     _chatbot_embedder = embedder
     _chatbot_collection = collection
 
-
 def get_chatbot_response(query: str) -> str:
-    """
-    Generate a chatbot reply to the given user query.
-
-    The response is generated by retrieving up to three relevant snippets
-    from the knowledge base using the MiniLM embeddings and then feeding
-    those snippets together with the user question into the conversational
-    model.  If no relevant information is found or the model generates
-    an empty response, a helpful fallback message is returned.
-
-    Parameters
-    ----------
-    query : str
-        The user's message.  Should be non‑empty and related to the
-        Codingo platform.
-
-    Returns
-    -------
-    str
-        The chatbot's reply, always a string.
-    """
-    # Basic validation of the query string
     if not query or not query.strip():
         return "Please type a question about the Codingo platform."
 
-    # Lazy initialisation of the vector store and Hugging Face model
     _init_vector_store()
     _init_hf_model()
-
-    # Unpack state
-    embedder = _chatbot_embedder  # type: ignore[assignment]
-    collection = _chatbot_collection  # type: ignore[assignment]
-    model = _hf_model  # type: ignore[assignment]
-    tokenizer = _hf_tokenizer  # type: ignore[assignment]
+    embedder = _chatbot_embedder
+    collection = _chatbot_collection
+    model = _hf_model
+    tokenizer = _hf_tokenizer
 
     import torch
 
-    # Embed the incoming query using the same sentence transformer
-    query_embedding = embedder.encode([query])[0]  # type: ignore[operator]
-    # Retrieve the three most similar documents from the vector store
+    query_embedding = embedder.encode([query])[0]
     results = collection.query(query_embeddings=[query_embedding.tolist()], n_results=3)
     retrieved_docs = results.get("documents", [[]])[0] if results else []
-
-    # Build a context string from the retrieved documents
     context = "\n".join(retrieved_docs[:3])
 
-    # Compose the system instruction.  The model is prompted as a
-    # persona called LUNA AI.  Keep responses concise and friendly, and
-    # redirect politely on irrelevant questions.
     system_instruction = (
         "You are LUNA AI, a helpful assistant for the Codingo recruitment "
         "platform. Use the provided context to answer questions about "
         "Codingo. If the question is not related to Codingo, politely "
         "redirect the conversation. Keep responses concise and friendly."
     )
+    prompt = f"{system_instruction}\n\nContext:\n{context}\n\nUser: {query}\nLUNA AI:"
+    inputs = tokenizer.encode(
+        prompt, return_tensors="pt", truncation=True, max_length=512, padding=True
+    ).to(model.device)
 
-    prompt = (
-        f"{system_instruction}\n\nContext:\n{context}\n\n"
-        f"User: {query}\nLUNA AI:"
-    )
-
-    # Tokenise the prompt and truncate to the maximum input length supported
-    inputs = tokenizer.encode(prompt, return_tensors="pt", truncation=True, max_length=512, padding=True)
-    inputs = inputs.to(model.device)
-
-    # Generate a continuation from the model
     with torch.no_grad():
         output_ids = model.generate(
             inputs,
-            max_length=inputs.shape[1] + 150,
+            max_new_tokens=150,
             num_beams=3,
             do_sample=True,
             temperature=0.7,
@@ -238,17 +133,14 @@ def get_chatbot_response(query: str) -> str:
             early_stopping=True,
         )
 
-    # Decode the output and strip the prompt from the beginning
     response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
     if "LUNA AI:" in response:
         response = response.split("LUNA AI:")[-1].strip()
     elif prompt in response:
         response = response.replace(prompt, "").strip()
 
-    # Fallback if the model didn't return anything useful
-    if not response:
-        return (
-            "I'm here to help you with questions about the Codingo platform. "
-            "What would you like to know?"
-        )
-    return response
+    return (
+        response
+        if response
+        else "I'm here to help you with questions about the Codingo platform. What would you like to know?"
+    )