Create app.py

app.py

@@ -0,0 +1,117 @@
+import streamlit as st
+from datasets import load_dataset
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from time import time
+import torch
+
+@st.cache(
+    allow_output_mutation=True,
+    hash_funcs={
+        AutoTokenizer: lambda x: None,
+        AutoModelForSeq2SeqLM: lambda x: None,
+    },
+    suppress_st_warning=True
+)
+def load_models(lan):
+    st_time = time()
+    tokenizer = AutoTokenizer.from_pretrained("Babelscape/mrebel-large", src_lang=_Tokens[lan], "tgt_lang": "tp_XX")
+    print("+++++ loading Model", time() - st_time)
+    model = AutoModelForSeq2SeqLM.from_pretrained("Babelscape/mrebel-large")
+    if torch.cuda.is_available():
+        _ = model.to("cuda:0") # comment if no GPU available
+    _ = model.eval()
+    print("+++++ loaded model", time() - st_time)
+    dataset = load_dataset('Babelscape/SREDFM', lan, split="validation", streaming=True)
+    dataset = [example for example in dataset.take(1001)]
+    return (tokenizer, model, dataset)
+
+def extract_triplets_typed(text):
+    triplets = []
+    relation = ''
+    text = text.strip()
+    current = 'x'
+    subject, relation, object_, object_type, subject_type = '','','','',''
+
+    for token in text.replace("<s>", "").replace("<pad>", "").replace("</s>", "").replace("tp_XX", "").replace("__en__", "").split():
+        if token == "<triplet>" or token == "<relation>":
+            current = 't'
+            if relation != '':
+                triplets.append({'head': subject.strip(), 'head_type': subject_type, 'type': relation.strip(),'tail': object_.strip(), 'tail_type': object_type})
+                relation = ''
+            subject = ''
+        elif token.startswith("<") and token.endswith(">"):
+            if current == 't' or current == 'o':
+                current = 's'
+                if relation != '':
+                    triplets.append({'head': subject.strip(), 'head_type': subject_type, 'type': relation.strip(),'tail': object_.strip(), 'tail_type': object_type})
+                object_ = ''
+                subject_type = token[1:-1]
+            else:
+                current = 'o'
+                object_type = token[1:-1]
+                relation = ''
+        else:
+            if current == 't':
+                subject += ' ' + token
+            elif current == 's':
+                object_ += ' ' + token
+            elif current == 'o':
+                relation += ' ' + token
+    if subject != '' and relation != '' and object_ != '' and object_type != '' and subject_type != '':
+        triplets.append({'head': subject.strip(), 'head_type': subject_type, 'type': relation.strip(),'tail': object_.strip(), 'tail_type': object_type})
+    return triplets
+
+st.markdown("""This is a demo for the Findings of EMNLP 2021 paper [REBEL: Relation Extraction By End-to-end Language generation](https://github.com/Babelscape/rebel/blob/main/docs/EMNLP_2021_REBEL__Camera_Ready_.pdf). The pre-trained model is able to extract triplets for up to 200 relation types from Wikidata or be used in downstream Relation Extraction task by fine-tuning. Find the model card [here](https://huggingface.co/Babelscape/rebel-large). Read more about it in the [paper](https://aclanthology.org/2021.findings-emnlp.204) and in the original [repository](https://github.com/Babelscape/rebel).""")
+
+lan = st.selectbox(
+    'Select a Language',
+    ('ar', 'ca', 'de', 'el', 'en', 'es', 'fr', 'hi', 'it', 'ja', 'ko', 'nl', 'pl', 'pt', 'ru', 'sv', 'vi', 'zh'))
+
+_Tokens = {'en': 'en_XX', 'de': 'de_DE', 'ca': 'ca_XX', 'ar': 'ar_AR', 'el': 'el_EL', 'it': 'it_IT', 'ja': 'ja_XX', 'ko': 'ko_KR', 'hi': 'hi_IN', 'pt': 'pt_XX', 'ru': 'ru_RU', 'pl': 'pl_PL', 'zh': 'zh_CN', 'fr': 'fr_XX', 'vi': 'vi_VN', 'sv':'sv_SE'}
+
+tokenizer, model, dataset = load_models(lan)
+
+agree = st.checkbox('Free input', False)
+if agree:
+    text = st.text_input('Input text', 'Els Red Hot Chili Peppers es van formar a Los Angeles per Kiedis, Flea, el guitarrista Hillel Slovak i el bateria Jack Irons.')
+    print(text)
+else:
+    dataset_example = st.slider('dataset id', 0, 1000, 0)
+    text = dataset[dataset_example]['context']
+length_penalty = st.slider('length_penalty', 0, 10, 0)
+num_beams = st.slider('num_beams', 1, 20, 3)
+num_return_sequences = st.slider('num_return_sequences', 1, num_beams, 2)
+
+gen_kwargs = {
+    "max_length": 256,
+    "length_penalty": length_penalty,
+    "num_beams": num_beams,
+    "num_return_sequences": num_return_sequences,
+    "forced_bos_token_id": None,
+}
+
+model_inputs = tokenizer(text, max_length=256, padding=True, truncation=True, return_tensors = 'pt')
+generated_tokens = model.generate(
+    model_inputs["input_ids"].to(model.device),
+    attention_mask=model_inputs["attention_mask"].to(model.device),
+    decoder_start_token_id = tokenizer.convert_tokens_to_ids("tp_XX"),
+    **gen_kwargs,
+)
+
+decoded_preds = tokenizer.batch_decode(generated_tokens, skip_special_tokens=False)
+st.title('Input text')
+
+st.write(text)
+
+if not agree:
+    st.title('Silver output')
+    st.write(dataset[dataset_example]['triplets'])
+    st.write(extract_triplets_typed(dataset[dataset_example]['triplets']))
+
+st.title('Prediction text')
+decoded_preds = [text.replace('<s>', '').replace('</s>', '').replace('<pad>', '') for text in decoded_preds]
+st.write(decoded_preds)
+
+for idx, sentence in enumerate(decoded_preds):
+    st.title(f'Prediction triplets sentence {idx}')
+    st.write(extract_triplets_typed(sentence))