Update chatbot.py

chatbot.py

@@ -10,7 +10,6 @@ import torch
 import faiss
 from sentence_transformers import SentenceTransformer
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import gradio as gr
 
 # ---------------------------------------------------------------------------
 # Defaults — edit these if you don't want to pass CLI flags every time
@@ -18,7 +17,7 @@ import gradio as gr
 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"
+DEFAULT_LLM_MODEL_PATH = "./qwen2.5-3b-instruct"
 
 RETRIEVAL_TOP_K = 4
 RETRIEVAL_MIN_SCORE = 0.55
@@ -99,21 +98,19 @@ def format_chunks_display(retrieved):
     lines = []
     for i, r in enumerate(retrieved, 1):
         preview = r["text"][:400] + ("..." if len(r["text"]) > 400 else "")
-        lines.append(f"**[{i}] score: {r['score']:.3f}**\nSource: `{r['source_url']}`\n\n> {preview}\n")
-    return "\n---\n".join(lines)
+        lines.append(f"[{i}] score: {r['score']:.3f}  source: {r['source_url']}\n  > {preview}")
+    return "\n".join(lines)
 
 # ---------------------------------------------------------------------------
 # Main
 # ---------------------------------------------------------------------------
 def parse_args():
-    p = argparse.ArgumentParser(description="HBL Internal Assistant (RAG chatbot)")
+    p = argparse.ArgumentParser(description="HBL Internal Assistant (RAG chatbot) — 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("--host", default="127.0.0.1", help="Bind address (default: localhost only)")
-    p.add_argument("--port", type=int, default=7860)
-    p.add_argument("--share", action="store_true", help="Create a public gradio.live tunnel (needs internet)")
+    p.add_argument("--show-chunks", action="store_true", help="Print retrieved chunks before each answer")
     return p.parse_args()
 
 def main():
@@ -127,7 +124,6 @@ def main():
     device = "cuda" if torch.cuda.is_available() else "cpu"
     print(f"Using device: {device}")
 
-    # Prevent CPU execution from saturating background threads and freezing the OS
     if device == "cpu":
         torch.set_num_threads(4)
         torch.set_num_interop_threads(4)
@@ -148,9 +144,6 @@ def main():
     print(f"Loading LLM from {args.llm_model}...")
     tokenizer = AutoTokenizer.from_pretrained(args.llm_model)
 
-    # FIX: on CPU, load in bfloat16 (half the RAM of float32) with
-    # low_cpu_mem_usage=True (avoids the ~2x transient spike during load).
-    # On CUDA, bfloat16 + device_map="auto" lets accelerate place weights on GPU directly.
     if device == "cuda":
         llm_model = AutoModelForCausalLM.from_pretrained(
             args.llm_model,
@@ -170,7 +163,7 @@ def main():
         )
         llm_model = llm_model.to("cpu")
 
-    print("LLM ready.")
+    print("LLM ready.\n")
 
     def retrieve(query, k=RETRIEVAL_TOP_K, min_score=RETRIEVAL_MIN_SCORE):
         t0 = time.time()
@@ -221,20 +214,44 @@ def main():
             traceback.print_exc()
             return f"⚠️ Internal error: {e}", "*Error occurred*"
 
-    chunks_output = gr.Markdown(label="Retrieved chunks")
-    demo = gr.ChatInterface(
-        fn=chatbot_respond,
-        type="messages",
-        additional_outputs=[chunks_output],
-        title="HBL Internal Assistant",
-        description="Ask about HBL products, accounts, fees, and policies.",
-        examples=[
-            "What are the requirements for HBL Asaan Account?",
-            "HBL debit card fees",
-            "How do I apply for HBL PersonalLoan?",
-        ],
-    )
-    demo.launch(server_name=args.host, server_port=args.port, share=args.share)
+    # ---------------------------------------------------------------------
+    # Terminal chat loop (replaces Gradio ChatInterface)
+    # ---------------------------------------------------------------------
+    history = []  # list of {"role": "user"/"assistant", "content": str}
+    print("=" * 60)
+    print("HBL Internal Assistant — terminal mode")
+    print("Type your question and press Enter.")
+    print("Commands: 'exit' or 'quit' to stop, 'reset' to clear history.")
+    print("=" * 60 + "\n")
+
+    while True:
+        try:
+            message = input("You: ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print("\nExiting.")
+            break
+
+        if not message:
+            continue
+        if message.lower() in ("exit", "quit"):
+            print("Exiting.")
+            break
+        if message.lower() == "reset":
+            history = []
+            print("(history cleared)\n")
+            continue
+
+        answer, chunks_display = chatbot_respond(message, history)
+
+        if args.show_chunks:
+            print("\n--- Retrieved chunks ---")
+            print(chunks_display)
+            print("------------------------\n")
+
+        print(f"\nAssistant: {answer}\n")
+
+        history.append({"role": "user", "content": message})
+        history.append({"role": "assistant", "content": answer})
 
 if __name__ == "__main__":
     main()