Update README.md

README.md

@@ -5,3 +5,106 @@ tags:
 - seq2seq
 license: cc-by-nc-sa-4.0
 ---
+To use the model with a pipeline:
+```python3
+from transformers import pipeline
+
+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
+
+triplet_extractor = pipeline('text2text-generation', model='Babelscape/rebel-large', tokenizer='Babelscape/rebel-large')
+extracted_text = triplet_extractor("Punta Cana is a resort town in the municipality of Higüey, in La Altagracia Province, the easternmost province of the Dominican Republic.Punta Cana is a resort town in the municipality of Higüey, in La Altagracia Province, the easternmost province of the Dominican Republic.")["generated_text"]
+
+extracted_triplets = extract_triplets(extracted_text)
+print(extracted_triplets)
+```
+
+Or using the transformers
+```python3
+
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+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
+
+# Load model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained("model/rebel-large")
+model = AutoModelForSeq2SeqLM.from_pretrained("model/rebel-large")
+gen_kwargs = {
+    "max_length": 256,
+    "length_penalty": 0,
+    "num_beams": 3,
+    "num_return_sequences": 3,
+}
+
+# Text to extract triplets from
+text = 'Punta Cana is a resort town in the municipality of Higüey, in La Altagracia Province, the easternmost province of the Dominican Republic.'
+
+# Tokenizer text
+model_inputs = tokenizer(text, max_length=256, padding=True, truncation=True, return_tensors = 'pt')
+
+# Generate
+generated_tokens = model.generate(
+    model_inputs["input_ids"].to(model.device),
+    attention_mask=model_inputs["attention_mask"].to(model.device),
+    **gen_kwargs,
+)
+
+# Extract text
+decoded_preds = tokenizer.batch_decode(generated_tokens, skip_special_tokens=False)
+
+# Extract triplets
+for idx, sentence in enumerate(decoded_preds):
+    print(f'Prediction triplets sentence {idx}')
+    print(extract_triplets(sentence))
+```