Update chatbot.py

chatbot.py

@@ -5,11 +5,10 @@ import pickle
 import argparse
 import traceback
 
-os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
 import torch
 import faiss
 from sentence_transformers import SentenceTransformer
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from llama_cpp import Llama
 
 # ---------------------------------------------------------------------------
 # Defaults — edit these if you don't want to pass CLI flags every time
@@ -17,7 +16,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 DEFAULT_FAISS_INDEX_PATH = "./hbl_site_index_COMPLETE.faiss"
 DEFAULT_CHUNKS_METADATA_PATH = "./hbl_site_metadata_COMPLETE.pkl"
 DEFAULT_EMBED_MODEL_PATH = "./bge-m3"
-DEFAULT_LLM_MODEL_PATH = "./qwen2.5-3b-instruct"
+DEFAULT_LLM_GGUF_PATH = "./qwen2.5-0.5b-instruct-q5_k_m.gguf"
 
 RETRIEVAL_TOP_K = 4
 RETRIEVAL_MIN_SCORE = 0.55
@@ -34,18 +33,18 @@ CONTEXT:
 {context}"""
 
 CODE_PATTERNS = [
-    r"```",                          # any fenced code block
-    r"\bdef\s+\w+\s*\(",             # python function def
-    r"\bimport\s+\w+",               # import statements
-    r"\bfunction\s+\w+\s*\(",        # js function
+    r"```",
+    r"\bdef\s+\w+\s*\(",
+    r"\bimport\s+\w+",
+    r"\bfunction\s+\w+\s*\(",
     r"\bconsole\.log\(",
     r"\bprint\(",
     r"\breturn\s+\w+",
 ]
 
 MATH_PATTERNS = [
-    r"^\s*-?\d+(\.\d+)?\s*[\+\-\*/x×]\s*-?\d+(\.\d+)?",   # "6+7", "12 * 3", "6 x 7"
-    r"\bwhat\s+is\s+\d+.{0,15}[\+\-\*/].{0,15}\d+",         # "what is 6 plus 7" style
+    r"^\s*-?\d+(\.\d+)?\s*[\+\-\*/x×]\s*-?\d+(\.\d+)?",
+    r"\bwhat\s+is\s+\d+.{0,15}[\+\-\*/].{0,15}\d+",
     r"\bcalculate\s+\d+.{0,15}\d+",
     r"\bsolve\s+(this|the)?\s*(equation|expression|problem)\b",
 ]
@@ -69,9 +68,7 @@ def is_writing_task(message):
 def strip_or_block(answer):
     if not contains_code(answer):
         return answer
-    # Remove fenced code blocks
     cleaned = re.sub(r"```.*?```", "\x00CODE_REMOVED\x00", answer, flags=re.DOTALL)
-    # Split into sentences/lines and drop ones that reference the removed code
     lines = cleaned.split("\n")
     result_lines = []
     reference_phrases = [
@@ -83,9 +80,9 @@ def strip_or_block(answer):
     for line in lines:
         low = line.lower()
         if "\x00CODE_REMOVED\x00" in line:
-            continue  # drop the marker line entirely
+            continue
         if any(p in low for p in reference_phrases):
-            continue  # drop sentences describing the removed code
+            continue
         result_lines.append(line)
     cleaned = "\n".join(result_lines).strip()
     cleaned += ("\n\n*(Note: I can explain loan interest calculations in plain language "
@@ -105,11 +102,16 @@ def format_chunks_display(retrieved):
 # Main
 # ---------------------------------------------------------------------------
 def parse_args():
-    p = argparse.ArgumentParser(description="HBL Internal Assistant (RAG chatbot) — terminal version")
+    p = argparse.ArgumentParser(description="HBL Internal Assistant (RAG chatbot) — llama.cpp terminal version")
     p.add_argument("--faiss-index", default=os.environ.get("FAISS_INDEX_PATH", DEFAULT_FAISS_INDEX_PATH))
     p.add_argument("--chunks-metadata", default=os.environ.get("CHUNKS_METADATA_PATH", DEFAULT_CHUNKS_METADATA_PATH))
     p.add_argument("--embed-model", default=os.environ.get("EMBED_MODEL_PATH", DEFAULT_EMBED_MODEL_PATH))
-    p.add_argument("--llm-model", default=os.environ.get("LLM_MODEL_PATH", DEFAULT_LLM_MODEL_PATH))
+    p.add_argument("--llm-gguf", default=os.environ.get("LLM_GGUF_PATH", DEFAULT_LLM_GGUF_PATH),
+                    help="Path to the .gguf model file")
+    p.add_argument("--n-threads", type=int, default=os.cpu_count(),
+                    help="CPU threads for llama.cpp to use (default: all logical cores)")
+    p.add_argument("--n-ctx", type=int, default=4096, help="Context window size")
+    p.add_argument("--max-tokens", type=int, default=300, help="Max tokens to generate per response")
     p.add_argument("--show-chunks", action="store_true", help="Print retrieved chunks before each answer")
     return p.parse_args()
 
@@ -119,14 +121,7 @@ def main():
     assert os.path.exists(args.faiss_index), f"FAISS index not found: {args.faiss_index}"
     assert os.path.exists(args.chunks_metadata), f"Chunks metadata not found: {args.chunks_metadata}"
     assert os.path.exists(args.embed_model), f"Embedding model folder not found: {args.embed_model}"
-    assert os.path.exists(args.llm_model), f"LLM model folder not found: {args.llm_model}"
-
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    print(f"Using device: {device}")
-
-    if device == "cpu":
-        torch.set_num_threads(4)
-        torch.set_num_interop_threads(4)
+    assert os.path.exists(args.llm_gguf), f"GGUF model file not found: {args.llm_gguf}"
 
     print("Loading FAISS index...")
     index = faiss.read_index(args.faiss_index)
@@ -138,32 +133,18 @@ def main():
     print(f"Loaded {index.ntotal} vectors, {len(chunks)} chunks.")
 
     print(f"Loading embedding model from {args.embed_model}...")
-    embed_model = SentenceTransformer(args.embed_model, device=device)
+    embed_model = SentenceTransformer(args.embed_model, device="cpu")
     print("Retrieval assets ready.")
 
-    print(f"Loading LLM from {args.llm_model}...")
-    tokenizer = AutoTokenizer.from_pretrained(args.llm_model)
-
-    if device == "cuda":
-        llm_model = AutoModelForCausalLM.from_pretrained(
-            args.llm_model,
-            torch_dtype=torch.bfloat16,
-            low_cpu_mem_usage=True,
-            local_files_only=True,
-            use_safetensors=True,
-            device_map="auto",
-        )
-    else:
-        llm_model = AutoModelForCausalLM.from_pretrained(
-            args.llm_model,
-            torch_dtype=torch.bfloat16,
-            low_cpu_mem_usage=True,
-            local_files_only=True,
-            use_safetensors=True,
-        )
-        llm_model = llm_model.to("cpu")
-
-    print("LLM ready.\n")
+    print(f"Loading LLM (GGUF) from {args.llm_gguf} with {args.n_threads} threads...")
+    t0 = time.time()
+    llm = Llama(
+        model_path=args.llm_gguf,
+        n_ctx=args.n_ctx,
+        n_threads=args.n_threads,
+        verbose=False,
+    )
+    print(f"LLM ready in {time.time() - t0:.2f}s.\n")
 
     def retrieve(query, k=RETRIEVAL_TOP_K, min_score=RETRIEVAL_MIN_SCORE):
         t0 = time.time()
@@ -178,22 +159,22 @@ def main():
         print(f"[timing] retrieve() {time.time() - t0:.2f}s, {len(results)} chunks")
         return results
 
-    def call_llm_with_history(system_prompt, history, current_message, max_new_tokens=400):
+    def call_llm_with_history(system_prompt, history, current_message, max_new_tokens=None):
+        max_new_tokens = max_new_tokens or args.max_tokens
         messages = [{"role": "system", "content": system_prompt}]
         messages += history
         messages.append({"role": "user", "content": current_message})
-        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        inputs = tokenizer(prompt, return_tensors="pt").to(device)
+
         t0 = time.time()
-        outputs = llm_model.generate(
-            **inputs,
-            max_new_tokens=max_new_tokens,
-            temperature=0.3,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id,
+        # llama-cpp-python's create_chat_completion applies the model's chat
+        # template internally, same role as tokenizer.apply_chat_template before.
+        result = llm.create_chat_completion(
+            messages=messages,
+            max_tokens=max_new_tokens,
+            temperature=0.0,   # greedy-ish; set >0 if you want variation back
         )
         print(f"[timing] generate_response() {time.time() - t0:.2f}s")
-        return tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True).strip()
+        return result["choices"][0]["message"]["content"].strip()
 
     def chatbot_respond(message, history):
         try:
@@ -215,11 +196,11 @@ def main():
             return f"⚠️ Internal error: {e}", "*Error occurred*"
 
     # ---------------------------------------------------------------------
-    # Terminal chat loop (replaces Gradio ChatInterface)
+    # Terminal chat loop
     # ---------------------------------------------------------------------
-    history = []  # list of {"role": "user"/"assistant", "content": str}
+    history = []
     print("=" * 60)
-    print("HBL Internal Assistant — terminal mode")
+    print("HBL Internal Assistant — terminal mode (llama.cpp)")
     print("Type your question and press Enter.")
     print("Commands: 'exit' or 'quit' to stop, 'reset' to clear history.")
     print("=" * 60 + "\n")