Update app.py

app.py

@@ -1,97 +1,53 @@
 import os
 import torch
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from peft import PeftModel
 
-BASE = "microsoft/phi-3-mini-4k-instruct"
-ADAPTER = "Znilsson/survivalai-phi3-lora"
-TOKEN = os.environ.get("HF_TOKEN")
+BASE    = "microsoft/phi-3-mini-4k-instruct"
+ADAPTER = "Znilsson/survivalai-phi3-lora"   # <-- replace if your adapter repo ID differs
+TOKEN   = os.environ.get("HF_TOKEN")
 
-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,
-)
+print("Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(BASE, trust_remote_code=True)
 
+print("Loading base model (fp16)...")
 model = AutoModelForCausalLM.from_pretrained(
     BASE,
-    quantization_config=quant_config,
-    device_map="cpu",                    # Spaces is CPU-only
+    dtype=torch.float16,
+    device_map="auto",
     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
-)
-
-# Do NOT merge_and_unload() on CPU in Spaces — it spikes memory too much
-# model = model.merge_and_unload()   # Comment this out for now
-
+print("Attaching + merging LoRA adapter...")
+model = PeftModel.from_pretrained(model, ADAPTER, token=TOKEN)
+model = model.merge_and_unload()
 model.eval()
-tokenizer = AutoTokenizer.from_pretrained(BASE, trust_remote_code=True)
-
-print("SurvivalAI is ready! (Running in 4-bit on CPU)")
-
-def respond(message, history):
-    # 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(
-        messages,
-        tokenize=True,
-        add_generation_prompt=True,
-        return_tensors="pt"
-    ).to(model.device)
+def chat(message, history):
+    prompt = ""
+    for user, assistant in history:
+        prompt += f"<|user|>\n{user}<|end|>\n<|assistant|>\n{assistant}<|end|>\n"
+    prompt += f"<|user|>\n{message}<|end|>\n<|assistant|>\n"
 
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     with torch.no_grad():
-        outputs = model.generate(
-            inputs,
-            max_new_tokens=512,          # Increased a bit for better survival answers
-            do_sample=True,
+        out = model.generate(
+            **inputs,
+            max_new_tokens=400,
             temperature=0.7,
             top_p=0.9,
-            repetition_penalty=1.1,
+            do_sample=True,
             pad_token_id=tokenizer.eos_token_id,
         )
+    resp = tokenizer.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
+    return resp.strip()
 
-    # 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(
-    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 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(),
+    fn=chat,
+    title="SurvivalAI",
+    description="Fine-tuned Phi-3-mini on survival & emergency preparedness corpus.",
 )
 
 if __name__ == "__main__":