Update app.py

app.py

@@ -9,7 +9,6 @@ import numpy as np
 
 # Download the punkt tokenizer
 nltk.download('punkt_tab')
-
 # Helper function to split text into topics using KMeans clustering
 def split_text_into_topics(text, n_topics):
     sentences = sent_tokenize(text)
@@ -42,29 +41,36 @@ def recursive_split(topic_dict, depth, max_depth, subtopics):
             
     return new_topic_dict
 
-# Plotting function to visualize the tree structure
-def plot_tree(tree, parent=None, graph=None, level=0):
-    if graph is None:
-        graph = nx.Graph()
-    
+# Function to convert the tree into edge data for Plotly visualization
+def get_edges(tree, parent=None, level=0):
+    edges = []
+    labels = []
+    pos = {}
+
     for key, value in tree.items():
         node_label = f'Topic {key}' if parent is None else f'Subtopic {key}'
-        graph.add_node(node_label, level=level)
-        if parent:
-            graph.add_edge(parent, node_label)
+        pos[node_label] = (level, len(labels))
+        labels.append(node_label)
         
+        if parent:
+            edges.append((parent, node_label))
+
         if isinstance(value, dict):
-            plot_tree(value, parent=node_label, graph=graph, level=level+1)
+            new_edges, new_labels, new_pos = get_edges(value, node_label, level+1)
+            edges += new_edges
+            labels += new_labels
+            pos.update(new_pos)
         else:
             for i, sentence in enumerate(value):
                 sentence_label = f"{node_label} - Sentence {i+1}"
-                graph.add_node(sentence_label, level=level+1)
-                graph.add_edge(node_label, sentence_label)
+                pos[sentence_label] = (level+1, len(labels))
+                labels.append(sentence_label)
+                edges.append((node_label, sentence_label))
     
-    return graph
+    return edges, labels, pos
 
 # Streamlit App layout
-st.title('Text Topic Tree Generator')
+st.title('Interactive Text Topic Tree Generator')
 
 # Upload file
 uploaded_file = st.file_uploader("Upload a text file", type="txt")
@@ -83,13 +89,57 @@ if uploaded_file is not None:
     # Recursively split the topics into subtopics
     full_tree = recursive_split(topic_dict, 0, max_depth, subtopics_per_topic)
     
-    # Create and display the tree graph
-    graph = plot_tree(full_tree)
+    # Get edges and positions for the plot
+    edges, labels, pos = get_edges(full_tree)
+    
+    # Plot the tree graph using Plotly
+    edge_x = []
+    edge_y = []
+    for edge in edges:
+        x0, y0 = pos[edge[0]]
+        x1, y1 = pos[edge[1]]
+        edge_x += [x0, x1, None]
+        edge_y += [y0, y1, None]
+    
+    node_x = [pos[label][0] for label in labels]
+    node_y = [pos[label][1] for label in labels]
+    
+    # Create edge trace
+    edge_trace = go.Scatter(
+        x=edge_x, y=edge_y,
+        line=dict(width=2, color='Gray'),
+        hoverinfo='none',
+        mode='lines'
+    )
+    
+    # Create node trace
+    node_trace = go.Scatter(
+        x=node_x, y=node_y,
+        mode='markers+text',
+        text=labels,
+        hoverinfo='text',
+        marker=dict(
+            showscale=True,
+            colorscale='YlGnBu',
+            size=20,
+            colorbar=dict(
+                thickness=15,
+                title='Depth',
+                xanchor='left',
+                titleside='right'
+            ),
+            line_width=2
+        )
+    )
+    
+    # Plot the figure
+    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),
+                        yaxis=dict(showgrid=False, zeroline=False)
+                    ))
     
-    # Plot the tree graph
-    pos = nx.spring_layout(graph)
-    levels = nx.get_node_attributes(graph, 'level')
-    plt.figure(figsize=(12, 8))
-    nx.draw(graph, pos, with_labels=True, node_size=3000, node_color="lightblue", font_size=10, font_weight="bold", labels={node: node for node in graph.nodes()})
-    plt.title("Tree Structure of Text Topics")
-    st.pyplot(plt)
+    st.plotly_chart(fig)