Update app.py

app.py

@@ -5,6 +5,7 @@ import subprocess
 from neo4j import GraphDatabase
 import matplotlib.pyplot as plt
 import networkx as nx
+from sklearn.feature_extraction.text import TfidfVectorizer
 
 # === Ensure spaCy model is installed ===
 def install_spacy_model():
@@ -14,10 +15,10 @@ 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
 nlp = spacy.load("en_core_web_sm")
 
-# Neo4j credentials
+# === Neo4j credentials ===
 uri = "neo4j+s://ff701b1c.databases.neo4j.io"
 username = "neo4j"
 password = "BfZM7YRKpFz1b_V7acAmOtaSQHPU9xK03rJlfPep88g"  
@@ -25,9 +26,18 @@ password = "BfZM7YRKpFz1b_V7acAmOtaSQHPU9xK03rJlfPep88g"
 # Connect to Neo4j
 driver = GraphDatabase.driver(uri, auth=(username, password))
 
-# Triple extraction function
+# === 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 with TF-IDF filtering and limit to 10 ===
 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_]
@@ -36,12 +46,15 @@ def extract_triples(text):
         for subj in subjects:
             for verb in verbs:
                 for obj in objects:
-                    triples.append((subj.text, verb.lemma_, obj.text))
-                    if len(triples) == 10:
-                        return triples 
+                    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))
+                        if len(triples) == 10:
+                            return triples
     return triples
 
-# Visualization function
+# === Visualization Function ===
 def show_graph(triples):
     G = nx.DiGraph()
     for s, p, o in triples:
@@ -55,16 +68,16 @@ def show_graph(triples):
     st.pyplot(plt)
 
 # === Streamlit UI ===
-st.title("🧠 Knowledge Graph Generator")
+st.title("🧠 Knowledge Graph Generator with TF-IDF Filtering")
 
 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")
+        st.write("### Extracted Triples (Top 10 filtered by TF-IDF):")
         for t in triples:
             st.write("🔗", t)
         show_graph(triples)
     else:
-        st.warning("Please enter some text.")
+        st.warning("Please enter some text.")