Add ploty for knowledge graph

app.py

@@ -50,6 +50,8 @@ st.set_page_config(
 
 import pandas as pd
 from final import *
+from pydantic import BaseModel
+import plotly.graph_objects as go
 
 # # Cache model loading
 # @st.cache_resource
@@ -78,6 +80,75 @@ def initialize_models():
     knowledge_graph = load_knowledge_graph()
     return nlp, tokenizer, model, knowledge_graph
 
+
+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)
+    
+    # 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 kg_builder.knowledge_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
+
 # Streamlit UI
 def main():
     st.title("Nexus NLP News Classifier")
@@ -132,6 +203,10 @@ def main():
                     df = pd.DataFrame(entities, columns=["Entity", "Type"])
                     st.dataframe(df)
 
+                 # Generate and display knowledge graph
+                fig = generate_knowledge_graph_viz(news_text, nlp, tokenizer, model)
+                st.plotly_chart(fig, use_container_width=True)
+
         else:
             st.warning("Please enter some text to analyze")
 

final.py

@@ -278,6 +278,7 @@ import google.generativeai as genai
 import json
 import os
 import dotenv
+import plotly.graph_objects as go
 
 # Load environment variables
 dotenv.load_dotenv()
@@ -303,6 +304,40 @@ def load_knowledge_graph():
             knowledge_graph.add_edge(u, v, **data)
     return knowledge_graph
 
+
+class KnowledgeGraphBuilder:
+    def __init__(self):
+        self.knowledge_graph = nx.DiGraph()
+        
+    def update_knowledge_graph(self, text, is_real, nlp):
+        entities = extract_entities(text, nlp)
+        for entity, entity_type in entities:
+            if not self.knowledge_graph.has_node(entity):
+                self.knowledge_graph.add_node(
+                    entity,
+                    type=entity_type,
+                    real_count=1 if is_real else 0,
+                    fake_count=0 if is_real else 1
+                )
+            else:
+                if is_real:
+                    self.knowledge_graph.nodes[entity]['real_count'] += 1
+                else:
+                    self.knowledge_graph.nodes[entity]['fake_count'] += 1
+
+        for i, (entity1, _) in enumerate(entities):
+            for entity2, _ in entities[i+1:]:
+                if not self.knowledge_graph.has_edge(entity1, entity2):
+                    self.knowledge_graph.add_edge(
+                        entity1,
+                        entity2,
+                        weight=1,
+                        is_real=is_real
+                    )
+                else:
+                    self.knowledge_graph[entity1][entity2]['weight'] += 1
+
+
 def setup_gemini():
     """Initialize Gemini model"""
     genai.configure(api_key=os.getenv("GEMINI_API"))

requirements.txt

@@ -11,4 +11,6 @@ uvicorn
 tiktoken
 sentencepiece
 timm
+plotly
+networkx
 https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.7.0/en_core_web_sm-3.7.0.tar.gz