Update README.md

README.md

@@ -70,23 +70,25 @@ pip install -q ./transformers
 ```
 
 ```python
-from transformers import AutoTokenizer, T5ForConditionalGeneration
+from transformers import MBart50TokenizerFast, MBartForConditionalGeneration
 
-ckpt = 'Narrativa/byt5-base-tweet-hate-detection'
+ckpt = 'mbart-large-50-finetuned-opus-en-pt-translation'
 
-tokenizer = AutoTokenizer.from_pretrained(ckpt)
-model = T5ForConditionalGeneration.from_pretrained(ckpt).to("cuda")
+tokenizer = MBart50TokenizerFast.from_pretrained(ckpt)
+model = MBartForConditionalGeneration.from_pretrained(ckpt).to("cuda")
 
-def classify_tweet(tweet):
+tokenizer.src_lang = 'en_XX'
 
-    inputs = tokenizer([tweet], padding='max_length', truncation=True, max_length=512, return_tensors='pt')
+def translate(text):
+
+    inputs = tokenizer(text, return_tensors='pt')
     input_ids = inputs.input_ids.to('cuda')
     attention_mask = inputs.attention_mask.to('cuda')
-    output = model.generate(input_ids, attention_mask=attention_mask)
+    output = model.generate(input_ids, attention_mask=attention_mask, forced_bos_token_id=tokenizer.lang_code_to_id['pt_XX'])
     return tokenizer.decode(output[0], skip_special_tokens=True)
     
     
-classify_tweet('here goes your tweet...')
+translate('here your English text to be translated to Portuguese...')
 ```
 
 Created by: [Narrativa](https://www.narrativa.com/)