Update app.py

app.py

@@ -56,7 +56,8 @@ from level_classifier_tool_2 import (
     build_phrase_index
 )
 from task_temp import rag_temp, rag_cls_temp, cls_temp, gen_temp
-from all_tools import classify_and_score, QuestionRetrieverTool
+from all_tools import classify_and_score, QuestionRetrieverTool, set_classifier_state, set_retrieval_index
+
 from phrases import BLOOMS_PHRASES, DOK_PHRASES
 from pathlib import Path
 # ------------------------ Prebuild embeddings once ------------------------
@@ -146,6 +147,8 @@ emb = HuggingFaceEmbeddings(
 storage_context = StorageContext.from_defaults(persist_dir=str(persist_dir))
 index = load_index_from_storage(storage_context, embed_model=emb)
 
+set_classifier_state(_backend, _BLOOM_INDEX, _DOK_INDEX)
+set_retrieval_index(index)
 # Datasets & GPU build code remains commented out...
 # @spaces.GPU(15)
 # def build_indexes_on_gpu(model="google/embeddinggemma-300m"):
@@ -156,7 +159,12 @@ index = load_index_from_storage(storage_context, embed_model=emb)
 #         encode_kwargs={"normalize_embeddings": True})
 #     idx = VectorStoreIndex.from_documents([Document(text=t) for t in texts], embed_model=emb)
 #     return idx
-
+TASK_TEMPLATES = {
+    "rag_temp": rag_temp,
+    "rag_cls_temp": rag_cls_temp,
+    "cls_temp": cls_temp,
+    "gen_temp": gen_temp,
+}
 # ------------------------ Agent setup with timeout ------------------------
 def make_agent(hf_token: str, model_id: str, provider: str, timeout: int, temperature: float, max_tokens: int):
     client = InferenceClient(
@@ -198,8 +206,8 @@ def run_pipeline(
         temperature=float(temperature),
         max_tokens=int(max_tokens),
     )
-
-    task = task_type.format(
+    template = TASK_TEMPLATES[task_type]
+    task = template.format(
         grade=grade,
         topic=topic,
         subject=subject,
@@ -251,7 +259,10 @@ with gr.Blocks() as demo:
         )
         subject = gr.Textbox(value="Math", label="Subject")
         task_type = gr.Dropdown(
-            choices=["rag_temp", "rag_cls_temp", "cls_temp", "gen_temp"],
+            choices=[("RAG Template", "rag_temp"),
+                     ("RAG+CLS Template", "rag_cls_temp"),
+                     ("Classification Template", "cls_temp"),
+                     ("Generation Template", "gen_temp")],
             label="task type"
         )