add model.py and app.py

app.py

@@ -0,0 +1 @@
+from model import *

model.py

@@ -0,0 +1,57 @@
+import sys, os
+
+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
+
+# Vit - as encoder
+from transformers import ViTFeatureExtractor
+from PIL import Image
+import requests
+import numpy as np
+
+# GPT2 / GPT2LM - as decoder
+from transformers import ViTFeatureExtractor, GPT2Tokenizer
+
+model_name_or_path = './outputs/ckpt_2/'
+flax_vit_gpt2_lm = FlaxViTGPT2LMForConditionalGeneration.from_pretrained(model_name_or_path)
+
+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 = 16
+num_beams = 4
+gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
+
+
+def predict(image):
+
+    image = Image.open(requests.get(url, stream=True).raw)
+    # batch dim is added automatically
+    encoder_inputs = feature_extractor(images=image, return_tensors="jax")
+    pixel_values = encoder_inputs.pixel_values
+
+    # generation
+    batch = {'pixel_values': pixel_values}
+    generation = flax_vit_gpt2_lm.generate(batch['pixel_values'], **gen_kwargs)
+
+    token_ids = np.array(generation.sequences)[0]
+    caption = tokenizer.decode(token_ids)
+
+    return caption, token_ids
+
+
+if __name__ == '__main__':
+
+
+    url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
+    image = Image.open(requests.get(url, stream=True).raw)
+    caption, token_ids = predict(image)
+
+    print(f'token_ids: {token_ids}')
+    print(f'caption: {caption}')

test_model.py

@@ -0,0 +1,74 @@
+import sys, os
+
+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
+
+# Vit - as encoder
+from transformers import ViTFeatureExtractor
+from PIL import Image
+import requests
+import numpy as np
+
+# GPT2 / GPT2LM - as decoder
+from transformers import ViTFeatureExtractor, GPT2Tokenizer
+
+model_name_or_path = './outputs/ckpt_2/'
+flax_vit_gpt2_lm = FlaxViTGPT2LMForConditionalGeneration.from_pretrained(model_name_or_path)
+
+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 = 16
+gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
+
+
+# encoder data
+url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
+image = Image.open(requests.get(url, stream=True).raw)
+# batch dim is added automatically
+encoder_inputs = feature_extractor(images=image, return_tensors="jax")
+pixel_values = encoder_inputs.pixel_values
+print(f'pixel_values.shape = {pixel_values.shape}')
+
+# decoder data
+sentence = 'mon chien est mignon'
+# IMPORTANT: For training/evaluation/attention_mask/loss
+sentence += ' ' + tokenizer.eos_token
+# batch dim is added automatically
+decoder_inputs = tokenizer(sentence, return_tensors="jax")
+print(decoder_inputs)
+print(f'input_ids.shape = {decoder_inputs.input_ids.shape}')
+
+# model data
+inputs = dict(decoder_inputs)
+inputs['pixel_values'] = pixel_values
+
+
+logits = flax_vit_gpt2_lm(**inputs)[0]
+preds = np.argmax(logits, axis=-1)
+print('=' * 60)
+print('Flax: Vit-GPT2-LM')
+print('predicted token ids:')
+print(preds)
+print('=' * 60)
+
+
+# Generation!
+batch = {'pixel_values': pixel_values}
+generation = flax_vit_gpt2_lm.generate(batch['pixel_values'], **gen_kwargs)
+print('generation:')
+print(generation)
+print('=' * 60)
+
+token_ids = np.array(generation.sequences)[0]
+caption = tokenizer.decode(token_ids)
+print(f'token_ids: {token_ids}')
+print(f'caption: {caption}')
+print('=' * 60)