Add print statements

modeling_cogvlm.py

@@ -456,6 +456,7 @@ class CogVLMModel(CogVLMPreTrainedModel):
             output_attentions: Optional[bool] = None,
             output_hidden_states: Optional[bool] = None,
             return_dict: Optional[bool] = None,
+            step: int = None,
     ) -> Union[Tuple, BaseModelOutputWithPast]:
         """take care of image_encode, token_type_ids, position_ids and (attention_mask = None is fine)"""
 
@@ -527,6 +528,7 @@ class CogVLMModel(CogVLMPreTrainedModel):
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
+            step=step,
         )
 
     def llm_forward(
@@ -541,6 +543,7 @@ class CogVLMModel(CogVLMPreTrainedModel):
             output_attentions: Optional[bool] = None,
             output_hidden_states: Optional[bool] = None,
             return_dict: Optional[bool] = None,
+            step: int = None,
     ) -> Union[Tuple, BaseModelOutputWithPast]:
         """largely copy from llama forward and adapt for cogvlm with `token_type_ids`"""
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
@@ -590,41 +593,42 @@ class CogVLMModel(CogVLMPreTrainedModel):
 
         hidden_states = inputs_embeds
 
-        # torch.save(hidden_states, "initial_hidden_states.pt")
-        # torch.save(attention_mask, "initial_attention_mask.pt")
-        # torch.save(token_type_ids, "initial_token_type_ids.pt")
-        # torch.save(position_ids, "initial_position_ids.pt")
-        
-        # from huggingface_hub import HfApi
-
-        # api = HfApi()
-        # api.upload_file(
-        #     path_or_fileobj="initial_hidden_states.pt",
-        #     path_in_repo="initial_hidden_states.pt",
-        #     repo_id="nielsr/test-cogvlm",
-        #     repo_type="dataset",
-        # )
-        # api = HfApi()
-        # api.upload_file(
-        #     path_or_fileobj="initial_attention_mask.pt",
-        #     path_in_repo="initial_attention_mask.pt",
-        #     repo_id="nielsr/test-cogvlm",
-        #     repo_type="dataset",
-        # )
-        # api = HfApi()
-        # api.upload_file(
-        #     path_or_fileobj="initial_token_type_ids.pt",
-        #     path_in_repo="initial_token_type_ids.pt",
-        #     repo_id="nielsr/test-cogvlm",
-        #     repo_type="dataset",
-        # )
-        # api = HfApi()
-        # api.upload_file(
-        #     path_or_fileobj="initial_position_ids.pt",
-        #     path_in_repo="initial_position_ids.pt",
-        #     repo_id="nielsr/test-cogvlm",
-        #     repo_type="dataset",
-        # )
+        if step == 1:
+            torch.save(hidden_states, "hidden_states_step_1.pt")
+            torch.save(attention_mask, "attention_mask_step_1.pt")
+            torch.save(token_type_ids, "token_type_ids_step_1.pt")
+            torch.save(position_ids, "position_ids_step_1.pt")
+            
+            from huggingface_hub import HfApi
+
+            api = HfApi()
+            api.upload_file(
+                path_or_fileobj="hidden_states_step_1.pt",
+                path_in_repo="hidden_states_step_1.pt",
+                repo_id="nielsr/test-cogvlm",
+                repo_type="dataset",
+            )
+            api = HfApi()
+            api.upload_file(
+                path_or_fileobj="attention_mask_step_1.pt",
+                path_in_repo="attention_mask_step_1.pt",
+                repo_id="nielsr/test-cogvlm",
+                repo_type="dataset",
+            )
+            api = HfApi()
+            api.upload_file(
+                path_or_fileobj="token_type_ids_step_1.pt",
+                path_in_repo="token_type_ids_step_1.pt",
+                repo_id="nielsr/test-cogvlm",
+                repo_type="dataset",
+            )
+            api = HfApi()
+            api.upload_file(
+                path_or_fileobj="position_ids_step_1.pt",
+                path_in_repo="position_ids_step_1.pt",
+                repo_id="nielsr/test-cogvlm",
+                repo_type="dataset",
+            )
 
         # decoder layers
         all_hidden_states = () if output_hidden_states else None
@@ -774,6 +778,7 @@ class CogVLMForCausalLM(CogVLMPreTrainedModel):
             output_hidden_states: Optional[bool] = None,
             return_dict: Optional[bool] = None,
             labels: Optional[torch.LongTensor] = None,
+            step: int = None,
     ) -> Union[Tuple, CausalLMOutputWithPast]:
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
@@ -794,6 +799,7 @@ class CogVLMForCausalLM(CogVLMPreTrainedModel):
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
+            step=step,
         )
 
         hidden_states = outputs[0]