Update README.md

README.md

@@ -50,6 +50,27 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
 tokenizer = AutoTokenizer.from_pretrained("tarudesu/ViHateT5-base-HSD")
 model = AutoModelForSeq2SeqLM.from_pretrained("tarudesu/ViHateT5-base-HSD")
+
+def generate_output(input_text, prefix):
+    # Add prefix
+    prefixed_input_text = prefix + ': ' + input_text
+
+    # Tokenize input text
+    input_ids = tokenizer.encode(prefixed_input_text, return_tensors="pt")
+
+    # Generate output
+    output_ids = model.generate(input_ids, max_length=256)
+
+    # Decode the generated output
+    output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+    return output_text
+
+sample = 'Tôi ghét bạn vl luôn!'
+prefix = 'hate-spans-detection' # Choose 1 from 3 prefixes ['hate-speech-detection', 'toxic-speech-detection', 'hate-spans-detection']
+
+result = generate_output(sample, prefix)
+print('Result: ', result)
 ```
 
 Please feel free to contact us by email luannt@uit.edu.vn if you have any further information!