Update README.md

README.md

@@ -8,7 +8,7 @@ datasets:
 # ByT5-base fine-tuned for Hate Speech Detection (on Tweets)
 [ByT5](https://huggingface.co/google/byt5-base) base fine-tuned on [tweets hate speech detection](https://huggingface.co/datasets/tweets_hate_speech_detection) dataset for **Sequence Classification** downstream task.
 
-# Details of ByT5 - Base
+# Details of ByT5 - Base 🧠
 
 ByT5 is a tokenizer-free version of [Google's T5](https://ai.googleblog.com/2020/02/exploring-transfer-learning-with-t5.html) and generally follows the architecture of [MT5](https://huggingface.co/google/mt5-base).
 ByT5 was only pre-trained on [mC4](https://www.tensorflow.org/datasets/catalog/c4#c4multilingual) excluding any supervised training with an average span-mask of 20 UTF-8 characters. Therefore, this model has to be fine-tuned before it is useable on a downstream task.
@@ -61,18 +61,18 @@ pip install -q transformers
 
 ```python
 from transformers import AutoTokenizer, T5ForConditionalGeneration
-tokenizer = AutoTokenizer.from_pretrained("Narrativa/byt5-base-tweet-hate-detection")
 
-model = T5ForConditionalGeneration.from_pretrained("Narrativa/byt5-base-tweet-hate-detection").to("cuda")
+ckpt = 'Narrativa/byt5-base-tweet-hate-detection'
 
-def classify_tweet(tweet):
+tokenizer = AutoTokenizer.from_pretrained(ckpt)
+model = T5ForConditionalGeneration.from_pretrained(ckpt).to("cuda")
 
-    inputs = tokenizer([tweet], padding="max_length", truncation=True, max_length=512, return_tensors="pt")
-    input_ids = inputs.input_ids.to("cuda")
-    attention_mask = inputs.attention_mask.to("cuda")
+def classify_tweet(tweet):
 
+    inputs = tokenizer([tweet], padding='max_length', truncation=True, max_length=512, 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)
-
     return tokenizer.decode(output[0], skip_special_tokens=True)