Update app.py

app.py

@@ -1,26 +1,22 @@
-# app.py — EduPrompt FastAPI backend (HF Spaces-safe)
+# app.py — EduPrompt FastAPI backend (final HF Spaces fix)
 
 from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import pipeline
-import time
-import os
-import asyncio
-
-# ---- Hugging Face cache: force writable dir on Spaces ----
-os.environ["HF_HOME"] = "/tmp"
-os.environ["HF_HUB_CACHE"] = "/tmp"
-os.environ["TRANSFORMERS_CACHE"] = "/tmp"
+from fastapi.middleware.cors import CORSMiddleware
+import os, time
 
-CACHE_DIR = "/tmp"  # single source of truth
+# ---- 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
 
-# ---- FastAPI + CORS ----
 app = FastAPI(title="EduPrompt API")
 
-from fastapi.middleware.cors import CORSMiddleware
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],   # set your domain(s) in production
+    allow_origins=["*"],
     allow_methods=["*"],
     allow_headers=["*"],
 )
@@ -29,131 +25,84 @@ app.add_middleware(
 def health():
     return {"ok": True, "service": "eduprompt-api"}
 
-# ---- Lazy models (loaded on first use), with cache_dir enforced ----
+# Lazy models (loaded on first use)
 _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.
-    """
+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 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
+    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)
 
 def get_models():
     global _summarizer, _rewriter, _proofreader, _code_explainer
     if _summarizer is None:
-        _summarizer = _make_pipeline("summarization", "t5-small")  # CPU-friendly
+        _summarizer = safe_pipeline("summarization", "t5-small")
     if _rewriter is None:
-        _rewriter = _make_pipeline("text2text-generation", "google/flan-t5-small")
+        _rewriter = safe_pipeline("text2text-generation", "google/flan-t5-small")
     if _proofreader is None:
-        _proofreader = _make_pipeline("text2text-generation", "google/flan-t5-small")
+        _proofreader = safe_pipeline("text2text-generation", "google/flan-t5-small")
     if _code_explainer is None:
-        _code_explainer = _make_pipeline("text2text-generation", "Salesforce/codet5p-220m")
+        _code_explainer = safe_pipeline("text2text-generation", "Salesforce/codet5p-220m")
     return _summarizer, _rewriter, _proofreader, _code_explainer
 
-# (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                 # summarize | rewrite | proofread | explain_code
-    input: str                # user text / code
-    params: dict | None = None
+    task: str
+    input: str
 
-# ---- Core endpoint ----
 @app.post("/run")
 async def run_task(data: InputData):
     start = time.time()
-    text = (data.input or "").strip()
-    task = (data.task or "").strip().lower()
-
+    task = data.task.strip().lower()
+    text = data.input.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 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"]
-            model_id = "t5-small"
+            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"
 
         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"]
-            model_id = "google/flan-t5-small"
+            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"
 
         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"]
-            model_id = "google/flan-t5-small"
-
-        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"]
-            model_id = "Salesforce/codet5p-220m"
+            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":
+            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}'."}
 
     except Exception as e:
         return {"error": f"inference_failed: {type(e).__name__}: {str(e)}"}
 
     latency = round((time.time() - start) * 1000, 2)
     return {
-        "enhancedPrompt": enhanced,
-        "output": out,
-        "model": model_id,
+        "enhancedPrompt": prompt,
+        "output": result,
+        "model": model_used,
         "latencyMs": latency
     }