Update modeling_prismatic.py

modeling_prismatic.py

@@ -1,12 +1,9 @@
 """
 modeling_prismatic.py
-
 Core HuggingFace-style PrismaticPreTrainedModel and PrismaticForConditionalGeneration class definitions, inheriting
 from the default `transformers.PretrainedModel`. Meant to be standalone and self-contained, but exactly replicate the
 logic in `prismatic.models.vlms.prismatic.py`.
-
 Note =>> for the time being, not adding the custom HF "docstring" formatting.
-
 References [LLaVa, IDEFICS-2]:
     => https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava/modeling_llava.py
     => https://github.com/huggingface/transformers/blob/main/src/transformers/models/idefics2/modeling_idefics2.py
@@ -411,7 +408,7 @@ class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
                 use_cache=use_cache,
                 output_attentions=True,
                 output_hidden_states=output_hidden_states,
-                return_dict=return_dict,
+                return_dict=True,
             )
 
         # === Otherwise =>> Assume Invalid! ===
@@ -437,9 +434,6 @@ class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
 
             return language_model_output
 
-        print("Forward")
-        
-        
         return PrismaticCausalLMOutputWithPast(
             loss=language_model_output.loss,
             logits=language_model_output.logits,
@@ -485,11 +479,6 @@ class PrismaticForConditionalGeneration(PrismaticPreTrainedModel):
             }
         )
 
-        # model_inputs["output_attentions"] = True
-
-        print("Prepare")
-        print(model_inputs.keys())
-
         return model_inputs
 
     # Defer to Language Model (all handle this differently, with different return types)
@@ -523,10 +512,8 @@ class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):
             )
 
         # Run VLA inference
-        print("=" * 100)
-        model_outputs = self.generate(input_ids, max_new_tokens=self.get_action_dim(unnorm_key), return_dict=True, **kwargs)
-        print(model_outputs.shape)
-        print("=" * 100)
+        model_outputs = self.generate(input_ids, max_new_tokens=self.get_action_dim(unnorm_key), **kwargs)
+        print(model_outputs.keys())
 
         # Extract predicted action tokens and translate into (normalized) continuous actions
         predicted_action_token_ids = generated_ids[0, -self.get_action_dim(unnorm_key) :].cpu().numpy()
@@ -544,7 +531,7 @@ class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):
             normalized_actions,
         )
 
-        return actions, generated_ids
+        return actions
 
     @staticmethod
     def _check_unnorm_key(norm_stats: Dict[str, Dict[str, Any]], unnorm_key: Optional[str]) -> str:
@@ -570,6 +557,5 @@ class OpenVLAForActionPrediction(PrismaticForConditionalGeneration):
     def get_action_stats(self, unnorm_key: Optional[str] = None) -> Dict[str, Any]:
         """Get all the logged statistics for the given dataset."""
         unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)
-        print("keys:", self.norm_stats.keys())
-        print("items:", self.norm_stats)
         return self.norm_stats[unnorm_key]["action"]
+