update test_model.py

tests/test_model.py

@@ -3,6 +3,9 @@ import sys, os
 current_path = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(current_path)
 
+from transformers import FlaxGPT2LMHeadModel as Orig_FlaxGPT2LMHeadModel
+from vit_gpt2.modeling_flax_gpt2 import FlaxGPT2LMHeadModel
+
 # Main model -  ViTGPT2LM
 from vit_gpt2.modeling_flax_vit_gpt2_lm import FlaxViTGPT2LMForConditionalGeneration
 
@@ -37,6 +40,8 @@ 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('=' * 60)
 print(f'pixel_values.shape = {pixel_values.shape}')
 
 # decoder data
@@ -68,11 +73,36 @@ 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('=' * 60)
 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}')
 
+orig_gpt2_lm = Orig_FlaxGPT2LMHeadModel.from_pretrained(text_model_name)
+gpt2_lm = FlaxGPT2LMHeadModel.from_pretrained(text_model_name)
+
+# Generation!
+num_beams = 1
+gen_kwargs = {"max_length": 6, "num_beams": num_beams}
+
+orig_gpt2_generated = orig_gpt2_lm.generate(decoder_input_ids[:, 0:3], **gen_kwargs)
+gpt2_generated = gpt2_lm.generate(decoder_input_ids[:, 0:3], **gen_kwargs)
+
+orig_token_ids = np.array(orig_gpt2_generated.sequences)[0]
+token_ids = np.array(gpt2_generated.sequences)[0]
+
+orig_caption = tokenizer.decode(orig_token_ids)
+caption = tokenizer.decode(token_ids)
+
+print('=' * 60)
+print(f'orig. GPT2 generated token ids: {orig_token_ids}')
+print(f'GPT2 generated token ids: {token_ids}')
+
+print('=' * 60)
+print(f'orig. GPT2 generated caption: {orig_caption}')
+print(f'GPT2 generated caption: {caption}')
+
 # model data
 model_inputs = {
     'pixel_values': pixel_values,
@@ -83,14 +113,14 @@ model_inputs = {
 }
 
 # Model call
-model_outputs = flax_vit_gpt2_lm(**model_inputs)
+model_outputs = model(**model_inputs)
 logits = model_outputs[0]
 preds = np.argmax(logits, axis=-1)
+
 print('=' * 60)
 print('Flax: Vit-GPT2-LM')
 print('predicted token ids:')
 print(preds)
-print('=' * 60)
 
 # encoder_last_hidden_state = model_outputs['encoder_last_hidden_state']
 # print(encoder_last_hidden_state)
@@ -103,10 +133,28 @@ num_beams = 1
 gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
 
 batch = {'pixel_values': pixel_values}
-generated = flax_vit_gpt2_lm.generate(batch['pixel_values'], **gen_kwargs)
+generated = model.generate(batch['pixel_values'], **gen_kwargs)
 token_ids = np.array(generated.sequences)[0]
-print(f'generated token ids: {token_ids}')
+
 print('=' * 60)
+print(f'generated token ids: {token_ids}')
+
 caption = tokenizer.decode(token_ids)
+
+print('=' * 60)
 print(f'generated caption: {caption}')
+
+# save
+os.makedirs('./model/', exist_ok=True)
+model.save_pretrained(save_directory='./model/')
+
+# load
+_model = FlaxViTGPT2LMForConditionalGeneration.from_pretrained('./model/')
+
+# check if the result is the same as before
+_generated = _model.generate(batch['pixel_values'], **gen_kwargs)
+_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}')