Create app.py

app.py

@@ -0,0 +1,56 @@
+# app.py
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+import os
+
+MODEL_ID = "Mayur74/tinyllama-finetune-mayur"   # or local folder name if you upload merged_model/
+USE_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+# If your model is private, the token should be set via Secrets (HF_SECRETS)
+HF_TOKEN = os.environ.get("HF_TOKEN")  # set this in Space Settings -> Secrets (recommended for private models)
+
+app = FastAPI(title="TinyLlama Inference")
+
+class GenRequest(BaseModel):
+    prompt: str
+    max_new_tokens: int = 128
+    temperature: float = 0.7
+
+# Load tokenizer & model once at startup
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_auth_token=HF_TOKEN)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.float16 if USE_DEVICE=="cuda" else torch.float32,
+        device_map="auto" if USE_DEVICE=="cuda" else None,
+        trust_remote_code=False,  # change if model needs trust_remote_code
+        use_auth_token=HF_TOKEN
+    )
+    # reduce cache if needed
+    model.eval()
+    return tokenizer, model
+
+tokenizer, model = load_model()
+
+@app.post("/generate")
+async def generate(req: GenRequest):
+    try:
+        inputs = tokenizer(req.prompt, return_tensors="pt", truncation=True)
+        if torch.cuda.is_available():
+            inputs = {k: v.to("cuda") for k, v in inputs.items()}
+
+        with torch.no_grad():
+            out = model.generate(
+                **inputs,
+                max_new_tokens=req.max_new_tokens,
+                temperature=req.temperature,
+                do_sample=True,
+                top_p=0.95,
+                pad_token_id=tokenizer.eos_token_id,
+            )
+        text = tokenizer.decode(out[0], skip_special_tokens=True)
+        return {"generated_text": text}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))