Create app.py

app.py

@@ -0,0 +1,311 @@
+import numpy as np
+import plotly.graph_objects as go
+import json
+import gradio as gr
+from nltk.corpus import words
+import nltk
+
+
+# load files w embeddings, attention scores, and tokens
+vocab_embeddings = np.load('vocab_embeddings.npy')
+with open('vocab_attention_scores.json', 'r') as f:
+    vocab_attention_scores = json.load(f)
+with open('vocab_tokens.json', 'r') as f:
+    vocab_tokens = json.load(f)
+
+# attention scores to numpy arrs
+b_gen_attention = np.array([score['B-GEN'] for score in vocab_attention_scores])
+i_gen_attention = np.array([score['I-GEN'] for score in vocab_attention_scores])
+b_unfair_attention = np.array([score['B-UNFAIR'] for score in vocab_attention_scores])
+i_unfair_attention = np.array([score['I-UNFAIR'] for score in vocab_attention_scores])
+b_stereo_attention = np.array([score['B-STEREO'] for score in vocab_attention_scores])
+i_stereo_attention = np.array([score['I-STEREO'] for score in vocab_attention_scores])
+o_attention = np.array([score['O'] for score in vocab_attention_scores])  # Use actual O scores
+
+# remove non-dict english words, but keep subwords ##
+nltk.download('words')
+english_words = set(words.words())
+
+filtered_indices = [i for i, token in enumerate(vocab_tokens) if token in english_words or token.startswith("##")]
+filtered_tokens = [vocab_tokens[i] for i in filtered_indices]
+
+b_gen_attention_filtered = b_gen_attention[filtered_indices]
+i_gen_attention_filtered = i_gen_attention[filtered_indices]
+b_unfair_attention_filtered = b_unfair_attention[filtered_indices]
+i_unfair_attention_filtered = i_unfair_attention[filtered_indices]
+b_stereo_attention_filtered = b_stereo_attention[filtered_indices]
+i_stereo_attention_filtered = i_stereo_attention[filtered_indices]
+o_attention_filtered = o_attention[filtered_indices]
+
+# plot top 500 O tokens for comparison
+top_500_o_indices = np.argsort(o_attention_filtered)[-500:]
+top_500_o_tokens = [filtered_tokens[i] for i in top_500_o_indices]
+o_attention_filtered_top_500 = o_attention_filtered[top_500_o_indices]
+
+# tool tip for tokens
+def create_hover_text(tokens, b_gen, i_gen, b_unfair, i_unfair, b_stereo, i_stereo, o_val):
+    hover_text = []
+    for i in range(len(tokens)):
+        hover_text.append(
+            f"Token: {tokens[i]}<br>"
+            f"B-GEN: {b_gen[i]:.3f}, I-GEN: {i_gen[i]:.3f}<br>"
+            f"B-UNFAIR: {b_unfair[i]:.3f}, I-UNFAIR: {i_unfair[i]:.3f}<br>"
+            f"B-STEREO: {b_stereo[i]:.3f}, I-STEREO: {i_stereo[i]:.3f}<br>"
+            f"O: {o_val[i]:.3f}"
+        )
+    return hover_text
+
+# ploting top 100 tokens for each entity
+def select_top_100(*data_arrays):
+    indices_list = []
+    for data in data_arrays:
+        if data is not None:
+            top_indices = np.argsort(data)[-100:] 
+            indices_list.append(top_indices)
+    
+    combined_indices = np.unique(np.concatenate(indices_list))
+
+    # filter based on combined indices
+    filtered_data = [data[combined_indices] if data is not None else None for data in data_arrays]
+    tokens_filtered = [filtered_tokens[i] for i in combined_indices]
+
+    return (*filtered_data, tokens_filtered)
+
+# plots for 1 2 and 3 D
+def create_plot(selected_dimensions):
+    # plot data
+    attention_map = {
+        'Generalization': b_gen_attention_filtered + i_gen_attention_filtered,
+        'Unfairness': b_unfair_attention_filtered + i_unfair_attention_filtered,
+        'Stereotype': b_stereo_attention_filtered + i_stereo_attention_filtered,
+    }
+
+    # init x, y, z so they can be moved around
+    x_data, y_data, z_data = None, None, None
+
+    # use selected dimentsions to order dimensions
+    if len(selected_dimensions) > 0:
+        x_data = attention_map[selected_dimensions[0]]
+    if len(selected_dimensions) > 1:
+        y_data = attention_map[selected_dimensions[1]] 
+    if len(selected_dimensions) > 2:
+        z_data = attention_map[selected_dimensions[2]]
+
+    # select top 100 dps for each selected dimension
+    x_data, y_data, z_data, tokens_filtered = select_top_100(x_data, y_data, z_data)
+
+    # filter the O tokens using the same dimensions
+    o_x = attention_map[selected_dimensions[0]][top_500_o_indices]
+    if len(selected_dimensions) > 1:
+        o_y = attention_map[selected_dimensions[1]][top_500_o_indices]
+    else:
+        o_y = np.zeros_like(o_x) 
+    if len(selected_dimensions) > 2:
+        o_z = attention_map[selected_dimensions[2]][top_500_o_indices]
+    else:
+        o_z = np.zeros_like(o_x)
+
+    # hover text for GUS tokens
+    classified_hover_text = create_hover_text(
+        tokens_filtered, 
+        b_gen_attention_filtered, i_gen_attention_filtered, 
+        b_unfair_attention_filtered, i_unfair_attention_filtered, 
+        b_stereo_attention_filtered, i_stereo_attention_filtered, 
+        o_attention_filtered
+    )
+
+    # hover text for O tokens
+    o_hover_text = create_hover_text(
+        top_500_o_tokens, 
+        b_gen_attention_filtered[top_500_o_indices], i_gen_attention_filtered[top_500_o_indices], 
+        b_unfair_attention_filtered[top_500_o_indices], i_unfair_attention_filtered[top_500_o_indices], 
+        b_stereo_attention_filtered[top_500_o_indices], i_stereo_attention_filtered[top_500_o_indices], 
+        o_attention_filtered_top_500
+    )
+
+
+    # plot
+    fig = go.Figure()
+
+    if x_data is not None and y_data is not None and z_data is not None:
+        # 3d scatter plot
+        fig.add_trace(go.Scatter3d(
+            x=x_data,
+            y=y_data,
+            z=z_data,
+            mode='markers',
+            marker=dict(
+                size=6,
+                color=x_data,  # color based on the x-axis data
+                colorscale='Viridis',
+                opacity=0.85,
+            ),
+            text=classified_hover_text, 
+            hoverinfo='text',
+            name='Classified Tokens'
+        ))
+        # add top 500 O tags to the plot too
+        fig.add_trace(go.Scatter3d(
+            x=o_x,
+            y=o_y,
+            z=o_z,
+            mode='markers',
+            marker=dict(
+                size=6,
+                color='grey',
+                opacity=0.5,
+            ),
+            text=o_hover_text,
+            hoverinfo='text',
+            name='O Tokens'
+        ))
+    elif x_data is not None and y_data is not None:
+        # 2d scatter plot
+        fig.add_trace(go.Scatter(
+            x=x_data,
+            y=y_data,
+            mode='markers',
+            marker=dict(
+                size=6,
+                color=x_data, # color based on the x-axis data
+                colorscale='Viridis',
+                opacity=0.85,
+            ),
+            text=classified_hover_text,
+            hoverinfo='text',
+            name='Classified Tokens'
+        ))
+        # add top 500 O tags to the plot too
+        fig.add_trace(go.Scatter(
+            x=o_x,
+            y=o_y,
+            mode='markers',
+            marker=dict(
+                size=6,
+                color='grey',
+                opacity=0.5,
+            ),
+            text=o_hover_text, 
+            hoverinfo='text',
+            name='O Tokens'
+        ))
+    elif x_data is not None:
+        # 1D scatter plot
+        fig.add_trace(go.Scatter(
+            x=x_data,
+            y=np.zeros_like(x_data),
+            mode='markers',
+            marker=dict(
+                size=6,
+                color=x_data, 
+                colorscale='Viridis',
+                opacity=0.85,
+            ),
+            text=classified_hover_text,
+            hoverinfo='text',
+            name='GUS Tokens'
+        ))
+        fig.add_trace(go.Scatter(
+            x=o_x,
+            y=np.zeros_like(o_x), 
+            mode='markers',
+            marker=dict(
+                size=6,
+                color='grey',
+                opacity=0.5,
+            ),
+            text=o_hover_text,
+            hoverinfo='text',
+            name='O Tokens'
+        ))
+
+    # update layout dynamically
+    if x_data is not None and y_data is not None and z_data is not None:
+        # 3D
+        fig.update_layout(
+            title="GUS-Net Entity Attentions Visualization",
+            scene=dict(
+                xaxis=dict(title=f"{selected_dimensions[0]} Attention"),
+                yaxis=dict(title=f"{selected_dimensions[1]} Attention"),
+                zaxis=dict(title=f"{selected_dimensions[2]} Attention"),
+            ),
+            margin=dict(l=0, r=0, b=0, t=40),
+        )
+    elif x_data is not None and y_data is not None:
+        # 2D
+        fig.update_layout(
+            title="GUS-Net Entity Attentions Visualization",
+            xaxis_title=f"{selected_dimensions[0]} Attention",
+            yaxis_title=f"{selected_dimensions[1]} Attention",
+            margin=dict(l=0, r=0, b=0, t=40),
+        )
+    elif x_data is not None:
+        # 1D
+        fig.update_layout(
+            title="GUS-Net Entity Attentions Visualization",
+            xaxis_title=f"{selected_dimensions[0]} Attention",
+            margin=dict(l=0, r=0, b=0, t=40),
+        )
+
+    return fig
+
+def get_top_tokens_for_entities(selected_dimensions):
+    entity_map = {
+        'Generalization': b_gen_attention_filtered + i_gen_attention_filtered,
+        'Unfairness': b_unfair_attention_filtered + i_unfair_attention_filtered,
+        'Stereotype': b_stereo_attention_filtered + i_stereo_attention_filtered,
+    }
+
+    top_tokens_info = {}
+    for dimension in selected_dimensions:
+        if dimension in entity_map:
+            attention_scores = entity_map[dimension]
+            top_indices = np.argsort(attention_scores)[-10:]  # top 10 tokens
+            top_tokens = [filtered_tokens[i] for i in top_indices]
+            top_scores = attention_scores[top_indices]
+            top_tokens_info[dimension] = list(zip(top_tokens, top_scores))
+
+    return top_tokens_info
+
+def update_gradio(selected_dimensions):
+    fig = create_plot(selected_dimensions)
+
+    top_tokens_info = get_top_tokens_for_entities(selected_dimensions)
+    
+    formatted_top_tokens = ""
+    for entity, tokens_info in top_tokens_info.items():
+        formatted_top_tokens += f"\nTop tokens for {entity}:\n"
+        for token, score in tokens_info:
+            formatted_top_tokens += f"Token: {token}, Attention Score: {score:.3f}\n"
+
+    return fig, formatted_top_tokens
+
+
+def render_gradio_interface():
+    with gr.Blocks() as interface:
+        with gr.Column():
+            dimensions_input = gr.CheckboxGroup(
+                choices=["Generalization", "Unfairness", "Stereotype"],
+                label="Select Dimensions to Plot",
+                value=["Generalization", "Unfairness", "Stereotype"]  # defaults to 3D
+            )
+
+            plot_output = gr.Plot(label="Token Attention Visualization")
+            top_tokens_output = gr.Textbox(label="Top Tokens for Each Entity Class", lines=10)
+
+            dimensions_input.change(
+                fn=update_gradio,
+                inputs=[dimensions_input],
+                outputs=[plot_output, top_tokens_output]
+            )
+
+            interface.load(
+                fn=lambda: update_gradio(["Generalization", "Unfairness", "Stereotype"]),
+                inputs=None, 
+                outputs=[plot_output, top_tokens_output]
+            )
+
+    return interface
+
+interface = render_gradio_interface()
+interface.launch()