Update app.py

app.py

@@ -3,64 +3,74 @@ import json
 import gradio as gr
 from huggingface_hub import InferenceClient
 from smolagents import CodeAgent, InferenceClientModel, tool
-import spaces
+
 from level_classifier_tool import (
     classify_levels_phrases,
     HFEmbeddingBackend,
     build_phrase_index
 )
+
+# ------------------------ Taxonomy phrases ------------------------
 BLOOMS_PHRASES = {
-    "Remember": [  
-        "define", "list", "recall", "identify", "state", "label", "name", "recognize", "find", "select", "match", "choose", "give", "write", "tell", "show"
+    "Remember": [
+        "define", "list", "recall", "identify", "state", "label", "name", "recognize", "find",
+        "select", "match", "choose", "give", "write", "tell", "show"
     ],
-    "Understand": [  
-        "classify", "interpret", "summarize", "explain", "estimate", "describe", "discuss", "predict", "paraphrase", "restate", "illustrate", "compare", "contrast", "report"
+    "Understand": [
+        "classify", "interpret", "summarize", "explain", "estimate", "describe", "discuss",
+        "predict", "paraphrase", "restate", "illustrate", "compare", "contrast", "report"
     ],
-    "Apply": [  
-        "apply", "solve", "use", "demonstrate", "calculate", "implement", "perform", "execute", "carry out", "practice", "employ", "sketch"
+    "Apply": [
+        "apply", "solve", "use", "demonstrate", "calculate", "implement", "perform",
+        "execute", "carry out", "practice", "employ", "sketch"
     ],
-    "Analyze": [  
-        "analyze", "differentiate", "organize", "structure", "break down", "distinguish", "dissect", "examine", "compare", "contrast", "attribute", "investigate"
+    "Analyze": [
+        "analyze", "differentiate", "organize", "structure", "break down", "distinguish",
+        "dissect", "examine", "compare", "contrast", "attribute", "investigate"
     ],
-    "Evaluate": [  
-        "evaluate", "judge", "critique", "assess", "defend", "argue", "select", "support", "appraise", "recommend", "conclude", "review"
+    "Evaluate": [
+        "evaluate", "judge", "critique", "assess", "defend", "argue", "select", "support",
+        "appraise", "recommend", "conclude", "review"
     ],
-    "Create": [  
-        "create", "design", "compose", "plan", "construct", "produce", "devise", "generate", "develop", "formulate", "invent", "build"
+    "Create": [
+        "create", "design", "compose", "plan", "construct", "produce", "devise", "generate",
+        "develop", "formulate", "invent", "build"
     ]
 }
 
 DOK_PHRASES = {
     "DOK1": [
         "define", "list", "recall", "compute", "identify", "state", "label", "how many",
-        "name", "recognize", "find", "determine", "select", "match", "choose", "give", 
+        "name", "recognize", "find", "determine", "select", "match", "choose", "give",
         "write", "tell", "show", "point out"
     ],
     "DOK2": [
         "classify", "interpret", "estimate", "organise", "summarise", "explain", "solve",
-        "categorize", "group", "compare", "contrast", "distinguish", "make observations", 
-        "collect data", "display data", "arrange", "sort", "paraphrase", "restate", "predict", 
+        "categorize", "group", "compare", "contrast", "distinguish", "make observations",
+        "collect data", "display data", "arrange", "sort", "paraphrase", "restate", "predict",
         "approximate", "demonstrate", "illustrate", "describe", "analyze data"
     ],
     "DOK3": [
         "justify", "analyze", "generalise", "compare", "construct", "investigate",
-        "support", "defend", "argue", "examine", "differentiate", "criticize", "debate", 
-        "test", "experiment", "hypothesize", "draw conclusions", "break down", "dissect", 
+        "support", "defend", "argue", "examine", "differentiate", "criticize", "debate",
+        "test", "experiment", "hypothesize", "draw conclusions", "break down", "dissect",
         "probe", "explore", "develop", "formulate"
     ],
     "DOK4": [
         "design", "synthesize", "model", "prove", "evaluate system", "critique", "create",
-        "compose", "plan", "invent", "devise", "generate", "build", "construct", "produce", 
-        "formulate", "improve", "revise", "assess", "appraise", "judge", "recommend", 
+        "compose", "plan", "invent", "devise", "generate", "build", "construct", "produce",
+        "formulate", "improve", "revise", "assess", "appraise", "judge", "recommend",
         "predict outcome", "simulate"
     ]
 }
 
-# Prebuild embeddings once
-_backend = HFEmbeddingBackend(model_name="sentence-transformers/all-MiniLM-L6-v2").to('cuda')
+# ------------------------ Prebuild embeddings once ------------------------
+_backend = HFEmbeddingBackend(model_name="sentence-transformers/all-MiniLM-L6-v2")
 _BLOOM_INDEX = build_phrase_index(_backend, BLOOMS_PHRASES)
 _DOK_INDEX = build_phrase_index(_backend, DOK_PHRASES)
 
+
+# ------------------------ Tool: classify and score ------------------------
 @tool
 def classify_and_score(
     question: str,
@@ -69,7 +79,7 @@ def classify_and_score(
     agg: str = "max"
 ) -> dict:
     """Classify a question against Bloom’s and DOK targets and return guidance.
-    
+
     Args:
         question: The question text to evaluate for cognitive demand.
         target_bloom: Target Bloom’s level or range. Accepts exact (e.g., "Analyze")
@@ -78,7 +88,7 @@ def classify_and_score(
             or span (e.g., "DOK2-DOK3").
         agg: Aggregation method over phrase similarities within a level
             (choices: "mean", "max", "topk_mean").
-    
+
     Returns:
         A dictionary with:
             ok: True if both Bloom’s and DOK match the targets.
@@ -97,21 +107,40 @@ def classify_and_score(
     )
 
     def _parse_target_bloom(t: str):
-        order = ["Remember","Understand","Apply","Analyze","Evaluate","Create"]
+        order = ["Remember", "Understand", "Apply", "Analyze", "Evaluate", "Create"]
         if t.endswith("+"):
             base = t[:-1]
+            if base not in order:
+                raise ValueError(f"Invalid Bloom target '{t}'")
             return set(order[order.index(base):])
+        if t not in order:
+            raise ValueError(f"Invalid Bloom target '{t}'")
         return {t}
 
     def _parse_target_dok(t: str):
-        order = ["DOK1","DOK2","DOK3","DOK4"]
+        order = ["DOK1", "DOK2", "DOK3", "DOK4"]
         if "-" in t:
             lo, hi = t.split("-")
-            return set(order[order.index(lo):order.index(hi)+1])
+            if lo not in order or hi not in order or order.index(lo) > order.index(hi):
+                raise ValueError(f"Invalid DOK range '{t}'")
+            return set(order[order.index(lo):order.index(hi) + 1])
+        if t not in order:
+            raise ValueError(f"Invalid DOK target '{t}'")
         return {t}
 
-    bloom_target_set = _parse_target_bloom(target_bloom)
-    dok_target_set = _parse_target_dok(target_dok)
+    try:
+        bloom_target_set = _parse_target_bloom(target_bloom)
+        dok_target_set = _parse_target_dok(target_dok)
+    except Exception as e:
+        return {
+            "ok": False,
+            "measured": {},
+            "feedback": (
+                f"Invalid targets: {e}. Use Bloom in "
+                "{Remember, Understand, Apply, Analyze, Evaluate, Create} "
+                "and DOK in {DOK1..DOK4} or ranges like 'DOK2-DOK3'."
+            ),
+        }
 
     bloom_best = res["blooms"]["best_level"]
     dok_best = res["dok"]["best_level"]
@@ -119,16 +148,19 @@ def classify_and_score(
     bloom_ok = bloom_best in bloom_target_set
     dok_ok = dok_best in dok_target_set
 
+    top_bloom_phrases = res["blooms"].get("top_phrases", {})
+    top_dok_phrases = res["dok"].get("top_phrases", {})
+
     feedback_parts = []
     if not bloom_ok:
         feedback_parts.append(
-            f"Shift Bloom’s from {bloom_best} toward {sorted(bloom_target_set)}. "
-            f"Top cues: {res['blooms']['top_phrases'].get(bloom_best, [])[:3]}"
+            f"Shift Bloom’s from {bloom_best} toward {sorted(list(bloom_target_set))}. "
+            f"Top cues: {top_bloom_phrases.get(bloom_best, [])[:3]}"
         )
     if not dok_ok:
         feedback_parts.append(
-            f"Shift DOK from {dok_best} toward {sorted(dok_target_set)}. "
-            f"Top cues: {res['dok']['top_phrases'].get(dok_best, [])[:3]}"
+            f"Shift DOK from {dok_best} toward {sorted(list(dok_target_set))}. "
+            f"Top cues: {top_dok_phrases.get(dok_best, [])[:3]}"
         )
 
     return {
@@ -144,25 +176,18 @@ def classify_and_score(
 
 
 # ------------------------ Agent setup with timeout ------------------------
-def make_agent(hf_token: str, model_id: str, provider: str, timeout: int, temperature: float, max_tokens: int):
+def make_agent(hf_token: str, model_id: str, timeout: int, temperature: float, max_tokens: int):
     client = InferenceClient(
         model=model_id,
-        provider=provider,
         timeout=timeout,
-        token=hf_token if hf_token else None,
+        token=hf_token or None,
     )
-
     model = InferenceClientModel(client=client)
     agent = CodeAgent(model=model, tools=[classify_and_score])
-    agent._ui_params = {"temperature": temperature, "max_tokens": max_tokens}  # attach for reference
+    # Not used by the agent core, but helpful for debugging/visibility
+    agent._ui_params = {"temperature": temperature, "max_tokens": max_tokens}
     return agent
-@spaces.GPU(duration=20)
-def load_model():
-    transformers_model = TransformersModel(
-        model_id='swiss-ai/Apertus-70B-Instruct-2509',
-        device_map="auto"
-    ).to("cuda")
-    return transformers_model
+
 
 # ------------------------ Agent task template -----------------------------
 TASK_TMPL = '''You generate {subject} question candidates for {grade} on "{topic}".
@@ -185,7 +210,33 @@ If you output JSON, ensure it is valid JSON (no trailing commas, use double quot
 '''
 
 
-# ------------------------ Gradio glue ------------------------------------
+# ------------------------ Utility: robust JSON extractor ------------------
+def extract_top_level_json(s: str) -> str:
+    """
+    Extract the first top-level JSON object from a string by tracking braces.
+    Returns the JSON string if found, else "".
+    """
+    start = s.find("{")
+    if start == -1:
+        return ""
+    depth = 0
+    for i in range(start, len(s)):
+        if s[i] == "{":
+            depth += 1
+        elif s[i] == "}":
+            depth -= 1
+            if depth == 0:
+                candidate = s[start:i + 1]
+                try:
+                    # validate
+                    json.loads(candidate)
+                    return candidate
+                except Exception:
+                    return ""
+    return ""
+
+
+# ------------------------ Pipeline ---------------------------------------
 def run_pipeline(
     hf_token,
     topic,
@@ -195,16 +246,14 @@ def run_pipeline(
     target_dok,
     attempts,
     model_id,
-    provider,
     timeout,
     temperature,
     max_tokens
 ):
-    # Build agent per run (or cache if you prefer)
+    # Build agent per run
     agent = make_agent(
-        hf_token=hf_token.strip(),
+        hf_token=(hf_token or "").strip(),
         model_id=model_id,
-        provider=provider,
         timeout=int(timeout),
         temperature=float(temperature),
         max_tokens=int(max_tokens),
@@ -221,25 +270,23 @@ def run_pipeline(
 
     # The agent will internally call the tool
     try:
-        result_text = agent.run(task, max_steps=int(attempts)*4)
+        result_text = agent.run(task, max_steps=int(attempts) * 4)
     except Exception as e:
         result_text = f"ERROR: {e}"
 
     # Try to extract final JSON
     final_json = ""
-    try:
-        # find JSON object in result_text (simple heuristic)
-        start = result_text.find("{")
-        end = result_text.rfind("}")
-        if start != -1 and end != -1 and end > start:
-            candidate = result_text[start:end+1]
+    candidate = extract_top_level_json(result_text or "")
+    if candidate:
+        try:
             final_json = json.dumps(json.loads(candidate), indent=2)
-    except Exception:
-        final_json = ""
+        except Exception:
+            final_json = ""
 
     return final_json, result_text
 
 
+# ------------------------ Gradio UI --------------------------------------
 with gr.Blocks() as demo:
     gr.Markdown("# Agent + Tool: Generate Questions to Target Difficulty")
     gr.Markdown(
@@ -248,54 +295,61 @@ with gr.Blocks() as demo:
     )
 
     with gr.Accordion("API Settings", open=False):
-        hf_token = gr.Textbox(label="Hugging Face Token (required if the endpoint needs auth)", type="password")
-        model_id = gr.Textbox(value=transformers_model, label="Model ID")
-        provider = gr.Textbox(value="novita", label="Provider")
+        hf_token = gr.Textbox(
+            label="Hugging Face Token (required if the endpoint needs auth)",
+            type="password"
+        )
+        model_id = gr.Textbox(
+            value="swiss-ai/Apertus-70B-Instruct-2509",
+            label="Model ID"
+        )
         timeout = gr.Slider(5, 120, value=30, step=1, label="Timeout (s)")
 
     with gr.Row():
         topic = gr.Textbox(value="Fractions", label="Topic")
         grade = gr.Dropdown(
-            choices=["Grade 1","Grade 2","Grade 3","Grade4","Grade 5","Grade 6","Grade 7","Grade 8","Grade 9",
-                     "Grade 10","Grade 11","Grade 12","Under Graduate","Post Graduate"],
+            choices=[
+                "Grade 1", "Grade 2", "Grade 3", "Grade 4", "Grade 5", "Grade 6",
+                "Grade 7", "Grade 8", "Grade 9",
+                "Grade 10", "Grade 11", "Grade 12",
+                "Under Graduate", "Post Graduate"
+            ],
             value="Grade 7",
             label="Grade"
         )
-        subject= gr.Textbox(value="Math", label="Subject")
+        subject = gr.Textbox(value="Math", label="Subject")
 
     with gr.Row():
         target_bloom = gr.Dropdown(
-            choices=["Remember","Understand","Apply","Analyze","Evaluate","Create"],
+            choices=["Remember", "Understand", "Apply", "Analyze", "Evaluate", "Create"],
             value="Analyze",
             label="Target Bloom’s"
         )
         target_dok = gr.Dropdown(
-            choices=["DOK1","DOK2","DOK3","DOK4","DOK1-DOK2","DOK2-DOK3","DOK3-DOK4"],
+            choices=["DOK1", "DOK2", "DOK3", "DOK4", "DOK1-DOK2", "DOK2-DOK3", "DOK3-DOK4"],
             value="DOK2-DOK3",
             label="Target Depth of Knowledge"
         )
         attempts = gr.Slider(1, 8, value=5, step=1, label="Max Attempts")
 
-    with gr.Accordion("⚙️ Generation Controls", open=False):
+    with gr.Accordion("Generation Controls", open=False):
         temperature = gr.Slider(0.0, 1.5, value=0.7, step=0.1, label="Temperature")
         max_tokens = gr.Slider(64, 1024, value=300, step=16, label="Max Tokens")
 
-    run_btn = gr.Button("Run Agent 🚀")
+    run_btn = gr.Button("Run Agent")
 
     final_json = gr.Code(label="Final Candidate (JSON if detected)", language="json")
     transcript = gr.Textbox(label="Agent Transcript", lines=18)
 
     run_btn.click(
         fn=run_pipeline,
-        inputs=[hf_token, topic, grade, subject, target_bloom, target_dok, attempts, model_id, provider, timeout, temperature, max_tokens],
+        inputs=[
+            hf_token, topic, grade, subject,
+            target_bloom, target_dok, attempts,
+            model_id, timeout, temperature, max_tokens
+        ],
         outputs=[final_json, transcript]
     )
 
 if __name__ == "__main__" or os.getenv("SYSTEM") == "spaces":
-    try:
-        load_model()  # triggers GPU allocation during startup
-    except Exception as e:
-        # don't crash the app if warmup fails; logs will show details
-        print("Warmup failed:", e)
     demo.launch()
-