Create app.py

app.py

@@ -0,0 +1,28 @@
+import os
+
+
+from fastapi import FastAPI
+from pydantic import BaseModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# 1. Load model & tokenizer once at startup
+MODEL_ID = "EQuIP-Queries/EQuIP_3B"
+# Specify cache_dir just in case
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
+
+# 2. Initialize FastAPI
+app = FastAPI()
+
+# 3. Define request schema
+class GenerateRequest(BaseModel):
+    prompt: str
+    max_new_tokens: int = 50
+
+# 4. Inference endpoint
+@app.post("/generate")
+async def generate(req: GenerateRequest):
+    inputs = tokenizer(req.prompt, return_tensors="pt")
+    ids = model.generate(**inputs, max_new_tokens=req.max_new_tokens)
+    text = tokenizer.decode(ids[0], skip_special_tokens=True)
+    return {"generated_text": text}