Upload modeling_e1.py with huggingface_hub

modeling_e1.py

@@ -377,6 +377,11 @@ class Pooler:
             attention_mask: Optional[torch.Tensor] = None,
             attentions: Optional[torch.Tensor] = None
         ) -> torch.Tensor:
+        if attention_mask is not None:
+            assert attention_mask.sum(dim=-1).min() > 0, (
+                "Pooler received samples with all-zero attention masks. "
+                "This causes NaN from division by zero. Filter empty inputs before pooling."
+            )
         final_emb: List[torch.Tensor] = []
         for pooling_type in self.pooling_types:
             final_emb.append(self.pooling_options[pooling_type](emb=emb, attention_mask=attention_mask, attentions=attentions))
@@ -2398,7 +2403,6 @@ class FAST_E1_ENCODER(E1PreTrainedModel, EmbeddingMixin):
     def set_input_embeddings(self, value: nn.Embedding) -> None:
         self.embed_tokens = value
 
-    @torch.inference_mode()
     def _embed(self, sequences: List[str], return_attention_mask: bool = False, **kwargs) -> torch.Tensor:
         batch = self.prep_tokens.get_batch_kwargs(sequences, device=self._device)
         last_hidden_state = self.forward(**batch, output_hidden_states=False, output_attentions=False).last_hidden_state
@@ -2602,7 +2606,6 @@ class E1Model(E1PreTrainedModel, EmbeddingMixin):
     def set_input_embeddings(self, value: nn.Embedding) -> None:
         self.model.set_input_embeddings(value)
 
-    @torch.inference_mode()
     def _embed(self, sequences: List[str], return_attention_mask: bool = False, **kwargs) -> torch.Tensor:
         return self.model._embed(sequences, return_attention_mask=return_attention_mask, **kwargs)
 
@@ -2656,7 +2659,6 @@ class E1ForMaskedLM(E1PreTrainedModel, EmbeddingMixin):
     def device_mesh(self) -> torch.distributed.device_mesh.DeviceMesh:
         return self.model.device_mesh
 
-    @torch.inference_mode()
     def _embed(self, sequences: List[str], return_attention_mask: bool = False, **kwargs) -> torch.Tensor:
         batch = self.prep_tokens.get_batch_kwargs(sequences, device=self._device)
         last_hidden_state = self.model(**batch, output_hidden_states=False, output_attentions=False).last_hidden_state
@@ -2778,7 +2780,6 @@ class E1ForSequenceClassification(E1PreTrainedModel, EmbeddingMixin):
     def device_mesh(self) -> torch.distributed.device_mesh.DeviceMesh:
         return self.model.device_mesh
 
-    @torch.inference_mode()
     def _embed(self, sequences: List[str], return_attention_mask: bool = False, **kwargs) -> torch.Tensor:
         batch = self.prep_tokens.get_batch_kwargs(sequences, device=self._device)
         last_hidden_state = self.model(**batch, output_hidden_states=False, output_attentions=False).last_hidden_state
@@ -2875,7 +2876,6 @@ class E1ForTokenClassification(E1PreTrainedModel, EmbeddingMixin):
     def device_mesh(self) -> torch.distributed.device_mesh.DeviceMesh:
         return self.model.device_mesh
 
-    @torch.inference_mode()
     def _embed(self, sequences: List[str], return_attention_mask: bool = False, **kwargs) -> torch.Tensor:
         batch = self.prep_tokens.get_batch_kwargs(sequences, device=self._device)
         last_hidden_state = self.model(**batch, output_hidden_states=False, output_attentions=False).last_hidden_state