add streaming

app.py

@@ -21,23 +21,25 @@ def load_models():
         _ = 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")
-    return (tokenizer, model, dataset)
+    dataset = load_dataset('Babelscape/rebel-dataset', split="validation", streaming=True)
+    return (tokenizer, model, dataset.take(1000))
 
 def extract_triplets(text):
     triplets = []
-    relation = ''
+    relation, subject, relation, object_ = '', '', '', ''
+    text = text.strip()
+    current = 'x'
     for token in text.split():
         if token == "<triplet>":
             current = 't'
             if relation != '':
-                triplets.append((subject, relation, object_))
+                triplets.append({'head': subject.strip(), 'type': relation.strip(),'tail': object_.strip()})
                 relation = ''
             subject = ''
         elif token == "<subj>":
             current = 's'
             if relation != '':
-                triplets.append((subject, relation, object_))
+                triplets.append({'head': subject.strip(), 'type': relation.strip(),'tail': object_.strip()})
             object_ = ''
         elif token == "<obj>":
             current = 'o'
@@ -49,10 +51,10 @@ def extract_triplets(text):
                 object_ += ' ' + token
             elif current == 'o':
                 relation += ' ' + token
-    triplets.append((subject, relation, object_))
+    if subject != '' and relation != '' and object_ != '':
+        triplets.append({'head': subject.strip(), 'type': relation.strip(),'tail': object_.strip()})
     return triplets
 
-
 tokenizer, model, dataset = load_models()
 
 agree = st.checkbox('Free input', False)