Upload modeling_fastesm.py with huggingface_hub

modeling_fastesm.py

@@ -556,10 +556,6 @@ class FastEsmPreTrainedModel(PreTrainedModel):
         Returns:
             Dictionary mapping sequences to embeddings, or None if sql=True
         """
-        sequences = list(set([seq[:max_len] for seq in sequences]))
-        sequences = sorted(sequences, key=len, reverse=True)
-        dataset = ProteinDataset(sequences)
-        dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers, collate_fn=self._collate_fn, shuffle=False)
         device = self.device
 
         def get_embeddings(residue_embeddings: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
@@ -570,6 +566,7 @@ class FastEsmPreTrainedModel(PreTrainedModel):
             else:
                 return residue_embeddings[:, 0, :]
 
+        sequences = list(set([seq[:max_len] for seq in sequences]))
         if sql:
             import sqlite3
             conn = sqlite3.connect(sql_db_path)
@@ -580,6 +577,9 @@ class FastEsmPreTrainedModel(PreTrainedModel):
             print(f"Found {len(already_embedded)} already embedded sequences in {sql_db_path}")
             print(f"Embedding {len(to_embed)} new sequences")
             if len(to_embed) > 0:
+                to_embed = sorted(to_embed, key=len, reverse=True)
+                dataset = ProteinDataset(to_embed)
+                dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers, collate_fn=self._collate_fn, shuffle=False)
                 with torch.no_grad():
                     for i, batch in tqdm(enumerate(dataloader), total=len(dataloader), desc='Embedding batches'):
                         seqs = sequences[i * batch_size:(i + 1) * batch_size]
@@ -598,6 +598,10 @@ class FastEsmPreTrainedModel(PreTrainedModel):
             conn.close()
             return None
             
+        sequences = list(set([seq[:max_len] for seq in sequences]))
+        sequences = sorted(sequences, key=len, reverse=True)
+        dataset = ProteinDataset(sequences)
+        dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers, collate_fn=self._collate_fn, shuffle=False)
         embeddings_dict = {}
         with torch.no_grad():
             for i, batch in tqdm(enumerate(dataloader), total=len(dataloader), desc='Embedding batches'):