Text_Summarization / utils.py
raj22rishi's picture
Upload 14 files
fb4a3c6 verified
raw
history blame contribute delete
No virus
353 Bytes
from transformers import AutoTokenizer
model_nm = 't5-base'
device = 'cuda'
def tokenize_for_inference(text):
tokenizer = AutoTokenizer.from_pretrained(model_nm)
model_inputs = tokenizer.encode(
text,
max_length = 512,
padding=True,
truncation=True,
return_tensors='pt'
)
return model_inputs.to(device)