maybe now it works?

app.py

@@ -1,12 +1,5 @@
-from logging import disable
-from pkg_resources import EggMetadata
 import streamlit as st
 import streamlit.components.v1 as components
-import networkx as nx
-import matplotlib.pyplot as plt
-from pyvis.network import Network
-from streamlit.state.session_state import SessionState
-from streamlit.type_util import Key
 import rebel
 import wikipedia
 from utils import clip_text
@@ -16,22 +9,23 @@ import os
 MAX_TOPICS = 3
 
 wiki_state_variables = {
-    'has_run_wiki':False,
+    'has_run_wiki': False,
     'wiki_suggestions': [],
-    'wiki_text' : [],
-    'nodes':[],
-    "topics":[],
-    "html_wiki":""
+    'wiki_text': [],
+    'nodes': [],
+    "topics": [],
+    "html_wiki": ""
 }
 
 free_text_state_variables = {
-    'has_run_free':False,
-    "html_free":""
+    'has_run_free': False,
+    "html_free": ""
 
 }
 
 BUTTON_COLUMS = 4
 
+
 def wiki_init_state_variables():
     for k in free_text_state_variables.keys():
         if k in st.session_state:
@@ -41,8 +35,10 @@ def wiki_init_state_variables():
         if k not in st.session_state:
             st.session_state[k] = v
 
+
 def wiki_generate_graph():
-    st.session_state["GRAPH_FILENAME"] = str(dt.now().timestamp()*1000) + ".html"
+    st.session_state["GRAPH_FILENAME"] = str(
+        dt.now().timestamp()*1000) + ".html"
 
     if 'wiki_text' not in st.session_state:
         return
@@ -52,20 +48,23 @@ def wiki_generate_graph():
     with st.spinner(text="Generating graph..."):
         texts = st.session_state['wiki_text']
         st.session_state['nodes'] = []
-        nodes = rebel.generate_knowledge_graph(texts, st.session_state["GRAPH_FILENAME"])
-        HtmlFile = open(st.session_state["GRAPH_FILENAME"], 'r', encoding='utf-8')
+        nodes = rebel.generate_knowledge_graph(
+            texts, st.session_state["GRAPH_FILENAME"])
+        HtmlFile = open(
+            st.session_state["GRAPH_FILENAME"], 'r', encoding='utf-8')
         source_code = HtmlFile.read()
         st.session_state["html_wiki"] = source_code
         os.remove(st.session_state["GRAPH_FILENAME"])
         for n in nodes:
             n = n.lower()
             if n not in st.session_state['topics']:
-                possible_topics = wikipedia.search(n, results = 2)
+                possible_topics = wikipedia.search(n, results=2)
                 st.session_state['nodes'].extend(possible_topics)
         st.session_state['nodes'] = list(set(st.session_state['nodes']))
         st.session_state['has_run_wiki'] = True
     st.success('Done!')
 
+
 def wiki_show_suggestion():
     st.session_state['wiki_suggestions'] = []
     with st.spinner(text="fetching wiki topics..."):
@@ -74,7 +73,9 @@ def wiki_show_suggestion():
             if (text is not None) and (text != ""):
                 subjects = text.split(",")[:MAX_TOPICS]
                 for subj in subjects:
-                    st.session_state['wiki_suggestions'] += wikipedia.search(subj, results = 3)
+                    st.session_state['wiki_suggestions'] += wikipedia.search(
+                        subj, results=3)
+
 
 def wiki_show_text(page_title):
     with st.spinner(text="fetching wiki page..."):
@@ -92,6 +93,7 @@ def wiki_show_text(page_title):
         except wikipedia.WikipediaException:
             st.session_state['wiki_suggestions'].remove(page_title)
 
+
 def wiki_add_text(term):
     if len(st.session_state['wiki_text']) > MAX_TOPICS:
         return
@@ -113,16 +115,20 @@ def wiki_add_text(term):
         print(e)
         st.session_state['nodes'].remove(term)
 
+
 def wiki_reset_session():
     for k in wiki_state_variables:
         del st.session_state[k]
 
+
 def free_reset_session():
     for k in free_text_state_variables:
         del st.session_state[k]
 
+
 def free_text_generate():
-    st.session_state["GRAPH_FILENAME"] = str(dt.now().timestamp()*1000) + ".html"
+    st.session_state["GRAPH_FILENAME"] = str(
+        dt.now().timestamp()*1000) + ".html"
     text = st.session_state['free_text'][0:100]
     rebel.generate_knowledge_graph([text], st.session_state["GRAPH_FILENAME"])
     HtmlFile = open(st.session_state["GRAPH_FILENAME"], 'r', encoding='utf-8')
@@ -131,9 +137,13 @@ def free_text_generate():
     os.remove(st.session_state["GRAPH_FILENAME"])
     st.session_state['has_run_free'] = True
 
+
 def free_text_layout():
-    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.")
-    st.button("Generate", on_click=free_text_generate, key="free_text_generate")
+    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.")
+    st.button("Generate", on_click=free_text_generate,
+              key="free_text_generate")
+
 
 def free_test_init_state_variables():
     for k in wiki_state_variables.keys():
@@ -144,47 +154,51 @@ def free_test_init_state_variables():
         if k not in st.session_state:
             st.session_state[k] = v
 
+
 st.title('RE:Belle')
 st.markdown(
-"""
+    """
 ### Building Beautiful Knowledge Graphs With REBEL
 """)
 st.selectbox(
-     'input method',
-     ('wikipedia', 'free text'),  key="input_method")
+    'input method',
+    ('wikipedia', 'free text'),  key="input_method")
 
 
 def show_wiki_hub_page():
     st.sidebar.button("Reset", on_click=wiki_reset_session, key="reset_key")
 
     st.sidebar.markdown(
-"""
+        """
 ## How To Create a Graph:
 - Enter wikipedia search terms, separated by comma's
 - Choose one or more of the suggested topics (max 3)
 - Click generate!
 """
-)
+    )
     cols = st.columns([8, 1])
     with cols[0]:
-        st.text_input("wikipedia search term", on_change=wiki_show_suggestion, key="text", value="graphs, are, awesome")
+        st.text_input("wikipedia search term", on_change=wiki_show_suggestion,
+                      key="text", value="graphs, are, awesome")
     with cols[1]:
         st.text('')
         st.text('')
-        st.button("Search", on_click=wiki_show_suggestion, key="show_suggestion_key")
+        st.button("Search", on_click=wiki_show_suggestion,
+                  key="show_suggestion_key")
 
     if len(st.session_state['wiki_suggestions']) != 0:
         num_buttons = len(st.session_state['wiki_suggestions'])
         num_cols = num_buttons if 0 < num_buttons < BUTTON_COLUMS else BUTTON_COLUMS
-        columns = st.columns([1] * num_cols )
+        columns = st.columns([1] * num_cols)
         for q in range(1 + num_buttons//num_cols):
             for i, (c, s) in enumerate(zip(columns, st.session_state['wiki_suggestions'][q*num_cols: (q+1)*num_cols])):
                 with c:
-                    st.button(s, on_click=wiki_show_text, args=(s,), key=str(i)+s+"wiki_suggestion")
+                    st.button(s, on_click=wiki_show_text, args=(
+                        s,), key=str(i)+s+"wiki_suggestion")
 
     if len(st.session_state['wiki_text']) != 0:
         for i, t in enumerate(st.session_state['wiki_text']):
-            new_expander = st.expander(label=t[:30] + "...", expanded=(i==0))
+            new_expander = st.expander(label=t[:30] + "...", expanded=(i == 0))
             with new_expander:
                 st.markdown(t)
 
@@ -209,23 +223,27 @@ def show_wiki_hub_page():
         for q in range(1 + num_buttons//num_cols):
             for i, (c, s) in enumerate(zip(columns, st.session_state["nodes"][q*num_cols: (q+1)*num_cols])):
                 with c:
-                    st.button(s, on_click=wiki_add_text, args=(s,), key=str(i)+s)
+                    st.button(s, on_click=wiki_add_text,
+                              args=(s,), key=str(i)+s)
+
 
 def show_free_text_hub_page():
-    st.sidebar.button("Reset", on_click=free_reset_session, key="free_reset_key")
+    st.sidebar.button("Reset", on_click=free_reset_session,
+                      key="free_reset_key")
     st.sidebar.markdown(
-"""
+        """
 ## How To Create a Graph:
 - Enter a text you'd like to see as a graph.
 - Click generate!
 """
-)
+    )
 
     free_text_layout()
 
     if st.session_state['has_run_free']:
         components.html(st.session_state["html_free"], width=720, height=600)
 
+
 if st.session_state['input_method'] == "wikipedia":
     wiki_init_state_variables()
     show_wiki_hub_page()
@@ -234,9 +252,8 @@ else:
     show_free_text_hub_page()
 
 
-
 st.sidebar.markdown(
-"""
+    """
 ## What This Is And Why We Built it
 
 This 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).
@@ -248,9 +265,9 @@ There is one problem though: building knowledge graphs from scratch is a time-co
 )
 
 st.sidebar.markdown(
-"""
+    """
 *Credits for the REBEL model go out to Pere-Lluís Huguet Cabot and Roberto Navigli.
 The code can be found [here](https://github.com/Babelscape/rebel),
 and the original paper [here](https://github.com/Babelscape/rebel/blob/main/docs/EMNLP_2021_REBEL__Camera_Ready_.pdf)*
 """
-)
+)

rebel.py

@@ -3,7 +3,6 @@ from transformers import pipeline
 from pyvis.network import Network
 from functools import lru_cache
 import spacy
-from spacy import displacy
 
 import streamlit as st
 
@@ -31,11 +30,11 @@ DEFAULT_LABEL_COLORS = {
 
 @st.experimental_singleton(max_entries=1)
 def get_pipeline():
-    triplet_extractor = pipeline('text2text-generation', model='Babelscape/rebel-large', tokenizer='Babelscape/rebel-large')
+    triplet_extractor = pipeline(
+        'text2text-generation', model='Babelscape/rebel-large', tokenizer='Babelscape/rebel-large')
     return triplet_extractor
 
 
-
 @st.experimental_singleton(max_entries=1)
 def load_spacy():
     nlp = spacy.load("en_core_web_sm")
@@ -47,13 +46,13 @@ def generate_knowledge_graph(texts: List[str], filename: str):
 
     doc = nlp("\n".join(texts).lower())
     NERs = [ent.text for ent in doc.ents]
-    NER_types =  [ent.label_ for ent in doc.ents]
+    NER_types = [ent.label_ for ent in doc.ents]
 
     triplets = []
     for triplet in texts:
         triplets.extend(generate_partial_graph(triplet))
-    heads = [ t["head"].lower() for t in triplets]
-    tails = [ t["tail"].lower() for t in triplets]
+    heads = [t["head"].lower() for t in triplets]
+    tails = [t["tail"].lower() for t in triplets]
 
     nodes = list(set(heads + tails))
     net = Network(directed=True, width="700px", height="700px")
@@ -73,7 +72,7 @@ def generate_knowledge_graph(texts: List[str], filename: str):
             net.add_node(n, shape="circle")
 
     unique_triplets = set()
-    stringify_trip = lambda x : x["tail"] + x["head"] + x["type"].lower()
+    def stringify_trip(x): return x["tail"] + x["head"] + x["type"].lower()
     for triplet in triplets:
         if stringify_trip(triplet) not in unique_triplets:
             net.add_edge(triplet["head"].lower(), triplet["tail"].lower(),
@@ -95,7 +94,8 @@ def generate_knowledge_graph(texts: List[str], filename: str):
 @lru_cache(maxsize=16)
 def generate_partial_graph(text: str):
     triplet_extractor = get_pipeline()
-    a = triplet_extractor(text, return_tensors=True, return_text=False)[0]["generated_token_ids"]["output_ids"]
+    a = triplet_extractor(text, return_tensors=True, return_text=False)[
+        0]["generated_token_ids"]["output_ids"]
     extracted_text = triplet_extractor.tokenizer.batch_decode(a)
     extracted_triplets = extract_triplets(extracted_text[0])
     return extracted_triplets
@@ -113,13 +113,15 @@ def extract_triplets(text):
         if token == "<triplet>":
             current = 't'
             if relation != '':
-                triplets.append({'head': subject.strip(), 'type': relation.strip(),'tail': object_.strip()})
+                triplets.append(
+                    {'head': subject.strip(), 'type': relation.strip(), 'tail': object_.strip()})
                 relation = ''
             subject = ''
         elif token == "<subj>":
             current = 's'
             if relation != '':
-                triplets.append({'head': subject.strip(), 'type': relation.strip(),'tail': object_.strip()})
+                triplets.append(
+                    {'head': subject.strip(), 'type': relation.strip(), 'tail': object_.strip()})
             object_ = ''
         elif token == "<obj>":
             current = 'o'
@@ -132,7 +134,7 @@ def extract_triplets(text):
             elif current == 'o':
                 relation += ' ' + token
     if subject != '' and relation != '' and object_ != '':
-        triplets.append({'head': subject.strip(), 'type': relation.strip(),'tail': object_.strip()})
+        triplets.append(
+            {'head': subject.strip(), 'type': relation.strip(), 'tail': object_.strip()})
 
     return triplets
-