Update functions.py

Added Knowledge Graph tab

functions.py

@@ -6,7 +6,7 @@ import plotly_express as px
 import nltk
 import plotly.graph_objects as go
 from optimum.onnxruntime import ORTModelForSequenceClassification
-from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification, AutoModelForTokenClassification
+from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification, AutoModelForTokenClassification, AutoModelForSeq2SeqLM
 from sentence_transformers import SentenceTransformer, CrossEncoder, util
 import streamlit as st
 import en_core_web_lg
@@ -31,6 +31,8 @@ margin-bottom: 2.5rem">{}</div> """
 def load_models():
     q_model = ORTModelForSequenceClassification.from_pretrained("nickmuchi/quantized-optimum-finbert-tone")
     ner_model = AutoModelForTokenClassification.from_pretrained("xlm-roberta-large-finetuned-conll03-english")
+    kg_model = AutoModelForSeq2SeqLM.from_pretrained("Babelscape/rebel-large")
+    kg_tokenizer = AutoTokenizer.from_pretrained("Babelscape/rebel-large")
     q_tokenizer = AutoTokenizer.from_pretrained("nickmuchi/quantized-optimum-finbert-tone")
     ner_tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-large-finetuned-conll03-english")
     sent_pipe = pipeline("text-classification",model=q_model, tokenizer=q_tokenizer)
@@ -38,7 +40,7 @@ def load_models():
     ner_pipe = pipeline("ner", model=ner_model, tokenizer=ner_tokenizer, grouped_entities=True)
     cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-12-v2')
     
-    return sent_pipe, sum_pipe, ner_pipe, cross_encoder
+    return sent_pipe, sum_pipe, ner_pipe, cross_encoder, kg_model, kg_tokenizer
 
 @st.experimental_singleton(suppress_st_warning=True)
 def load_asr_model(asr_model_name):
@@ -358,7 +360,324 @@ def make_spans(text,results):
 def fin_ext(text):
     results = remote_clx(sent_tokenizer(text))
     return make_spans(text,results)
-    
+
+## Knowledge Graphs code
+
+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
+
+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
+            
+def save_network_html(kb, filename="network.html"):
+    # create network
+    net = Network(directed=True, width="700px", height="700px", bgcolor="#eeeeee")
+
+    # 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)
+            
+
 nlp = get_spacy()    
-sent_pipe, sum_pipe, ner_pipe, cross_encoder  = load_models()
+sent_pipe, sum_pipe, ner_pipe, cross_encoder, kg_model, kg_tokenizer  = load_models()
 sbert = load_sbert('all-MiniLM-L12-v2')