Update app.py

app.py

@@ -1,25 +1,36 @@
-from fastapi import FastAPI,Query
+from fastapi import FastAPI, Query, HTTPException
 from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
-import os
 from pydantic import BaseModel
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import HTMLResponse
 from fastapi.staticfiles import StaticFiles
+import os
+import torch
 
-
-# ✅ Force Hugging Face cache to /tmp (writable in Spaces)
+# ✅ Hugging Face cache directory
 os.environ["HF_HOME"] = "/tmp"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp"
 
+# -----------------------
+# Model Setup
+# -----------------------
+model_id = "LLM360/K2-Think"
 
-model_id = "rabiyulfahim/qa_python_gpt2"
-
+# Load tokenizer and model
+print("Loading tokenizer and model...")
 tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir="/tmp")
-model = AutoModelForCausalLM.from_pretrained(model_id, cache_dir="/tmp")
-
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    cache_dir="/tmp",
+    device_map="auto",       # Automatically select GPU/CPU
+    torch_dtype=torch.float16
+)
+print("Model loaded successfully!")
 
-app = FastAPI(title="QA GPT2 API UI", description="Serving HuggingFace model with FastAPI")
+# -----------------------
+# FastAPI Setup
+# -----------------------
+app = FastAPI(title="K2-Think QA API", description="Serving K2-Think Hugging Face model with FastAPI")
 
 app.add_middleware(
     CORSMiddleware,
@@ -28,28 +39,25 @@ app.add_middleware(
     allow_methods=["*"],
     allow_headers=["*"],
 )
-# Request schema
+
+# Mount static folder
+app.mount("/static", StaticFiles(directory="static"), name="static")
+
+# -----------------------
+# Request Schema
+# -----------------------
 class QueryRequest(BaseModel):
     question: str
     max_new_tokens: int = 50
     temperature: float = 0.7
     top_p: float = 0.9
 
-
+# -----------------------
+# Endpoints
+# -----------------------
 @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}
-
-
-# Mount static folder
-app.mount("/static", StaticFiles(directory="static"), name="static")
+    return {"message": "Welcome to K2-Think QA API 🚀"}
 
 @app.get("/ui", response_class=HTMLResponse)
 def serve_ui():
@@ -57,55 +65,42 @@ def serve_ui():
     with open(html_path, "r", encoding="utf-8") as f:
         return HTMLResponse(f.read())
 
-
-# Health check endpoint
 @app.get("/health")
 def health():
     return {"status": "ok"}
 
-# Inference endpoint
+@app.get("/ask")
+def ask(question: str = Query(...), max_new_tokens: int = Query(50)):
+    try:
+        inputs = tokenizer(question, return_tensors="pt")
+        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
+        )
+        answer = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
+        return {"question": question, "answer": answer}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
 @app.post("/predict")
 def predict(request: QueryRequest):
-    inputs = tokenizer(request.question, return_tensors="pt")
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=request.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
-    )
-
-    answer = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
-    return {
-        "question": request.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
-    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
-    )
-
-    # Decode output
-    answer = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
-
-    return {
-        "question": question,
-        "answer": answer
-    }
+    try:
+        inputs = tokenizer(request.question, return_tensors="pt")
+        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,
+            return_dict_in_generate=True
+        )
+        answer = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
+        return {"question": request.question, "answer": answer}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))