update test_model.py

tests/test_model.py

@@ -6,28 +6,30 @@ sys.path.append(current_path)
 # Main model -  ViTGPT2LM
 from vit_gpt2.modeling_flax_vit_gpt2_lm import FlaxViTGPT2LMForConditionalGeneration
 
-# Vit - as encoder
+# ViT - as encoder
 from transformers import ViTFeatureExtractor
 from PIL import Image
 import requests
 import numpy as np
+import jax
+import jax.numpy as jnp
 
-# GPT2 / GPT2LM - as decoder
-from transformers import ViTFeatureExtractor, GPT2Tokenizer
+# GPT2+LM - as decoder
+from transformers import GPT2Tokenizer
 
-model_name_or_path = './outputs/ckpt_2/'
-flax_vit_gpt2_lm = FlaxViTGPT2LMForConditionalGeneration.from_pretrained(model_name_or_path)
+max_length = 8
 
-vit_model_name = 'google/vit-base-patch16-224-in21k'
-feature_extractor = ViTFeatureExtractor.from_pretrained(vit_model_name)
+vision_model_name = 'google/vit-base-patch16-224-in21k'
+text_model_name = 'asi/gpt-fr-cased-small'
 
-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}
+flax_vit_gpt2_lm = FlaxViTGPT2LMForConditionalGeneration.from_vision_text_pretrained(
+    vision_pretrained_model_name_or_path=vision_model_name,
+    text_pretrained_model_name_or_path=text_model_name
+)
+model = flax_vit_gpt2_lm
 
+feature_extractor = ViTFeatureExtractor.from_pretrained(vision_model_name)
+tokenizer = GPT2Tokenizer.from_pretrained(text_model_name)
 
 # encoder data
 url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
@@ -42,16 +44,47 @@ 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}')
+# Setup the tokenizer for targets
+with tokenizer.as_target_tokenizer():
+    labels = tokenizer(sentence, max_length=max_length, padding="max_length", truncation=True, return_tensors="np")
+
+def shift_tokens_right(input_ids: jnp.ndarray, pad_token_id: int, decoder_start_token_id: int) -> jnp.ndarray:
+    """
+    Shift input ids one token to the right.
+    """
+    shifted_input_ids = jnp.roll(input_ids, 1, axis=-1)
+    shifted_input_ids = jax.ops.index_update(shifted_input_ids, (..., 0), decoder_start_token_id)
+    # replace possible -100 values in labels by `pad_token_id`
+    shifted_input_ids = jnp.where(shifted_input_ids == -100, pad_token_id, shifted_input_ids)
+
+    return shifted_input_ids
+
+decoder_input_ids = shift_tokens_right(
+    jnp.array(labels["input_ids"]),
+    model.config.text_config.pad_token_id,
+    model.config.decoder_start_token_id
+)
+decoder_input_ids = np.asarray(decoder_input_ids)
+# We need decoder_attention_mask so we can ignore pad tokens from loss
+decoder_attention_mask = labels["attention_mask"]
+
+print(f'decoder_inputs = {decoder_input_ids}')
+print(f'decoder_input_ids.shape = {decoder_input_ids.shape}')
+print(f'decoder_attention_mask = {decoder_attention_mask}')
+print(f'decoder_attention_mask.shape = {decoder_attention_mask.shape}')
 
 # model data
-inputs = dict(decoder_inputs)
-inputs['pixel_values'] = pixel_values
-
-
-logits = flax_vit_gpt2_lm(**inputs)[0]
+model_inputs = {
+    'pixel_values': pixel_values,
+    'attention_mask': None,
+    'decoder_input_ids': decoder_input_ids,
+    'decoder_attention_mask': decoder_attention_mask,
+    'decoder_position_ids': None,
+}
+
+# Model call
+model_outputs = flax_vit_gpt2_lm(**model_inputs)
+logits = model_outputs[0]
 preds = np.argmax(logits, axis=-1)
 print('=' * 60)
 print('Flax: Vit-GPT2-LM')
@@ -59,16 +92,21 @@ print('predicted token ids:')
 print(preds)
 print('=' * 60)
 
+# encoder_last_hidden_state = model_outputs['encoder_last_hidden_state']
+# print(encoder_last_hidden_state)
+# encoder_kwargs = {}
+# encoder_outputs = flax_vit_gpt2_lm.encode(pixel_values, return_dict=True, **encoder_kwargs)
+# print(encoder_outputs['last_hidden_state'])
 
 # Generation!
+num_beams = 1
+gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
+
 batch = {'pixel_values': pixel_values}
-generation = flax_vit_gpt2_lm.generate(batch['pixel_values'], **gen_kwargs)
-print('generation:')
-print(generation)
+generated = flax_vit_gpt2_lm.generate(batch['pixel_values'], **gen_kwargs)
+token_ids = np.array(generated.sequences)[0]
+print(f'generated token ids: {token_ids}')
 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(f'generated caption: {caption}')
 print('=' * 60)