Update app.py

app.py

@@ -1,23 +1,26 @@
-# app.py — EduPrompt FastAPI backend
+# app.py — EduPrompt FastAPI backend (HF Spaces-safe)
 
 from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import pipeline
 import time
 import os
+import asyncio
 
-# --- Hugging Face cache fix for Spaces (permission-safe) ---
+# ---- Hugging Face cache: force writable dir on Spaces ----
 os.environ["HF_HOME"] = "/tmp"
 os.environ["HF_HUB_CACHE"] = "/tmp"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp"
 
-# --- FastAPI app + CORS ---
+CACHE_DIR = "/tmp"  # single source of truth
+
+# ---- FastAPI + CORS ----
 app = FastAPI(title="EduPrompt API")
 
 from fastapi.middleware.cors import CORSMiddleware
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
+    allow_origins=["*"],   # set your domain(s) in production
     allow_methods=["*"],
     allow_headers=["*"],
 )
@@ -26,31 +29,67 @@ app.add_middleware(
 def health():
     return {"ok": True, "service": "eduprompt-api"}
 
-# --- Lazy model loaders ---
+# ---- Lazy models (loaded on first use), with cache_dir enforced ----
 _summarizer = None
 _rewriter = None
 _proofreader = None
 _code_explainer = None
 
+def _make_pipeline(task: str, model_id: str):
+    """
+    Create a HF pipeline that always caches to /tmp (writable on Spaces).
+    Retries once on cache-related OSError.
+    """
+    try:
+        return pipeline(task, model=model_id, cache_dir=CACHE_DIR)
+    except OSError as e:
+        # Rare HF cache race; wait briefly and retry once
+        if "/.cache" in str(e) or "PermissionError" in str(e):
+            time.sleep(1.5)
+            os.environ["HF_HOME"] = CACHE_DIR
+            os.environ["HF_HUB_CACHE"] = CACHE_DIR
+            os.environ["TRANSFORMERS_CACHE"] = CACHE_DIR
+            return pipeline(task, model=model_id, cache_dir=CACHE_DIR)
+        raise
+
 def get_models():
     global _summarizer, _rewriter, _proofreader, _code_explainer
     if _summarizer is None:
-        _summarizer = pipeline("summarization", model="t5-small")
+        _summarizer = _make_pipeline("summarization", "t5-small")  # CPU-friendly
     if _rewriter is None:
-        _rewriter = pipeline("text2text-generation", model="google/flan-t5-small")
+        _rewriter = _make_pipeline("text2text-generation", "google/flan-t5-small")
     if _proofreader is None:
-        _proofreader = pipeline("text2text-generation", model="google/flan-t5-small")
+        _proofreader = _make_pipeline("text2text-generation", "google/flan-t5-small")
     if _code_explainer is None:
-        _code_explainer = pipeline("text2text-generation", model="Salesforce/codet5p-220m")
+        _code_explainer = _make_pipeline("text2text-generation", "Salesforce/codet5p-220m")
     return _summarizer, _rewriter, _proofreader, _code_explainer
 
-# --- Request schema ---
+# (Optional) tiny warmup to trigger downloads after first request to /run
+@app.on_event("startup")
+async def _post_start_note():
+    # We don't download at startup to keep boot fast; models load on first call.
+    # Leaving this here in case you ever want to warm them:
+    # asyncio.create_task(_warm_once())
+    pass
+
+async def _warm_once():
+    try:
+        s, r, p, c = get_models()
+        _ = s("warm up", max_length=10, min_length=5, do_sample=False)
+        _ = r("rewrite: warm up", max_new_tokens=8)
+        _ = p("proofread: warm up", max_new_tokens=8)
+        _ = c("explain: print(1)", max_new_tokens=12)
+    except Exception:
+        # Ignore warm errors; real requests will still retry/load.
+        pass
+
+# ---- Request schema ----
 class InputData(BaseModel):
-    task: str
-    input: str
+    task: str                 # summarize | rewrite | proofread | explain_code
+    input: str                # user text / code
     params: dict | None = None
 
-# --- Core endpoint ---
+# ---- Core endpoint ----
 @app.post("/run")
 async def run_task(data: InputData):
     start = time.time()
@@ -62,28 +101,50 @@ async def run_task(data: InputData):
     if task not in {"summarize", "rewrite", "proofread", "explain_code"}:
         return {"error": f"Unsupported task '{task}'."}
 
-    enhanced = ""
+    # Load models (lazy, cached to /tmp)
     try:
         summarizer, rewriter, proofreader, code_explainer = get_models()
+    except Exception as e:
+        return {"error": f"model_load_failed: {type(e).__name__}: {str(e)}"}
 
+    enhanced = ""
+    try:
         if task == "summarize":
             enhanced = f"You are an expert explainer. Summarize clearly and concisely:\n{text}"
-            out = summarizer(enhanced, max_length=120, min_length=30, do_sample=False, truncation=True)[0]["summary_text"]
+            out = summarizer(
+                enhanced,
+                max_length=120,
+                min_length=30,
+                do_sample=False,
+                truncation=True
+            )[0]["summary_text"]
             model_id = "t5-small"
 
         elif task == "rewrite":
             enhanced = f"You are a writing assistant. Rewrite this text for clarity and tone:\n{text}"
-            out = rewriter(enhanced, max_new_tokens=150, truncation=True)[0]["generated_text"]
+            out = rewriter(
+                enhanced,
+                max_new_tokens=150,
+                truncation=True
+            )[0]["generated_text"]
             model_id = "google/flan-t5-small"
 
         elif task == "proofread":
             enhanced = f"You are a grammar and style editor. Correct and improve this text:\n{text}"
-            out = proofreader(enhanced, max_new_tokens=150, truncation=True)[0]["generated_text"]
+            out = proofreader(
+                enhanced,
+                max_new_tokens=150,
+                truncation=True
+            )[0]["generated_text"]
             model_id = "google/flan-t5-small"
 
-        else:
+        else:  # explain_code
             enhanced = f"You are a programming tutor. Explain what this code does in simple language:\n{text}"
-            out = code_explainer(enhanced, max_new_tokens=200, truncation=True)[0]["generated_text"]
+            out = code_explainer(
+                enhanced,
+                max_new_tokens=200,
+                truncation=True
+            )[0]["generated_text"]
             model_id = "Salesforce/codet5p-220m"
 
     except Exception as e: