Update README.md

README.md

@@ -21,11 +21,11 @@ language:
 widget:
 - text: >-
     The Red Hot Chili Peppers were formed in Los Angeles by Kiedis, Flea, guitarist Hillel Slovak and drummer Jack Irons.
-    example_title: English
+  example_title: English
 inference:
   parameters:
-    src_lang: __en__
-    tgt_lang: <triplet>
+    src_lang: en
+    tgt_lang: en
 tags:
 - seq2seq
 - relation-extraction
@@ -61,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=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.
+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.", src_lang="en", 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):
@@ -145,11 +145,8 @@ 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 ("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")
+tokenizer = AutoTokenizer.from_pretrained("Babelscape/mrebel-base", src_lang="en", tgt_lang="en") 
+# Here we set English ("en") as source language. To change the source language swap the first token of the input for your desired language or change to supported language. 
 model = AutoModelForSeq2SeqLM.from_pretrained("Babelscape/mrebel-base")
 gen_kwargs = {
     "max_length": 256,
@@ -169,7 +166,6 @@ 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 = tokenizer.convert_tokens_to_ids("tp_XX"),
     **gen_kwargs,
 )