models/t5_small_medium_title_generation.py

import nltk
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch

def t5model(prompt: str) -> str:
 tokenizer = AutoTokenizer.from_pretrained("fabiochiu/t5-small-medium-title-generation")
 model = AutoModelForSeq2SeqLM.from_pretrained("fabiochiu/t5-small-medium-title-generation", device_map="cuda:0", torch_dtype=torch.float16)
 inputs = tokenizer(
 ["summarize:" + prompt],
 return_tensors="pt",
 max_length=1024,
 truncation=True
 )
 
 # Move the inputs tensor to the same device as the model tensor
 inputs = {k: v.to(model.device) for k, v in inputs.items()}
 
 outputs = model.generate(
 **inputs,
 num_beams=8,
 do_sample=True,
 min_length=8,
 max_length=15
 )
 
 decoded_output = tokenizer.batch_decode(
 outputs, skip_special_tokens=True
 )[0]
 
 return decoded_output