Add inputs_embeds argument

modeling_hf_nomic_bert.py

@@ -977,14 +977,18 @@ class NomicBertEmbeddings(nn.Module):
         if self.type_vocab_size > 0:
             self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
 
-    def forward(self, input_ids, position_ids=None, token_type_ids=None):
+    def forward(self, input_ids=None, position_ids=None, token_type_ids=None, inputs_embeds=None):
         """
         input_ids: (batch, seqlen)
         position_ids: (batch, seqlen)
         token_type_ids: (batch, seqlen)
         """
         batch_size, seqlen = input_ids.shape
-        embeddings = self.word_embeddings(input_ids)
+
+        if inputs_embeds is None:
+            embeddings = self.word_embeddings(input_ids)
+        else:
+            embeddings = inputs_embeds
 
         if self.type_vocab_size > 0:
             if token_type_ids is None:
@@ -1680,10 +1684,18 @@ class NomicBertModel(NomicBertPreTrainedModel):
         token_type_ids=None,
         return_dict=None,
         matryoshka_dim=None,
+        inputs_embeds=None,
     ):
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
         if token_type_ids is None:
             token_type_ids = torch.zeros_like(input_ids)
-        hidden_states = self.embeddings(input_ids, position_ids=position_ids, token_type_ids=token_type_ids)
+        hidden_states = self.embeddings(
+            input_ids=input_ids, 
+            position_ids=position_ids, 
+            token_type_ids=token_type_ids,
+            inputs_embeds=inputs_embeds,
+        )
         hidden_states = self.emb_ln(hidden_states)
         hidden_states = self.emb_drop(hidden_states)