jaibadachiya/knowledge_graph

1# app.py2 3import streamlit as st4import spacy5import subprocess6from neo4j import GraphDatabase7import matplotlib.pyplot as plt8import networkx as nx9from sklearn.feature_extraction.text import TfidfVectorizer10 11# === Ensure spaCy model is installed ===12def install_spacy_model():13    try:14        spacy.load("en_core_web_sm")15    except OSError:16        subprocess.run(["python", "-m", "spacy", "download", "en_core_web_sm"])17        spacy.load("en_core_web_sm")18 19install_spacy_model()20 21# Load spaCy model22nlp = spacy.load("en_core_web_sm")23 24# === Neo4j credentials ===25NEO4J_URI = "neo4j+s://ff701b1c.databases.neo4j.io"26NEO4J_USERNAME = "neo4j"27NEO4J_PASSWORD = "BfZM7YRKpFz1b_V7acAmOtaSQHPU9xK03rJlfPep88g"  28 29def get_neo4j_driver():30    try:31        driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USERNAME, NEO4J_PASSWORD))32        return driver33    except Exception as e:34        st.error(f"Failed to connect to Neo4j: {e}")35        return None36 37# === TF-IDF Filtering (Optional for noise reduction) ===38def compute_tfidf_keywords(text: str, top_n=100):39    vectorizer = TfidfVectorizer(stop_words='english')40    X = vectorizer.fit_transform([text])41    scores = zip(vectorizer.get_feature_names_out(), X.toarray()[0])42    sorted_scores = sorted(scores, key=lambda x: x[1], reverse=True)43    return {word for word, _ in sorted_scores[:top_n]}44 45# === Triple Extraction ===46def extract_triples(text, use_tfidf=False):47    doc = nlp(text)48    tfidf_keywords = compute_tfidf_keywords(text) if use_tfidf else None49    triples = []50 51    for sent in doc.sents:52        subject = ""53        obj = ""54        verb = ""55 56        noun_chunks = list(sent.noun_chunks)57        root = [token for token in sent if token.dep_ == "ROOT"]58        if root:59            verb = root[0].lemma_60 61        for chunk in noun_chunks:62            if chunk.root.dep_ in ("nsubj", "nsubjpass") and not subject:63                subject = chunk.text64            elif chunk.root.dep_ in ("dobj", "pobj", "attr") and not obj:65                obj = chunk.text66 67        if subject and verb and obj:68            if tfidf_keywords:69                if subject.lower() in tfidf_keywords or obj.lower() in tfidf_keywords:70                    triples.append((subject.strip(), verb.strip(), obj.strip()))71            else:72                triples.append((subject.strip(), verb.strip(), obj.strip()))73 74    return triples75 76# === Visualization Function ===77def show_graph(triples):78    if not triples:79        st.warning("No triples found to visualize.")80        return81 82    G = nx.DiGraph()83    for s, p, o in triples:84        G.add_node(s)85        G.add_node(o)86        G.add_edge(s, o, label=p)87 88    pos = nx.spring_layout(G, seed=42)  # fixed layout89    plt.figure(figsize=(10, 8))90    nx.draw(G, pos, with_labels=True, node_color='skyblue', node_size=2000, font_size=10, edge_color='gray')91    nx.draw_networkx_edge_labels(G, pos, edge_labels={(u, v): d['label'] for u, v, d in G.edges(data=True)})92    st.pyplot(plt.gcf())93    plt.clf()94 95# === Streamlit UI ===96st.title("🧠 Knowledge Graph Generator")97 98text_input = st.text_area("Paste your text here:", height=200)99use_tfidf = st.checkbox("Use TF-IDF filtering (Optional: Recommended for large texts)")100 101if st.button("Generate Graph"):102    if text_input:103        all_triples = extract_triples(text_input, use_tfidf=use_tfidf)104 105        if all_triples:106            st.subheader("🔗 Extracted Triples:")107            for triple in all_triples:108                st.markdown(f"- **({triple[0]} → {triple[1]} → {triple[2]})**")109 110            show_graph(all_triples)111        else:112            st.warning("No valid triples could be extracted. Try different text.")113    else:114        st.warning("Please enter some text.")115