Duplicate from fabiochiu/text-to-kb

Co-authored-by: Fabio Chiusano <fabiochiu@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,27 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1 @@
+myvenv

README.md

@@ -0,0 +1,14 @@
+---
+title: Text To Kb
+emoji: 📉
+colorFrom: pink
+colorTo: blue
+sdk: streamlit
+sdk_version: 1.9.0
+app_file: app.py
+pinned: false
+license: mit
+duplicated_from: fabiochiu/text-to-kb
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces#reference

app.py

@@ -0,0 +1,197 @@
+import streamlit as st
+import streamlit.components.v1 as components
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+import utils
+from kb import KB
+
+texts = {
+    "Napoleon": "Napoleon Bonaparte (born Napoleone di Buonaparte; 15 August 1769 – 5 May 1821), and later known by his regnal name Napoleon I, was a French military and political leader who rose to prominence during the French Revolution and led several successful campaigns during the Revolutionary Wars. He was the de facto leader of the French Republic as First Consul from 1799 to 1804. As Napoleon I, he was Emperor of the French from 1804 until 1814 and again in 1815. Napoleon's political and cultural legacy has endured, and he has been one of the most celebrated and controversial leaders in world history.",
+    "Kobe Bryant": "Kobe Bean Bryant (August 23, 1978 – January 26, 2020) was an American professional basketball player. A shooting guard, he spent his entire 20-year career with the Los Angeles Lakers in the National Basketball Association (NBA). Widely regarded as one of the greatest basketball players of all time, Bryant won five NBA championships, was an 18-time All-Star, a 15-time member of the All-NBA Team, a 12-time member of the All-Defensive Team, the 2008 NBA Most Valuable Player (MVP), and a two-time NBA Finals MVP. Bryant also led the NBA in scoring twice, and ranks fourth in league all-time regular season and postseason scoring. He was posthumously voted into the Naismith Memorial Basketball Hall of Fame in 2020 and named to the NBA 75th Anniversary Team in 2021.",
+    "Google": "Originally known as BackRub. Google is a search engine that started development in 1996 by Sergey Brin and Larry Page as a research project at Stanford University to find files on the Internet. Larry and Sergey later decided the name of their search engine needed to change and chose Google, which is inspired from the term googol. The company is headquartered in Mountain View, California."
+}
+
+urls = {
+    "Crypto": "https://www.investopedia.com/terms/c/cryptocurrency.asp",
+    "Jhonny Depp": "https://www.britannica.com/biography/Johnny-Depp",
+    "Rome": "https://www.timeout.com/rome/things-to-do/best-things-to-do-in-rome"
+}
+
+st.header("Extracting a Knowledge Base from text")
+
+# sidebar
+with st.sidebar:
+    st.markdown("_Read the accompanying article [Building a Knowledge Base from Texts: a Full Practical Example](https://medium.com/nlplanet/building-a-knowledge-base-from-texts-a-full-practical-example-8dbbffb912fa)_")
+    st.header("What is a Knowledge Base")
+    st.markdown("A [**Knowledge Base (KB)**](https://en.wikipedia.org/wiki/Knowledge_base) is information stored in structured data, ready to be used for analysis or inference. Usually a KB is stored as a graph (i.e. a [**Knowledge Graph**](https://www.ibm.com/cloud/learn/knowledge-graph)), where nodes are **entities** and edges are **relations** between entities.")
+    st.markdown("_For example, from the text \"Fabio lives in Italy\" we can extract the relation triplet <Fabio, lives in, Italy>, where \"Fabio\" and \"Italy\" are entities._")
+    st.header("How to build a Knowledge Graph")
+    st.markdown("To build a Knowledge Graph from text, we typically need to perform two steps:\n- Extract entities, a.k.a. **Named Entity Recognition (NER)**, i.e. the nodes.\n- Extract relations between the entities, a.k.a. **Relation Classification (RC)**, i.e. the edges.\nRecently, end-to-end approaches have been proposed to tackle both tasks simultaneously. This task is usually referred to as **Relation Extraction (RE)**. In this demo, an end-to-end model called [**REBEL**](https://github.com/Babelscape/rebel/blob/main/docs/EMNLP_2021_REBEL__Camera_Ready_.pdf) is used, trained by [Babelscape](https://babelscape.com/).")
+    st.header("How REBEL works")
+    st.markdown("REBEL is a **text2text** model obtained by fine-tuning [**BART**](https://huggingface.co/docs/transformers/model_doc/bart) for translating a raw input sentence containing entities and implicit relations into a set of triplets that explicitly refer to those relations. You can find [REBEL in the Hugging Face Hub](https://huggingface.co/Babelscape/rebel-large).")
+    st.header("Further steps")
+    st.markdown("Even though they are not visualized, the knowledge graph saves information about the provenience of each relation (e.g. from which articles it has been extracted and other metadata), along with Wikipedia data about each entity.")
+    st.markdown("Other libraries used:\n- [wikipedia](https://pypi.org/project/wikipedia/): For validating extracted entities checking if they have a corresponding Wikipedia page.\n- [newspaper](https://github.com/codelucas/newspaper): For parsing articles from URLs.\n- [pyvis](https://pyvis.readthedocs.io/en/latest/index.html): For graphs visualizations.\n- [GoogleNews](https://github.com/Iceloof/GoogleNews): For reading Google News latest articles about a topic.")
+    st.header("Considerations")
+    st.markdown("If you look closely at the extracted knowledge graphs, some extracted relations are false. Indeed, relation extraction models are still far from perfect and require further steps in the pipeline to build reliable knowledge graphs. Consider this demo as a starting step!")
+
+# Loading the model
+st_model_load = st.text('Loading NER model... It may take a while.')
+
+@st.cache(allow_output_mutation=True)
+def load_model():
+    print("Loading model...")
+    tokenizer = AutoTokenizer.from_pretrained("Babelscape/rebel-large")
+    model = AutoModelForSeq2SeqLM.from_pretrained("Babelscape/rebel-large")
+    print("Model loaded!")
+    return tokenizer, model
+
+tokenizer, model = load_model()
+st.success('Model loaded!')
+st_model_load.text("")
+
+# Choose from where to generate the KB
+options = [
+    "Text",
+    "Article at URL",
+    "Multiple news articles"
+]
+if 'option' not in st.session_state:
+    st.session_state.option = options[0]
+option = st.selectbox('Build a Knowledge Base from:', options, index=options.index(st.session_state.option))
+
+text_option, text = None, None
+url_option, url = None, None
+news_option = None
+
+if option == "Text":
+    text_options = [
+        "Napoleon",
+        "Kobe Bryant",
+        "Google",
+        "Free text"
+    ]
+    if 'text_option' not in st.session_state or st.session_state.text_option is None:
+        st.session_state.text_option = text_options[0]
+    text_option = st.selectbox('Choose text option:', text_options, index=text_options.index(st.session_state.text_option))
+
+    disabled = False
+    if text_option != "Free text":
+        disabled = True
+        text = texts[text_option]
+    else:
+        if 'text' not in st.session_state:
+            st.session_state.text = ""
+        text = st.session_state.text
+    text = st.text_area('Text:', value=text, height=300, disabled=disabled, max_chars=10000)
+elif option == "Article at URL":
+    url_options = [
+        "Crypto",
+        "Jhonny Depp",
+        "Rome",
+        "Free URL"
+    ]
+    if 'url_option' not in st.session_state or st.session_state.url_option is None:
+        st.session_state.url_option = url_options[0]
+    url_option = st.selectbox('Choose URL option:', url_options, index=url_options.index(st.session_state.url_option))
+
+    disabled = False
+    if url_option != "Free URL":
+        disabled = True
+        url = urls[url_option]
+    else:
+        if 'url' not in st.session_state:
+            st.session_state.url = ""
+        url = st.session_state.url
+    url = st.text_input('URL:', value=url, disabled=disabled)
+else:
+    news_options = [
+        "Google",
+        "Apple",
+        "Elon Musk",
+        "Kobe Bryant"
+    ]
+    if 'news_option' not in st.session_state or st.session_state.news_option is None:
+        st.session_state.news_option = news_options[0]
+    news_option = st.selectbox('Use articles about:', news_options, index=news_options.index(st.session_state.news_option))
+
+def generate_kb():
+    st.session_state.option = option
+    st.session_state.text_option = text_option
+    st.session_state.text = text
+    st.session_state.url_option = url_option
+    st.session_state.url = url
+    st.session_state.news_option = news_option
+
+    # load correct kb
+    if option == "Text":
+        if text_option == "Napoleon":
+            kb = utils.load_kb("networks/network_1_napoleon.p")
+        elif text_option == "Kobe Bryant":
+            kb = utils.load_kb("networks/network_1_bryant.p")
+        elif text_option == "Google":
+            kb = utils.load_kb("networks/network_1_google.p")
+        else:
+            kb = utils.from_text_to_kb(text, model, tokenizer, "", verbose=True)
+    elif option == "Article at URL":
+        if url_option == "Crypto":
+            kb = utils.load_kb("networks/network_2_crypto.p")
+        elif url_option == "Jhonny Depp":
+            kb = utils.load_kb("networks/network_2_depp.p")
+        elif url_option == "Rome":
+            kb = utils.load_kb("networks/network_2_rome.p")
+        else:
+            try:
+                kb = utils.from_url_to_kb(url, model, tokenizer)
+            except Exception as e:
+                print("Couldn't extract article from URL")
+                st.session_state.error_url = "Couldn't extract article from URL"
+                return
+    else:
+        if news_option == "Google":
+            kb = utils.load_kb("networks/network_3_google.p")
+        elif news_option == "Apple":
+            kb = utils.load_kb("networks/network_3_apple.p")
+        elif news_option == "Elon Musk":
+            kb = utils.load_kb("networks/network_3_musk.p")
+        elif news_option == "Kobe Bryant":
+            kb = utils.load_kb("networks/network_3_bryant.p")
+
+    # save chart
+    utils.save_network_html(kb, filename="networks/network.html")
+    st.session_state.kb_chart = "networks/network.html"
+    st.session_state.kb_text = kb.get_textual_representation()
+    st.session_state.error_url = None
+
+
+st.session_state.option = option
+st.session_state.text_option = text_option
+st.session_state.text = text
+st.session_state.url_option = url_option
+st.session_state.url = url
+st.session_state.news_option = news_option
+
+button_text = "Show KB"
+if (option == "Text" and text_option == "Free text") or (option == "Article at URL" and url_option == "Free URL"):
+    button_text = "Generate KB"
+
+# generate KB button
+st.button(button_text, on_click=generate_kb)
+
+# kb chart session state
+if 'kb_chart' not in st.session_state:
+    st.session_state.kb_chart = None
+if 'kb_text' not in st.session_state:
+    st.session_state.kb_text = None
+if 'error_url' not in st.session_state:
+    st.session_state.error_url = None
+
+# show graph
+if st.session_state.error_url:
+    st.markdown(st.session_state.error_url)
+elif st.session_state.kb_chart:
+    with st.container():
+        st.subheader("Generated KB")
+        st.markdown("*You can interact with the graph and zoom.*")
+        html_source_code = open(st.session_state.kb_chart, 'r', encoding='utf-8').read()
+        components.html(html_source_code, width=700, height=700)
+        st.markdown(st.session_state.kb_text)

kb.py

@@ -0,0 +1,99 @@
+import wikipedia
+
+class KB():
+    def __init__(self):
+        self.entities = {} # { entity_title: {...} }
+        self.relations = [] # [ head: entity_title, type: ..., tail: entity_title,
+          # meta: { article_url: { spans: [...] } } ]
+        self.sources = {} # { article_url: {...} }
+
+    def merge_with_kb(self, kb2):
+        for r in kb2.relations:
+            article_url = list(r["meta"].keys())[0]
+            source_data = kb2.sources[article_url]
+            self.add_relation(r, source_data["article_title"],
+                              source_data["article_publish_date"])
+
+    def are_relations_equal(self, r1, r2):
+        return all(r1[attr] == r2[attr] for attr in ["head", "type", "tail"])
+
+    def exists_relation(self, r1):
+        return any(self.are_relations_equal(r1, r2) for r2 in self.relations)
+
+    def merge_relations(self, r2):
+        r1 = [r for r in self.relations
+              if self.are_relations_equal(r2, r)][0]
+
+        # if different article
+        article_url = list(r2["meta"].keys())[0]
+        if article_url not in r1["meta"]:
+            r1["meta"][article_url] = r2["meta"][article_url]
+
+        # if existing article
+        else:
+            spans_to_add = [span for span in r2["meta"][article_url]["spans"]
+                            if span not in r1["meta"][article_url]["spans"]]
+            r1["meta"][article_url]["spans"] += spans_to_add
+
+    def get_wikipedia_data(self, candidate_entity):
+        try:
+            page = wikipedia.page(candidate_entity, auto_suggest=False)
+            entity_data = {
+                "title": page.title,
+                "url": page.url,
+                "summary": page.summary
+            }
+            return entity_data
+        except:
+            return None
+
+    def add_entity(self, e):
+        self.entities[e["title"]] = {k:v for k,v in e.items() if k != "title"}
+
+    def add_relation(self, r, article_title, article_publish_date):
+        # check on wikipedia
+        candidate_entities = [r["head"], r["tail"]]
+        entities = [self.get_wikipedia_data(ent) for ent in candidate_entities]
+
+        # if one entity does not exist, stop
+        if any(ent is None for ent in entities):
+            return
+
+        # manage new entities
+        for e in entities:
+            self.add_entity(e)
+
+        # rename relation entities with their wikipedia titles
+        r["head"] = entities[0]["title"]
+        r["tail"] = entities[1]["title"]
+
+        # add source if not in kb
+        article_url = list(r["meta"].keys())[0]
+        if article_url not in self.sources:
+            self.sources[article_url] = {
+                "article_title": article_title,
+                "article_publish_date": article_publish_date
+            }
+
+        # manage new relation
+        if not self.exists_relation(r):
+            self.relations.append(r)
+        else:
+            self.merge_relations(r)
+
+    def get_textual_representation(self):
+        res = ""
+        res += "### Entities\n"
+        for e in self.entities.items():
+            # shorten summary
+            e_temp = (e[0], {k:(v[:100] + "..." if k == "summary" else v) for k,v in e[1].items()})
+            res += f"- {e_temp}\n"
+        res += "\n"
+        res += "### Relations\n"
+        for r in self.relations:
+            res += f"- {r}\n"
+        res += "\n"
+        res += "### Sources\n"
+        for s in self.sources.items():
+            res += f"- {s}\n"
+        return res

networks/network.html

@@ -0,0 +1,116 @@
+<html>
+<head>
+<link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/vis-network@latest/styles/vis-network.css" type="text/css" />
+<script type="text/javascript" src="https://cdn.jsdelivr.net/npm/vis-network@latest/dist/vis-network.min.js"> </script>
+<center>
+<h1></h1>
+</center>
+
+<!-- <link rel="stylesheet" href="../node_modules/vis/dist/vis.min.css" type="text/css" />
+<script type="text/javascript" src="../node_modules/vis/dist/vis.js"> </script>-->
+
+<style type="text/css">
+
+        #mynetwork {
+            width: 700px;
+            height: 700px;
+            background-color: #ffffff;
+            border: 1px solid lightgray;
+            position: relative;
+            float: left;
+        }
+
+        
+
+        
+
+        
+</style>
+
+</head>
+
+<body>
+<div id = "mynetwork"></div>
+
+
+<script type="text/javascript">
+
+    // initialize global variables.
+    var edges;
+    var nodes;
+    var network; 
+    var container;
+    var options, data;
+
+    
+    // This method is responsible for drawing the graph, returns the drawn network
+    function drawGraph() {
+        var container = document.getElementById('mynetwork');
+        
+        
+
+        // parsing and collecting nodes and edges from the python
+        nodes = new vis.DataSet([{"color": "#00FF00", "id": "Napoleon", "label": "Napoleon", "shape": "circle"}, {"color": "#00FF00", "id": "French Revolution", "label": "French Revolution", "shape": "circle"}, {"color": "#00FF00", "id": "France", "label": "France", "shape": "circle"}]);
+        edges = new vis.DataSet([{"arrows": "to", "from": "Napoleon", "label": "participant in", "title": "participant in", "to": "French Revolution"}, {"arrows": "to", "from": "French Revolution", "label": "participant", "title": "participant", "to": "Napoleon"}, {"arrows": "to", "from": "French Revolution", "label": "country", "title": "country", "to": "France"}]);
+
+        // adding nodes and edges to the graph
+        data = {nodes: nodes, edges: edges};
+
+        var options = {
+    "configure": {
+        "enabled": false
+    },
+    "edges": {
+        "color": {
+            "inherit": true
+        },
+        "smooth": {
+            "enabled": true,
+            "type": "dynamic"
+        }
+    },
+    "interaction": {
+        "dragNodes": true,
+        "hideEdgesOnDrag": false,
+        "hideNodesOnDrag": false
+    },
+    "physics": {
+        "enabled": true,
+        "repulsion": {
+            "centralGravity": 0.2,
+            "damping": 0.09,
+            "nodeDistance": 200,
+            "springConstant": 0.05,
+            "springLength": 200
+        },
+        "solver": "repulsion",
+        "stabilization": {
+            "enabled": true,
+            "fit": true,
+            "iterations": 1000,
+            "onlyDynamicEdges": false,
+            "updateInterval": 50
+        }
+    }
+};
+        
+        
+
+        
+
+        network = new vis.Network(container, data, options);
+	 
+        
+
+
+        
+
+        return network;
+
+    }
+
+    drawGraph();
+
+</script>
+</body>
+</html>

requirements.txt

@@ -0,0 +1,6 @@
+torch
+transformers
+pyvis
+GoogleNews
+newspaper3k
+wikipedia

utils.py

@@ -0,0 +1,198 @@
+from pyvis.network import Network
+from GoogleNews import GoogleNews
+from newspaper import Article, ArticleException
+import math
+import torch
+from kb import KB
+import pickle
+
+def extract_relations_from_model_output(text):
+    relations = []
+    relation, subject, relation, object_ = '', '', '', ''
+    text = text.strip()
+    current = 'x'
+    text_replaced = text.replace("<s>", "").replace("<pad>", "").replace("</s>", "")
+    for token in text_replaced.split():
+        if token == "<triplet>":
+            current = 't'
+            if relation != '':
+                relations.append({
+                    'head': subject.strip(),
+                    'type': relation.strip(),
+                    'tail': object_.strip()
+                })
+                relation = ''
+            subject = ''
+        elif token == "<subj>":
+            current = 's'
+            if relation != '':
+                relations.append({
+                    'head': subject.strip(),
+                    'type': relation.strip(),
+                    'tail': object_.strip()
+                })
+            object_ = ''
+        elif token == "<obj>":
+            current = 'o'
+            relation = ''
+        else:
+            if current == 't':
+                subject += ' ' + token
+            elif current == 's':
+                object_ += ' ' + token
+            elif current == 'o':
+                relation += ' ' + token
+    if subject != '' and relation != '' and object_ != '':
+        relations.append({
+            'head': subject.strip(),
+            'type': relation.strip(),
+            'tail': object_.strip()
+        })
+    return relations
+
+def from_text_to_kb(text, model, tokenizer, article_url, span_length=128, article_title=None,
+                    article_publish_date=None, verbose=False):
+    # tokenize whole text
+    inputs = tokenizer([text], return_tensors="pt")
+
+    # compute span boundaries
+    num_tokens = len(inputs["input_ids"][0])
+    if verbose:
+        print(f"Input has {num_tokens} tokens")
+    num_spans = math.ceil(num_tokens / span_length)
+    if verbose:
+        print(f"Input has {num_spans} spans")
+    overlap = math.ceil((num_spans * span_length - num_tokens) / 
+                        max(num_spans - 1, 1))
+    spans_boundaries = []
+    start = 0
+    for i in range(num_spans):
+        spans_boundaries.append([start + span_length * i,
+                                 start + span_length * (i + 1)])
+        start -= overlap
+    if verbose:
+        print(f"Span boundaries are {spans_boundaries}")
+
+    # transform input with spans
+    tensor_ids = [inputs["input_ids"][0][boundary[0]:boundary[1]]
+                  for boundary in spans_boundaries]
+    tensor_masks = [inputs["attention_mask"][0][boundary[0]:boundary[1]]
+                    for boundary in spans_boundaries]
+    inputs = {
+        "input_ids": torch.stack(tensor_ids),
+        "attention_mask": torch.stack(tensor_masks)
+    }
+
+    # generate relations
+    num_return_sequences = 3
+    gen_kwargs = {
+        "max_length": 256,
+        "length_penalty": 0,
+        "num_beams": 3,
+        "num_return_sequences": num_return_sequences
+    }
+    generated_tokens = model.generate(
+        **inputs,
+        **gen_kwargs,
+    )
+
+    # decode relations
+    decoded_preds = tokenizer.batch_decode(generated_tokens,
+                                           skip_special_tokens=False)
+
+    # create kb
+    kb = KB()
+    i = 0
+    for sentence_pred in decoded_preds:
+        current_span_index = i // num_return_sequences
+        relations = extract_relations_from_model_output(sentence_pred)
+        for relation in relations:
+            relation["meta"] = {
+                article_url: {
+                    "spans": [spans_boundaries[current_span_index]]
+                }
+            }
+            kb.add_relation(relation, article_title, article_publish_date)
+        i += 1
+
+    return kb
+
+def get_article(url):
+    article = Article(url)
+    article.download()
+    article.parse()
+    return article
+
+def from_url_to_kb(url, model, tokenizer):
+    article = get_article(url)
+    config = {
+        "article_title": article.title,
+        "article_publish_date": article.publish_date
+    }
+    kb = from_text_to_kb(article.text, model, tokenizer, article.url, **config)
+    return kb
+
+def get_news_links(query, lang="en", region="US", pages=1):
+    googlenews = GoogleNews(lang=lang, region=region)
+    googlenews.search(query)
+    all_urls = []
+    for page in range(pages):
+        googlenews.get_page(page)
+        all_urls += googlenews.get_links()
+    return list(set(all_urls))
+
+def from_urls_to_kb(urls, model, tokenizer, verbose=False):
+    kb = KB()
+    if verbose:
+        print(f"{len(urls)} links to visit")
+    for url in urls:
+        if verbose:
+            print(f"Visiting {url}...")
+        try:
+            kb_url = from_url_to_kb(url, model, tokenizer)
+            kb.merge_with_kb(kb_url)
+        except ArticleException:
+            if verbose:
+                print(f"  Couldn't download article at url {url}")
+    return kb
+
+def save_network_html(kb, filename="network.html"):
+    # create network
+    net = Network(directed=True, width="700px", height="700px")
+
+    # nodes
+    color_entity = "#00FF00"
+    for e in kb.entities:
+        net.add_node(e, shape="circle", color=color_entity)
+
+    # edges
+    for r in kb.relations:
+        net.add_edge(r["head"], r["tail"],
+                    title=r["type"], label=r["type"])
+
+    # save network
+    net.repulsion(
+        node_distance=200,
+        central_gravity=0.2,
+        spring_length=200,
+        spring_strength=0.05,
+        damping=0.09
+    )
+    net.set_edge_smooth('dynamic')
+    net.show(filename)
+
+def save_kb(kb, filename):
+    with open(filename, "wb") as f:
+        pickle.dump(kb, f)
+
+class CustomUnpickler(pickle.Unpickler):
+    def find_class(self, module, name):
+        if name == 'KB':
+            return KB
+        return super().find_class(module, name)
+
+def load_kb(filename):
+    res = None
+    with open(filename, "rb") as f:
+        res = CustomUnpickler(f).load()
+    return res