Update app.py

app.py

@@ -1,14 +1,5 @@
-# app.py — EduPrompt API (final: per-task load, Spaces-safe caches, smart retries)
-
-import os, time
-from fastapi import FastAPI
-from pydantic import BaseModel
-from fastapi.middleware.cors import CORSMiddleware
-from transformers import pipeline
-
-# =========================
-# Hard-force ALL caches to /tmp (writable on Spaces)
-# =========================
+import os
+os.environ["HOME"] = "/tmp"
 BASE = "/tmp"
 os.environ["HF_HOME"] = f"{BASE}/hf"
 os.environ["HF_HUB_CACHE"] = f"{BASE}/hf"
@@ -28,20 +19,22 @@ for d in (
 ):
     os.makedirs(d, exist_ok=True)
 
-# =========================
-# FastAPI app + CORS
-# =========================
+import time
+from fastapi import FastAPI
+from pydantic import BaseModel
+from fastapi.middleware.cors import CORSMiddleware
+from transformers import pipeline
+
 app = FastAPI(title="EduPrompt API")
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],   # tighten in prod
+    allow_origins=["*"],
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
 @app.get("/")
 def health():
-    # prove /tmp is writable and show cache path
     writable = True
     try:
         with open(f"{BASE}/eduprompt_write_test.txt", "w") as f:
@@ -53,45 +46,33 @@ def health():
         "service": "eduprompt-api",
         "tmpWritable": writable,
         "TRANSFORMERS_CACHE": os.environ["TRANSFORMERS_CACHE"],
+        "HOME": os.environ["HOME"],
     }
 
-# =========================
-# Lazy singletons (loaded per task)
-# =========================
 _summarizer = None
 _rewriter = None
 _proofreader = None
 _code_explainer = None
 
 def _model_cache_dir(model_id: str) -> str:
-    # each model gets its own directory to avoid lock fights
     p = os.path.join(os.environ["TRANSFORMERS_CACHE"], model_id.replace("/", "_"))
     os.makedirs(p, exist_ok=True)
     return p
 
 def safe_pipeline(task: str, model_id: str):
-    """
-    Build a pipeline that caches to /tmp per model.
-    Some pipelines reject 'cache_dir' -> retry without it.
-    Also handles rare permission/lock races by a short retry.
-    """
     cache_dir = _model_cache_dir(model_id)
     print(f"[init] task={task} model={model_id} cache={cache_dir}")
-    # Try with cache_dir
     try:
         return pipeline(task, model=model_id, cache_dir=cache_dir,
                         trust_remote_code=True, device=-1)
     except ValueError as e:
-        # Some models complain: "model_kwargs not used: ['cache_dir']"
         if "cache_dir" in str(e):
             print(f"[init] {model_id} rejects cache_dir, retrying without it")
             return pipeline(task, model=model_id, trust_remote_code=True, device=-1)
         raise
     except OSError as e:
-        # Permission/lock race — wait and retry once
         print(f"[init] OSError on {model_id}: {e}; retrying once")
         time.sleep(1.5)
-        # Re-assert env (some libs re-read)
         os.environ["HF_HOME"] = f"{BASE}/hf"
         os.environ["HF_HUB_CACHE"] = f"{BASE}/hf"
         os.environ["TRANSFORMERS_CACHE"] = f"{BASE}/hf/transformers"
@@ -107,9 +88,6 @@ def safe_pipeline(task: str, model_id: str):
             raise
 
 def get_model(task: str):
-    """
-    Load ONLY the model needed for this task.
-    """
     global _summarizer, _rewriter, _proofreader, _code_explainer
     if task == "summarize":
         if _summarizer is None:
@@ -129,24 +107,20 @@ def get_model(task: str):
         return _code_explainer, "Salesforce/codet5p-220m"
     raise ValueError(f"Unsupported task '{task}'")
 
-# =========================
-# Request schema
-# =========================
 class InputData(BaseModel):
-    task: str                  # summarize | rewrite | proofread | explain_code
+    task: str
     input: str
     params: dict | None = None
 
 def _clean_params(params: dict | None):
-    # Block params that some pipelines reject in generate/forward
     forbidden = {"cache_dir"}
     return {k: v for k, v in (params or {}).items() if k not in forbidden}
 
-# =========================
-# Core endpoint
-# =========================
 @app.post("/run")
 async def run_task(data: InputData):
+    print("TRANSFORMERS_CACHE:", os.environ.get("TRANSFORMERS_CACHE"))
+    print("HOME:", os.environ.get("HOME"))
+    print("Current user:", os.getuid() if hasattr(os, "getuid") else "unknown")
     t0 = time.time()
     task = (data.task or "").strip().lower()
     text = (data.input or "").strip()
@@ -156,7 +130,6 @@ async def run_task(data: InputData):
     if task not in {"summarize", "rewrite", "proofread", "explain_code"}:
         return {"error": f"Unsupported task '{task}'."}
 
-    # load only what we need
     try:
         model, model_used = get_model(task)
     except Exception as e:
@@ -183,7 +156,6 @@ async def run_task(data: InputData):
             out = model(prompt, max_new_tokens=200, truncation=True, **params)[0]["generated_text"]
 
     except Exception as e:
-        # print full stack to logs for debugging; return friendly message to client
         import traceback
         print(traceback.format_exc())
         return {"error": f"inference_failed: {type(e).__name__}: {str(e)}"}
@@ -193,4 +165,4 @@ async def run_task(data: InputData):
         "output": out,
         "model": model_used,
         "latencyMs": round((time.time() - t0) * 1000, 2),
-    }
+    }