File size: 4,290 Bytes
2816512 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
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():
st_time = time()
tokenizer = AutoTokenizer.from_pretrained("Babelscape/rebel-large")
print("+++++ loading Model", time() - st_time)
model = AutoModelForSeq2SeqLM.from_pretrained("Babelscape/rebel-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/rebel-dataset', split="validation", streaming=True)
dataset = [example for example in dataset.take(1001)]
return (tokenizer, model, dataset)
def extract_triplets(text):
triplets = []
relation, subject, relation, object_ = '', '', '', ''
text = text.strip()
current = 'x'
for token in text.split():
if token == "<triplet>":
current = 't'
if relation != '':
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()})
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_ != '':
triplets.append({'head': subject.strip(), 'type': relation.strip(),'tail': object_.strip()})
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).""")
tokenizer, model, dataset = load_models()
agree = st.checkbox('Free input', False)
if agree:
text = st.text_input('Input text', 'Punta Cana is a resort town in the municipality of Higüey, in La Altagracia Province, the easternmost province of the Dominican Republic.')
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,
}
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),
**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(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(sentence)) |