Update app.py

app.py

@@ -1,51 +1,98 @@
 import os
-import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from peft import PeftModel
 
 BASE = "microsoft/phi-3-mini-4k-instruct"
 ADAPTER = "Znilsson/survivalai-phi3-lora"
 TOKEN = os.environ.get("HF_TOKEN")
 
-print("Loading base model (first load ~3 min)...")
-tokenizer = AutoTokenizer.from_pretrained(BASE, trust_remote_code=True)
+print("Loading base model (this may take 2-4 minutes on first run)...")
+
+# 4-bit quantization config (huge memory saver)
+quant_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_use_double_quant=True,
+)
+
 model = AutoModelForCausalLM.from_pretrained(
-    BASE, torch_dtype=torch.float32, trust_remote_code=True, low_cpu_mem_usage=True
+    BASE,
+    quantization_config=quant_config,
+    device_map="cpu",                    # Spaces is CPU-only
+    trust_remote_code=True,
+    torch_dtype=torch.float16,           # Avoid deprecation warning
+    attn_implementation="eager",         # Bypass flash-attn / window_size issues
+    low_cpu_mem_usage=True,
+)
+
+print("Attaching LoRA adapter (SurvivalAI fine-tune)...")
+model = PeftModel.from_pretrained(
+    model, 
+    ADAPTER, 
+    token=TOKEN,
+    is_trainable=False
 )
-print("Attaching LoRA adapter...")
-model = PeftModel.from_pretrained(model, ADAPTER, token=TOKEN)
-model = model.merge_and_unload()
+
+# Do NOT merge_and_unload() on CPU in Spaces — it spikes memory too much
+# model = model.merge_and_unload()   # Comment this out for now
+
 model.eval()
-print("Ready.")
+tokenizer = AutoTokenizer.from_pretrained(BASE, trust_remote_code=True)
+
+print("SurvivalAI is ready! (Running in 4-bit on CPU)")
 
 def respond(message, history):
-    msgs = [{"role": "user", "content": message}]
+    # Build full conversation for proper context
+    messages = []
+    for user_msg, assistant_msg in history or []:
+        messages.append({"role": "user", "content": user_msg})
+        if assistant_msg:
+            messages.append({"role": "assistant", "content": assistant_msg})
+    messages.append({"role": "user", "content": message})
+
+    # Apply Phi-3 chat template
     inputs = tokenizer.apply_chat_template(
-        msgs, tokenize=True, add_generation_prompt=True, return_tensors="pt"
-    )
+        messages,
+        tokenize=True,
+        add_generation_prompt=True,
+        return_tensors="pt"
+    ).to(model.device)
+
     with torch.no_grad():
-        out = model.generate(
+        outputs = model.generate(
             inputs,
-            max_new_tokens=400,
+            max_new_tokens=512,          # Increased a bit for better survival answers
             do_sample=True,
             temperature=0.7,
             top_p=0.9,
             repetition_penalty=1.1,
+            pad_token_id=tokenizer.eos_token_id,
         )
-    response = tokenizer.decode(out[0][inputs.shape[1]:], skip_special_tokens=True)
-    return response
 
+    # Decode only the new tokens
+    response = tokenizer.decode(
+        outputs[0][inputs.shape[1]:], 
+        skip_special_tokens=True
+    )
+    return response.strip()
+
+# Gradio interface
 demo = gr.ChatInterface(
-    respond,
-    title="SurvivalAI — Phi-3 LoRA demo",
-    description="Fine-tuned on ~150k survival/preparedness QA pairs. Slow on free CPU (~20s/response).",
+    fn=respond,
+    title="🌲 SurvivalAI — Phi-3 LoRA (Survival / Preparedness Expert)",
+    description="Fine-tuned on survival knowledge from Survivor Library, Army manuals, FEMA, Grokipedia, etc. "
+                "Running quantized on CPU — responses may take 15–60 seconds. Offline-capable foundation for our handheld version.",
     examples=[
         "How do I purify water from a stream with nothing but a pot?",
-        "My friend is hypothermic. What do I do?",
-        "List three edible wild plants in temperate forests.",
+        "My friend is hypothermic. What are the immediate steps?",
+        "List three edible wild plants in temperate forests and how to identify them safely.",
+        "How do I build a basic debris shelter in a forest?",
     ],
+    theme=gr.themes.Soft(),
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()