Update app.py

app.py

@@ -1,19 +1,16 @@
 import os
 import streamlit as st
+import torch
 from transformers import (
-    pipeline, AutoTokenizer, AutoModelForCausalLM
+    pipeline,
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    BitsAndBytesConfig,
 )
 from peft import PeftModel
 from supabase import create_client
 from sentence_transformers import SentenceTransformer
 
-# Try to import bitsandbytes for 4-bit; fall back if missing
-try:
-    from transformers import BitsAndBytesConfig
-    BNB_AVAILABLE = True
-except ImportError:
-    BNB_AVAILABLE = False
-
 # ── Supabase setup ─────────────────────────────────────────────────────────
 SUPA_URL = os.getenv("SUPABASE_URL")
 SUPA_KEY = os.getenv("SUPABASE_SERVICE_ROLE_KEY")
@@ -28,14 +25,14 @@ embedder = get_embedder()
 
 @st.cache_data(show_spinner=False)
 def fetch_mems(query, k=3):
-    vec = embedder.encode(query).astype('float32').tolist()
+    vec = embedder.encode(query).astype("float32").tolist()
     return supabase.rpc(
-        "match_memories",
+        "match_memories", 
         {"query_embedding": vec, "match_count": k}
     ).execute().data
 
 def add_mem(speaker, text):
-    vec = embedder.encode(text).astype('float32').tolist()
+    vec = embedder.encode(text).astype("float32").tolist()
     supabase.table("memories").insert({
         "speaker": speaker, "text": text, "embedding": vec
     }).execute()
@@ -44,6 +41,7 @@ def add_mem(speaker, text):
 @st.cache_resource(show_spinner=False)
 def load_generator():
     REPO = "sourize/phi2-memory-lora"
+
     # 1) Tokenizer
     tokenizer = AutoTokenizer.from_pretrained(
         REPO, trust_remote_code=True, padding_side="left"
@@ -51,8 +49,9 @@ def load_generator():
     if tokenizer.pad_token_id is None:
         tokenizer.add_special_tokens({"pad_token": "[PAD]"})
 
-    # 2) Base model load (with or without 4-bit)
-    if BNB_AVAILABLE:
+    # 2) Decide quantization vs. fp16/fp32
+    use_4bit = torch.cuda.is_available()
+    if use_4bit:
         bnb_config = BitsAndBytesConfig(
             load_in_4bit=True,
             bnb_4bit_quant_type="nf4",
@@ -66,10 +65,12 @@ def load_generator():
             device_map="auto"
         )
     else:
+        # CPU or no CUDA: use fp16 if available, else fp32
+        dtype = torch.float16 if torch.cuda.is_available() or torch.cuda.device_count()>0 else torch.float32
         base = AutoModelForCausalLM.from_pretrained(
             "microsoft/phi-2",
             trust_remote_code=True,
-            torch_dtype="auto",
+            torch_dtype=dtype,
             device_map="auto"
         )