Update README.md

README.md

@@ -21,14 +21,12 @@ This model was trained on a lfqa dataset. The model provide long-form answers to
 
 ```python
 import torch
-from transformers import AutoTokenizer, AutoModel, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
 model_name = "efederici/it5-small-lfqa"
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-model = model.to(device)
+tokenizer = AutoTokenizer.from_pretrained("efederici/it5-small-lfqa", from_flax=True)
+model = AutoModelForSeq2SeqLM.from_pretrained("efederici/it5-small-lfqa", from_flax=True)
 
 query = "con chi si era messo in contatto elon musk?"
 
@@ -42,8 +40,8 @@ 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"].to(device),
-                                           attention_mask=model_input["attention_mask"].to(device),
+output = model.generate(input_ids=model_input["input_ids"],
+                                           attention_mask=model_input["attention_mask"],
                                            min_length=10,
                                            max_length=256,
                                            do_sample=False,