Update app.py

app.py

@@ -15,7 +15,7 @@ def install_spacy_model():
         subprocess.run(["python", "-m", "spacy", "download", "en_core_web_sm"])
 install_spacy_model()
 
-# Load the spaCy model
+# Load spaCy model
 nlp = spacy.load("en_core_web_sm")
 
 # === Neo4j credentials ===
@@ -27,29 +27,38 @@ password = "BfZM7YRKpFz1b_V7acAmOtaSQHPU9xK03rJlfPep88g"
 driver = GraphDatabase.driver(uri, auth=(username, password))
 
 # === TF-IDF Filtering ===
-def compute_tfidf_keywords(text: str, top_n=60):
+def compute_tfidf_keywords(text: str, top_n=100):
     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 with TF-IDF filtering and limit to 10 ===
+# === Enhanced Triple Extraction ===
 def extract_triples(text):
     doc = nlp(text)
     tfidf_keywords = compute_tfidf_keywords(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:
-                for obj in objects:
-                    if (subj.text.lower() in tfidf_keywords or 
-                        verb.lemma_.lower() in tfidf_keywords or 
-                        obj.text.lower() in tfidf_keywords):
-                        triples.append((subj.text, verb.lemma_, obj.text))
+        subject = ""
+        obj = ""
+        verb = ""
+
+        noun_chunks = list(sent.noun_chunks)
+        root = [token for token in sent if token.dep_ == "ROOT"]
+        if root:
+            verb = root[0].lemma_
+
+        for chunk in noun_chunks:
+            if chunk.root.dep_ in ("nsubj", "nsubjpass"):
+                subject = chunk.text
+            elif chunk.root.dep_ in ("dobj", "pobj", "attr"):
+                obj = chunk.text
+
+        if subject and verb and obj:
+            triples.append((subject.strip(), verb.strip(), obj.strip()))
+
     return triples
 
 # === Visualization Function ===
@@ -62,20 +71,24 @@ def show_graph(triples):
     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)})
+    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 with TF-IDF Filtering")
+st.title("🧠 Knowledge Graph Generator (Enhanced with TF-IDF & Chunking)")
 
 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 (Top 10 filtered by TF-IDF):")
-        for t in triples:
+        all_triples = extract_triples(text_input)
+
+        # Display only the first 10
+        st.write("### Extracted Triples (showing top 10)")
+        for t in all_triples[:10]:
             st.write("🔗", t)
-        show_graph(triples)
+
+        # Visualize all triples
+        show_graph(all_triples)
     else:
         st.warning("Please enter some text.")