add app

app.py

@@ -0,0 +1,99 @@
+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="train[:1%]")
+    return (tokenizer, model, dataset)
+
+def extract_triplets(text):
+    triplets = []
+    relation = ''
+    for token in text.split():
+        if token == "<triplet>":
+            current = 't'
+            if relation != '':
+                triplets.append((subject, relation, object_))
+                relation = ''
+            subject = ''
+        elif token == "<subj>":
+            current = 's'
+            if relation != '':
+                triplets.append((subject, relation, object_))
+            object_ = ''
+        elif token == "<obj>":
+            current = 'o'
+            relation = ''
+        else:
+            if current == 't':
+                subject += ' ' + token
+            elif current == 's':
+                object_ += ' ' + token
+            elif current == 'o':
+                relation += ' ' + token
+    triplets.append((subject, relation, object_))
+    return triplets
+
+
+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))