Updates

model.py

@@ -4,7 +4,10 @@ import random
 from huggingface_hub import hf_hub_download
 
 
-from transformers import AutoTokenizer, TFEncoderDecoderModel
+from transformers import BertTokenizer, EncoderDecoderModel
+
+model = EncoderDecoderModel.from_pretrained("patrickvonplaten/bert2bert-cnn_dailymail-fp16")
+tokenizer = BertTokenizer.from_pretrained("patrickvonplaten/bert2bert-cnn_dailymail-fp16")
 
 loc = "ydshieh/bert2bert-cnn_dailymail-fp16"
 
@@ -23,14 +26,13 @@ os.makedirs(model_dir, exist_ok=True)
 #     file_path = hf_hub_download("ydshieh/bert2bert-cnn_dailymail-fp16", f"ckpt_epoch_3_step_6900/{fn}")
 #     shutil.copyfile(file_path, os.path.join(model_dir, fn))
 
-model = TFEncoderDecoderModel.from_pretrained(loc)
-tokenizer = AutoTokenizer.from_pretrained(loc)
 
 
 def predict(article):
 
-    input_ids = tokenizer(article, return_tensors="tf").input_ids
+    input_ids = tokenizer(article, return_tensors="pt").input_ids
     output_ids = model.generate(input_ids)
+
     summary = tokenizer.decode(output_ids[0], skip_special_tokens=True)
     return summary