app.py

import os
from fastapi import FastAPI
from pydantic import BaseModel
from typing import Optional

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] = 64

@app.on_event("startup")
async def load_model():
    global generator
    from transformers import pipeline
    generator = pipeline("text-generation", model=MODEL_ID, cache_dir=CACHE_DIR)

@app.get("/health")
async def health():
    return {"status": "ok", "model": MODEL_ID, "cache": CACHE_DIR}

@app.post("/generate")
async def generate(req: GenerateRequest):
    if generator is None:
        return {"error": "model not loaded yet"}
    out = generator(req.prompt, max_length=req.max_length, num_return_sequences=1)
    return {"generated_text": out[0]["generated_text"]}