update.

README.md

@@ -30,7 +30,7 @@ Please checkout the details in our [github repository](https://github.com/stonyb
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from digit_tokenization import enable_digit_tokenization # digit_tokenization.py from https://github.com/stonybrooknlp/teabreac
 
-model_name = "poet-large-iirc-retrieved"
+model_name = "StonyBrookNLP/poet-large-iirc-retrieved"
 tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False) # Fast doesn't work with digit tokenization
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 enable_digit_tokenization(tokenizer)
@@ -43,7 +43,7 @@ input_ids = tokenizer(
     input_texts, return_tensors="pt",
     truncation=True, max_length=800,
     add_special_tokens=True, padding=True,
-)
+)["input_ids"]
 generated_ids = model.generate(input_ids, min_length=1,  max_length=50)
 generated_predictions = tokenizer.batch_decode(generated_ids, skip_special_tokens=False)
 generated_predictions = [