consciousAI/question-answering-generative-t5-v1-base-s-q-c

from transformers import (
    AutoModelForSeq2SeqLM,
    AutoTokenizer
)

def _generate(query, context, model, device):
    
    FT_MODEL = AutoModelForSeq2SeqLM.from_pretrained(model).to(device)
    FT_MODEL_TOKENIZER = AutoTokenizer.from_pretrained(model)
    input_text = "question: " + query + "</s> question_context: " + context
    
    input_tokenized = FT_MODEL_TOKENIZER.encode(input_text, return_tensors='pt', truncation=True, padding='max_length', max_length=1024).to(device)
    _tok_count_assessment = FT_MODEL_TOKENIZER.encode(input_text, return_tensors='pt', truncation=True).to(device)

    summary_ids = FT_MODEL.generate(input_tokenized, 
                                       max_length=30, 
                                       min_length=5, 
                                       num_beams=2,
                                       early_stopping=True,
                                   )
    output = [FT_MODEL_TOKENIZER.decode(id, clean_up_tokenization_spaces=True, skip_special_tokens=True) for id in summary_ids] 
    
    return str(output[0])

device = [0 if torch.cuda.is_available() else 'cpu'][0]
_generate(query, context, model="consciousAI/t5-v1-base-s-q-c-multi-task-qgen-v2", device=device)