BART-Squad2

Aquí hay una manera rápida de hacer que la respuesta a preguntas funcione localmente:

from transformers import AutoTokenizer, AutoModelForQuestionAnswering

# Cargar el tokenizador y el modelo
 tokenizador = AutoTokenizer.from_pretrained('Primer/bart-squad2')
modelo = AutoModelForQuestionAnswering.from_pretrained('Primer/bart-squad2')
modelo.to('cuda')
modelo.eval()

# Definir función para responder preguntas
def answer(question, text):
    seq = '' + question + ' ' + text + ' '
    tokens = tokenizador.encode_plus(seq, return_tensors='pt', padding='max_length', max_length=1024)
    input_ids = tokens['input_ids'].to('cuda')
    attention_mask = tokens['attention_mask'].to('cuda')
    start, end, _ = modelo(input_ids, attention_mask=attention_mask)
    start_idx = int(start.argmax().int())
    end_idx = int(end.argmax().int())
    print(tokenizador.decode(input_ids[0, start_idx:end_idx]).strip())
    
# Ejemplo de uso
question = 'Where does Tom live?'
contexto = 'Tom is an engineer in San Francisco.'
answer(question, contexto)
# San Francisco