SandhyaMadhunagula/GNN-Knowledge-Graph

1from flask import Flask, render_template, request, redirect, url_for, session2import networkx as nx3from pyvis.network import Network4import os, re, pickle5from dotenv import load_dotenv6from PyPDF2 import PdfReader7from docx import Document8from transformers import AutoTokenizer, AutoModelForSeq2SeqLM9import torch10import csv11from flask import Response12import io13 14app = Flask(__name__)15app.secret_key = "secret_key_for_session"16 17model_name = "Babelscape/rebel-large" 18device = "cuda" if torch.cuda.is_available() else "cpu"19 20 21 22load_dotenv() # This loads the variables from .env23HF_TOKEN = os.getenv("HF_TOKEN")24rebel_tokenizer = AutoTokenizer.from_pretrained(model_name, token=HF_TOKEN)25 26#rebel_tokenizer = AutoTokenizer.from_pretrained(model_name, token=HF_TOKEN)27rebel_model = AutoModelForSeq2SeqLM.from_pretrained(model_name, token=HF_TOKEN, low_cpu_mem_usage=True).to(device)28 29 30DB_FILE = "graph_database.pkl"31 32def save_db(graph):33    with open(DB_FILE, "wb") as f:34        pickle.dump(graph, f)35 36def load_db():37    if os.path.exists(DB_FILE):38        try:39            with open(DB_FILE, "rb") as f:40                return pickle.load(f)41        except: return nx.DiGraph()42    return nx.DiGraph()43 44G = load_db()45 46def extract_triples(text):47    inputs = rebel_tokenizer(text, return_tensors="pt", truncation=True, max_length=256).to(device)48    gen_kwargs = {"max_length": 128, "length_penalty": 0, "num_beams": 1, "num_return_sequences": 1}49    generated_tokens = rebel_model.generate(**inputs, **gen_kwargs)50    decoded = rebel_tokenizer.batch_decode(generated_tokens, skip_special_tokens=False)[0]51 52    triples = []53    current_subject, current_relation, current_object = "", "", ""54    current_state = ""55    56    # ADD THESE TWO LINES TO FIX THE "FIRST WORD" PROBLEM57    clean_decoded = decoded.replace("<s>", "").replace("</s>", "")58    clean_decoded = clean_decoded.replace("<triplet>", " <triplet> ").replace("<subj>", " <subj> ").replace("<obj>", " <obj> ")59 60    # CHANGE THIS LOOP TO USE clean_decoded61    for token in clean_decoded.split():62        if token == "<triplet>":63            current_state = "s"64            if current_subject and current_relation and current_object:65                triples.append((current_subject.strip(), current_relation.strip(), current_object.strip()))66            current_subject, current_relation, current_object = "", "", ""67        elif token == "<subj>": current_state = "o"68        elif token == "<obj>": current_state = "r"69        else:70            if current_state == "s": current_subject += " " + token71            elif current_state == "o": current_object += " " + token72            elif current_state == "r": current_relation += " " + token73    74    if current_subject and current_relation and current_object:75        triples.append((current_subject.strip(), current_relation.strip(), current_object.strip()))76    return triples77 78def visualize_graph():79    # Use absolute paths for the cloud environment80    base_path = os.path.dirname(os.path.abspath(__file__))81    static_path = os.path.join(base_path, 'static')82    83    if not os.path.exists(static_path):84        os.makedirs(static_path)85 86    net = Network(height="600px", width="100%", directed=True, bgcolor="#ffffff", font_color="black", cdn_resources='remote')87    net.force_atlas_2based(gravity=-50, central_gravity=0.01, spring_length=150, damping=0.4)88    89    for node in G.nodes():90        net.add_node(node, label=node, color="#00d2ff", size=25, shadow={'enabled': True, 'color': 'rgba(0,210,255,0.6)', 'size': 10})91    for source, target, data in G.edges(data=True):92        net.add_edge(source, target, label=data.get("label", ""), color="#a29bfe")93    94    # Save using the absolute path95    save_path = os.path.join(static_path, "graph.html")96    net.save_graph(save_path)97 98@app.route("/", methods=["GET", "POST"])99def index():100    global G101    answer = None102    user_query = ""103    text = session.get('user_text', "") 104 105    if request.method == "POST":106        # 1. HANDLE FILE UPLOAD OR TEXT BOX107        if "file" in request.files and request.files["file"].filename != "":108            file = request.files["file"]109            ext = file.filename.split('.')[-1].lower()110            if ext == "pdf":111                reader = PdfReader(file)112                text = " ".join([page.extract_text() for page in reader.pages])113            elif ext == "docx":114                text = " ".join([p.text for p in Document(file).paragraphs])115            elif ext == "txt":116                text = file.read().decode("utf-8")117        elif "text" in request.form and request.form["text"].strip():118            text = request.form["text"]119       120        # 2. PROCESS DATA (Only if we have new text)121        if text and "query" not in request.form:122            session['user_text'] = text123            sentences = [s.strip() for s in re.split(r'[\n.!?]', text) if len(s.strip()) > 10]124            print(f"--- 🚀 AI is extracting from {len(sentences)} sentences ---")125            for i, sent in enumerate(sentences):126                print(f"📄 Processing {i+1}/{len(sentences)}...")127                for s, r, o in extract_triples(sent):128                    G.add_edge(s.title().strip(), o.title().strip(), label=r.strip())129            save_db(G)130            visualize_graph()131 132        # 3. HANDLE SEARCH QUERY133        if "query" in request.form:134            user_query = request.form["query"].strip()135            keywords = [w.lower() for w in user_query.split() if len(w) > 3]136            results = []137            for node in G.nodes():138                if any(k in node.lower() for k in keywords):139                    for n in G.successors(node):140                        results.append(f"<b>{node}</b> {G[node][n]['label']} <b>{n}</b>")141                    for p in G.predecessors(node):142                        results.append(f"<b>{p}</b> {G[p][node]['label']} <b>{node}</b>")143            answer = " • " + "<br> • ".join(list(set(results))[:8]) if results else f"Nothing found for '{user_query}'."144 145    db_triples = [{"s": s, "r": d['label'], "o": t} for s, t, d in G.edges(data=True)]146    return render_template("index.html", answer=answer, graph=os.path.exists("static/graph.html"), user_query=user_query, user_text=text, db_triples=db_triples)147 148@app.route("/export_csv")149def export_csv():150    # 1. Create a string buffer to hold CSV data151    output = io.StringIO()152    writer = csv.writer(output)153    154    # 2. Write the Header155    writer.writerow(['Subject', 'Relationship', 'Object'])156    157    # 3. Write the Data from the Graph G158    for s, t, d in G.edges(data=True):159        writer.writerow([s, d.get('label', ''), t])160    161    # 4. Prepare the response for download162    output.seek(0)163    return Response(164        output,165        mimetype="text/csv",166        headers={"Content-disposition": "attachment; filename=knowledge_graph.csv"}167    )168@app.route("/clear")169def clear_db():170    global G171    G = nx.DiGraph()172    session.clear()173    if os.path.exists(DB_FILE): os.remove(DB_FILE)174    if os.path.exists("static/graph.html"): os.remove("static/graph.html")175    return redirect(url_for('index'))176 177#if __name__ == "__main__":178 #   app.run(debug=True)179if __name__ == "__main__":180    # 0.0.0.0 makes it accessible to the internet181    app.run(host="0.0.0.0", port=7860)