Spaces:
Paused
Paused
| import os | |
| # --- Environment Setup --- | |
| os.environ["HF_HOME"] = "/app/cache" | |
| os.environ["TRANSFORMERS_CACHE"] = "/app/cache" | |
| os.environ["SENTENCE_TRANSFORMERS_HOME"] = "/app/cache" | |
| os.environ["CHROMA_DB_DIR"] = "/app/chroma_db" | |
| os.environ["XDG_CONFIG_HOME"] = "/app/.streamlit" | |
| import streamlit as st | |
| st.set_page_config(layout="wide") # ✅ Must come first before any other st.* calls | |
| # --- Libraries --- | |
| import networkx as nx | |
| import spacy | |
| import subprocess | |
| from docx import Document | |
| import pdfplumber | |
| from pyvis.network import Network | |
| from sentence_transformers import SentenceTransformer | |
| import chromadb | |
| from chromadb.config import Settings | |
| from transformers import pipeline | |
| import tempfile | |
| import streamlit.components.v1 as components | |
| # --- Setup --- | |
| def ensure_spacy_model(): | |
| nlp = spacy.load("en_core_web_sm") | |
| nlp.max_length = 2_000_000 | |
| return nlp | |
| nlp = ensure_spacy_model() | |
| embedder = SentenceTransformer("all-MiniLM-L6-v2") | |
| os.makedirs("/app/chroma_db", exist_ok=True) | |
| client = chromadb.PersistentClient(path="/app/chroma_db") | |
| collection = client.get_or_create_collection(name="graphrag") | |
| llm = pipeline("text2text-generation", model="google/flan-t5-large") | |
| # --- Utils --- | |
| def extract_text(upload): | |
| if upload is None: | |
| return "" | |
| ext = os.path.splitext(upload.name)[1].lower() | |
| with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp: | |
| tmp.write(upload.read()) | |
| tmp_path = tmp.name | |
| if ext == ".txt": | |
| return open(tmp_path, "r", encoding="utf-8").read() | |
| elif ext == ".pdf": | |
| with pdfplumber.open(tmp_path) as pdf: | |
| return "\n".join(page.extract_text() for page in pdf.pages if page.extract_text()) | |
| elif ext == ".docx": | |
| doc = Document(tmp_path) | |
| return "\n".join(p.text for p in doc.paragraphs) | |
| return "" | |
| def build_graph(text): | |
| doc = nlp(text) | |
| G = nx.Graph() | |
| entity_data = [(ent.text.strip(), ent.label_) for ent in doc.ents if ent.text.strip()] | |
| unique_entities = list(set(entity_data)) | |
| entity_texts = [ent[0] for ent in unique_entities] | |
| embeddings = embedder.encode(entity_texts, convert_to_numpy=True).tolist() | |
| for (text, label), emb in zip(unique_entities, embeddings): | |
| G.add_node(text, label=label, embedding=emb) | |
| unique_ids = [f"{ent}_{i}" for i, ent in enumerate(entity_texts)] | |
| collection.add(documents=entity_texts, embeddings=embeddings, ids=unique_ids) | |
| for sent in doc.sents: | |
| ents = [ent.text.strip() for ent in sent.ents if ent.text.strip() in G.nodes] | |
| for i in range(len(ents)): | |
| for j in range(i + 1, len(ents)): | |
| G.add_edge(ents[i], ents[j]) | |
| return G | |
| def retrieve_relevant_entities(query, top_k=5): | |
| query_emb = embedder.encode([query], convert_to_numpy=True).tolist() | |
| results = collection.query(query_embeddings=query_emb, n_results=top_k) | |
| return [match for match in results['ids'][0] if match in collection.peek()['ids']] | |
| def expand_graph_context(graph, entities, depth=1): | |
| subgraph_nodes = set() | |
| for entity in entities: | |
| if entity in graph: | |
| neighbors = nx.single_source_shortest_path_length(graph, entity, cutoff=depth).keys() | |
| subgraph_nodes.update(neighbors) | |
| return graph.subgraph(subgraph_nodes) | |
| def format_graph_context(subgraph): | |
| lines = [] | |
| for node in subgraph.nodes(data=True): | |
| lines.append(f"Entity: {node[0]}, Type: {node[1].get('label', 'N/A')}") | |
| for edge in subgraph.edges(): | |
| lines.append(f"Connection: {edge[0]} ↔ {edge[1]}") | |
| return "\n".join(lines) | |
| def generate_answer(question, context_text): | |
| prompt = f"""You are a compliance assistant. Use the following knowledge graph context to answer the question. | |
| {context_text} | |
| Question: {question} | |
| Answer:""" | |
| return llm(prompt, max_length=512, do_sample=False)[0]['generated_text'] | |
| def plot_graph(G, name="graph"): | |
| import shutil | |
| output_dir = "/app/tmp" | |
| lib_dir = os.path.join(output_dir, "lib") | |
| os.makedirs(output_dir, exist_ok=True) | |
| original_cwd = os.getcwd() | |
| os.chdir(output_dir) | |
| try: | |
| net = Network( | |
| height="600px", | |
| width="100%", | |
| bgcolor="#222222", | |
| font_color="white", | |
| cdn_resources="local" | |
| ) | |
| for node, data in G.nodes(data=True): | |
| net.add_node(node, label=f"{node}\n[{data.get('label', '')}]") | |
| for src, dst in G.edges(): | |
| net.add_edge(src, dst) | |
| net.repulsion(node_distance=150, spring_length=200) | |
| net.show_buttons(filter_=['physics']) | |
| html_path = f"{name}.html" | |
| net.save_graph(html_path) | |
| return os.path.join(output_dir, html_path) | |
| finally: | |
| os.chdir(original_cwd) | |
| # --- Streamlit App --- | |
| st.title("\U0001F4D8 GraphRAG Agent") | |
| uploaded_file = st.file_uploader("Upload document", type=["pdf", "txt", "docx"]) | |
| if uploaded_file: | |
| with st.spinner("Processing document..."): | |
| text = extract_text(uploaded_file) | |
| G = build_graph(text) | |
| html_file = plot_graph(G) | |
| st.subheader("Entity Graph") | |
| with open(html_file, 'r', encoding='utf-8') as f: | |
| components.html(f.read(), height=650, scrolling=True) | |
| st.download_button("Download Graph HTML", data=open(html_file, 'rb'), file_name="graph.html") | |
| question = st.text_input("Ask a question (Graph-RAG):") | |
| if question: | |
| with st.spinner("Querying..."): | |
| relevant_entities = retrieve_relevant_entities(question) | |
| subgraph = expand_graph_context(G, relevant_entities) | |
| context_text = format_graph_context(subgraph) | |
| answer = generate_answer(question, context_text) | |
| st.markdown(f"**Answer:** {answer}") | |