Spaces:
Runtime error
Runtime error
| import streamlit as st | |
| from datasets import load_dataset | |
| from transformers import AutoModelForSeq2SeqLM, AutoTokenizer | |
| from time import time | |
| import torch | |
| 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)) |