Update app.py

app.py

@@ -1,39 +1,50 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel
 import torch
 
 app = FastAPI()
 
-# ---- Load your HF model repo ----
-MODEL_REPO = "hello-ram/unsolth_gpt.20"
+# 1. Base model
+BASE_MODEL = "gpt2"
+
+# 2. LoRA adapter repo
+LORA_REPO = "hello-ram/unsolth_gpt.20"
+
 
 print("Loading tokenizer...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
 
-print("Loading model...")
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_REPO,
+print("Loading base model...")
+base_model = AutoModelForCausalLM.from_pretrained(
+    BASE_MODEL,
     torch_dtype=torch.float16,
+    device_map="auto",
+)
+
+print("Applying LoRA adapter...")
+model = PeftModel.from_pretrained(
+    base_model,
+    LORA_REPO,
     device_map="auto"
 )
 
-# ---------- ROUTES -------------
+model.eval()
+
 
 @app.get("/")
 async def root():
-    return {
-        "message": "🚀 FastAPI MPT Model Running on Hugging Face Spaces",
-        "endpoints": ["/", "/status", "/generate"]
-    }
+    return {"msg": "LoRA model running", "endpoints": ["/status", "/generate"]}
+
 
 @app.get("/status")
 async def status():
     return {
         "status": "ok",
-        "model": MODEL_REPO,
-        "device": str(model.device),
-        "torch_dtype": str(model.dtype)
+        "base_model": BASE_MODEL,
+        "lora_model": LORA_REPO,
+        "device": str(model.device)
     }
 
 
@@ -45,11 +56,12 @@ class InputText(BaseModel):
 async def generate_text(data: InputText):
     inputs = tokenizer(data.text, return_tensors="pt").to(model.device)
 
-    output = model.generate(
-        **inputs,
-        max_new_tokens=200,
-        temperature=0.7
-    )
+    with torch.no_grad():
+        output = model.generate(
+            **inputs,
+            max_new_tokens=200,
+            temperature=0.7
+        )
 
-    generated = tokenizer.decode(output[0], skip_special_tokens=True)
-    return {"response": generated}
+    text = tokenizer.decode(output[0], skip_special_tokens=True)
+    return {"response": text}