Update README.md

README.md

@@ -1,6 +1,8 @@
 ---
 language:
 - en
+widget:
+- text: "generate question: <hl> 42 <hl> is the answer to life, the universe and everything. </s>"
 tags:
 - seq2seq
 license: cc-by-nc-sa-4.0
@@ -9,6 +11,11 @@ To use the model with a pipeline:
 ```python3
 from transformers import pipeline
 
+triplet_extractor = pipeline('text2text-generation', model='Babelscape/rebel-large', tokenizer='Babelscape/rebel-large')
+# We need to use the tokenizer manually since we need special tokens.
+extracted_text = triplet_extractor.tokenizer.decode(triplet_extractor("Punta Cana is a resort town in the municipality of Higuey, in La Altagracia Province, the eastern most province of the Dominican Republic", return_tensors=True, return_text=False)[0]["generated_token_ids"])
+print(extracted_text)
+# Function to parse the generated text and extract the triplets
 def extract_triplets(text):
     triplets = []
     relation = ''
@@ -36,10 +43,6 @@ def extract_triplets(text):
                 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)
 ```
@@ -78,8 +81,8 @@ def extract_triplets(text):
     return triplets
 
 # Load model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained("model/rebel-large")
-model = AutoModelForSeq2SeqLM.from_pretrained("model/rebel-large")
+tokenizer = AutoTokenizer.from_pretrained("Babelscape/rebel-large")
+model = AutoModelForSeq2SeqLM.from_pretrained("Babelscape/rebel-large")
 gen_kwargs = {
     "max_length": 256,
     "length_penalty": 0,