Update app.py

app.py

@@ -1,65 +1,98 @@
-from fastapi import FastAPI, Request
+# app.py — EduPrompt FastAPI backend
+
+from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import pipeline
 import time
+import os
+
+# --- Hugging Face cache fix for Spaces (permission-safe) ---
+os.environ["HF_HOME"] = "/tmp"
+os.environ["HF_HUB_CACHE"] = "/tmp"
+os.environ["TRANSFORMERS_CACHE"] = "/tmp"
 
+# --- FastAPI app + CORS ---
 app = FastAPI(title="EduPrompt API")
 
 from fastapi.middleware.cors import CORSMiddleware
-
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"], allow_methods=["*"], allow_headers=["*"]
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
 )
 
 @app.get("/")
 def health():
     return {"ok": True, "service": "eduprompt-api"}
 
+# --- Lazy model loaders ---
+_summarizer = None
+_rewriter = None
+_proofreader = None
+_code_explainer = None
 
-# Models
-summarizer = pipeline("summarization", model="t5-small")
-rewriter = pipeline("text2text-generation", model="google/flan-t5-small")
-proofreader = pipeline("text2text-generation", model="google/flan-t5-small")
-code_explainer = pipeline("text2text-generation", model="Salesforce/codet5p-220m")
+def get_models():
+    global _summarizer, _rewriter, _proofreader, _code_explainer
+    if _summarizer is None:
+        _summarizer = pipeline("summarization", model="t5-small")
+    if _rewriter is None:
+        _rewriter = pipeline("text2text-generation", model="google/flan-t5-small")
+    if _proofreader is None:
+        _proofreader = pipeline("text2text-generation", model="google/flan-t5-small")
+    if _code_explainer is None:
+        _code_explainer = pipeline("text2text-generation", model="Salesforce/codet5p-220m")
+    return _summarizer, _rewriter, _proofreader, _code_explainer
 
-# Request schema
+# --- Request schema ---
 class InputData(BaseModel):
     task: str
     input: str
     params: dict | None = None
 
+# --- Core endpoint ---
 @app.post("/run")
 async def run_task(data: InputData):
     start = time.time()
-    text = data.input.strip()
-    task = data.task.lower()
+    text = (data.input or "").strip()
+    task = (data.task or "").strip().lower()
+
+    if not text:
+        return {"error": "Empty input text."}
+    if task not in {"summarize", "rewrite", "proofread", "explain_code"}:
+        return {"error": f"Unsupported task '{task}'."}
 
-    # Build enhanced prompt (simple version)
     enhanced = ""
-    if task == "summarize":
-        enhanced = f"You are an expert explainer. Summarize clearly and concisely:\n{text}"
-        result = summarizer(enhanced, max_length=120, min_length=30, do_sample=False)[0]["summary_text"]
+    try:
+        summarizer, rewriter, proofreader, code_explainer = get_models()
+
+        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"
 
-    elif task == "rewrite":
-        enhanced = f"You are a writing assistant. Rewrite this text for clarity and tone:\n{text}"
-        result = rewriter(enhanced, max_new_tokens=150)[0]["generated_text"]
+        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"
 
-    elif task == "proofread":
-        enhanced = f"You are a grammar and style editor. Correct and improve this text:\n{text}"
-        result = proofreader(enhanced, max_new_tokens=150)[0]["generated_text"]
+        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"
 
-    elif task == "explain_code":
-        enhanced = f"You are a programming tutor. Explain what this code does in simple language:\n{text}"
-        result = code_explainer(enhanced, max_new_tokens=200)[0]["generated_text"]
+        else:
+            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"
 
-    else:
-        return {"error": "Unsupported 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": result,
-        "model": task,
+        "output": out,
+        "model": model_id,
         "latencyMs": latency
     }