Update app.py

app.py

@@ -1,28 +1,31 @@
-from fastapi import FastAPI, Request
+from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
-import os
 import uuid
+import os
 
+# FastAPI app setup
 app = FastAPI()
-import os; os.environ["HF_HOME"] = "/tmp/huggingface"
 
-model_id = "hugging-quants/Meta-Llama-3.1-8B-Instruct-GPTQ-INT4"
+# Use HF cache location that's safe in HF Spaces
+os.environ["HF_HOME"] = "/data/huggingface"
+
+# Use a CPU-compatible model (non-GPTQ)
+model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
 hf_token = os.environ.get("HF_TOKEN")
 
+# Load model and tokenizer (no GPU-specific args)
 tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    torch_dtype=torch.float16,
-    device_map="auto",
-    low_cpu_mem_usage=True,
     token=hf_token
-)
+).to("cpu")
 
-# Store per-session system prompts
+# In-memory store for system prompts per session
 session_prompts = {}
 
+# Request body models
 class SystemPrompt(BaseModel):
     prompt: str
 
@@ -44,14 +47,12 @@ def chat(message: UserMessage):
 
     full_prompt = f"<|system|>\n{system}\n<|user|>\n{message.message}\n<|assistant|>\n"
 
-    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
+    inputs = tokenizer(full_prompt, return_tensors="pt").to("cpu")
     outputs = model.generate(
         **inputs,
         max_new_tokens=200,
         pad_token_id=tokenizer.eos_token_id,
     )
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    # Strip input part to isolate model's answer
     answer = response.replace(full_prompt.strip(), "").strip()
     return {"response": answer}