Commit

app.py

@@ -1,6 +1,7 @@
 import os
 import streamlit as st
 import torch
+import logging
 from transformers import (
     pipeline,
     AutoTokenizer,
@@ -8,90 +9,86 @@ from transformers import (
     BitsAndBytesConfig,
 )
 from peft import PeftModel
-import logging
 
 # ── Configuration ──────────────────────────────────────────────────────────
 BASE_MODEL     = "microsoft/phi-2"
 ADAPTER_REPO   = "sourize/phi2-memory-lora"
 CONTEXT_TURNS  = 6
-MAX_NEW_TOKENS = 32
+MAX_NEW_TOKENS = 128
 OFFLOAD_DIR    = "offload"
 
 SYSTEM = (
-    "You are a helpful assistant for DeepTalks with base Phi-2 "
+    "You are a helpful assistant for DeepTalks with base Phi-2\n"
     "fine-tuned by Sourish for domain support.\n"
     "Answer **only** using the conversation context below.\n"
     "Do NOT output any lines beginning with 'User:' or 'Assistant:'.\n"
-    "If you don't know, say 'I don't know.'\n"
+    "If you don't know, say \"I don't know.\"\n"
 )
 
-
+# ── Model + Pipeline Loader ─────────────────────────────────────────────────
 @st.cache_resource(show_spinner=False)
 def load_pipeline():
     # 1) Tokenizer
     tokenizer = AutoTokenizer.from_pretrained(
-        BASE_MODEL,
-        trust_remote_code=True,
-        padding_side="left",
+        BASE_MODEL, trust_remote_code=True, padding_side="left"
     )
     if tokenizer.pad_token_id is None:
         tokenizer.add_special_tokens({"pad_token": "[PAD]"})
 
-    # 2) Base model (4-bit quant on GPU, else FP16/FP32)
+    # 2) Choose quantization config
     if torch.cuda.is_available():
-        bnb_config = BitsAndBytesConfig(
+        quant_config = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_quant_type="nf4",
             bnb_4bit_compute_dtype="float16",
             low_cpu_mem_usage=True,
         )
-        base = AutoModelForCausalLM.from_pretrained(
-            BASE_MODEL,
-            trust_remote_code=True,
-            quantization_config=bnb_config,
-            device_map="auto",
-            offload_folder=OFFLOAD_DIR,
-            offload_state_dict=True,
-        )
     else:
-        dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-        base = AutoModelForCausalLM.from_pretrained(
-            BASE_MODEL,
-            trust_remote_code=True,
-            torch_dtype=dtype,
-            device_map="auto",
-            offload_folder=OFFLOAD_DIR,
-            offload_state_dict=True,
+        quant_config = BitsAndBytesConfig(
+            load_in_8bit=True,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
         )
 
-    # 3) Resize embeddings
-    base.resize_token_embeddings(len(tokenizer))
+    # 3) Load base model
+    base = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL,
+        trust_remote_code=True,
+        quantization_config=quant_config,
+        device_map="auto",
+        offload_folder=OFFLOAD_DIR,
+        offload_state_dict=True,
+        torch_dtype=None  # auto
+    )
 
-    # 4) Overlay LoRA adapter from Hugging Face Hub
+    # 4) Resize embeddings & overlay LoRA
+    base.resize_token_embeddings(len(tokenizer))
     model = PeftModel.from_pretrained(
         base,
         ADAPTER_REPO,
         trust_remote_code=True,
         device_map="auto",
-        torch_dtype="auto",
         offload_folder=OFFLOAD_DIR,
         offload_state_dict=True,
+        torch_dtype="auto",
     )
     model.eval()
 
-    # 5) Build generation pipeline (greedy for speed)
+    # 5) Build sampler pipeline
     gen = pipeline(
         "text-generation",
         model=model,
         tokenizer=tokenizer,
         device_map="auto",
         max_new_tokens=MAX_NEW_TOKENS,
-        do_sample=False,      # greedy decoding
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9,
         use_cache=True,
         return_full_text=False,
     )
 
-    logging.info(f"Loaded pipeline with model {BASE_MODEL} + adapter {ADAPTER_REPO}")
+    logging.info("Pipeline loaded.")
     return gen
 
 generator = load_pipeline()
@@ -101,24 +98,25 @@ st.set_page_config(layout="centered")
 st.title("🧠 DeepTalks")
 st.subheader("Your personal AI Companion", divider='grey')
 
-# initialize chat history
+# initialize history
 if "history" not in st.session_state:
-    st.session_state.history = []  # List of (role, text)
+    st.session_state.history = []  # list of (role, text)
 
-# render chat history
+# render past messages
 for role, text in st.session_state.history:
-    st.chat_message("user" if role == "You" else "assistant").write(text)
+    st.chat_message("user" if role=="You" else "assistant").write(text)
 
 # user input
 user_input = st.chat_input("Your message…")
 if user_input:
-    # show & store user turn
+    # show & store user
     st.chat_message("user").write(user_input)
     st.session_state.history.append(("You", user_input))
 
-    # build context from last CONTEXT_TURNS
+    # build clean context from last turns (texts only)
     recent = st.session_state.history[-CONTEXT_TURNS*2:]
-    context = "\n".join(t for _, t in recent)
+    context = "\n".join(text for _, text in recent)
+
     prompt = f"""{SYSTEM}
 
 Context:
@@ -127,22 +125,20 @@ Context:
 User: {user_input}
 Assistant:"""
 
-    # generate reply
+    # generate with spinner
     with st.spinner("Thinking…"):
         try:
-            # pipeline was set up with `return_full_text=False`, so we get just the reply
             reply = generator(prompt)[0]["generated_text"].strip()
-            # strip any accidental echoes
+            # strip stray markers
             for marker in ["User:", "Assistant:"]:
                 if marker in reply:
                     reply = reply.split(marker)[0].strip()
-            # if it somehow ends up empty, backstop with an apology
             if not reply:
-                reply = "I'm sorry, I didn't catch that. Could you rephrase?"
+                reply = "I’m sorry, I didn’t catch that. Could you rephrase?"
         except Exception as e:
-            reply = "I'm sorry, something went wrong."
-            st.error(f"Generation error: {e}")
+            reply = "I’m sorry, something went wrong."
+            st.error(f"Error: {e}")
 
-    # show & store assistant turn
+    # show & store assistant
     st.chat_message("assistant").write(reply)
     st.session_state.history.append(("Bot", reply))