Update app.py

app.py

@@ -1,16 +1,20 @@
-from fastapi import FastAPI
-from pydantic import BaseModel
-from fastapi.middleware.cors import CORSMiddleware
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import unsloth  # MUST be imported before transformers
+from unsloth import FastLanguageModel
+from transformers import AutoTokenizer
 import torch
 import os
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
 
-# Make HF Spaces writable
+# Optional: HF cache
 os.environ["HF_HOME"] = "/tmp"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp"
 
+# -------------------------------
 # FastAPI app
-app = FastAPI(title="Unsloth GPT API")
+# -------------------------------
+app = FastAPI(title="Unsolth GPT OSS API")
 
 app.add_middleware(
     CORSMiddleware,
@@ -19,20 +23,46 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+# -------------------------------
 # Model variables
+# -------------------------------
 model = None
 tokenizer = None
-model_id = "hello-ram/unsolth_gpt.20"
 
+# Paths
+base_model_name = "unsloth/gpt-oss-20b"  # Pretrained GPT-OSS base
+lora_model_path = "./finetuned_model"     # Your LoRA weights in the Space repo
+
+# -------------------------------
 # Load model lazily
+# -------------------------------
 def load_model():
     global model, tokenizer
     if model is None or tokenizer is None:
-        tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir="/tmp")
-        model = AutoModelForCausalLM.from_pretrained(model_id, cache_dir="/tmp")
+        tokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)
+
+        base_model = FastLanguageModel.from_pretrained(
+            base_model_name, trust_remote_code=True
+        )
+
+        model = FastLanguageModel.get_peft_model(
+            base_model,
+            r=8,
+            target_modules=[
+                "q_proj", "k_proj", "v_proj", "o_proj",
+                "gate_proj", "up_proj", "down_proj"
+            ],
+            lora_alpha=16,
+            lora_dropout=0,
+            bias="none",
+            use_gradient_checkpointing="unsloth",
+            state_dict=torch.load(os.path.join(lora_model_path, "adapter_model.safetensors"))
+        )
         model.eval()
 
+# -------------------------------
 # Input schema
+# -------------------------------
 class QueryRequest(BaseModel):
     question: str
     max_new_tokens: int = 64
@@ -40,26 +70,33 @@ class QueryRequest(BaseModel):
     top_p: float = 0.9
     reasoning_effort: str = "medium"
 
+# -------------------------------
 # Health check
+# -------------------------------
 @app.get("/")
 def health():
     return {"status": "ok"}
 
-# Main predict endpoint
+# -------------------------------
+# Prediction endpoint
+# -------------------------------
 @app.post("/predict")
 def predict(req: QueryRequest):
     load_model()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
 
-    input_ids = tokenizer(req.question, return_tensors="pt").input_ids
+    inputs = tokenizer.apply_chat_template(
+        [{"role": "user", "content": req.question}],
+        add_generation_prompt=True,
+        return_tensors="pt",
+        return_dict=True,
+        reasoning_effort=req.reasoning_effort
+    ).to(device)
 
     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,
+            **inputs,
+            max_new_tokens=req.max_new_tokens
         )
 
     answer = tokenizer.decode(output[0], skip_special_tokens=True)