Commit

app.py

@@ -25,7 +25,6 @@ SYSTEM = (
     "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
@@ -35,7 +34,7 @@ def load_pipeline():
     if tokenizer.pad_token_id is None:
         tokenizer.add_special_tokens({"pad_token": "[PAD]"})
 
-    # 2) Choose quantization config
+    # 2) Base model: 4-bit on CUDA, plain FP16/FP32 on CPU
     if torch.cuda.is_available():
         quant_config = BitsAndBytesConfig(
             load_in_4bit=True,
@@ -43,43 +42,40 @@ def load_pipeline():
             bnb_4bit_compute_dtype="float16",
             low_cpu_mem_usage=True,
         )
+        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,
+        )
     else:
-        quant_config = BitsAndBytesConfig(
-            load_in_8bit=True,
-            llm_int8_threshold=6.0,
-            llm_int8_has_fp16_weight=False,
+        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="cpu",           # force CPU
         )
 
-    # 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) Resize embeddings & overlay LoRA
+    # 3) Resize + LoRA overlay
     base.resize_token_embeddings(len(tokenizer))
     model = PeftModel.from_pretrained(
         base,
         ADAPTER_REPO,
         trust_remote_code=True,
-        device_map="auto",
-        offload_folder=OFFLOAD_DIR,
-        offload_state_dict=True,
-        torch_dtype="auto",
+        device_map="auto" if torch.cuda.is_available() else None,
+        torch_dtype=None,
     )
     model.eval()
 
-    # 5) Build sampler pipeline
+    # 4) Build generation pipeline
     gen = pipeline(
         "text-generation",
         model=model,
         tokenizer=tokenizer,
-        device_map="auto",
+        device_map="auto" if torch.cuda.is_available() else None,
         max_new_tokens=MAX_NEW_TOKENS,
         do_sample=True,
         temperature=0.7,
@@ -98,26 +94,20 @@ st.set_page_config(layout="centered")
 st.title("🧠 DeepTalks")
 st.subheader("Your personal AI Companion", divider='grey')
 
-# initialize history
 if "history" not in st.session_state:
-    st.session_state.history = []  # list of (role, text)
+    st.session_state.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
     st.chat_message("user").write(user_input)
     st.session_state.history.append(("You", user_input))
 
-    # build clean context from last turns (texts only)
     recent = st.session_state.history[-CONTEXT_TURNS*2:]
-    context = "\n".join(text for _, text in recent)
-
-    prompt = f"""{SYSTEM}
+    context = "\n".join(t for _, t in recent)
+    prompt  = f"""{SYSTEM}
 
 Context:
 {context}
@@ -125,11 +115,9 @@ Context:
 User: {user_input}
 Assistant:"""
 
-    # generate with spinner
     with st.spinner("Thinking…"):
         try:
             reply = generator(prompt)[0]["generated_text"].strip()
-            # strip stray markers
             for marker in ["User:", "Assistant:"]:
                 if marker in reply:
                     reply = reply.split(marker)[0].strip()
@@ -139,6 +127,5 @@ Assistant:"""
             reply = "I’m sorry, something went wrong."
             st.error(f"Error: {e}")
 
-    # show & store assistant
     st.chat_message("assistant").write(reply)
     st.session_state.history.append(("Bot", reply))