Update README.md

README.md

@@ -24,8 +24,7 @@ widget:
     example_title: English
 inference:
   parameters:
-    decoder_start_token_id: 250058
-    src_lang: en_XX
+    src_lang: __en__
     tgt_lang: <triplet>
 tags:
 - seq2seq
@@ -51,9 +50,9 @@ mREBEL is introduced in the ACL 2023 paper [RED^{FM}: a Filtered and Multilingua
         url = "https://arxiv.org/abs/2306.09802",
     }
 
-The original repository for the paper can be found [here](https://github.com/Babelscape/rebel)
+The original repository for the paper can be found [here](https://github.com/Babelscape/rebel#REDFM)
 
-Be aware that the inference widget at the right does not output special tokens, which are necessary to distinguish the subject, object and relation types. For a demo of REBEL and its pre-training dataset check the [Spaces demo](https://huggingface.co/spaces/Babelscape/rebel-demo).
+Be aware that the inference widget at the right does not output special tokens, which are necessary to distinguish the subject, object and relation types. For a demo of mREBEL and its pre-training dataset check the [Spaces demo](https://huggingface.co/spaces/Babelscape/mrebel-demo).
 
 ## Pipeline usage
 
@@ -62,7 +61,7 @@ from transformers import pipeline
 
 triplet_extractor = pipeline('translation_xx_to_yy', model='Babelscape/mrebel-base', tokenizer='Babelscape/mrebel-base')
 # We need to use the tokenizer manually since we need special tokens.
-extracted_text = triplet_extractor.tokenizer.batch_decode([triplet_extractor("The Red Hot Chili Peppers were formed in Los Angeles by Kiedis, Flea, guitarist Hillel Slovak and drummer Jack Irons.", decoder_start_token_id=250058, src_lang="en_XX", tgt_lang="<triplet>", return_tensors=True, return_text=False)[0]["translation_token_ids"]]) # change en_XX for the language of the source.
+extracted_text = triplet_extractor.tokenizer.batch_decode([triplet_extractor("The Red Hot Chili Peppers were formed in Los Angeles by Kiedis, Flea, guitarist Hillel Slovak and drummer Jack Irons.", decoder_start_token_id=tokenizer.convert_tokens_to_ids("tp_XX"), src_lang="__en__", tgt_lang="<triplet>", return_tensors=True, return_text=False)[0]["translation_token_ids"]]) # change __en__ for the language of the source.
 print(extracted_text[0])
 # Function to parse the generated text and extract the triplets
 def extract_triplets_typed(text):
@@ -146,7 +145,11 @@ def extract_triplets_typed(text):
     return triplets
 
 # Load model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained("Babelscape/mrebel-base", src_lang="en_XX", "tgt_lang": "tp_XX") # Here we set English as source language. To change the source language just change it here or swap the first token of the input for your desired language
+tokenizer = AutoTokenizer.from_pretrained("Babelscape/mrebel-base", src_lang="en_XX", tgt_lang="tp_XX") 
+# Here we set English ("en_XX") as source language. To change the source language swap the first token of the input for your desired language or change to supported language. For catalan ("ca_XX") or greek ("el_EL") (not included in mBART pretraining) you need a workaround:
+# tokenizer._src_lang = "ca_XX"
+# tokenizer.cur_lang_code_id = tokenizer.convert_tokens_to_ids("ca_XX")
+# tokenizer.set_src_lang_special_tokens("ca_XX")
 model = AutoModelForSeq2SeqLM.from_pretrained("Babelscape/mrebel-base")
 gen_kwargs = {
     "max_length": 256,
@@ -166,7 +169,7 @@ model_inputs = tokenizer(text, max_length=256, padding=True, truncation=True, re
 generated_tokens = model.generate(
     model_inputs["input_ids"].to(model.device),
     attention_mask=model_inputs["attention_mask"].to(model.device),
-    decoder_start_token_id = self.tokenizer.convert_tokens_to_ids("tp_XX"),
+    decoder_start_token_id = tokenizer.convert_tokens_to_ids("tp_XX"),
     **gen_kwargs,
 )