Update app.py

app.py

@@ -1,108 +1,119 @@
-# app.py — EduPrompt FastAPI backend (final HF Spaces fix)
+# app.py — EduPrompt API (per-task lazy load + cache-safe on Spaces)
 
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import pipeline
 from fastapi.middleware.cors import CORSMiddleware
+from transformers import pipeline
 import os, time
 
-# ---- Fix for Hugging Face Spaces cache permissions ----
-CACHE_DIR = "/tmp"
-os.environ["HF_HOME"] = CACHE_DIR
-os.environ["HF_HUB_CACHE"] = CACHE_DIR
-os.environ["TRANSFORMERS_CACHE"] = CACHE_DIR
+# ---------- Force ALL caches to /tmp ----------
+BASE = "/tmp"
+os.environ["HF_HOME"] = f"{BASE}/hf"
+os.environ["HF_HUB_CACHE"] = f"{BASE}/hf"
+os.environ["HUGGINGFACE_HUB_CACHE"] = f"{BASE}/hf"
+os.environ["TRANSFORMERS_CACHE"] = f"{BASE}/hf/transformers"
+os.environ["XDG_CACHE_HOME"] = f"{BASE}/xdg"
+os.environ["TORCH_HOME"] = f"{BASE}/torch"
+os.environ["SENTENCEPIECE_CACHE"] = f"{BASE}/sp"
+for d in [
+    os.environ["HF_HOME"], os.environ["HF_HUB_CACHE"], os.environ["HUGGINGFACE_HUB_CACHE"],
+    os.environ["TRANSFORMERS_CACHE"], os.environ["XDG_CACHE_HOME"],
+    os.environ["TORCH_HOME"], os.environ["SENTENCEPIECE_CACHE"]
+]:
+    os.makedirs(d, exist_ok=True)
 
 app = FastAPI(title="EduPrompt API")
-
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
+    allow_origins=["*"],  # tighten in prod
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
 @app.get("/")
 def health():
-    return {"ok": True, "service": "eduprompt-api"}
-
-# Lazy models (loaded on first use)
+    # prove /tmp is writable
+    try:
+        with open(f"{BASE}/eduprompt_write_test.txt", "w") as f:
+            f.write("ok")
+        writable = True
+    except Exception:
+        writable = False
+    return {"ok": True, "service": "eduprompt-api", "tmpWritable": writable}
+
+# ---------- lazy singletons ----------
 _summarizer = None
 _rewriter = None
 _proofreader = None
 _code_explainer = None
 
 def safe_pipeline(task: str, model_id: str):
-    """Always download/cache models inside /tmp (Spaces-safe)."""
-    try:
-        return pipeline(task, model=model_id, cache_dir=CACHE_DIR)
-    except OSError:
-        # Force reset and retry once if cache issue
-        time.sleep(1)
-        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)
+    """Cache every model in its own /tmp subdir. CPU-only."""
+    model_cache = os.path.join(os.environ["TRANSFORMERS_CACHE"], model_id.replace("/", "_"))
+    os.makedirs(model_cache, exist_ok=True)
+    return pipeline(task, model=model_id, cache_dir=model_cache, trust_remote_code=True, device=-1)
 
-def get_models():
+def get_model(task: str):
+    """Load ONLY the model needed for this task."""
     global _summarizer, _rewriter, _proofreader, _code_explainer
-    if _summarizer is None:
-        _summarizer = safe_pipeline("summarization", "t5-small")
-    if _rewriter is None:
-        _rewriter = safe_pipeline("text2text-generation", "google/flan-t5-small")
-    if _proofreader is None:
-        _proofreader = safe_pipeline("text2text-generation", "google/flan-t5-small")
-    if _code_explainer is None:
-        _code_explainer = safe_pipeline("text2text-generation", "Salesforce/codet5p-220m")
-    return _summarizer, _rewriter, _proofreader, _code_explainer
+    if task == "summarize":
+        if _summarizer is None:
+            _summarizer = safe_pipeline("summarization", "t5-small")
+        return _summarizer, "t5-small"
+    if task == "rewrite":
+        if _rewriter is None:
+            _rewriter = safe_pipeline("text2text-generation", "google/flan-t5-small")
+        return _rewriter, "google/flan-t5-small"
+    if task == "proofread":
+        if _proofreader is None:
+            _proofreader = safe_pipeline("text2text-generation", "google/flan-t5-small")
+        return _proofreader, "google/flan-t5-small"
+    if task == "explain_code":
+        if _code_explainer is None:
+            _code_explainer = safe_pipeline("text2text-generation", "Salesforce/codet5p-220m")
+        return _code_explainer, "Salesforce/codet5p-220m"
+    raise ValueError(f"Unsupported task '{task}'")
 
 class InputData(BaseModel):
-    task: str
+    task: str                  # summarize | rewrite | proofread | explain_code
     input: str
+    params: dict | None = None
 
 @app.post("/run")
 async def run_task(data: InputData):
     start = time.time()
-    task = data.task.strip().lower()
-    text = data.input.strip()
+    task = (data.task or "").strip().lower()
+    text = (data.input or "").strip()
     if not text:
         return {"error": "Empty input text."}
+    if task not in {"summarize", "rewrite", "proofread", "explain_code"}:
+        return {"error": f"Unsupported task '{task}'."}
 
+    # Load only what we need
     try:
-        summarizer, rewriter, proofreader, code_explainer = get_models()
+        model, model_used = get_model(task)
     except Exception as e:
         return {"error": f"model_load_failed: {type(e).__name__}: {str(e)}"}
 
     try:
         if task == "summarize":
             prompt = f"You are an expert explainer. Summarize clearly and concisely:\n{text}"
-            result = summarizer(prompt, max_length=120, min_length=30, truncation=True)[0]["summary_text"]
-            model_used = "t5-small"
-
+            output = model(prompt, max_length=120, min_length=30, truncation=True, do_sample=False)[0]["summary_text"]
         elif task == "rewrite":
-            prompt = f"Rewrite this text for clarity and tone:\n{text}"
-            result = rewriter(prompt, max_new_tokens=150, truncation=True)[0]["generated_text"]
-            model_used = "google/flan-t5-small"
-
+            prompt = f"You are a writing assistant. Rewrite this text for clarity and tone:\n{text}"
+            output = model(prompt, max_new_tokens=150, truncation=True)[0]["generated_text"]
         elif task == "proofread":
             prompt = f"Correct and improve grammar and style:\n{text}"
-            result = proofreader(prompt, max_new_tokens=150, truncation=True)[0]["generated_text"]
-            model_used = "google/flan-t5-small"
-
-        elif task == "explain_code":
+            output = model(prompt, max_new_tokens=150, truncation=True)[0]["generated_text"]
+        else:  # explain_code
             prompt = f"Explain what this code does in simple language:\n{text}"
-            result = code_explainer(prompt, max_new_tokens=200, truncation=True)[0]["generated_text"]
-            model_used = "Salesforce/codet5p-220m"
-
-        else:
-            return {"error": f"Unsupported task '{task}'."}
-
+            output = model(prompt, max_new_tokens=200, truncation=True)[0]["generated_text"]
     except Exception as e:
         return {"error": f"inference_failed: {type(e).__name__}: {str(e)}"}
 
-    latency = round((time.time() - start) * 1000, 2)
     return {
         "enhancedPrompt": prompt,
-        "output": result,
+        "output": output,
         "model": model_used,
-        "latencyMs": latency
-    }
+        "latencyMs": round((time.time() - start) * 1000, 2),
+    }