Enhanced the knowledge graph

app.py

@@ -32,6 +32,92 @@ def initialize_models():
 class NewsInput(BaseModel):
     text: str
 
+# def generate_knowledge_graph_viz(text, nlp, tokenizer, model):
+#     kg_builder = KnowledgeGraphBuilder()
+    
+#     # Get prediction
+#     prediction, _ = predict_with_model(text, tokenizer, model)
+#     is_fake = prediction == "FAKE"
+    
+#     # Update knowledge graph
+#     kg_builder.update_knowledge_graph(text, not is_fake, nlp)
+
+#     # Randomly select subset of edges (e.g. 10% of edges)
+#     edges = list(kg_builder.knowledge_graph.edges())
+#     selected_edges = random.sample(edges, k=int(len(edges) * 0.3))
+    
+#     # Create a new graph with selected edges
+#     selected_graph = nx.DiGraph()
+#     selected_graph.add_nodes_from(kg_builder.knowledge_graph.nodes(data=True))
+#     selected_graph.add_edges_from(selected_edges)
+    
+#     pos = nx.spring_layout(selected_graph)
+    
+#     edge_trace = go.Scatter(
+#         x=[], y=[],
+#         line=dict(
+#             width=2, 
+#             color='rgba(255,0,0,0.7)' if is_fake else 'rgba(0,255,0,0.7)'
+#         ),
+#         hoverinfo='none',
+#         mode='lines'
+#     )
+    
+#     # Create visualization
+#     pos = nx.spring_layout(kg_builder.knowledge_graph)
+    
+#     edge_trace = go.Scatter(
+#         x=[], y=[],
+#         line=dict(
+#             width=2,
+#             color='rgba(255,0,0,0.7)' if is_fake else 'rgba(0,255,0,0.7)'
+#         ),
+#         hoverinfo='none',
+#         mode='lines'
+#     )
+    
+#     node_trace = go.Scatter(
+#         x=[], y=[],
+#         mode='markers+text',
+#         hoverinfo='text',
+#         textposition='top center',
+#         marker=dict(
+#             size=15,
+#             color='white',
+#             line=dict(width=2, color='black')
+#         ),
+#         text=[]
+#     )
+    
+#     # Add edges
+#     for edge in selected_graph.edges():
+#         x0, y0 = pos[edge[0]]
+#         x1, y1 = pos[edge[1]]
+#         edge_trace['x'] += (x0, x1, None)
+#         edge_trace['y'] += (y0, y1, None)
+    
+#     # Add nodes
+#     for node in kg_builder.knowledge_graph.nodes():
+#         x, y = pos[node]
+#         node_trace['x'] += (x,)
+#         node_trace['y'] += (y,)
+#         node_trace['text'] += (node,)
+    
+#     fig = go.Figure(
+#         data=[edge_trace, node_trace],
+#         layout=go.Layout(
+#             showlegend=False,
+#             hovermode='closest',
+#             margin=dict(b=0,l=0,r=0,t=0),
+#             xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
+#             yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
+#             plot_bgcolor='rgba(0,0,0,0)',
+#             paper_bgcolor='rgba(0,0,0,0)'
+#         )
+#     )
+    
+#     return fig
+
 def generate_knowledge_graph_viz(text, nlp, tokenizer, model):
     kg_builder = KnowledgeGraphBuilder()
     
@@ -42,18 +128,29 @@ def generate_knowledge_graph_viz(text, nlp, tokenizer, model):
     # Update knowledge graph
     kg_builder.update_knowledge_graph(text, not is_fake, nlp)
 
-    # Randomly select subset of edges (e.g. 10% of edges)
-    edges = list(kg_builder.knowledge_graph.edges())
-    selected_edges = random.sample(edges, k=int(len(edges) * 0.3))
+    # Get all edges from the knowledge graph
+    all_edges = list(kg_builder.knowledge_graph.edges())
+    total_edges = len(all_edges)
+    
+    # Select only 50% of edges to display
+    display_edge_count = int(total_edges * 0.5)
+    display_edges = random.sample(all_edges, k=min(display_edge_count, total_edges))
+    
+    # Determine how many edges should be the opposite color (15% of displayed edges)
+    opposite_color_count = int(len(display_edges) * 0.15)
+    
+    # Randomly select which edges will have the opposite color
+    opposite_color_edges = set(random.sample(display_edges, k=opposite_color_count))
     
     # Create a new graph with selected edges
     selected_graph = nx.DiGraph()
     selected_graph.add_nodes_from(kg_builder.knowledge_graph.nodes(data=True))
-    selected_graph.add_edges_from(selected_edges)
+    selected_graph.add_edges_from(display_edges)
     
     pos = nx.spring_layout(selected_graph)
     
-    edge_trace = go.Scatter(
+    # Create two edge traces - one for dominant color, one for opposite color
+    dominant_edge_trace = go.Scatter(
         x=[], y=[],
         line=dict(
             width=2, 
@@ -63,14 +160,136 @@ def generate_knowledge_graph_viz(text, nlp, tokenizer, model):
         mode='lines'
     )
     
-    # Create visualization
-    pos = nx.spring_layout(kg_builder.knowledge_graph)
+    opposite_edge_trace = go.Scatter(
+        x=[], y=[],
+        line=dict(
+            width=2, 
+            color='rgba(0,255,0,0.7)' if is_fake else 'rgba(255,0,0,0.7)'
+        ),
+        hoverinfo='none',
+        mode='lines'
+    )
     
-    edge_trace = go.Scatter(
+    node_trace = go.Scatter(
+        x=[], y=[],
+        mode='markers+text',
+        hoverinfo='text',
+        textposition='top center',
+        marker=dict(
+            size=15,
+            color='white',
+            line=dict(width=2, color='black')
+        ),
+        text=[]
+    )
+    
+    # Add edges with appropriate colors
+    for edge in display_edges:
+        x0, y0 = pos[edge[0]]
+        x1, y1 = pos[edge[1]]
+        
+        if edge in opposite_color_edges:
+            opposite_edge_trace['x'] += (x0, x1, None)
+            opposite_edge_trace['y'] += (y0, y1, None)
+        else:
+            dominant_edge_trace['x'] += (x0, x1, None)
+            dominant_edge_trace['y'] += (y0, y1, None)
+    
+    # Add nodes
+    for node in selected_graph.nodes():
+        x, y = pos[node]
+        node_trace['x'] += (x,)
+        node_trace['y'] += (y,)
+        node_trace['text'] += (node,)
+    
+    fig = go.Figure(
+        data=[dominant_edge_trace, opposite_edge_trace, node_trace],
+        layout=go.Layout(
+            showlegend=False,
+            hovermode='closest',
+            margin=dict(b=0,l=0,r=0,t=0),
+            xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
+            yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
+            plot_bgcolor='rgba(0,0,0,0)',
+            paper_bgcolor='rgba(0,0,0,0)'
+        )
+    )
+    
+    return fig
+
+def generate_knowledge_graph_viz(text, nlp, tokenizer, model):
+    kg_builder = KnowledgeGraphBuilder()
+    
+    # Get prediction
+    prediction, _ = predict_with_model(text, tokenizer, model)
+    is_fake = prediction == "FAKE"
+    
+    # Update knowledge graph
+    kg_builder.update_knowledge_graph(text, not is_fake, nlp)
+
+    # Get all edges from the knowledge graph
+    all_edges = list(kg_builder.knowledge_graph.edges())
+    total_edges = len(all_edges)
+    
+    # Select only 60% of edges to display (0.3 + 0.15 + 0.15)
+    display_edge_count = int(total_edges * 0.6)
+    display_edges = random.sample(all_edges, k=min(display_edge_count, total_edges))
+    
+    # Determine edge counts for each color
+    primary_color_count = int(total_edges * 0.3)  # 30% primary color (green for real, red for fake)
+    opposite_color_count = int(total_edges * 0.15)  # 15% opposite color
+    orange_color_count = int(total_edges * 0.15)  # 15% orange
+    
+    # Ensure we don't exceed the number of display edges
+    total_colored = primary_color_count + opposite_color_count + orange_color_count
+    if total_colored > len(display_edges):
+        ratio = len(display_edges) / total_colored
+        primary_color_count = int(primary_color_count * ratio)
+        opposite_color_count = int(opposite_color_count * ratio)
+        orange_color_count = int(orange_color_count * ratio)
+    
+    # Shuffle display edges to ensure random distribution
+    random.shuffle(display_edges)
+    
+    # Assign colors to edges
+    primary_color_edges = set(display_edges[:primary_color_count])
+    opposite_color_edges = set(display_edges[primary_color_count:primary_color_count+opposite_color_count])
+    orange_color_edges = set(display_edges[primary_color_count+opposite_color_count:
+                                          primary_color_count+opposite_color_count+orange_color_count])
+    
+    # Create a new graph with selected edges
+    selected_graph = nx.DiGraph()
+    selected_graph.add_nodes_from(kg_builder.knowledge_graph.nodes(data=True))
+    selected_graph.add_edges_from(display_edges)
+    
+    pos = nx.spring_layout(selected_graph)
+    
+    # Create three edge traces - primary, opposite, and orange
+    primary_edge_trace = go.Scatter(
         x=[], y=[],
         line=dict(
-            width=2,
-            color='rgba(255,0,0,0.7)' if is_fake else 'rgba(0,255,0,0.7)'
+            width=2, 
+            color='rgba(255,0,0,0.7)' if is_fake else 'rgba(0,255,0,0.7)'  # Red if fake, green if real
+        ),
+        hoverinfo='none',
+        mode='lines'
+    )
+    
+    opposite_edge_trace = go.Scatter(
+        x=[], y=[],
+        line=dict(
+            width=2, 
+            color='rgba(0,255,0,0.7)' if is_fake else 'rgba(255,0,0,0.7)'  # Green if fake, red if real
+        ),
+        hoverinfo='none',
+        mode='lines'
+    )
+    
+    orange_edge_trace = go.Scatter(
+        x=[], y=[],
+        line=dict(
+            width=2, 
+            color='rgba(255,165,0,0.7)'  # Orange
         ),
         hoverinfo='none',
         mode='lines'
@@ -89,22 +308,30 @@ def generate_knowledge_graph_viz(text, nlp, tokenizer, model):
         text=[]
     )
     
-    # Add edges
-    for edge in selected_graph.edges():
+    # Add edges with appropriate colors
+    for edge in display_edges:
         x0, y0 = pos[edge[0]]
         x1, y1 = pos[edge[1]]
-        edge_trace['x'] += (x0, x1, None)
-        edge_trace['y'] += (y0, y1, None)
+        
+        if edge in primary_color_edges:
+            primary_edge_trace['x'] += (x0, x1, None)
+            primary_edge_trace['y'] += (y0, y1, None)
+        elif edge in opposite_color_edges:
+            opposite_edge_trace['x'] += (x0, x1, None)
+            opposite_edge_trace['y'] += (y0, y1, None)
+        elif edge in orange_color_edges:
+            orange_edge_trace['x'] += (x0, x1, None)
+            orange_edge_trace['y'] += (y0, y1, None)
     
     # Add nodes
-    for node in kg_builder.knowledge_graph.nodes():
+    for node in selected_graph.nodes():
         x, y = pos[node]
         node_trace['x'] += (x,)
         node_trace['y'] += (y,)
         node_trace['text'] += (node,)
     
     fig = go.Figure(
-        data=[edge_trace, node_trace],
+        data=[primary_edge_trace, opposite_edge_trace, orange_edge_trace, node_trace],
         layout=go.Layout(
             showlegend=False,
             hovermode='closest',