Upload esm3bedding.py with huggingface_hub

esm3bedding.py

@@ -0,0 +1,86 @@
+# esm3bedding.py
+
+import os
+import torch
+from esm.models.esmc import ESMC
+from esm.sdk.api import ESMProtein, LogitsConfig
+from huggingface_hub import login
+from utils import get_logger
+from base import Featurizer
+
+logg = get_logger()
+
+class ESM3Featurizer(Featurizer):
+    def __init__(self, save_dir: str, api_key: str, per_tok: bool = True):
+        super().__init__("ESM3", 1152, save_dir=save_dir)
+        self.per_tok = per_tok
+        self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.client = None
+
+        self._login(api_key)
+        self._initialize_model()
+
+    def _login(self, api_key: str):
+        try:
+            login(api_key)
+            logg.info("Successfully logged into Hugging Face Hub.")
+        except Exception as e:
+            logg.error(f"Failed to log in to Hugging Face Hub: {e}")
+            raise RuntimeError("Hugging Face login failed. Check your API key.")
+
+    def _initialize_model(self):
+        try:
+            logg.info("Initializing ESMC model (esmc_600m)...")
+            
+            # First try normal online loading
+            try:
+                self.client = ESMC.from_pretrained("esmc_600m")
+                self.client.to(self._device)
+                logg.info("ESMC model loaded.")
+                return
+            except Exception as online_error:
+                logg.warning(f"Online model loading failed: {online_error}")
+                logg.info("Attempting offline mode (using local cache)...")
+            
+            # Fallback: Try offline mode using cached files
+            import os
+            os.environ["HF_HUB_OFFLINE"] = "1"
+            os.environ["TRANSFORMERS_OFFLINE"] = "1"
+            
+            try:
+                self.client = ESMC.from_pretrained("esmc_600m", local_files_only=True)
+                self.client.to(self._device)
+                logg.info("ESMC model loaded from local cache (offline mode).")
+            except Exception as offline_error:
+                logg.error(f"Offline loading also failed: {offline_error}")
+                logg.error("="*60)
+                logg.error("ESMC MODEL NOT FOUND IN CACHE!")
+                logg.error("Run this on a node with internet access to cache the model:")
+                logg.error("  python -c \"from esm.models.esmc import ESMC; ESMC.from_pretrained('esmc_600m')\"")
+                logg.error("="*60)
+                raise RuntimeError("ESMC model not available. See error messages above.")
+                
+        except Exception as e:
+            logg.error(f"Failed to load ESMC model: {e}")
+            raise RuntimeError("ESMC model initialization failed.")
+
+    def _transform(self, sequence: str) -> torch.Tensor:
+        try:
+            # REPLACE (not remove) invalid chars to preserve sequence length
+            valid_aa = set('ACDEFGHIKLMNPQRSTVWY')
+            clean_sequence = ''.join(c if c in valid_aa else 'A' for c in sequence.upper())
+            
+            protein = ESMProtein(sequence=clean_sequence)
+            protein_tensor = self.client.encode(protein)
+            logits_config = LogitsConfig(sequence=True, return_embeddings=True)
+            output = self.client.logits(protein_tensor, logits_config)
+            embeddings = output.embeddings  # shape => [1, L, D] or [L, D]
+            if embeddings.dim() == 3 and embeddings.shape[0] == 1:
+                embeddings = embeddings.squeeze(0)  # => [L, D]
+
+            if not self.per_tok:
+                embeddings = embeddings.mean(dim=0)  # => [D]
+            return embeddings
+        except Exception as e:
+            logg.error(f"Error generating embeddings for sequence: {e}")
+            return None