addressing review comments

Signed-off-by: Peter St. John <pstjohn@nvidia.com>

Author: pstjohn

models/esm2/src/esm/collator.py

@@ -171,9 +171,17 @@ def __call__(self, features, return_tensors=None):
         batch = self.flattening_collator(features, return_tensors)
 
         special_tokens_mask = batch.pop("special_tokens_mask", None)
-        batch["input_ids"], batch["labels"] = self.mlm_collator.torch_mask_tokens(
-            batch["input_ids"], special_tokens_mask=special_tokens_mask
-        )
+
+        if return_tensors == "pt":
+            batch["input_ids"], batch["labels"] = self.mlm_collator.torch_mask_tokens(
+                batch["input_ids"], special_tokens_mask=special_tokens_mask
+            )
+        elif return_tensors == "np":
+            batch["input_ids"], batch["labels"] = self.mlm_collator.numpy_mask_tokens(
+                batch["input_ids"], special_tokens_mask=special_tokens_mask
+            )
+        else:
+            raise ValueError(f'return_tensors must be one of ("pt", "np"), {return_tensors=} not suported')
 
         return batch
 

models/esm2/src/esm/modeling_esm_te.py

@@ -176,33 +176,39 @@ def forward(
                 raise ValueError(
                     "cu_seq_lens_q, cu_seq_lens_k, max_length_q, and max_length_k must be provided when using THD inputs."
                 )
+            assert hidden_states.dim() == 3 and hidden_states.size(0) == 1, (
+                "THD expects embeddings shaped [1, total_tokens, hidden_size]."
+            )
+            hidden_states = hidden_states.squeeze(0)
 
         elif self.config.attn_input_format == "bshd":
             if any(x is not None for x in [cu_seq_lens_q, cu_seq_lens_k, max_length_q, max_length_k]):
                 raise ValueError(
                     "cu_seq_lens_q, cu_seq_lens_k, max_length_q, and max_length_k are not allowed when using BSHD inputs."
                 )
 
-        if self.config.attn_input_format == "bshd" and self.te_rope_emb is not None:
-            te_rope_emb = self.te_rope_emb.to(
-                device=hidden_states.device, dtype=hidden_states.dtype, non_blocking=True
-            )
-            seq_len = hidden_states.shape[1]
-            if te_rope_emb.size(0) < seq_len:
-                raise RuntimeError(
-                    f"ROPE length {te_rope_emb.size(0)} < input seq length {seq_len}. "
-                    f"Increase max_position_embeddings."
+        te_rope_emb = None
+        if self.config.position_embedding_type == "rotary":
+            if self.config.attn_input_format == "bshd":
+                te_rope_emb = self.te_rope_emb.to(
+                    device=hidden_states.device, dtype=hidden_states.dtype, non_blocking=True
+                )
+                seq_len = hidden_states.shape[1]
+                if te_rope_emb.size(0) < seq_len:
+                    raise RuntimeError(
+                        f"ROPE length {te_rope_emb.size(0)} < input seq length {seq_len}. "
+                        f"Increase max_position_embeddings."
+                    )
+                te_rope_emb = te_rope_emb[:seq_len]
+
+            elif self.config.attn_input_format == "thd":
+                assert cu_seq_lens_q is not None
+                te_rope_emb = self.rotary_embeddings(max_seq_len=cu_seq_lens_q[-1]).to(
+                    device=hidden_states.device, dtype=hidden_states.dtype, non_blocking=True
                 )
-            te_rope_emb = te_rope_emb[:seq_len]
 
-        elif self.config.attn_input_format == "thd":
-            assert cu_seq_lens_q is not None
-            te_rope_emb = self.rotary_embeddings(max_seq_len=cu_seq_lens_q[-1]).to(
-                device=hidden_states.device, dtype=hidden_states.dtype, non_blocking=True
-            )
-            hidden_states = hidden_states.squeeze(0)
-        else:
-            te_rope_emb = None
+            else:
+                raise ValueError(f"Unsupported attention input format: {self.config.attn_input_format}")
 
         for layer_module in self.layers:
             if output_hidden_states: