Update app.py

app.py

@@ -1,106 +1,117 @@
-from fastapi import FastAPI,Query
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
-import os
+from fastapi import FastAPI
 from pydantic import BaseModel
 from fastapi.middleware.cors import CORSMiddleware
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+import torch
+import os
 
-
-# ✅ Force Hugging Face cache to /tmp (writable in Spaces)
+# ─────────────────────────────────────────────
+# Make HF Spaces writable (required)
+# ─────────────────────────────────────────────
 os.environ["HF_HOME"] = "/tmp"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp"
 
-
+# ─────────────────────────────────────────────
+# Load model & tokenizer
+# ─────────────────────────────────────────────
 model_id = "hello-ram/unsolth_gpt.20"
-#helloram
+
 tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir="/tmp")
 model = AutoModelForCausalLM.from_pretrained(model_id, cache_dir="/tmp")
+model.eval()     # VERY IMPORTANT on CPU
 
-
-app = FastAPI(title="QA GPT2 API", description="Serving HuggingFace model with FastAPI")
+# ─────────────────────────────────────────────
+# FastAPI config
+# ─────────────────────────────────────────────
+app = FastAPI(title="Unsloth GPT API", description="LoRA + Chat Template + Reasoning API")
 
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
-    allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
-# Request schema
+
+# POST input schema
 class QueryRequest(BaseModel):
     question: str
-    max_new_tokens: int = 50
+    max_new_tokens: int = 64
     temperature: float = 0.7
     top_p: float = 0.9
+    reasoning_effort: str = "medium"  # low / medium / high
 
 
+# ─────────────────────────────────────────────
+# Home
+# ─────────────────────────────────────────────
 @app.get("/")
 def home():
-    return {"message": "Welcome to QA GPT2 API 🚀"}
-
-@app.get("/ask")
-def ask(question: str, max_new_tokens: int = 50):
-    inputs = tokenizer(question, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=max_new_tokens)
-    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return {"question": question, "answer": answer}
-
+    return {"message": "Unsloth GPT API Running Successfully 🚀"}
 
 
-# Health check endpoint
+# ─────────────────────────────────────────────
+# HEALTH CHECK
+# ─────────────────────────────────────────────
 @app.get("/health")
 def health():
     return {"status": "ok"}
 
-# Inference endpoint
+
+# ─────────────────────────────────────────────
+# MAIN PREDICTION ENDPOINT (FULL CHAT TEMPLATE)
+# ─────────────────────────────────────────────
 @app.post("/predict")
-def predict(request: QueryRequest):
+def predict(req: QueryRequest):
+
+    # Build correct Unsloth chat template input
     messages = [
-        {"role": "system", "content": "reasoning language: english\n\nYou are a helpful assistant."},
-        {"role": "user", "content": request.question}
+        {"role": "system", "content": "reasoning language: english\nYou are a helpful assistant."},
+        {"role": "user", "content": req.question},
     ]
 
-    inputs = tokenizer.apply_chat_template(
+    # Apply Unsloth chat template
+    input_ids = tokenizer.apply_chat_template(
         messages,
         add_generation_prompt=True,
+        reasoning_effort=req.reasoning_effort,
         return_tensors="pt"
-    ).to(model.device)
-
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=request.max_new_tokens,
-        do_sample=True,
-        temperature=request.temperature,
-        top_p=request.top_p,
-        pad_token_id=tokenizer.eos_token_id
     )
 
-    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return {"question": request.question, "answer": answer}
+    # Generate
+    with torch.no_grad():
+        output = model.generate(
+            input_ids=input_ids,
+            max_new_tokens=req.max_new_tokens,
+            do_sample=True,
+            temperature=req.temperature,
+            top_p=req.top_p,
+            pad_token_id=tokenizer.eos_token_id,
+        )
 
+    # Decode
+    answer = tokenizer.decode(output[0], skip_special_tokens=True)
 
+    return {
+        "question": req.question,
+        "answer": answer
+    }
 
 
-@app.get("/answers")
-def predict(question: str = Query(..., description="The question to ask"), max_new_tokens: int = Query(50, description="Max new tokens to generate")):
-    # Tokenize the input question
+# ─────────────────────────────────────────────
+# SIMPLE NON-CHAT ENDPOINT (/ask)
+# ─────────────────────────────────────────────
+@app.get("/ask")
+def ask(question: str, max_new_tokens: int = 50):
+
     inputs = tokenizer(question, return_tensors="pt")
 
-    # Generate output from model
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
-        temperature=0.7,
-        top_p=0.9,
-        pad_token_id=tokenizer.eos_token_id,
-        return_dict_in_generate=True
-    )
+    with torch.no_grad():
+        output = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            pad_token_id=tokenizer.eos_token_id,
+        )
 
-    # Decode output
-    answer = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
+    answer = tokenizer.decode(output[0], skip_special_tokens=True)
 
-    return {
-        "question": question,
-        "answer": answer
-    }
+    return {"question": question, "answer": answer}