Update app.py

app.py

@@ -1,25 +1,18 @@
-from fastapi import FastAPI, Request
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-app = FastAPI()
-
-# Load model and tokenizer
-model_id = "meta-llama/Llama-3.1-8B"
-
-tokenizer = AutoTokenizer.from_pretrained(model_id)
+model_id = "hugging-quants/Meta-Llama-3.1-8B-Instruct-GPTQ-INT4"
 
+tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=True)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    device_map="auto"  # Will auto-detect if CUDA or CPU
+    device_map="auto",
+    torch_dtype=torch.float16,
+    low_cpu_mem_usage=True,
+    use_auth_token=True  # ⬅️ ensures gated model access
 )
 
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=256)
-
-@app.post("/chat")
-async def chat(request: Request):
-    data = await request.json()
-    prompt = data.get("prompt", "")
-    output = pipe(prompt)[0]['generated_text']
-    return {"response": output}
+def generate(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    outputs = model.generate(**inputs, max_new_tokens=128)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)