Duplicate from Babelscape/rebel-demo

Co-authored-by: Mamta Narang <mamta@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,27 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,38 @@
+---
+title: Rebel Demo
+emoji: 🌍
+colorFrom: purple
+colorTo: pink
+sdk: streamlit
+app_file: app.py
+pinned: false
+duplicated_from: Babelscape/rebel-demo
+---
+
+# Configuration
+
+`title`: _string_  
+Display title for the Space
+
+`emoji`: _string_  
+Space emoji (emoji-only character allowed)
+
+`colorFrom`: _string_  
+Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
+
+`colorTo`: _string_  
+Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
+
+`sdk`: _string_  
+Can be either `gradio` or `streamlit`
+
+`sdk_version` : _string_  
+Only applicable for `streamlit` SDK.  
+See [doc](https://hf.co/docs/hub/spaces) for more info on supported versions.
+
+`app_file`: _string_  
+Path to your main application file (which contains either `gradio` or `streamlit` Python code).  
+Path is relative to the root of the repository.
+
+`pinned`: _boolean_  
+Whether the Space stays on top of your list.

app.py

@@ -0,0 +1,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))

requirements.txt

@@ -0,0 +1,3 @@
+torch
+streamlit
+transformers