Update README.md

README.md

@@ -46,7 +46,7 @@ Here is an example of generation for Wikipedia page retrieval for open-domain fa
 
 ```python
 import pickle
-from transformers import BartTokenizer, BartForConditionalGeneration
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
 # OPTIONAL: load the prefix tree (trie), you need to additionally download
 # https://huggingface.co/facebook/genre-kilt/blob/main/trie.py and 
@@ -55,8 +55,8 @@ from transformers import BartTokenizer, BartForConditionalGeneration
 # with open("kilt_titles_trie_dict.pkl", "rb") as f:
 #     trie = Trie.load_from_dict(pickle.load(f))
 
-tokenizer = BartTokenizer.from_pretrained("facebook/genre-kilt")
-model = BartForConditionalGeneration.from_pretrained("facebook/genre-kilt").eval()
+tokenizer = AutoTokenizer.from_pretrained("facebook/genre-kilt")
+model = AutoModelForSeq2SeqLM.from_pretrained("facebook/genre-kilt").eval()
 
 sentences = ["Einstein was a German physicist."]