Spaces:
Build error
Build error
Duplicate from ml6team/Knowledge-graphs
Browse filesCo-authored-by: Julien Chaumond <[email protected]>
- .gitattributes +27 -0
- .gitignore +2 -0
- .vscode/launch.json +16 -0
- .vscode/settings.json +7 -0
- README.md +39 -0
- __pycache__/app.cpython-38.pyc +0 -0
- __pycache__/rebel.cpython-38.pyc +0 -0
- __pycache__/utils.cpython-38.pyc +0 -0
- app.py +256 -0
- rebel.py +122 -0
- requirements.txt +124 -0
- utils.py +6 -0
.gitattributes
ADDED
@@ -0,0 +1,27 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
*.7z filter=lfs diff=lfs merge=lfs -text
|
2 |
+
*.arrow filter=lfs diff=lfs merge=lfs -text
|
3 |
+
*.bin filter=lfs diff=lfs merge=lfs -text
|
4 |
+
*.bin.* filter=lfs diff=lfs merge=lfs -text
|
5 |
+
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
6 |
+
*.ftz filter=lfs diff=lfs merge=lfs -text
|
7 |
+
*.gz filter=lfs diff=lfs merge=lfs -text
|
8 |
+
*.h5 filter=lfs diff=lfs merge=lfs -text
|
9 |
+
*.joblib filter=lfs diff=lfs merge=lfs -text
|
10 |
+
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
11 |
+
*.model filter=lfs diff=lfs merge=lfs -text
|
12 |
+
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
13 |
+
*.onnx filter=lfs diff=lfs merge=lfs -text
|
14 |
+
*.ot filter=lfs diff=lfs merge=lfs -text
|
15 |
+
*.parquet filter=lfs diff=lfs merge=lfs -text
|
16 |
+
*.pb filter=lfs diff=lfs merge=lfs -text
|
17 |
+
*.pt filter=lfs diff=lfs merge=lfs -text
|
18 |
+
*.pth filter=lfs diff=lfs merge=lfs -text
|
19 |
+
*.rar filter=lfs diff=lfs merge=lfs -text
|
20 |
+
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
21 |
+
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
22 |
+
*.tflite filter=lfs diff=lfs merge=lfs -text
|
23 |
+
*.tgz filter=lfs diff=lfs merge=lfs -text
|
24 |
+
*.xz filter=lfs diff=lfs merge=lfs -text
|
25 |
+
*.zip filter=lfs diff=lfs merge=lfs -text
|
26 |
+
*.zstandard filter=lfs diff=lfs merge=lfs -text
|
27 |
+
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
.gitignore
ADDED
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
1 |
+
venv
|
2 |
+
test.html
|
.vscode/launch.json
ADDED
@@ -0,0 +1,16 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
// Use IntelliSense to learn about possible attributes.
|
3 |
+
// Hover to view descriptions of existing attributes.
|
4 |
+
// For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
|
5 |
+
"version": "0.2.0",
|
6 |
+
"configurations": [
|
7 |
+
{
|
8 |
+
"name": "Python: Current File",
|
9 |
+
"type": "python",
|
10 |
+
"request": "launch",
|
11 |
+
"program": "${file}",
|
12 |
+
"console": "integratedTerminal",
|
13 |
+
"justMyCode": false
|
14 |
+
}
|
15 |
+
]
|
16 |
+
}
|
.vscode/settings.json
ADDED
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
{
|
2 |
+
"workbench.colorCustomizations": {
|
3 |
+
"activityBar.background": "#09323E",
|
4 |
+
"titleBar.activeBackground": "#0C4656",
|
5 |
+
"titleBar.activeForeground": "#F6FCFE"
|
6 |
+
}
|
7 |
+
}
|
README.md
ADDED
@@ -0,0 +1,39 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
---
|
2 |
+
title: RE:Belle
|
3 |
+
emoji: 🌐
|
4 |
+
colorFrom: red
|
5 |
+
colorTo: blue
|
6 |
+
sdk: streamlit
|
7 |
+
sdk_version: 1.2.0
|
8 |
+
app_file: app.py
|
9 |
+
pinned: false
|
10 |
+
duplicated_from: ml6team/Knowledge-graphs
|
11 |
+
---
|
12 |
+
|
13 |
+
# Configuration
|
14 |
+
|
15 |
+
`title`: _string_
|
16 |
+
Display title for the Space
|
17 |
+
|
18 |
+
`emoji`: _string_
|
19 |
+
Space emoji (emoji-only character allowed)
|
20 |
+
|
21 |
+
`colorFrom`: _string_
|
22 |
+
Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
|
23 |
+
|
24 |
+
`colorTo`: _string_
|
25 |
+
Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
|
26 |
+
|
27 |
+
`sdk`: _string_
|
28 |
+
Can be either `gradio`, `streamlit`, or `static`
|
29 |
+
|
30 |
+
`sdk_version` : _string_
|
31 |
+
Only applicable for `streamlit` SDK.
|
32 |
+
See [doc](https://hf.co/docs/hub/spaces) for more info on supported versions.
|
33 |
+
|
34 |
+
`app_file`: _string_
|
35 |
+
Path to your main application file (which contains either `gradio` or `streamlit` Python code, or `static` html code).
|
36 |
+
Path is relative to the root of the repository.
|
37 |
+
|
38 |
+
`pinned`: _boolean_
|
39 |
+
Whether the Space stays on top of your list.
|
__pycache__/app.cpython-38.pyc
ADDED
Binary file (5.05 kB). View file
|
|
__pycache__/rebel.cpython-38.pyc
ADDED
Binary file (3.7 kB). View file
|
|
__pycache__/utils.cpython-38.pyc
ADDED
Binary file (314 Bytes). View file
|
|
app.py
ADDED
@@ -0,0 +1,256 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from logging import disable
|
2 |
+
from pkg_resources import EggMetadata
|
3 |
+
import streamlit as st
|
4 |
+
import streamlit.components.v1 as components
|
5 |
+
import networkx as nx
|
6 |
+
import matplotlib.pyplot as plt
|
7 |
+
from pyvis.network import Network
|
8 |
+
from streamlit.state.session_state import SessionState
|
9 |
+
from streamlit.type_util import Key
|
10 |
+
import rebel
|
11 |
+
import wikipedia
|
12 |
+
from utils import clip_text
|
13 |
+
from datetime import datetime as dt
|
14 |
+
import os
|
15 |
+
|
16 |
+
MAX_TOPICS = 3
|
17 |
+
|
18 |
+
wiki_state_variables = {
|
19 |
+
'has_run_wiki':False,
|
20 |
+
'wiki_suggestions': [],
|
21 |
+
'wiki_text' : [],
|
22 |
+
'nodes':[],
|
23 |
+
"topics":[],
|
24 |
+
"html_wiki":""
|
25 |
+
}
|
26 |
+
|
27 |
+
free_text_state_variables = {
|
28 |
+
'has_run_free':False,
|
29 |
+
"html_free":""
|
30 |
+
|
31 |
+
}
|
32 |
+
|
33 |
+
BUTTON_COLUMS = 4
|
34 |
+
|
35 |
+
def 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] = v
|
43 |
+
|
44 |
+
def 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 |
+
return
|
49 |
+
if len(st.session_state['wiki_text']) == 0:
|
50 |
+
st.error("please enter a topic and select a wiki page first")
|
51 |
+
return
|
52 |
+
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_code
|
59 |
+
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'] = True
|
67 |
+
st.success('Done!')
|
68 |
+
|
69 |
+
def 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.text
|
74 |
+
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 |
+
|
79 |
+
def 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 |
+
|
95 |
+
def wiki_add_text(term):
|
96 |
+
if len(st.session_state['wiki_text']) > MAX_TOPICS:
|
97 |
+
return
|
98 |
+
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 |
+
|
116 |
+
def wiki_reset_session():
|
117 |
+
for k in wiki_state_variables:
|
118 |
+
del st.session_state[k]
|
119 |
+
|
120 |
+
def free_reset_session():
|
121 |
+
for k in free_text_state_variables:
|
122 |
+
del st.session_state[k]
|
123 |
+
|
124 |
+
def 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_code
|
131 |
+
os.remove(st.session_state["GRAPH_FILENAME"])
|
132 |
+
st.session_state['has_run_free'] = True
|
133 |
+
|
134 |
+
def 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 |
+
|
138 |
+
def 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] = v
|
146 |
+
|
147 |
+
st.title('RE:Belle')
|
148 |
+
st.markdown(
|
149 |
+
"""
|
150 |
+
### Building Beautiful Knowledge Graphs With REBEL
|
151 |
+
""")
|
152 |
+
st.selectbox(
|
153 |
+
'input method',
|
154 |
+
('wikipedia', 'free text'), key="input_method")
|
155 |
+
|
156 |
+
|
157 |
+
def show_wiki_hub_page():
|
158 |
+
st.sidebar.button("Reset", on_click=wiki_reset_session, key="reset_key")
|
159 |
+
|
160 |
+
st.sidebar.markdown(
|
161 |
+
"""
|
162 |
+
## How To Create a Graph:
|
163 |
+
- Enter wikipedia search terms, separated by comma's
|
164 |
+
- Choose one or more of the suggested topics (max 3)
|
165 |
+
- Click generate!
|
166 |
+
"""
|
167 |
+
)
|
168 |
+
cols = st.columns([8, 1])
|
169 |
+
with cols[0]:
|
170 |
+
st.text_input("wikipedia search term", on_change=wiki_show_suggestion, key="text", value="graphs, are, awesome")
|
171 |
+
with cols[1]:
|
172 |
+
st.text('')
|
173 |
+
st.text('')
|
174 |
+
st.button("Search", on_click=wiki_show_suggestion, key="show_suggestion_key")
|
175 |
+
|
176 |
+
if len(st.session_state['wiki_suggestions']) != 0:
|
177 |
+
num_buttons = len(st.session_state['wiki_suggestions'])
|
178 |
+
num_cols = num_buttons if 0 < num_buttons < BUTTON_COLUMS else BUTTON_COLUMS
|
179 |
+
columns = st.columns([1] * num_cols )
|
180 |
+
for q in range(1 + num_buttons//num_cols):
|
181 |
+
for i, (c, s) in enumerate(zip(columns, st.session_state['wiki_suggestions'][q*num_cols: (q+1)*num_cols])):
|
182 |
+
with c:
|
183 |
+
st.button(s, on_click=wiki_show_text, args=(s,), key=str(i)+s+"wiki_suggestion")
|
184 |
+
|
185 |
+
if len(st.session_state['wiki_text']) != 0:
|
186 |
+
for i, t in enumerate(st.session_state['wiki_text']):
|
187 |
+
new_expander = st.expander(label=t[:30] + "...", expanded=(i==0))
|
188 |
+
with new_expander:
|
189 |
+
st.markdown(t)
|
190 |
+
|
191 |
+
if len(st.session_state['wiki_text']) > 0:
|
192 |
+
st.button("Generate", on_click=wiki_generate_graph, key="gen_graph")
|
193 |
+
st.sidebar.markdown(
|
194 |
+
"""
|
195 |
+
## How to expand the graph
|
196 |
+
- Click a button below the graph to expand that node
|
197 |
+
(Only nodes that have wiki pages will be expanded)
|
198 |
+
- Hit the Generate button again to expand your graph!
|
199 |
+
"""
|
200 |
+
)
|
201 |
+
|
202 |
+
if st.session_state['has_run_wiki']:
|
203 |
+
|
204 |
+
components.html(st.session_state["html_wiki"], width=720, height=600)
|
205 |
+
num_buttons = len(st.session_state["nodes"])
|
206 |
+
num_cols = num_buttons if 0 < num_buttons < BUTTON_COLUMS else BUTTON_COLUMS
|
207 |
+
columns = st.columns([1] * num_cols + [1])
|
208 |
+
|
209 |
+
for q in range(1 + num_buttons//num_cols):
|
210 |
+
for i, (c, s) in enumerate(zip(columns, st.session_state["nodes"][q*num_cols: (q+1)*num_cols])):
|
211 |
+
with c:
|
212 |
+
st.button(s, on_click=wiki_add_text, args=(s,), key=str(i)+s)
|
213 |
+
|
214 |
+
def show_free_text_hub_page():
|
215 |
+
st.sidebar.button("Reset", on_click=free_reset_session, key="free_reset_key")
|
216 |
+
st.sidebar.markdown(
|
217 |
+
"""
|
218 |
+
## How To Create a Graph:
|
219 |
+
- Enter a text you'd like to see as a graph.
|
220 |
+
- Click generate!
|
221 |
+
"""
|
222 |
+
)
|
223 |
+
|
224 |
+
free_text_layout()
|
225 |
+
|
226 |
+
if st.session_state['has_run_free']:
|
227 |
+
components.html(st.session_state["html_free"], width=720, height=600)
|
228 |
+
|
229 |
+
if st.session_state['input_method'] == "wikipedia":
|
230 |
+
wiki_init_state_variables()
|
231 |
+
show_wiki_hub_page()
|
232 |
+
else:
|
233 |
+
free_test_init_state_variables()
|
234 |
+
show_free_text_hub_page()
|
235 |
+
|
236 |
+
|
237 |
+
|
238 |
+
st.sidebar.markdown(
|
239 |
+
"""
|
240 |
+
## What This Is And Why We Built it
|
241 |
+
|
242 |
+
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).
|
243 |
+
|
244 |
+
Knowledge 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).
|
245 |
+
|
246 |
+
There 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.
|
247 |
+
"""
|
248 |
+
)
|
249 |
+
|
250 |
+
st.sidebar.markdown(
|
251 |
+
"""
|
252 |
+
*Credits for the REBEL model go out to Pere-Lluís Huguet Cabot and Roberto Navigli.
|
253 |
+
The code can be found [here](https://github.com/Babelscape/rebel),
|
254 |
+
and the original paper [here](https://github.com/Babelscape/rebel/blob/main/docs/EMNLP_2021_REBEL__Camera_Ready_.pdf)*
|
255 |
+
"""
|
256 |
+
)
|
rebel.py
ADDED
@@ -0,0 +1,122 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from typing import List
|
2 |
+
from transformers import pipeline
|
3 |
+
from pyvis.network import Network
|
4 |
+
from functools import lru_cache
|
5 |
+
import spacy
|
6 |
+
from spacy import displacy
|
7 |
+
|
8 |
+
|
9 |
+
DEFAULT_LABEL_COLORS = {
|
10 |
+
"ORG": "#7aecec",
|
11 |
+
"PRODUCT": "#bfeeb7",
|
12 |
+
"GPE": "#feca74",
|
13 |
+
"LOC": "#ff9561",
|
14 |
+
"PERSON": "#aa9cfc",
|
15 |
+
"NORP": "#c887fb",
|
16 |
+
"FACILITY": "#9cc9cc",
|
17 |
+
"EVENT": "#ffeb80",
|
18 |
+
"LAW": "#ff8197",
|
19 |
+
"LANGUAGE": "#ff8197",
|
20 |
+
"WORK_OF_ART": "#f0d0ff",
|
21 |
+
"DATE": "#bfe1d9",
|
22 |
+
"TIME": "#bfe1d9",
|
23 |
+
"MONEY": "#e4e7d2",
|
24 |
+
"QUANTITY": "#e4e7d2",
|
25 |
+
"ORDINAL": "#e4e7d2",
|
26 |
+
"CARDINAL": "#e4e7d2",
|
27 |
+
"PERCENT": "#e4e7d2",
|
28 |
+
}
|
29 |
+
|
30 |
+
def generate_knowledge_graph(texts: List[str], filename: str):
|
31 |
+
nlp = spacy.load("en_core_web_sm")
|
32 |
+
doc = nlp("\n".join(texts).lower())
|
33 |
+
NERs = [ent.text for ent in doc.ents]
|
34 |
+
NER_types = [ent.label_ for ent in doc.ents]
|
35 |
+
|
36 |
+
triplets = []
|
37 |
+
for triplet in texts:
|
38 |
+
triplets.extend(generate_partial_graph(triplet))
|
39 |
+
heads = [ t["head"].lower() for t in triplets]
|
40 |
+
tails = [ t["tail"].lower() for t in triplets]
|
41 |
+
|
42 |
+
nodes = list(set(heads + tails))
|
43 |
+
net = Network(directed=True, width="700px", height="700px")
|
44 |
+
|
45 |
+
for n in nodes:
|
46 |
+
if n in NERs:
|
47 |
+
NER_type = NER_types[NERs.index(n)]
|
48 |
+
if NER_type in NER_types:
|
49 |
+
if NER_type in DEFAULT_LABEL_COLORS.keys():
|
50 |
+
color = DEFAULT_LABEL_COLORS[NER_type]
|
51 |
+
else:
|
52 |
+
color = "#666666"
|
53 |
+
net.add_node(n, title=NER_type, shape="circle", color=color)
|
54 |
+
else:
|
55 |
+
net.add_node(n, shape="circle")
|
56 |
+
else:
|
57 |
+
net.add_node(n, shape="circle")
|
58 |
+
|
59 |
+
unique_triplets = set()
|
60 |
+
stringify_trip = lambda x : x["tail"] + x["head"] + x["type"].lower()
|
61 |
+
for triplet in triplets:
|
62 |
+
if stringify_trip(triplet) not in unique_triplets:
|
63 |
+
net.add_edge(triplet["head"].lower(), triplet["tail"].lower(),
|
64 |
+
title=triplet["type"], label=triplet["type"])
|
65 |
+
unique_triplets.add(stringify_trip(triplet))
|
66 |
+
|
67 |
+
net.repulsion(
|
68 |
+
node_distance=200,
|
69 |
+
central_gravity=0.2,
|
70 |
+
spring_length=200,
|
71 |
+
spring_strength=0.05,
|
72 |
+
damping=0.09
|
73 |
+
)
|
74 |
+
net.set_edge_smooth('dynamic')
|
75 |
+
net.show(filename)
|
76 |
+
return nodes
|
77 |
+
|
78 |
+
|
79 |
+
@lru_cache(maxsize=16)
|
80 |
+
def generate_partial_graph(text: str):
|
81 |
+
triplet_extractor = pipeline('text2text-generation', model='Babelscape/rebel-large', tokenizer='Babelscape/rebel-large')
|
82 |
+
a = triplet_extractor(text, return_tensors=True, return_text=False)[0]["generated_token_ids"]["output_ids"]
|
83 |
+
extracted_text = triplet_extractor.tokenizer.batch_decode(a)
|
84 |
+
extracted_triplets = extract_triplets(extracted_text[0])
|
85 |
+
return extracted_triplets
|
86 |
+
|
87 |
+
|
88 |
+
def extract_triplets(text):
|
89 |
+
"""
|
90 |
+
Function to parse the generated text and extract the triplets
|
91 |
+
"""
|
92 |
+
triplets = []
|
93 |
+
relation, subject, relation, object_ = '', '', '', ''
|
94 |
+
text = text.strip()
|
95 |
+
current = 'x'
|
96 |
+
for token in text.replace("<s>", "").replace("<pad>", "").replace("</s>", "").split():
|
97 |
+
if token == "<triplet>":
|
98 |
+
current = 't'
|
99 |
+
if relation != '':
|
100 |
+
triplets.append({'head': subject.strip(), 'type': relation.strip(),'tail': object_.strip()})
|
101 |
+
relation = ''
|
102 |
+
subject = ''
|
103 |
+
elif token == "<subj>":
|
104 |
+
current = 's'
|
105 |
+
if relation != '':
|
106 |
+
triplets.append({'head': subject.strip(), 'type': relation.strip(),'tail': object_.strip()})
|
107 |
+
object_ = ''
|
108 |
+
elif token == "<obj>":
|
109 |
+
current = 'o'
|
110 |
+
relation = ''
|
111 |
+
else:
|
112 |
+
if current == 't':
|
113 |
+
subject += ' ' + token
|
114 |
+
elif current == 's':
|
115 |
+
object_ += ' ' + token
|
116 |
+
elif current == 'o':
|
117 |
+
relation += ' ' + token
|
118 |
+
if subject != '' and relation != '' and object_ != '':
|
119 |
+
triplets.append({'head': subject.strip(), 'type': relation.strip(),'tail': object_.strip()})
|
120 |
+
|
121 |
+
return triplets
|
122 |
+
|
requirements.txt
ADDED
@@ -0,0 +1,124 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
altair==4.2.0
|
2 |
+
argon2-cffi==21.3.0
|
3 |
+
argon2-cffi-bindings==21.2.0
|
4 |
+
astor==0.8.1
|
5 |
+
attrs==21.4.0
|
6 |
+
backcall==0.2.0
|
7 |
+
backports.zoneinfo==0.2.1
|
8 |
+
base58==2.1.1
|
9 |
+
beautifulsoup4==4.10.0
|
10 |
+
bleach==4.1.0
|
11 |
+
blinker==1.4
|
12 |
+
blis==0.7.5
|
13 |
+
cachetools==5.0.0
|
14 |
+
catalogue==2.0.6
|
15 |
+
certifi==2021.10.8
|
16 |
+
cffi==1.15.0
|
17 |
+
charset-normalizer==2.0.10
|
18 |
+
click==7.1.2
|
19 |
+
cycler==0.11.0
|
20 |
+
cymem==2.0.6
|
21 |
+
debugpy==1.5.1
|
22 |
+
decorator==5.1.1
|
23 |
+
defusedxml==0.7.1
|
24 |
+
entrypoints==0.3
|
25 |
+
filelock==3.4.2
|
26 |
+
fonttools==4.28.5
|
27 |
+
gitdb==4.0.9
|
28 |
+
GitPython==3.1.26
|
29 |
+
got==0.0.1
|
30 |
+
huggingface-hub==0.2.1
|
31 |
+
idna==3.3
|
32 |
+
importlib-resources==5.4.0
|
33 |
+
ipykernel==6.6.1
|
34 |
+
ipython==7.31.0
|
35 |
+
ipython-genutils==0.2.0
|
36 |
+
ipywidgets==7.6.5
|
37 |
+
jedi==0.18.1
|
38 |
+
Jinja2==3.0.3
|
39 |
+
joblib==1.1.0
|
40 |
+
jsonpickle==2.0.0
|
41 |
+
jsonschema==4.3.3
|
42 |
+
jupyter-client==7.1.0
|
43 |
+
jupyter-core==4.9.1
|
44 |
+
jupyterlab-pygments==0.1.2
|
45 |
+
jupyterlab-widgets==1.0.2
|
46 |
+
kiwisolver==1.3.2
|
47 |
+
langcodes==3.3.0
|
48 |
+
MarkupSafe==2.0.1
|
49 |
+
matplotlib==3.5.1
|
50 |
+
matplotlib-inline==0.1.3
|
51 |
+
mistune==0.8.4
|
52 |
+
murmurhash==1.0.6
|
53 |
+
nbclient==0.5.9
|
54 |
+
nbconvert==6.4.0
|
55 |
+
nbformat==5.1.3
|
56 |
+
nest-asyncio==1.5.4
|
57 |
+
networkx==2.6.3
|
58 |
+
notebook==6.4.6
|
59 |
+
numpy==1.22.0
|
60 |
+
packaging==21.3
|
61 |
+
pandas==1.3.5
|
62 |
+
pandocfilters==1.5.0
|
63 |
+
parso==0.8.3
|
64 |
+
pathy==0.6.1
|
65 |
+
pexpect==4.8.0
|
66 |
+
pickleshare==0.7.5
|
67 |
+
Pillow==9.0.0
|
68 |
+
preshed==3.0.6
|
69 |
+
prometheus-client==0.12.0
|
70 |
+
prompt-toolkit==3.0.24
|
71 |
+
protobuf==3.19.3
|
72 |
+
ptyprocess==0.7.0
|
73 |
+
pyarrow==6.0.1
|
74 |
+
pycparser==2.21
|
75 |
+
pydantic==1.8.2
|
76 |
+
pydeck==0.7.1
|
77 |
+
Pygments==2.11.2
|
78 |
+
Pympler==1.0.1
|
79 |
+
pyparsing==3.0.6
|
80 |
+
pyrsistent==0.18.0
|
81 |
+
python-dateutil==2.8.2
|
82 |
+
pytz==2021.3
|
83 |
+
pytz-deprecation-shim==0.1.0.post0
|
84 |
+
pyvis==0.1.9
|
85 |
+
PyYAML==6.0
|
86 |
+
pyzmq==22.3.0
|
87 |
+
regex==2021.11.10
|
88 |
+
requests==2.27.1
|
89 |
+
sacremoses==0.0.47
|
90 |
+
Send2Trash==1.8.0
|
91 |
+
six==1.16.0
|
92 |
+
smart-open==5.2.1
|
93 |
+
smmap==5.0.0
|
94 |
+
soupsieve==2.3.1
|
95 |
+
spacy==3.2.1
|
96 |
+
spacy-legacy==3.0.8
|
97 |
+
spacy-loggers==1.0.1
|
98 |
+
srsly==2.4.2
|
99 |
+
streamlit==1.3.1
|
100 |
+
terminado==0.12.1
|
101 |
+
testpath==0.5.0
|
102 |
+
thinc==8.0.13
|
103 |
+
tokenizers==0.10.3
|
104 |
+
toml==0.10.2
|
105 |
+
toolz==0.11.2
|
106 |
+
torch==1.10.1
|
107 |
+
tornado==6.1
|
108 |
+
tqdm==4.62.3
|
109 |
+
traitlets==5.1.1
|
110 |
+
transformers==4.15.0
|
111 |
+
typer==0.4.0
|
112 |
+
typing-extensions==4.0.1
|
113 |
+
tzdata==2021.5
|
114 |
+
tzlocal==4.1
|
115 |
+
urllib3==1.26.8
|
116 |
+
validators==0.18.2
|
117 |
+
wasabi==0.9.0
|
118 |
+
watchdog==2.1.6
|
119 |
+
wcwidth==0.2.5
|
120 |
+
webencodings==0.5.1
|
121 |
+
widgetsnbextension==3.5.2
|
122 |
+
wikipedia==1.4.0
|
123 |
+
zipp==3.7.0
|
124 |
+
https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.2.0/en_core_web_sm-3.2.0.tar.gz#egg=en_core_web_sm
|
utils.py
ADDED
@@ -0,0 +1,6 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
|
2 |
+
def clip_text(t, lenght = 4):
|
3 |
+
t_sub = t.replace("...", "dotdotdot")
|
4 |
+
t_clipped = ".".join(t_sub.split(".")[:lenght]) + "."
|
5 |
+
t_reverted = t_clipped.replace("dotdotdot", "...")
|
6 |
+
return t_reverted
|