Update app.py

app.py

@@ -13,6 +13,7 @@ from captum.attr import LayerIntegratedGradients, TokenReferenceBase
 from captum.attr import visualization as viz
 from huggingface_hub import InferenceClient
 from datetime import datetime
+import uuid
 
 
 HF = 'hf'
@@ -45,19 +46,30 @@ def load_drafts():
             return {}
     return {}
 
-def update_cache_row(user, pol_a, pol_b, a_list, idx, b_text, rel, inter, just):
-    """Fires automatically on keystrokes/clicks to save progress"""
+def update_cache_row(user, session_id, dom_a, pol_a, dom_b, pol_b, tar_col, ctx_col, a_list, idx, b_text, rel, inter, just):
+    """Fires automatically on keystrokes/clicks to save progress and workspace state"""
     if not user or not a_list or idx >= len(a_list) or not b_text: return
     curr_a = a_list[idx]
     
     drafts = load_drafts()
-    if user not in drafts: drafts[user] = {}
+    # Upgraded structure to hold workspace settings AND row data
+    if user not in drafts: drafts[user] = {"workspace": {}, "rows": {}}
+    
+    # Save the active workspace so we can restore it on reload
+    drafts[user]["workspace"] = {
+        "session_id": session_id,
+        "dom_a": dom_a, "pol_a": pol_a,
+        "dom_b": dom_b, "pol_b": pol_b,
+        "tar_col": tar_col, "ctx_col": ctx_col
+    }
     
     cache_key = f"{pol_a}|{pol_b}|{curr_a}"
-    if cache_key not in drafts[user]: drafts[user][cache_key] = {}
+    if cache_key not in drafts[user]["rows"]: drafts[user]["rows"][cache_key] = {}
     
-    # Store the exact state of this specific row
-    drafts[user][cache_key][b_text] = {"rel": rel, "inter": inter, "just": just}
+    # Store the exact state of this specific row with the unique session tag
+    drafts[user]["rows"][cache_key][b_text] = {
+        "rel": rel, "inter": inter, "just": just, "session_id": session_id
+    }
     
     with open(DRAFT_FILE, 'w') as f:
         json.dump(drafts, f)
@@ -313,6 +325,7 @@ def get_worksheet_by_number(spreadsheet, worksheet_number, format=True):
             
         if 'Policy' in df.columns:
             df['Policy'] = df['Policy'].ffill()
+        
             
     return df
     
@@ -497,6 +510,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
     
     hf_df_state = gr.State()
     user_tag_state = gr.State()
+    session_id_state = gr.State(lambda: str(uuid.uuid4().hex[:12]))
     
     target_a_list_state = gr.State([]) 
     pending_tasks_state = gr.State({})
@@ -666,8 +680,12 @@ with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
     for i in range(MAX_ROWS):
         _, b_text, rel_radio, _, inter_dd, just_box, _, _, _, _, _ = eval_rows[i]
         
-        # Gather the exact state needed to cache this row
-        inputs_to_cache = [user_tag_state, policy_a_dd, policy_b_dd, target_a_list_state, current_index_state, b_text, rel_radio, inter_dd, just_box]
+        # Gather the exact state needed to cache this row AND the workspace config
+        inputs_to_cache = [
+            user_tag_state, session_id_state, 
+            domain_a_dd, policy_a_dd, domain_b_dd, policy_b_dd, target_col_dd, context_col_dd,
+            target_a_list_state, current_index_state, b_text, rel_radio, inter_dd, just_box
+        ]
         
         # Trigger cache save silently in the background on any change
         rel_radio.change(fn=update_cache_row, inputs=inputs_to_cache)
@@ -678,12 +696,50 @@ with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
     def authenticate(email):
         user_tag, msg = get_or_create_user(email)
         if not user_tag:
-            return gr.update(value=f"<font color='red'>{msg}</font>"), gr.update(visible=True), gr.update(visible=False), None, None
+            return (gr.update(value=f"<font color='red'>{msg}</font>"), gr.update(visible=True), gr.update(visible=False), None, None,
+                    gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update())
             
         hf_df = load_hf_dataset()
-        return gr.update(value=f"{msg} Loaded {len(hf_df)} annotations."), gr.update(visible=False), gr.update(visible=True), user_tag, hf_df
-
-    login_btn.click(fn=authenticate, inputs=[email_box], outputs=[login_status, login_box, app_box, user_tag_state, hf_df_state])
+        
+        # Check Cache for Session Recovery
+        drafts = load_drafts()
+        user_data = drafts.get(user_tag, {})
+        ws = user_data.get("workspace", {})
+        
+        # If we found a saved workspace, restore the dropdowns!
+        if ws.get("pol_a") and ws.get("pol_b"):
+            msg += f" Restored your previous session workspace."
+            return (
+                gr.update(value=f"{msg} Loaded {len(hf_df)} annotations."), 
+                gr.update(visible=False), 
+                gr.update(visible=True), 
+                user_tag, 
+                hf_df,
+                gr.update(value=ws["dom_a"]),
+                gr.update(choices=get_policy_list(ws["dom_a"]), value=ws["pol_a"]),
+                gr.update(value=ws["dom_b"]),
+                gr.update(choices=get_policy_list(ws["dom_b"]), value=ws["pol_b"]),
+                gr.update(value=ws["tar_col"]),
+                gr.update(value=ws["ctx_col"])
+            )
+        else:
+            return (
+                gr.update(value=f"{msg} Loaded {len(hf_df)} annotations."), 
+                gr.update(visible=False), 
+                gr.update(visible=True), 
+                user_tag, 
+                hf_df,
+                gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update()
+            )
+            
+    login_btn.click(
+        fn=authenticate, 
+        inputs=[email_box], 
+        outputs=[
+            login_status, login_box, app_box, user_tag_state, hf_df_state,
+            domain_a_dd, policy_a_dd, domain_b_dd, policy_b_dd, target_col_dd, context_col_dd # <-- ADDED THESE
+        ]
+    )
     domain_a_dd.change(fn=lambda d: gr.update(choices=get_policy_list(d), value=None), inputs=domain_a_dd, outputs=policy_a_dd)
     domain_b_dd.change(fn=lambda d: gr.update(choices=get_policy_list(d), value=None), inputs=domain_b_dd, outputs=policy_b_dd)
 
@@ -716,7 +772,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
         # Pull Drafts for this specific user and Target A
         drafts = load_drafts()
         cache_key = f"{pol_a}|{pol_b}|{curr_a_eng}"
-        user_draft = drafts.get(user_tag, {}).get(cache_key, {})
+        user_draft = drafts.get(user_tag, {}).get("rows", {}).get(cache_key, {})
         
         # Run model predictions for this batch
         preds = get_model_predictions(curr_a_eng, bs_to_eval_eng)
@@ -937,8 +993,10 @@ with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
                 # CLEAR CACHE ON SUCCESSFUL SAVE
                 drafts = load_drafts()
                 cache_key = f"{pol_a}|{pol_b}|{current_a_eng}"
-                if user_tag in drafts and cache_key in drafts[user_tag]:
-                    del drafts[user_tag][cache_key]
+                
+                # Check inside the "rows" sub-dictionary
+                if user_tag in drafts and "rows" in drafts[user_tag] and cache_key in drafts[user_tag]["rows"]:
+                    del drafts[user_tag]["rows"][cache_key]
                     with open(DRAFT_FILE, 'w') as f:
                         json.dump(drafts, f)