Upload app.py

app.py

@@ -8,7 +8,7 @@ import gradio as gr
 from gradio_toggle import Toggle
 from transformers import TextIteratorStreamer
 from model import load_model
-# from scheduler import load_scheduler
+from scheduler import load_scheduler
 from schemas import UserRequest, SteeringOutput, CONFIG
 
 logging.basicConfig(level=logging.INFO, format='%(asctime)s %(name)s %(levelname)s:%(message)s')
@@ -18,7 +18,7 @@ model_name = "DeepSeek-R1-Distill-Qwen-7B"
 examples = pd.read_csv("assets/examples.csv")
 
 instances = {}
-# scheduler = load_scheduler()
+scheduler = load_scheduler()
 model = load_model()
 
 
@@ -191,9 +191,9 @@ def cleanup_instance(request: gr.Request):
     session_id = request.session_hash
 
     if session_id in instances:
-        # for data in instances[session_id]:
-        #     if isinstance(data, SteeringOutput):
-        #         scheduler.append(data.model_dump())
+        for data in instances[session_id]:
+            if isinstance(data, SteeringOutput):
+                scheduler.append(data.model_dump())
                 
         del instances[session_id]