Create app.py

app.py

@@ -0,0 +1,30 @@
+from os import getenv
+from fastapi import FastAPI
+from pydantic import BaseModel
+from typing import Optional
+from transformers import pipeline
+
+MODEL_ID = getenv("MODEL_ID", "gpt2")            # set by env if you want another model
+GEN_KWARGS = {"max_length": 64, "num_return_sequences": 1}
+
+app = FastAPI(title="FastAPI Hugging Face Space")
+
+# load pipeline once on startup
+generator = pipeline("text-generation", model=MODEL_ID)
+
+class GenerateRequest(BaseModel):
+    prompt: str
+    max_length: Optional[int] = None
+
+@app.get("/health")
+async def health():
+    return {"status": "ok", "model": MODEL_ID}
+
+@app.post("/generate")
+async def generate(req: GenerateRequest):
+    kwargs = GEN_KWARGS.copy()
+    if req.max_length:
+        kwargs["max_length"] = req.max_length
+    out = generator(req.prompt, **kwargs)
+    # pipeline returns a list with dicts containing "generated_text"
+    return {"generated_text": out[0]["generated_text"]}