Update app.py

app.py

@@ -3,26 +3,22 @@ from fastapi import FastAPI
 from pydantic import BaseModel
 from typing import Optional
 
-MODEL_ID = os.getenv("MODEL_ID", "gpt2")
-CACHE_DIR = os.getenv("TRANSFORMERS_CACHE", "/app/.cache")
-
-# ensure cache dir exists
+MODEL_ID = os.getenv("MODEL_ID", "distilgpt2")  # smaller model for CPU
+CACHE_DIR = os.getenv("HF_HOME", "/app/.cache")
 os.makedirs(CACHE_DIR, exist_ok=True)
 
 app = FastAPI(title="FastAPI Hugging Face Space")
-
 generator = None
 
 class GenerateRequest(BaseModel):
     prompt: str
-    max_length: Optional[int] = None
+    max_length: Optional[int] = 64
 
 @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)
+    generator = pipeline("text-generation", model=MODEL_ID, cache_dir=CACHE_DIR)
 
 @app.get("/health")
 async def health():
@@ -32,6 +28,5 @@ async def health():
 async def generate(req: GenerateRequest):
     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)
+    out = generator(req.prompt, max_length=req.max_length, num_return_sequences=1)
     return {"generated_text": out[0]["generated_text"]}