Add a test against pytorch's GPT2

tests/test_load.py

@@ -1,48 +0,0 @@
-import sys, os
-
-current_path = os.path.dirname(os.path.abspath(__file__))
-sys.path.append(current_path)
-
-# Vit - as encoder
-from transformers import ViTFeatureExtractor
-from PIL import Image
-import requests
-import numpy as np
-
-url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
-image = Image.open(requests.get(url, stream=True).raw)
-
-feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')
-encoder_inputs = feature_extractor(images=image, return_tensors="jax")
-pixel_values = encoder_inputs.pixel_values
-
-# GPT2 / GPT2LM - as decoder
-from transformers import ViTFeatureExtractor, GPT2Tokenizer
-
-name = 'asi/gpt-fr-cased-small'
-tokenizer = GPT2Tokenizer.from_pretrained(name)
-decoder_inputs = tokenizer("mon chien est mignon", return_tensors="jax")
-
-inputs = dict(decoder_inputs)
-inputs['pixel_values'] = pixel_values
-print(inputs)
-
-
-
-
-
-# With the LM head in GPT2LM
-from vit_gpt2.modeling_flax_vit_gpt2_lm import FlaxViTGPT2LMForConditionalGeneration
-flax_vit_gpt2_lm = FlaxViTGPT2LMForConditionalGeneration.from_pretrained(
-    '.',
-)
-
-logits = flax_vit_gpt2_lm(**inputs)[0]
-preds = np.argmax(logits, axis=-1)
-print('=' * 60)
-print('Flax: Vit + modified GPT2LM')
-print(preds)
-
-# flax_vit_gpt2_lm.save_pretrained('.')
-
-del flax_vit_gpt2_lm

tests/test_model.py

@@ -22,6 +22,9 @@ from transformers import GPT2Tokenizer
 
 max_length = 8
 
+# ================================================================================
+# Models preparation
+
 vision_model_name = 'google/vit-base-patch16-224-in21k'
 text_model_name = 'asi/gpt-fr-cased-small'
 
@@ -34,6 +37,9 @@ flax_vit_gpt2_lm = FlaxViTGPT2LMForConditionalGeneration.from_vision_text_pretra
 )
 model = flax_vit_gpt2_lm
 
+# ================================================================================
+# Inputs preparation
+
 feature_extractor = ViTFeatureExtractor.from_pretrained(vision_model_name)
 tokenizer = GPT2Tokenizer.from_pretrained(text_model_name)
 
@@ -56,6 +62,7 @@ sentence += ' ' + tokenizer.eos_token
 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.
@@ -82,6 +89,9 @@ 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}')
 
+# ================================================================================
+# Check `FlaxGPT2LMHeadModel` has the same results in the new version (when no `encoder_outputs` is provided).
+
 orig_gpt2_lm = Orig_FlaxGPT2LMHeadModel.from_pretrained(text_model_name)
 gpt2_lm = FlaxGPT2LMHeadModel.from_pretrained(text_model_name)
 
@@ -108,6 +118,8 @@ print(f'GPT2 generated caption: {caption}')
 
 assert list(orig_token_ids) == list(token_ids)
 
+# ================================================================================
+
 # model data
 model_inputs = {
     'pixel_values': pixel_values,
@@ -117,6 +129,9 @@ model_inputs = {
     'decoder_position_ids': None,
 }
 
+# ================================================================================
+# Check `model.__call__()`
+
 # Model call
 model_outputs = model(**model_inputs)
 logits = model_outputs[0]
@@ -133,6 +148,9 @@ print(preds)
 # encoder_outputs = flax_vit_gpt2_lm.encode(pixel_values, return_dict=True, **encoder_kwargs)
 # print(encoder_outputs['last_hidden_state'])
 
+# ================================================================================
+# Check generation
+
 # Generation!
 num_beams = 1
 gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
@@ -149,6 +167,9 @@ caption = tokenizer.decode(token_ids)
 print('=' * 60)
 print(f'generated caption: {caption}')
 
+# ================================================================================
+# Check save & load
+
 # save
 os.makedirs('./model/', exist_ok=True)
 model.save_pretrained(save_directory='./model/')
@@ -163,3 +184,29 @@ _token_ids = np.array(_generated.sequences)[0]
 print('=' * 60)
 print(f'new generated token ids: {_token_ids}')
 print(f'token_ids == new_token_ids: {token_ids == _token_ids}')
+
+# ================================================================================
+# Check PyTorch version's output - it should be the same as above
+
+import torch
+from transformers import ViTModel, GPT2Config, GPT2LMHeadModel
+
+vision_model_pt = ViTModel.from_pretrained(vision_model_name)
+config = GPT2Config.from_pretrained(text_model_name)
+config.is_encoder_decoder = True
+config.add_cross_attention = True
+text_model_pt = GPT2LMHeadModel.from_pretrained(text_model_name, config=config)
+
+encoder_inputs_pt = feature_extractor(images=image, return_tensors="pt")
+vision_model_pt_outputs = vision_model_pt(**encoder_inputs)
+
+generated = text_model_pt.generate(encoder_outputs=vision_model_pt_outputs, **gen_kwargs)
+token_ids = np.array(generated.sequences)[0]
+
+print('=' * 60)
+print(f'Pytorch\'s GPT2 LM generated token ids: {token_ids}')
+
+caption = tokenizer.decode(token_ids)
+
+print('=' * 60)
+print(f'Pytorch\'s GPT2 LM generated caption: {caption}')

tests/test_save.py

@@ -1,48 +0,0 @@
-import sys, os
-
-current_path = os.path.dirname(os.path.abspath(__file__))
-sys.path.append(current_path)
-
-# Vit - as encoder
-from transformers import ViTFeatureExtractor
-from PIL import Image
-import requests
-import numpy as np
-
-url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
-image = Image.open(requests.get(url, stream=True).raw)
-
-feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')
-encoder_inputs = feature_extractor(images=image, return_tensors="jax")
-pixel_values = encoder_inputs.pixel_values
-
-# GPT2 / GPT2LM - as decoder
-from transformers import ViTFeatureExtractor, GPT2Tokenizer
-
-name = 'asi/gpt-fr-cased-small'
-tokenizer = GPT2Tokenizer.from_pretrained(name)
-decoder_inputs = tokenizer("mon chien est mignon", return_tensors="jax")
-
-inputs = dict(decoder_inputs)
-inputs['pixel_values'] = pixel_values
-print(inputs)
-
-
-
-
-
-# With the LM head in GPT2LM
-from vit_gpt2.modeling_flax_vit_gpt2_lm import FlaxViTGPT2LMForConditionalGeneration
-flax_vit_gpt2_lm = FlaxViTGPT2LMForConditionalGeneration.from_vit_gpt2_pretrained(
-    'google/vit-base-patch16-224-in21k', 'asi/gpt-fr-cased-small'
-)
-
-logits = flax_vit_gpt2_lm(**inputs)[0]
-preds = np.argmax(logits, axis=-1)
-print('=' * 60)
-print('Flax: Vit + modified GPT2LM')
-print(preds)
-
-flax_vit_gpt2_lm.save_pretrained('.')
-
-del flax_vit_gpt2_lm