Update test_model.py

tests/test_model.py

@@ -95,7 +95,7 @@ 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!
+# generation!
 num_beams = 1
 gen_kwargs = {"max_length": 6, "num_beams": num_beams}
 
@@ -138,20 +138,27 @@ logits = model_outputs[0]
 preds = np.argmax(logits, axis=-1)
 
 print('=' * 60)
-print('Flax: Vit-GPT2-LM')
-print('predicted token ids:')
+print('Flax ViT-GPT2-LM - predicted token ids:')
 print(preds)
 
-# 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'])
+encoder_last_hidden_state = model_outputs['encoder_last_hidden_state']
+print('=' * 60)
+print("encoder_last_hidden_state given by model.__call__():")
+print(encoder_last_hidden_state)
+
+encoder_outputs = model.encode(pixel_values, return_dict=True)
+print('=' * 60)
+print("encoder's last_hidden_state given by model.encode():")
+print(encoder_outputs['last_hidden_state'])
+
+total_diff = np.sum(np.abs(encoder_outputs['last_hidden_state'] - encoder_last_hidden_state))
+print('=' * 60)
+print(f"total difference: {total_diff}")
 
 # ================================================================================
-# Check generation
+# Check model generation
 
-# Generation!
+# generation
 num_beams = 1
 gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
 
@@ -215,17 +222,19 @@ logits = text_model_pt_outputs[0]
 preds = np.argmax(logits.detach().numpy(), axis=-1)
 
 print('=' * 60)
-print('PyTroch: Vit --> GPT2-LM')
+print('PyTroch: ViT --> GPT2-LM')
 print('predicted token ids:')
 print(preds)
 
-#generated = text_model_pt.generate(encoder_outputs=vision_model_pt_outputs, **gen_kwargs)
-#token_ids = np.array(generated.sequences)[0]
+model_logits = np.array(model_outputs.logits)
+text_model_pt_logits = text_model_pt_outputs.logits.detach().cpu().numpy()
+total_diff = np.sum(np.abs(model_logits - text_model_pt_logits))
 
-#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}')
+print('=' * 60)
+print("model_logits:")
+print(model_logits)
+print('=' * 60)
+print("text_model_pt_logits:")
+print(text_model_pt_logits)
+print('=' * 60)
+print(f"total difference between logits: {total_diff}")