Upload modeling_esm_plusplus.py with huggingface_hub

modeling_esm_plusplus.py

@@ -711,7 +711,7 @@ class EmbeddingMixin:
                         seqs = to_embed[i * batch_size:(i + 1) * batch_size]
                         input_ids, attention_mask = batch['input_ids'].to(device), batch['attention_mask'].to(device)
                         residue_embeddings = self._embed(input_ids, attention_mask).float() # sql requires float32
-                        embeddings = get_embeddings(residue_embeddings, attention_mask).cpu()
+                        embeddings = get_embeddings(residue_embeddings, attention_mask)
                         for seq, emb, mask in zip(seqs, embeddings, attention_mask):
                             if full_embeddings:
                                 emb = emb[mask.bool()].reshape(-1, hidden_size)
@@ -743,11 +743,11 @@ class EmbeddingMixin:
                     seqs = to_embed[i * batch_size:(i + 1) * batch_size]
                     input_ids, attention_mask = batch['input_ids'].to(device), batch['attention_mask'].to(device)
                     residue_embeddings = self._embed(input_ids, attention_mask)
-                    embeddings = get_embeddings(residue_embeddings, attention_mask).to(embed_dtype).cpu()
+                    embeddings = get_embeddings(residue_embeddings, attention_mask).to(embed_dtype)
                     for seq, emb, mask in zip(seqs, embeddings, attention_mask):
                         if full_embeddings:
                             emb = emb[mask.bool()].reshape(-1, hidden_size)
-                        embeddings_dict[seq] = emb
+                        embeddings_dict[seq] = emb.cpu()
 
         if save:
             torch.save(embeddings_dict, save_path)