Update app.py

app.py

@@ -1,12 +1,17 @@
-# app.py
 import streamlit as st
 import spacy
 import subprocess
 from neo4j import GraphDatabase
 import matplotlib.pyplot as plt
 import networkx as nx
+import logging
+from sklearn.feature_extraction.text import TfidfVectorizer
+import re
 
-# === Ensure spaCy model is installed ===
+# Set up logging configuration
+logging.basicConfig(level=logging.INFO)
+
+# Ensure spaCy model is installed
 def install_spacy_model():
     try:
         spacy.load("en_core_web_sm")
@@ -14,55 +19,118 @@ def install_spacy_model():
         subprocess.run(["python", "-m", "spacy", "download", "en_core_web_sm"])
 install_spacy_model()
 
-# Load the model after ensuring it's installed
+# Load the spaCy model after ensuring it's installed
 nlp = spacy.load("en_core_web_sm")
 
 # Neo4j credentials
 uri = "neo4j+s://ff701b1c.databases.neo4j.io"
 username = "neo4j"
-password = "BfZM7YRKpFz1b_V7acAmOtaSQHPU9xK03rJlfPep88g"  
-
-# Connect to Neo4j
-driver = GraphDatabase.driver(uri, auth=(username, password))
-
-# Triple extraction function
-def extract_triples(text):
-    doc = nlp(text)
-    triples = []
-    for sent in doc.sents:
-        subjects = [tok for tok in sent if "subj" in tok.dep_]
-        verbs = [tok for tok in sent if tok.pos_ == "VERB"]
-        objects = [tok for tok in sent if "obj" in tok.dep_]
-        for subj in subjects:
-            for verb in verbs:
+password = "BfZM7YRKpFz1b_V7acAmOtaSQHPU9xK03rJlfPep88g"
+
+driver = None
+try:
+    driver = GraphDatabase.driver(uri, auth=(username, password))
+    logging.info("✅ Connected to Neo4j!")
+
+    def close_driver():
+        if driver:
+            driver.close()
+            logging.info("🔒 Neo4j driver closed.")
+
+    def create_entity(tx, name: str):
+        tx.run("MERGE (e:Entity {name: $name})", name=name)
+
+    def create_relationship(tx, subj: str, pred: str, obj: str):
+        tx.run("""
+            MERGE (a:Entity {name: $subj})
+            MERGE (b:Entity {name: $obj})
+            MERGE (a)-[:RELATION {name: $pred}]->(b)
+        """, subj=subj, pred=pred, obj=obj)
+
+    # Text Processing
+    def load_and_clean_text(file_path: str) -> str:
+        with open(file_path, 'r', encoding='utf-8') as file:
+            text = file.read()
+        text = re.sub(r'\n+', ' ', text)
+        return re.sub(r'\s+', ' ', text).strip().lower()
+
+    # TF-IDF Filtering
+    def compute_tfidf_keywords(text: str, top_n=60):
+        vectorizer = TfidfVectorizer(stop_words='english')
+        X = vectorizer.fit_transform([text])
+        scores = zip(vectorizer.get_feature_names_out(), X.toarray()[0])
+        sorted_scores = sorted(scores, key=lambda x: x[1], reverse=True)
+        return {word for word, _ in sorted_scores[:top_n]}
+
+    # Triple Extraction
+    def get_full_phrase(token) -> str:
+        return ' '.join(tok.text for tok in token.subtree if tok.dep_ != 'punct').strip()
+
+    def extract_rich_triples(doc, tfidf_keywords) -> list:
+        triples = []
+        for sent in doc.sents:
+            subjects = [tok for tok in sent if "subj" in tok.dep_]
+            objects = [tok for tok in sent if "obj" in tok.dep_]
+            verbs = [tok for tok in sent if tok.pos_ == "VERB"]
+            for subj in subjects:
                 for obj in objects:
-                    triples.append((subj.text, verb.lemma_, obj.text))
-    return triples
-
-# Visualization function
-def show_graph(triples):
-    G = nx.DiGraph()
-    for s, p, o in triples:
-        G.add_node(s)
-        G.add_node(o)
-        G.add_edge(s, o, label=p)
-    pos = nx.spring_layout(G)
-    plt.figure(figsize=(10, 8))
-    nx.draw(G, pos, with_labels=True, node_color='skyblue', node_size=2000, font_size=10)
-    nx.draw_networkx_edge_labels(G, pos, edge_labels={(u,v):d['label'] for u,v,d in G.edges(data=True)})
-    st.pyplot(plt)
-
-# === Streamlit UI ===
-st.title("🧠 Knowledge Graph Generator")
-
-text_input = st.text_area("Paste your text here", height=200)
-
-if st.button("Generate Graph"):
-    if text_input:
-        triples = extract_triples(text_input)
-        st.write("### Extracted Triples")
-        for t in triples:
-            st.write("🔗", t)
-        show_graph(triples)
-    else:
-        st.warning("Please enter some text.")
+                    for verb in verbs:
+                        s = get_full_phrase(subj)
+                        o = get_full_phrase(obj)
+                        if s.lower() in tfidf_keywords or o.lower() in tfidf_keywords:
+                            triples.append((s, verb.lemma_, o))
+        return triples
+
+    # Graph Visualization
+    def visualize_knowledge_graph(triples: list, title: str = "Knowledge Graph"):
+        G = nx.DiGraph()
+        for subj, pred, obj in triples:
+            G.add_node(subj, label='Subject')
+            G.add_node(obj, label='Object')
+            G.add_edge(subj, obj, label=pred)
+
+        pos = nx.spring_layout(G, k=1.2, seed=42)
+        node_colors = ['skyblue' if G.nodes[n]['label'] == 'Subject' else 'lightgreen' for n in G.nodes]
+
+        plt.figure(figsize=(16, 16))
+        nx.draw(G, pos, with_labels=True, node_size=1200, node_color=node_colors,
+                font_size=10, font_weight='bold', edge_color='gray', alpha=0.8)
+        nx.draw_networkx_edge_labels(G, pos, edge_labels={(u, v): d['label'] for u, v, d in G.edges(data=True)},
+                                     font_size=8, font_color='red')
+        plt.title(title, fontsize=20)
+        plt.show()
+
+    # Streamlit UI
+    st.title("🧠 Knowledge Graph Generator")
+
+    text_input = st.text_area("Paste your text here", height=200)
+
+    if st.button("Generate Graph"):
+        if text_input:
+            triples = extract_rich_triples(nlp(text_input), compute_tfidf_keywords(text_input))
+            logging.info(f"🧠 Extracted {len(triples)} filtered triples.")
+            for t in triples[:10]:
+                st.write("🔗", t)
+
+            # Push to Neo4j
+            with driver.session() as session:
+                for subj, pred, obj in triples:
+                    session.execute_write(create_entity, subj)
+                    session.execute_write(create_entity, obj)
+                    session.execute_write(create_relationship, subj, pred, obj)
+
+            logging.info("📡 Triples successfully stored in Neo4j.")
+            st.success("Triples successfully stored in Neo4j.")
+
+            # Final Visualization
+            visualize_knowledge_graph(triples, title="Filtered Knowledge Graph (TF-IDF)")
+        else:
+            st.warning("Please enter some text.")
+
+except Exception as e:
+    logging.error(f"❌ An error occurred: {e}", exc_info=True)
+    st.error("An error occurred. Please check the logs.")
+
+finally:
+    if driver:
+        close_driver()