Update app.py

app.py

@@ -44,9 +44,10 @@ def custom_forward(input_ids, attention_mask):
 lig = LayerIntegratedGradients(custom_forward, model.roberta.embeddings.word_embeddings)
 
 llm_client = InferenceClient("Qwen/Qwen3-8B", token=HF_TOKEN)
+
 def generate_row_explanation(a_list, idx, text_b):
     if not a_list or idx >= len(a_list) or not text_b:
-        return "", ""
+        return "", "", "", ""
         
     policy_a = clean_policy_text(a_list[idx])
     policy_b = clean_policy_text(text_b)
@@ -77,7 +78,12 @@ def generate_row_explanation(a_list, idx, text_b):
     attributions = attributions.cpu().detach().numpy()
 
     tokens = tokenizer.convert_ids_to_tokens(input_ids[0])
-    score_list = [f"'{t.replace('Ġ', '').strip()}': {s:.3f}" for t, s in zip(tokens, attributions) if t.replace('Ġ', '').strip()]
+    
+    # NEW: Create a pure JSON dictionary of the attributions to save to the CSV
+    ig_dict = {t.replace('Ġ', '').strip(): float(s) for t, s in zip(tokens, attributions) if t.replace('Ġ', '').strip()}
+    ig_json_str = json.dumps(ig_dict)
+
+    score_list = [f"'{k}': {v:.3f}" for k, v in ig_dict.items()]
     formatted_scores = ", ".join(score_list)
 
     # 2. Call Qwen LLM
@@ -98,11 +104,13 @@ Write a highly analytical, 2 to 3 sentence explanation of the model's reasoning.
             think_content = match.group(1).strip()
             final_answer = raw_output.replace(match.group(0), '').strip()
             html_out = f"""<details style="margin-bottom: 12px; padding: 10px; background-color: #f3f4f6; border-radius: 6px; border: 1px solid #e5e7eb;"><summary style="cursor: pointer; font-weight: bold; color: #4b5563; outline: none;">🧠 Click to peek into the AI's thought process</summary><div style="margin-top: 10px; font-size: 0.9em; color: #6b7280; white-space: pre-wrap;">{think_content}</div></details>"""
-            return html_out, final_answer
             
-        return "", raw_output
+            # Return: UI HTML, UI TextBox, Hidden AI Text, Hidden IG JSON
+            return html_out, final_answer, raw_output, ig_json_str
+            
+        return "", raw_output, raw_output, ig_json_str
     except Exception as e:
-        return "", f"⚠️ Explainability Error: {str(e)}"
+        return "", f"⚠️ Explainability Error: {str(e)}", "", ""
         
 
 def bucket_score(score):
@@ -364,6 +372,7 @@ def load_hf_dataset():
             "Target_Column", "Target_A_Row", "Target_B_Row",
             "Context_Column", "Context_A_Chunk", "Context_B_Chunk",
             "Model_Coarse_Label", "Model_Drill_Down_Label", "Model_Confidences", # New Columns
+            "AI_Justification", "IG_JSON",
             "Coherence_Label", "Drill_Down_Label", "Justification", "AnnotatorUsername"
         ])
 
@@ -516,6 +525,8 @@ with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
                         m_coarse_st = gr.State("")
                         m_drill_st = gr.State("")
                         m_conf_st = gr.State("")
+                        m_ai_just_st = gr.State("") 
+                        m_ig_json_st = gr.State("")
                         
                         with gr.Row(equal_height=True):
                             b_text = gr.Textbox(label=f"Target B", interactive=False, scale=4, min_width=200, lines=3, max_lines=8)
@@ -535,11 +546,13 @@ with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
                                 explain_btn.click(
                                     fn=generate_row_explanation,
                                     inputs=[target_a_list_state, current_index_state, b_text],
-                                    outputs=[explain_html, just_box]
-    )
+                                    outputs=[explain_html, just_box, m_ai_just_st, m_ig_json_st] # <-- ADD OUTPUTS
+                                )
+                                
                         
                         # Tracking 9 items per row now
-                        eval_rows.append((row_container, b_text, rel_radio, conf_md, inter_dd, just_box, m_coarse_st, m_drill_st, m_conf_st))
+                        # eval_rows.append((row_container, b_text, rel_radio, conf_md, inter_dd, just_box, m_coarse_st, m_drill_st, m_conf_st))
+                        eval_rows.append((row_container, b_text, rel_radio, conf_md, inter_dd, just_box, m_coarse_st, m_drill_st, m_conf_st, m_ai_just_st, m_ig_json_st))
 
                 with gr.Row():
                     skip_btn = gr.Button("Skip Target A", size="lg")
@@ -614,7 +627,8 @@ with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
         updates = []
         
         # 9 components per row to reset
-        empty_row = [gr.update(visible=False), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), "", "", ""]
+        # empty_row = [gr.update(visible=False), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), "", "", ""]
+        empty_row = [gr.update(visible=False), gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), "", "", "", "", ""]
         
         if not a_list:
             prog_txt = t_text("**Progress:** No unannotated items found.", lang)
@@ -646,7 +660,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
                     gr.update(value=""),      # conf_md
                     gr.update(choices=[], value=None),      
                     gr.update(value=""),      # just_box
-                    "", "", ""                # Reset the 3 hidden model states
+                    "", "", "", "", ""             # Reset the 5 hidden model states
                 ])
             else:
                 updates.extend(empty_row)
@@ -656,11 +670,11 @@ with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
     def load_workspace(dom_a, pol_a, dom_b, pol_b, tar_col, ctx_col, hf_df, user_tag, lang):
         if not pol_a or not pol_b:
             err = t_text("Error: Select both policies.", lang)
-            return [gr.update(value=err)] + [gr.skip()] * (14 + MAX_ROWS*9)
+            return [gr.update(value=err)] + [gr.skip()] * (14 + MAX_ROWS*11)
             
         if tar_col == ctx_col:
             err = t_text("Error: Target and Context cannot be the same.", lang)
-            return [gr.update(value=err)] + [gr.skip()] * (14 + MAX_ROWS*9)
+            return [gr.update(value=err)] + [gr.skip()] * (14 + MAX_ROWS*11)
             
         df_a = DOMAIN_MAP[dom_a]
         df_b = DOMAIN_MAP[dom_b]
@@ -724,14 +738,16 @@ with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
             b_val_eng = b_eng_list[i] 
             
             # Row data length is now 8 elements: [b_text, rel_radio, conf_md, inter_dd, just_box, m_coarse, m_drill, m_conf]
-            rel = row_data[i*8 + 1]
-            inter = row_data[i*8 + 3]
-            just = row_data[i*8 + 4]
+            rel = row_data[i*10 + 1]
+            inter = row_data[i*10 + 3]
+            just = row_data[i*10 + 4]
             
             # Extract the independent model predictions from hidden states
-            model_coarse = row_data[i*8 + 5]
-            model_drill = row_data[i*8 + 6]
-            model_conf = row_data[i*8 + 7]
+            model_coarse = row_data[i*10 + 5]
+            model_drill = row_data[i*10 + 6]
+            model_conf = row_data[i*10 + 7]
+            ai_just = row_data[i*10 + 8]   # 
+            ig_json = row_data[i*10 + 9]
             
             has_rel = bool(rel)
             has_inter = bool(inter)
@@ -758,11 +774,14 @@ with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
                     "Model_Coarse_Label": model_coarse,     # Log model prediction
                     "Model_Drill_Down_Label": model_drill,  # Log model drill
                     "Model_Confidences": model_conf,        # Log model JSON confidence
-                    "Coherence_Label": rel,                 # Log User prediction
+                    "AI_Justification": ai_just,            # Log pure AI Thoughts
+                    "IG_JSON": ig_json,                     # Log Captum Gradients
+                    "Coherence_Label": rel,
                     "Drill_Down_Label": inter,              # Log User prediction
                     "Justification": just.strip(),
                     "AnnotatorUsername": user_tag
                 })
+                
         
         if new_rows:
             new_df = pd.DataFrame(new_rows)
@@ -819,9 +838,13 @@ with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as demo:
     first_pass_outputs = []
     
     # Unpack 9 items per row
-    for container, b, r, c_md, inter, j, m_co, m_dr, m_cf in eval_rows:
-        row_outputs.extend([container, b, r, c_md, inter, j, m_co, m_dr, m_cf])
-        row_inputs.extend([b, r, c_md, inter, j, m_co, m_dr, m_cf])
+    # for container, b, r, c_md, inter, j, m_co, m_dr, m_cf in eval_rows:
+    #     row_outputs.extend([container, b, r, c_md, inter, j, m_co, m_dr, m_cf])
+    #     row_inputs.extend([b, r, c_md, inter, j, m_co, m_dr, m_cf])
+    #     first_pass_outputs.extend([r, c_md, inter, m_co, m_dr, m_cf])
+    for container, b, r, c_md, inter, j, m_co, m_dr, m_cf, m_ai_j, m_ig_j in eval_rows:
+        row_outputs.extend([container, b, r, c_md, inter, j, m_co, m_dr, m_cf, m_ai_j, m_ig_j])
+        row_inputs.extend([b, r, c_md, inter, j, m_co, m_dr, m_cf, m_ai_j, m_ig_j])
         first_pass_outputs.extend([r, c_md, inter, m_co, m_dr, m_cf])
 
     lang_selector.change(