add full support for inputs_embeds (#10)

- add full support for inputs_embeds (2fd43c9a0641a75fa975f3257d97e5c55b3fa940)


Co-authored-by: Jack Morris <jxm@users.noreply.huggingface.co>

modeling_hf_nomic_bert.py

@@ -983,22 +983,21 @@ class NomicBertEmbeddings(nn.Module):
         position_ids: (batch, seqlen)
         token_type_ids: (batch, seqlen)
         """
-        batch_size, seqlen = input_ids.shape
-
         if inputs_embeds is None:
             embeddings = self.word_embeddings(input_ids)
         else:
             embeddings = inputs_embeds
-
+        batch_size, seqlen, _ = embeddings.shape
+        
         if self.type_vocab_size > 0:
             if token_type_ids is None:
-                token_type_ids = torch.zeros(seqlen, dtype=torch.long, device=input_ids.device)
+                token_type_ids = torch.zeros(seqlen, dtype=torch.long, device=embeddings.device)
             token_type_embeddings = self.token_type_embeddings(token_type_ids)
             embeddings = embeddings + token_type_embeddings
 
         if self.max_position_embeddings > 0:
             if position_ids is None:
-                position_ids = torch.arange(seqlen, dtype=torch.long, device=input_ids.device)
+                position_ids = torch.arange(seqlen, dtype=torch.long, device=embeddings.device)
             position_embeddings = self.position_embeddings(position_ids)
             embeddings = embeddings + position_embeddings
         return embeddings
@@ -1688,8 +1687,6 @@ class NomicBertModel(NomicBertPreTrainedModel):
     ):
         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=input_ids, 
             position_ids=position_ids, 
@@ -1699,7 +1696,7 @@ class NomicBertModel(NomicBertPreTrainedModel):
         hidden_states = self.emb_ln(hidden_states)
         hidden_states = self.emb_drop(hidden_states)
 
-        attention_mask = self.get_extended_attention_mask(attention_mask, input_ids.shape)
+        attention_mask = self.get_extended_attention_mask(attention_mask, hidden_states.shape[:-1])
         sequence_output = self.encoder(hidden_states, attention_mask=attention_mask, return_dict=return_dict)
 
         pooled_output = self.pooler(sequence_output) if self.pooler is not None else None