import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

tokenizer = AutoTokenizer.from_pretrained("efederici/it5-efficient-small-lfqa")
model = AutoModelForSeq2SeqLM.from_pretrained("efederici/it5-efficient-small-lfqa")

query = "con chi si era messo in contatto elon musk?"

# concatenated texts/document text
doc = """
La notizia dell’acquisizione da parte di Elon Musk del 9,2 per cento delle azioni di Twitter e del suo successivo ingresso nel consiglio di amministrazione della società hanno attirato grandi attenzioni, non solo da parte degli analisti finanziari, ma anche di chi si occupa di social media e del modo in cui viene impiegata la piattaforma da centinaia di milioni di persone in tutto il mondo. Musk, che ha un grande seguito su Twitter, in passato aveva più volte criticato il social network, accusandolo di non tutelare a sufficienza le libertà di espressione, anche in casi limite come l’assalto al Congresso degli Stati Uniti del 2021.
Alcune settimane fa, Musk si era messo in contatto con Parag Agrawal, CEO di Twitter da fine novembre 2021, e con il suo predecessore e cofondatore della società, Jack Dorsey, annunciando di avere avviato l’acquisizione di alcune quote dell’azienda e di essere disponibile per discutere di soluzioni per migliorarla. Secondo fonti del New York Times, dopo i primi contatti, Agrawal aveva proposto a Musk di avere un ruolo più attivo oltre a quello di azionista, offrendogli la possibilità di entrare nel consiglio di amministrazione.
"""

query_and_docs = f"Domanda: {query} Contesto: {doc}"

model_input = tokenizer(query_and_docs, truncation=True, padding=True, return_tensors="pt")

output = model.generate(
  input_ids=model_input["input_ids"],
  attention_mask=model_input["attention_mask"],
  min_length=10,
  max_length=256,
  do_sample=False, 
  early_stopping=True,
  num_beams=8,
  temperature=1.0,
  top_k=None,
  top_p=None,
  no_repeat_ngram_size=3,
  num_return_sequences=1
)
                                           
tokenizer.batch_decode(output, skip_special_tokens=True, clean_up_tokenization_spaces=True)

The model will predict: 'Elon Musk si era messo in contatto con Parag Agrawal, CEO di Twitter da fine novembre 2021 e con il suo predecessore e cofondatore della società, Jack Dorsey, annunciando di avere avviato l’acquisizione di alcune quote dell’azienda e di essere disponibile per discutere soluzioni per migliorarla.'