Commit

app.py

@@ -8,108 +8,80 @@ from transformers import (
     BitsAndBytesConfig,
 )
 from peft import LoraConfig, get_peft_model
-from supabase import create_client
-from sentence_transformers import SentenceTransformer
 from safetensors.torch import load_file as safe_load
 
-# ── Supabase setup ─────────────────────────────────────────────────────────
-SUPA_URL = os.getenv("SUPABASE_URL")
-SUPA_KEY = os.getenv("SUPABASE_SERVICE_ROLE_KEY")
-supabase = create_client(SUPA_URL, SUPA_KEY)
+# ── Configuration ──────────────────────────────────────────────────────────
+MODEL_REPO      = "models/phi2-deeptalk-lora"
+BASE_MODEL      = "microsoft/phi-2"
+CONTEXT_TURNS   = 7        # how many past messages to include
+MAX_NEW_TOKENS  = 32       # shorter = faster
+TEMPERATURE     = 0.0      # 0.0 = greedy
+TOP_P           = 1.0      # disable nucleus sampling
+DEVICE_MAP      = "auto"
 
-# ── Embedder & memory RPC ──────────────────────────────────────────────────
-@st.cache_resource(show_spinner=False)
-def get_embedder():
-    return SentenceTransformer("paraphrase-MiniLM-L3-v2")
-embedder = get_embedder()
-
-@st.cache_data(show_spinner=False)
-def fetch_mems(query, k=3):
-    vec = embedder.encode(query).astype("float32").tolist()
-    return supabase.rpc(
-        "match_memories",
-        {"query_embedding": vec, "match_count": k}
-    ).execute().data
-
-def add_mem(speaker, text):
-    vec = embedder.encode(text).astype("float32").tolist()
-    supabase.table("memories").insert({
-        "speaker": speaker,
-        "text": text,
-        "embedding": vec
-    }).execute()
-
-# ── Model + tokenizer (adapter locally, tokenizer remote) ─────────────────
+SYSTEM = (
+    "You are a helpful assistant for DeepTalks with a base model Phi-2 "
+    "fine-tuned by Sourish for domain-specific support.\n"
+    "Base replies **only** on the context below. "
+    "If you don't know, say “I don't know.”\n"
+)
+
+# ── Model Loader ───────────────────────────────────────────────────────────
 @st.cache_resource(show_spinner=False)
 def load_generator():
-    base_dir    = os.path.dirname(__file__)
-    LOCAL_REPO  = os.path.join(base_dir, "models", "phi2-deeptalk-lora")
-    OFFLOAD_DIR = os.path.join(base_dir, "offload")
-    os.makedirs(OFFLOAD_DIR, exist_ok=True)
-
-    # 1) Tokenizer from official Phi-2
+    # 1) Tokenizer (always from HuggingFace cache)
     tokenizer = AutoTokenizer.from_pretrained(
-        "microsoft/phi-2",
+        BASE_MODEL,
         trust_remote_code=True,
         padding_side="left",
-        local_files_only=False  # allow remote fetch (cached)
     )
     if tokenizer.pad_token_id is None:
-        tokenizer.add_special_tokens({"pad_token": "[PAD]"})
+        tokenizer.add_special_tokens({"pad_token":"[PAD]"})
 
-    # 2) Load base model (quantized on GPU if available)
+    # 2) Base model in 4-bit or fp16/32
     if torch.cuda.is_available():
-        bnb_config = BitsAndBytesConfig(
+        bnb = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_quant_type="nf4",
             bnb_4bit_compute_dtype="float16",
             low_cpu_mem_usage=True,
         )
         base = AutoModelForCausalLM.from_pretrained(
-            "microsoft/phi-2",
+            BASE_MODEL,
             trust_remote_code=True,
-            quantization_config=bnb_config,
-            device_map="auto",
-            offload_folder=OFFLOAD_DIR,
-            offload_state_dict=True,
+            quantization_config=bnb,
+            device_map=DEVICE_MAP,
         )
     else:
         dtype = torch.float16 if torch.cuda.is_available() else torch.float32
         base = AutoModelForCausalLM.from_pretrained(
-            "microsoft/phi-2",
+            BASE_MODEL,
             trust_remote_code=True,
             torch_dtype=dtype,
-            device_map="auto",
-            offload_folder=OFFLOAD_DIR,
-            offload_state_dict=True,
+            device_map=DEVICE_MAP,
         )
 
-    # 3) Resize embeddings to match tokenizer
+    # 3) Resize & wrap LoRA
     base.resize_token_embeddings(len(tokenizer))
+    peft_config = LoraConfig.from_pretrained(MODEL_REPO, local_files_only=True)
+    model       = get_peft_model(base, peft_config)
 
-    # 4) Load LoRA config & wrap base
-    peft_config = LoraConfig.from_pretrained(
-        LOCAL_REPO,
-        local_files_only=True
-    )
-    model = get_peft_model(base, peft_config)
-
-    # 5) Manually load adapter weights
-    adapter_path = os.path.join(LOCAL_REPO, "adapter_model.safetensors")
-    state_dict = safe_load(adapter_path)
+    # 4) Load adapter weights (.safetensors)
+    adapter_file = os.path.join(MODEL_REPO, "adapter_model.safetensors")
+    state_dict   = safe_load(adapter_file)
     model.load_state_dict(state_dict, strict=False)
     model.eval()
 
-    # 6) Build generation pipeline
+    # 5) Build pipeline (greedy for speed)
     gen = pipeline(
         "text-generation",
         model=model,
         tokenizer=tokenizer,
-        device_map="auto",
-        max_new_tokens=32,
-        do_sample=True,
-        temperature=0.2,
-        top_p=0.8,
+        device_map=DEVICE_MAP,
+        max_new_tokens=MAX_NEW_TOKENS,
+        do_sample=False,
+        temperature=TEMPERATURE,
+        top_p=TOP_P,
         use_cache=True,
         return_full_text=False,
     )
@@ -117,49 +89,34 @@ def load_generator():
 
 tokenizer, generator = load_generator()
 
-# ── System prompt ───────────────────────────────────────────────────────────
-SYSTEM = (
-    "You are a helpful assistant for DeepTalks with a base model as Phi-2 "
-    "and fine tuned by Sourish for my domain specific role.\n"
-    "My domain is assisting you within my expertise by listening to you, "
-    "understanding you & supporting you.\n"
-    "Answer **only** using the information in the memory below.\n"
-    "If the answer is not in memory, reply: \"I don't know.\"\n"
-    "Do NOT repeat any lines beginning with 'User:'.\n"
-)
-
 # ── Streamlit UI ──────────────────────────────────────────────────────────
 st.set_page_config(layout="centered")
 st.title("🧠 Memory-Aware Phi-2 Chat")
 
+# initialize history
 if "history" not in st.session_state:
-    st.session_state.history = []  # list of (role, message)
+    st.session_state.history = []  # list of (role, text)
 
-# Render existing history
-for role, msg in st.session_state.history:
-    st.chat_message("user" if role == "You" else "assistant").write(msg)
+# render existing
+for role, text in st.session_state.history:
+    st.chat_message("user" if role=="You" else "assistant").write(text)
 
-# Input at bottom
+# user input
 user_input = st.chat_input("Type your message...")
 
 if user_input:
-    # Show & store user
+    # show user
     st.chat_message("user").write(user_input)
     st.session_state.history.append(("You", user_input))
-    add_mem("user", user_input)
-
-    # Fetch memories & build prompt
-    mems = fetch_mems(user_input, k=3)
-    mem_block = "\n".join(m["text"] for m in mems)
-    prompt = f"""{SYSTEM}
 
-Memory:
-{mem_block}
+    # build context from last turns
+    recent = st.session_state.history[-CONTEXT_TURNS*2:]  # each turn = 2 entries
+    ctx = "\n".join(f"{'User' if r=='You' else 'Assistant'}: {t}"
+                    for r,t in recent)
 
-User: {user_input}
-Assistant:"""
+    prompt = f"{SYSTEM}\nContext:\n{ctx}\nUser: {user_input}\nAssistant:"
 
-    # Generate with spinner
+    # generate
     with st.spinner("Thinking..."):
         try:
             out = generator(prompt)[0]["generated_text"].strip()
@@ -167,7 +124,6 @@ Assistant:"""
             out = "Sorry, I encountered an error."
             st.error(f"Generation error: {e}")
 
-    # Show & store assistant
+    # show bot
     st.chat_message("assistant").write(out)
     st.session_state.history.append(("Bot", out))
-    add_mem("assistant", out)