File size: 375 Bytes
7e9fae4 |
1 2 3 4 5 6 7 8 9 10 11 12 |
from transformers import T5Tokenizer, T5ForConditionalGeneration
tokenizer = T5Tokenizer.from_pretrained("t5-small")
model = T5ForConditionalGeneration.from_pretrained("t5-small")
def infer_t5(input):
input_ids = tokenizer(input, return_tensors="pt").input_ids
outputs = model.generate(input_ids)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
|