Add features (show_solution, store_user_solution) & Fix bug (mha)

api/main.py

@@ -3,13 +3,14 @@
 from pydantic import BaseModel
 import sys
 import os
+from pathlib import Path
 
 # Add the root directory to PYTHONPATH so we can import torch_judge
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
 from torch_judge.tasks import TASKS, get_task
 from torch_judge.web_engine import execute_code
-from api.parser import get_all_templates
+from api.parser import get_all_solutions, get_all_templates
 
 app = FastAPI(title="TorchCode UI Backend")
 
@@ -24,10 +25,21 @@
 
 # Load templates on startup
 TEMPLATES = get_all_templates()
+SOLUTIONS = get_all_solutions()
+ROOT_DIR = Path(__file__).resolve().parent.parent
+USER_CODE_DIR = ROOT_DIR / "user_solutions"
 
 class SubmitRequest(BaseModel):
     code: str
 
+class CodeRequest(BaseModel):
+    code: str
+
+def get_user_code_path(task_id: str) -> Path:
+    if task_id not in TASKS:
+        raise HTTPException(status_code=404, detail="Task not found")
+    return USER_CODE_DIR / f"{task_id}.py"
+
 @app.get("/api/tasks")
 def list_tasks():
     tasks_list = []
@@ -55,6 +67,39 @@ def get_task_details(task_id: str):
         "initial_code": template.get("initial_code", "# Write your code here.")
     }
 
+@app.get("/api/tasks/{task_id}/solution")
+def get_task_solution(task_id: str):
+    task = get_task(task_id)
+    if not task:
+        raise HTTPException(status_code=404, detail="Task not found")
+
+    solution = SOLUTIONS.get(task_id)
+    if not solution:
+        raise HTTPException(status_code=404, detail="Solution not found")
+
+    return solution
+
+@app.get("/api/code/{task_id}")
+def get_saved_code(task_id: str):
+    code_path = get_user_code_path(task_id)
+    if not code_path.exists():
+        return {"code": None, "saved": False}
+    return {"code": code_path.read_text(encoding="utf-8"), "saved": True}
+
+@app.put("/api/code/{task_id}")
+def save_code(task_id: str, request: CodeRequest):
+    code_path = get_user_code_path(task_id)
+    USER_CODE_DIR.mkdir(parents=True, exist_ok=True)
+    code_path.write_text(request.code, encoding="utf-8")
+    return {"saved": True, "path": str(code_path.relative_to(ROOT_DIR))}
+
+@app.delete("/api/code/{task_id}")
+def delete_saved_code(task_id: str):
+    code_path = get_user_code_path(task_id)
+    if code_path.exists():
+        code_path.unlink()
+    return {"deleted": True}
+
 @app.post("/api/submit/{task_id}")
 def submit_code(task_id: str, request: SubmitRequest):
     result = execute_code(task_id, request.code)

api/parser.py

@@ -45,6 +45,46 @@ def parse_notebook_template(filepath: str) -> dict:
             "initial_code": "# Error loading template code."
         }
 
+def parse_notebook_solution(filepath: str) -> dict:
+    """
+    Parses a Jupyter Notebook solution to extract explanatory text and the
+    reference solution code.
+    """
+    try:
+        with open(filepath, 'r', encoding='utf-8') as f:
+            nb = json.load(f)
+
+        solution_parts = []
+
+        for cell in nb.get("cells", []):
+            source = cell.get("source", [])
+            source_str = "".join(source).strip()
+
+            if not source_str:
+                continue
+
+            if cell["cell_type"] == "markdown":
+                filtered_source = [
+                    line for line in source
+                    if "![Open In Colab]" not in line
+                ]
+                markdown = "".join(filtered_source).strip()
+                if markdown:
+                    solution_parts.append(markdown)
+
+            elif cell["cell_type"] == "code" and ("# ✅ SOLUTION" in source_str or "# SOLUTION" in source_str):
+                solution_parts.append(f"```python\n{source_str}\n```")
+                break
+
+        return {
+            "solution": "\n\n".join(solution_parts) or "Solution not found in notebook."
+        }
+    except Exception as e:
+        print(f"Error parsing {filepath}: {e}")
+        return {
+            "solution": "Error loading solution."
+        }
+
 def get_all_templates(templates_dir: str = "../templates") -> dict:
     """
     Returns a dictionary mapping task_ids to their extracted template data.
@@ -69,6 +109,26 @@ def get_all_templates(templates_dir: str = "../templates") -> dict:
 
     return templates
 
+def get_all_solutions(solutions_dir: str = "../solutions") -> dict:
+    """
+    Returns a dictionary mapping task_ids to their extracted solution data.
+    task_id is inferred from the filename (e.g., '01_relu_solution.ipynb' -> 'relu').
+    """
+    solutions = {}
+
+    base_dir = os.path.dirname(os.path.abspath(__file__))
+    solutions_path = os.path.join(base_dir, solutions_dir)
+
+    for filepath in glob.glob(os.path.join(solutions_path, "*.ipynb")):
+        filename = os.path.basename(filepath)
+
+        match = re.match(r"^\d+_(.+)_solution\.ipynb$", filename)
+        if match:
+            task_id = match.group(1)
+            solutions[task_id] = parse_notebook_solution(filepath)
+
+    return solutions
+
 if __name__ == "__main__":
     # Test parser
     res = get_all_templates()

api/requirements.txt

@@ -1,3 +1,4 @@
 fastapi
 uvicorn
 pydantic
+torch

solutions/06_multihead_attention_solution.ipynb

@@ -100,7 +100,7 @@
    "source": [
     "# Run judge\n",
     "from torch_judge import check\n",
-    "check(\"mha\")"
+    "check(\"multihead_attention\")"
    ]
   }
  ],

templates/06_multihead_attention.ipynb

@@ -112,7 +112,7 @@
    "source": [
     "# ✅ SUBMIT\n",
     "from torch_judge import check\n",
-    "check(\"mha\")"
+    "check(\"multihead_attention\")"
    ]
   }
  ],