Update app.py

app.py

@@ -1,30 +1,37 @@
-from os import getenv
+import os
 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}
+MODEL_ID = os.getenv("MODEL_ID", "gpt2")
+CACHE_DIR = os.getenv("TRANSFORMERS_CACHE", "/app/.cache")
+
+# ensure cache dir exists
+os.makedirs(CACHE_DIR, exist_ok=True)
 
 app = FastAPI(title="FastAPI Hugging Face Space")
 
-# load pipeline once on startup
-generator = pipeline("text-generation", model=MODEL_ID)
+generator = None
 
 class GenerateRequest(BaseModel):
     prompt: str
     max_length: Optional[int] = None
 
+@app.on_event("startup")
+async def load_model():
+    global generator
+    # import here so transformers uses the configured cache
+    from transformers import pipeline
+    generator = pipeline("text-generation", model=MODEL_ID)
+
 @app.get("/health")
 async def health():
-    return {"status": "ok", "model": MODEL_ID}
+    return {"status": "ok", "model": MODEL_ID, "cache": CACHE_DIR}
 
 @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"
+    if generator is None:
+        return {"error": "model not loaded yet"}
+    max_len = req.max_length or 64
+    out = generator(req.prompt, max_length=max_len, num_return_sequences=1)
     return {"generated_text": out[0]["generated_text"]}