update model

app.py

@@ -1,19 +1,17 @@
 import streamlit as st
-from PIL import Image
-import numpy as np
 
 
 # Designing the interface
-st.title("🖼️ French Image Captioning Demo 📝")
+st.title("🖼️ Image Captioning Demo 📝")
 st.write("[Yih-Dar SHIEH](https://huggingface.co/ydshieh)")
 
 st.sidebar.markdown(
     """
-    An image captioning model [ViT-GPT2](https://huggingface.co/flax-community/vit-gpt2) by combining the ViT model and a French GPT2 model.
+    An image captioning model [ViT-GPT2](https://huggingface.co/flax-community/vit-gpt2) by combining the ViT model with the GPT2 model.
     [Part of the [Huggingface JAX/Flax event](https://discuss.huggingface.co/t/open-to-the-community-community-week-using-jax-flax-for-nlp-cv/).]\n
-    The GPT2 model source code is modified so it can accept an encoder's output.
-    The pretained weights of both models are loaded, with a set of randomly initialized cross-attention weigths.
-    The model is trained on 65000 images from the COCO dataset for about 1500 steps (batch_size=256), with the original English cpationis being translated to French for training purpose.
+    The encoder (ViT) and decoder (GPT2) are combined using Hugging Face transformers' `FlaxVisionEncoderDecoderModel`.
+    The pretrained weights of both models are loaded, with a set of randomly initialized cross-attention weights.
+    The model is trained on the COCO 2017 dataset for about 6900 steps (batch_size=256).
     """
 )
 
@@ -21,6 +19,7 @@ st.sidebar.markdown(
 #show = st.image(image, use_column_width=True)
 #show.image(image, 'Preloaded Image', use_column_width=True)
 
+
 with st.spinner('Loading and compiling ViT-GPT2 model ...'):
 
     from model import *
@@ -43,16 +42,21 @@ show.image(image, '\n\nSelected Image', width=480)
 # For newline
 st.sidebar.write('\n')
 
+
 with st.spinner('Generating image caption ...'):
 
     caption = predict(image)
+
+    caption_en = caption
+    st.header(f'**Prediction (in English) **{caption_en}')
     
-    caption_en = translator.translate(caption, src='fr', dest='en').text
-    st.header(f'**Prediction (in French) **{caption}')
-    st.header(f'**English Translation**: {caption_en}')
+    # caption_en = translator.translate(caption, src='fr', dest='en').text
+    # st.header(f'**Prediction (in French) **{caption}')
+    # st.header(f'**English Translation**: {caption_en}')
+
 
 st.sidebar.header("ViT-GPT2 predicts:")
-st.sidebar.write(f"**French**: {caption}")
-st.sidebar.write(f"**English Translation**: {caption_en}")
+st.sidebar.write(f"**English**: {caption}")
+
 
-image.close()
+image.close()

model.py

@@ -1,83 +1,67 @@
-import os, sys, shutil
-import numpy as np
+import os, shutil
 from PIL import Image
-
 import jax
-from transformers import ViTFeatureExtractor
-from transformers import GPT2Tokenizer
+from transformers import FlaxVisionEncoderDecoderModel, ViTFeatureExtractor, AutoTokenizer
 from huggingface_hub import hf_hub_download
 
 from googletrans import Translator
 translator = Translator()
 
-current_path = os.path.dirname(os.path.abspath(__file__))
-sys.path.append(current_path)
-
-# Main model -  ViTGPT2LM
-from vit_gpt2.modeling_flax_vit_gpt2_lm import FlaxViTGPT2LMForConditionalGeneration
 
 # create target model directory
 model_dir = './models/'
 os.makedirs(model_dir, exist_ok=True)
-# copy config file
-filepath = hf_hub_download("flax-community/vit-gpt2", "checkpoints/ckpt_5/config.json")
-shutil.copyfile(filepath, os.path.join(model_dir, 'config.json'))
-# copy model file
-filepath = hf_hub_download("flax-community/vit-gpt2", "checkpoints/ckpt_5/flax_model.msgpack")
-shutil.copyfile(filepath, os.path.join(model_dir, 'flax_model.msgpack'))
-
-flax_vit_gpt2_lm = FlaxViTGPT2LMForConditionalGeneration.from_pretrained(model_dir)
-
-vit_model_name = 'google/vit-base-patch16-224-in21k'
-feature_extractor = ViTFeatureExtractor.from_pretrained(vit_model_name)
 
-gpt2_model_name = 'asi/gpt-fr-cased-small'
-tokenizer = GPT2Tokenizer.from_pretrained(gpt2_model_name)
-
-max_length = 32
-num_beams = 8
+files_to_download = [
+    "config.json",
+    "flax_model.msgpack",
+    "merges.txt",
+    "special_tokens_map.json",
+    "tokenizer.json",
+    "tokenizer_config.json",
+    "vocab.json",
+]
+
+# copy files from checkpoint hub:
+for fn in files_to_download:
+    file_path = hf_hub_download("ydshieh/vit-gpt2-coco-en", f"ckpt_epoch_3_step_6900/{fn}")
+    shutil.copyfile(file_path, os.path.join(model_dir, fn))
+
+model = FlaxVisionEncoderDecoderModel.from_pretrained(model_dir)
+feature_extractor = ViTFeatureExtractor.from_pretrained(model_dir)
+tokenizer = AutoTokenizer.from_pretrained(model_dir)
+
+max_length = 16
+num_beams = 4
 gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
 
 
 @jax.jit
-def predict_fn(pixel_values):
+def generate(pixel_values):
+    output_ids = model.generate(pixel_values, **gen_kwargs).sequences
+    return output_ids
 
-    return flax_vit_gpt2_lm.generate(pixel_values, **gen_kwargs)
 
 def predict(image):
 
-    # batch dim is added automatically
-    encoder_inputs = feature_extractor(images=image, return_tensors="jax")
-    pixel_values = encoder_inputs.pixel_values
-
-    # generation
-    generation = predict_fn(pixel_values)
+    pixel_values = feature_extractor(images=image, return_tensors="np").pixel_values
+    output_ids = generate(pixel_values)
+    preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
+    preds = [pred.strip() for pred in preds]
 
-    token_ids = np.array(generation.sequences)[0]
-    caption = tokenizer.decode(token_ids)
-    caption = caption.replace('<s>', '').replace('</s>', '').replace('<pad>', '')
-    caption = caption.replace("à l'arrière-plan", '').replace("Une photo en noir et blanc d'", '').replace("Une photo noire et blanche d'", '').replace("en arrière-plan", '').replace("Un gros plan d'", '').replace("un gros plan d'", '').replace("Une image d'", '')
-    while '  ' in caption:
-        caption = caption.replace('  ', ' ')
-    caption = caption.strip()
-    if caption:
-        caption = caption[0].upper() + caption[1:]
+    return preds[0]
 
-    return caption
 
-def compile():
+def _compile():
 
     image_path = 'samples/val_000000039769.jpg'
     image = Image.open(image_path)
-
     caption = predict(image)
     image.close()
 
-def predict_dummy(image):
-    
-    return 'dummy caption!'
 
-compile()
+_compile()
+
 
 sample_dir = './samples/'
 sample_fns = tuple([f"{int(f.replace('COCO_val2014_', '').replace('.jpg', ''))}.jpg" for f in os.listdir(sample_dir) if f.startswith('COCO_val2014_')])