Update app.py

app.py

@@ -2,6 +2,8 @@ import gradio as gr
 import pandas as pd
 from openai import OpenAI
 import os
+import plotly.graph_objects as go
+import networkx as nx
 
 # Initialize NVIDIA API client
 client = OpenAI(
@@ -16,6 +18,9 @@ relations_df = pd.read_csv("extracted_entities_relations_countries.csv")
 # Get list of entities for dropdown
 entity_list = validated_entities_df['entity'].dropna().unique().tolist()
 
+# Get all unique countries from the data
+all_countries_in_data = relations_df['country'].dropna().unique().tolist()
+
 def get_entity_triples(entity):
     """Extract all triples associated with the selected entity"""
     # Filter rows where entity appears in head or tail
@@ -29,7 +34,7 @@ def get_entity_triples(entity):
     # Extract unique triples
     triples = []
     for _, row in filtered_df.iterrows():
-        triple = f"{row['head']} | {row['relation']} | {row['tail']}"
+        triple = (row['head'], row['relation'], row['tail'])
         if triple not in triples:
             triples.append(triple)
     
@@ -41,12 +46,171 @@ def get_entity_triples(entity):
     
     return triples, sentences, countries
 
+def create_knowledge_graph(entity, triples):
+    """Create interactive knowledge graph visualization using plotly"""
+    if not triples:
+        return None
+    
+    # Create directed graph
+    G = nx.DiGraph()
+    
+    # Add edges (triples)
+    for head, relation, tail in triples:
+        G.add_edge(head, tail, label=relation)
+    
+    # Generate layout
+    pos = nx.spring_layout(G, k=2, iterations=50, seed=42)
+    
+    # Create edge traces
+    edge_traces = []
+    edge_labels = []
+    
+    for edge in G.edges(data=True):
+        x0, y0 = pos[edge[0]]
+        x1, y1 = pos[edge[1]]
+        
+        # Edge line
+        edge_trace = go.Scatter(
+            x=[x0, x1, None],
+            y=[y0, y1, None],
+            mode='lines',
+            line=dict(width=2, color='#888'),
+            hoverinfo='none',
+            showlegend=False
+        )
+        edge_traces.append(edge_trace)
+        
+        # Edge label (relation)
+        edge_labels.append(go.Scatter(
+            x=[(x0 + x1) / 2],
+            y=[(y0 + y1) / 2],
+            mode='text',
+            text=[edge[2]['label']],
+            textposition='middle center',
+            textfont=dict(size=8, color='#666'),
+            hoverinfo='text',
+            hovertext=edge[2]['label'],
+            showlegend=False
+        ))
+    
+    # Create node trace
+    node_x = []
+    node_y = []
+    node_text = []
+    node_colors = []
+    node_sizes = []
+    
+    for node in G.nodes():
+        x, y = pos[node]
+        node_x.append(x)
+        node_y.append(y)
+        node_text.append(node)
+        
+        # Highlight the main entity
+        if node == entity:
+            node_colors.append('#FF6B6B')
+            node_sizes.append(30)
+        else:
+            node_colors.append('#4ECDC4')
+            node_sizes.append(20)
+    
+    node_trace = go.Scatter(
+        x=node_x,
+        y=node_y,
+        mode='markers+text',
+        text=node_text,
+        textposition='top center',
+        textfont=dict(size=10, color='#000'),
+        marker=dict(
+            size=node_sizes,
+            color=node_colors,
+            line=dict(width=2, color='#fff')
+        ),
+        hoverinfo='text',
+        hovertext=node_text,
+        showlegend=False
+    )
+    
+    # Create figure
+    fig = go.Figure(data=edge_traces + edge_labels + [node_trace])
+    
+    fig.update_layout(
+        title=dict(
+            text=f"Knowledge Graph for: {entity}",
+            x=0.5,
+            xanchor='center',
+            font=dict(size=16)
+        ),
+        showlegend=False,
+        hovermode='closest',
+        margin=dict(b=20, l=5, r=5, t=40),
+        xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
+        yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
+        plot_bgcolor='#f9f9f9',
+        height=600
+    )
+    
+    return fig
+
+def create_africa_map(countries):
+    """Create Africa map with highlighted countries"""
+    if not countries:
+        return None
+    
+    # Use all countries from the dataset for the base map
+    locations = []
+    z_values = []
+    hover_text = []
+    
+    for country in all_countries_in_data:
+        locations.append(country)
+        if country in countries:
+            z_values.append(1)
+            hover_text.append(f"{country} (Selected Entity)")
+        else:
+            z_values.append(0)
+            hover_text.append(country)
+    
+    fig = go.Figure(data=go.Choropleth(
+        locations=locations,
+        locationmode='country names',
+        z=z_values,
+        text=hover_text,
+        colorscale=[[0, '#E8E8E8'], [1, '#FF6B6B']],
+        showscale=False,
+        marker_line_color='white',
+        marker_line_width=1,
+        hovertemplate='<b>%{text}</b><extra></extra>'
+    ))
+    
+    fig.update_geos(
+        scope='africa',
+        showframe=True,
+        showcoastlines=True,
+        projection_type='natural earth',
+        bgcolor='#f0f0f0'
+    )
+    
+    fig.update_layout(
+        title=dict(
+            text=f"Countries with Entity Mentions ({len(countries)} of {len(all_countries_in_data)} total)",
+            x=0.5,
+            xanchor='center',
+            font=dict(size=16)
+        ),
+        margin=dict(l=0, r=0, t=40, b=0),
+        height=500,
+        geo=dict(bgcolor='#f9f9f9')
+    )
+    
+    return fig
+
 def generate_entity_description(entity, triples):
     """Generate entity description using Llama 405B based on triples"""
     if not triples:
         return "No triples found for this entity."
     
-    triples_text = "\n".join([f"- {triple}" for triple in triples])
+    triples_text = "\n".join([f"- {head} | {rel} | {tail}" for head, rel, tail in triples])
     
     prompt = f"""Based on the following knowledge graph triples, provide a comprehensive description of "{entity}":
 
@@ -107,16 +271,19 @@ Please create a well-structured paragraph that synthesizes this information spec
 def process_entity(entity):
     """Main function to process selected entity"""
     if not entity:
-        return "Please select an entity.", "", "", ""
+        return None, None, "", "", ""
     
     # Get triples, sentences, and countries
     triples, sentences, countries = get_entity_triples(entity)
     
     if not triples:
-        return "No data found for this entity.", "", "", ""
+        return None, None, "No data found for this entity.", "", ""
+    
+    # Create knowledge graph
+    kg_fig = create_knowledge_graph(entity, triples)
     
-    # Format triples for display
-    triples_display = "\n".join(triples)
+    # Create Africa map
+    map_fig = create_africa_map(countries)
     
     # Generate entity description
     description = generate_entity_description(entity, triples)
@@ -127,15 +294,15 @@ def process_entity(entity):
     # Format metadata
     metadata = f"**Associated Countries:** {', '.join(countries)}\n\n**Number of Triples:** {len(triples)}\n\n**Number of Sentences:** {len(sentences)}"
     
-    return triples_display, description, country_paragraphs, metadata
+    return kg_fig, map_fig, description, country_paragraphs, metadata
 
 # Create Gradio interface
-with gr.Blocks(title="Entity Knowledge Graph Explorer") as demo:
+with gr.Blocks(title="Entity Knowledge Graph Explorer", theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
         # 🔍 Entity Knowledge Graph Explorer
         
-        Select an entity to explore its knowledge graph triples, generate AI descriptions, 
+        Select an entity to explore its knowledge graph, generate AI descriptions, 
         and view country-specific insights using Llama 405B.
         """
     )
@@ -153,13 +320,15 @@ with gr.Blocks(title="Entity Knowledge Graph Explorer") as demo:
             gr.Markdown("### Metadata")
             metadata_output = gr.Markdown()
         
-        with gr.Column(scale=2):
-            gr.Markdown("### Knowledge Graph Triples")
-            triples_output = gr.Textbox(
-                label="Extracted Triples (Head | Relation | Tail)",
-                lines=10,
-                max_lines=15
-            )
+        with gr.Column(scale=3):
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown("### Knowledge Graph Visualization")
+                    kg_plot = gr.Plot(label="Knowledge Graph")
+                
+                with gr.Column():
+                    gr.Markdown("### Geographic Distribution")
+                    map_plot = gr.Plot(label="Africa Map")
             
             gr.Markdown("### AI-Generated Entity Description")
             description_output = gr.Markdown()
@@ -171,14 +340,14 @@ with gr.Blocks(title="Entity Knowledge Graph Explorer") as demo:
     search_btn.click(
         fn=process_entity,
         inputs=[entity_dropdown],
-        outputs=[triples_output, description_output, country_output, metadata_output]
+        outputs=[kg_plot, map_plot, description_output, country_output, metadata_output]
     )
     
     # Also trigger on dropdown change
     entity_dropdown.change(
         fn=process_entity,
         inputs=[entity_dropdown],
-        outputs=[triples_output, description_output, country_output, metadata_output]
+        outputs=[kg_plot, map_plot, description_output, country_output, metadata_output]
     )
     
     gr.Markdown(