Update app.py

app.py

@@ -4,7 +4,7 @@ import zipfile
 import tempfile
 import base64
 from PIL import Image
-from transformers import VisionEncoderDecoderModel, ViTFeatureExtractor, AutoTokenizer
+from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
 import pandas as pd
 from nltk.corpus import wordnet
 import spacy
@@ -28,8 +28,16 @@ nlp = spacy.load("en_core_web_sm")
 # Load the pre-trained model for image captioning
 model_name = "NourFakih/Vit-GPT2-COCO2017Flickr-85k-09"
 model = VisionEncoderDecoderModel.from_pretrained(model_name)
-feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
+feature_extractor = ViTImageProcessor.from_pretrained(model_name)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
+# GPT2 only has bos/eos tokens but not decoder_start/pad tokens
+tokenizer.pad_token = tokenizer.eos_token
+
+# update the model config
+model.config.eos_token_id = tokenizer.eos_token_id
+model.config.decoder_start_token_id = tokenizer.bos_token_id
+model.config.pad_token_id = tokenizer.pad_token_id
+
 
 def generate_caption(image):
     pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values