Views
No views yet
1from transformers import T5ForConditionalGeneration, AutoTokenizer
2
3tokenizer = AutoTokenizer.from_pretrained("google/byt5-small")
4
5model = T5ForConditionalGeneration.from_pretrained("marianna13/byt5-small-NSFW-image-urls")
6
7def get_label(text):
8 input_ids = tokenizer(text, return_tensors="pt", padding=True).input_ids
9 outputs = model.generate(input_ids)
10
11 label = tokenizer.batch_decode(outputs, skip_special_tokens=True)
12
13 return label
14