awacke1/REBEL-Knowledge-Graph-Generator

1from logging import disable2from pkg_resources import EggMetadata3import streamlit as st4import streamlit.components.v1 as components5import networkx as nx6import matplotlib.pyplot as plt7from pyvis.network import Network8from streamlit.state.session_state import SessionState9from streamlit.type_util import Key10import rebel11import wikipedia12from utils import clip_text13from datetime import datetime as dt14import os15 16MAX_TOPICS = 317 18wiki_state_variables = {19    'has_run_wiki':False,20    'wiki_suggestions': [],21    'wiki_text' : [],22    'nodes':[],23    "topics":[],24    "html_wiki":""25}26 27free_text_state_variables = {28    'has_run_free':False,29    "html_free":""30 31}32 33BUTTON_COLUMS = 434 35def wiki_init_state_variables():36    for k in free_text_state_variables.keys():37        if k in st.session_state:38            del st.session_state[k]39 40    for k, v in wiki_state_variables.items():41        if k not in st.session_state:42            st.session_state[k] = v43 44def wiki_generate_graph():45    st.session_state["GRAPH_FILENAME"] = str(dt.now().timestamp()*1000) + ".html"46 47    if 'wiki_text' not in st.session_state:48        return49    if len(st.session_state['wiki_text']) == 0:50        st.error("please enter a topic and select a wiki page first")51        return52    with st.spinner(text="Generating graph..."):53        texts = st.session_state['wiki_text']54        st.session_state['nodes'] = []55        nodes = rebel.generate_knowledge_graph(texts, st.session_state["GRAPH_FILENAME"])56        HtmlFile = open(st.session_state["GRAPH_FILENAME"], 'r', encoding='utf-8')57        source_code = HtmlFile.read()58        st.session_state["html_wiki"] = source_code59        os.remove(st.session_state["GRAPH_FILENAME"])60        for n in nodes:61            n = n.lower()62            if n not in st.session_state['topics']:63                possible_topics = wikipedia.search(n, results = 2)64                st.session_state['nodes'].extend(possible_topics)65        st.session_state['nodes'] = list(set(st.session_state['nodes']))66        st.session_state['has_run_wiki'] = True67    st.success('Done!')68 69def wiki_show_suggestion():70    st.session_state['wiki_suggestions'] = []71    with st.spinner(text="fetching wiki topics..."):72        if st.session_state['input_method'] == "wikipedia":73            text = st.session_state.text74            if (text is not None) and (text != ""):75                subjects = text.split(",")[:MAX_TOPICS]76                for subj in subjects:77                    st.session_state['wiki_suggestions'] += wikipedia.search(subj, results = 3)78 79def wiki_show_text(page_title):80    with st.spinner(text="fetching wiki page..."):81        try:82            page = wikipedia.page(title=page_title, auto_suggest=False)83            st.session_state['wiki_text'].append(clip_text(page.summary))84            st.session_state['topics'].append(page_title.lower())85            st.session_state['wiki_suggestions'].remove(page_title)86 87        except wikipedia.DisambiguationError as e:88            with st.spinner(text="Woops, ambigious term, recalculating options..."):89                st.session_state['wiki_suggestions'].remove(page_title)90                temp = st.session_state['wiki_suggestions'] + e.options[:3]91                st.session_state['wiki_suggestions'] = list(set(temp))92        except wikipedia.WikipediaException:93            st.session_state['wiki_suggestions'].remove(page_title)94 95def wiki_add_text(term):96    if len(st.session_state['wiki_text']) > MAX_TOPICS:97        return98    try:99        page = wikipedia.page(title=term, auto_suggest=False)100        extra_text = clip_text(page.summary)101 102        st.session_state['wiki_text'].append(extra_text)103        st.session_state['topics'].append(term.lower())104        st.session_state['nodes'].remove(term)105 106    except wikipedia.DisambiguationError as e:107        print(e)108        with st.spinner(text="Woops, ambigious term, recalculating options..."):109            st.session_state['nodes'].remove(term)110            temp = st.session_state['nodes'] + e.options[:3]111            st.session_state['nodes'] = list(set(temp))112    except wikipedia.WikipediaException as e:113        print(e)114        st.session_state['nodes'].remove(term)115 116def wiki_reset_session():117    for k in wiki_state_variables:118        del st.session_state[k]119 120def free_reset_session():121    for k in free_text_state_variables:122        del st.session_state[k]123 124def free_text_generate():125    st.session_state["GRAPH_FILENAME"] = str(dt.now().timestamp()*1000) + ".html"126    text = st.session_state['free_text'][0:100]127    rebel.generate_knowledge_graph([text], st.session_state["GRAPH_FILENAME"])128    HtmlFile = open(st.session_state["GRAPH_FILENAME"], 'r', encoding='utf-8')129    source_code = HtmlFile.read()130    st.session_state["html_free"] = source_code131    os.remove(st.session_state["GRAPH_FILENAME"])132    st.session_state['has_run_free'] = True133 134def free_text_layout():135    st.text_area("Free text", key="free_text", height=5, value="Tardigrades, known colloquially as water bears or moss piglets, are a phylum of eight-legged segmented micro-animals.")136    st.button("Generate", on_click=free_text_generate, key="free_text_generate")137 138def free_test_init_state_variables():139    for k in wiki_state_variables.keys():140        if k in st.session_state:141            del st.session_state[k]142 143    for k, v in free_text_state_variables.items():144        if k not in st.session_state:145            st.session_state[k] = v146 147st.title('RE:Belle')148st.markdown(149"""150### Building Beautiful Knowledge Graphs With REBEL151""")152st.selectbox(153     'input method',154     ('wikipedia', 'free text'),  key="input_method")155 156 157def show_wiki_hub_page():158    # st.sidebar.button("Reset", on_click=wiki_reset_session, key="reset_key")159 160    cols = st.columns([8, 1])161    with cols[0]:162        st.text_input("wikipedia search term", on_change=wiki_show_suggestion, key="text", value="graphs, are, awesome")163    with cols[1]:164        st.text('')165        st.text('')166        st.button("Search", on_click=wiki_show_suggestion, key="show_suggestion_key")167 168    if len(st.session_state['wiki_suggestions']) != 0:169        num_buttons = len(st.session_state['wiki_suggestions'])170        num_cols = num_buttons if 0 < num_buttons < BUTTON_COLUMS else BUTTON_COLUMS171        columns = st.columns([1] * num_cols )172        for q in range(1 + num_buttons//num_cols):173            for i, (c, s) in enumerate(zip(columns, st.session_state['wiki_suggestions'][q*num_cols: (q+1)*num_cols])):174                with c:175                    st.button(s, on_click=wiki_show_text, args=(s,), key=str(i)+s+"wiki_suggestion")176 177    if len(st.session_state['wiki_text']) != 0:178        for i, t in enumerate(st.session_state['wiki_text']):179            new_expander = st.expander(label=t[:30] + "...", expanded=(i==0))180            with new_expander:181                st.markdown(t)182 183    if len(st.session_state['wiki_text']) > 0:184        st.button("Generate", on_click=wiki_generate_graph, key="gen_graph")185 186    if st.session_state['has_run_wiki']:187 188        components.html(st.session_state["html_wiki"], width=720, height=600)189        num_buttons = len(st.session_state["nodes"])190        num_cols = num_buttons if 0 < num_buttons < BUTTON_COLUMS else BUTTON_COLUMS191        columns = st.columns([1] * num_cols + [1])192 193        for q in range(1 + num_buttons//num_cols):194            for i, (c, s) in enumerate(zip(columns, st.session_state["nodes"][q*num_cols: (q+1)*num_cols])):195                with c:196                    st.button(s, on_click=wiki_add_text, args=(s,), key=str(i)+s)197 198def show_free_text_hub_page():199    free_text_layout()200    if st.session_state['has_run_free']:201        components.html(st.session_state["html_free"], width=720, height=600)202 203if st.session_state['input_method'] == "wikipedia":204    wiki_init_state_variables()205    show_wiki_hub_page()206else:207    free_test_init_state_variables()208    show_free_text_hub_page()209 210 211 212# st.sidebar.markdown(213"""214## What This Is And Why We Built it215 216This space shows how a transformer network can be used to convert *human* text into a computer-queryable format: a **knowledge graph**. Knowledge graphs are graphs where each node (or *vertex* if you're fancy) represent a concept/person/thing and each edge the link between those concepts. If you'd like to know more, you can read [this blogpost](https://www.ml6.eu/knowhow/knowledge-graphs-an-introduction-and-business-applications).217 218Knowledge graphs aren't just cool to look at, they are an extremely versatile way of storing data, and are used in machine learning to perform tasks like fraud detection. You can read more about the applications of knowledge graphs in ML in [this blogpost](https://blog.ml6.eu/how-are-knowledge-graphs-and-machine-learning-related-ff6f5c1760b5).219 220There is one problem though: building knowledge graphs from scratch is a time-consuming and tedious task, so it would be a lot easier if we could leverage machine learning to **create** them from existing texts. This demo shows how a model named **REBEL** has been trained to do just that: it reads summaries from Wikipedia (or any other text you input), and generates a graph containing the information it distills from the text.221"""222)223 224# st.sidebar.markdown(225"""226*Credits for the REBEL model go out to Pere-Lluís Huguet Cabot and Roberto Navigli.227The code can be found [here](https://github.com/Babelscape/rebel),228and the original paper [here](https://github.com/Babelscape/rebel/blob/main/docs/EMNLP_2021_REBEL__Camera_Ready_.pdf)*229"""230)