Upload modeling_esm_plusplus.py with huggingface_hub

modeling_esm_plusplus.py

@@ -762,8 +762,7 @@ class EmbeddingMixin:
                         for seq, emb, mask in zip(seqs, embeddings, attention_mask):
                             if full_embeddings:
                                 emb = emb[mask.bool()].reshape(-1, hidden_size)
-                            c.execute("INSERT OR REPLACE INTO embeddings VALUES (?, ?)", 
-                                    (seq, emb.cpu().numpy().tobytes()))
+                            c.execute("INSERT OR REPLACE INTO embeddings VALUES (?, ?)", (seq, emb.cpu().numpy().tobytes()))
                         
                         if (i + 1) % 100 == 0:
                             conn.commit()
@@ -979,7 +978,12 @@ class ESMplusplusForSequenceClassification(ESMplusplusForMaskedLM, EmbeddingMixi
         self.mse = nn.MSELoss()
         self.ce = nn.CrossEntropyLoss()
         self.bce = nn.BCEWithLogitsLoss()
-        self.pooler = Pooler(['cls','mean'])
+        # if kwargs has pooling_types, use them, otherwise use ['cls', 'mean']
+        if 'pooling_types' in kwargs and isinstance(kwargs['pooling_types'], List[str]) and len(kwargs['pooling_types']) > 0:
+            pooling_types = kwargs['pooling_types']
+        else:
+            pooling_types = ['cls', 'mean']
+        self.pooler = Pooler(pooling_types)
         self.init_weights()
 
     def _embed(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor: