do not use esm from transformers (#343)

* do not use esm from transformers

* minor

Author: arogozhnikov

chai_lab/data/dataset/embeddings/esm.py

@@ -6,37 +6,43 @@
 from contextlib import contextmanager
 
 import torch
-from transformers import logging as tr_logging
 
 from chai_lab.data.dataset.embeddings.embedding_context import EmbeddingContext
 from chai_lab.data.dataset.structure.chain import Chain
 from chai_lab.data.parsing.structure.entity_type import EntityType
-from chai_lab.utils.paths import downloads_path
-from chai_lab.utils.tensor_utils import move_data_to_device
+from chai_lab.utils.paths import download_if_not_exists, downloads_path
 from chai_lab.utils.typing import typecheck
 
 _esm_model: list = []  # persistent in-process container
 
 os.register_at_fork(after_in_child=lambda: _esm_model.clear())
 
 
-# unfortunately huggingface complains on pooler layer in ESM being non-initialized.
-# Did not find a way to filter specifically that logging message :/
-tr_logging.set_verbosity_error()
+ESM_URL = "https://chaiassets.com/chai1-inference-depencencies/esm2/traced_sdpa_esm2_t36_3B_UR50D_fp16.pt"
+
 
 esm_cache_folder = downloads_path.joinpath("esm")
 
 
 @contextmanager
-def esm_model(model_name: str, device):
+def esm_model(device):
     """Context transiently keeps ESM model on specified device."""
-    from transformers import EsmModel
+
+    local_esm_path = downloads_path.joinpath(
+        "esm/traced_sdpa_esm2_t36_3B_UR50D_fp16.pt"
+    )
+    download_if_not_exists(ESM_URL, local_esm_path)
 
     if len(_esm_model) == 0:
         # lazy loading of the model
-        _esm_model.append(
-            EsmModel.from_pretrained(model_name, cache_dir=esm_cache_folder)
-        )
+        if device != torch.device("cuda:0"):
+            # load on cpu first, then move to device
+            model = torch.jit.load(local_esm_path, map_location="cpu").to(device)
+        else:
+            # skip loading on CPU.
+            model = torch.jit.load(local_esm_path).to(device)
+
+        _esm_model.append(model)
 
     [model] = _esm_model
     model.to(device)
@@ -45,28 +51,82 @@ def esm_model(model_name: str, device):
     model.to("cpu")  # move model back to CPU when done
 
 
+token_map = {
+    "<cls>": 0,
+    "<pad>": 1,
+    "<eos>": 2,
+    "<unk>": 3,
+    "L": 4,
+    "A": 5,
+    "G": 6,
+    "V": 7,
+    "S": 8,
+    "E": 9,
+    "R": 10,
+    "T": 11,
+    "I": 12,
+    "D": 13,
+    "P": 14,
+    "K": 15,
+    "Q": 16,
+    "N": 17,
+    "F": 18,
+    "Y": 19,
+    "M": 20,
+    "H": 21,
+    "W": 22,
+    "C": 23,
+    "X": 24,
+    "B": 25,
+    "U": 26,
+    "Z": 27,
+    "O": 28,
+    ".": 29,
+    "-": 30,
+    "<null_1>": 31,
+    "<mask>": 32,
+}
+
+
+class DumbTokenizer:
+    def __init__(self, token_map: dict[str, int]):
+        self.token_map = token_map
+
+    def tokenize(self, text: str) -> list[int]:
+        tokens = []
+        i = 0
+        while i < len(text):
+            for token in self.token_map:
+                if text.startswith(token, i):
+                    tokens.append(self.token_map[token])
+                    i += len(token)
+                    break
+            else:
+                raise RuntimeError("Unknown token: " + text[i:])
+        return tokens
+
+
+esm_tokenizer = DumbTokenizer(token_map=token_map)
+
+
 def _get_esm_contexts_for_sequences(
     prot_sequences: set[str], device
 ) -> dict[str, EmbeddingContext]:
     if len(prot_sequences) == 0:
         return {}  # skip loading ESM
 
-    # local import, requires huggingface transformers
-    from transformers import EsmTokenizer
-
-    model_name = "facebook/esm2_t36_3B_UR50D"
-    tokenizer = EsmTokenizer.from_pretrained(model_name, cache_dir=esm_cache_folder)
-
     seq2embedding_context = {}
 
     with torch.no_grad():
-        with esm_model(model_name=model_name, device=device) as model:
+        with esm_model(device=device) as model:
             for seq in prot_sequences:
-                inputs = tokenizer(seq, return_tensors="pt")
-                inputs = move_data_to_device(dict(**inputs), device=device)
-                outputs = model(**inputs)
+                # add bos/eos, tokenize
+                token_ids = torch.asarray(esm_tokenizer.tokenize(f"<cls>{seq}<eos>"))
+                token_ids = token_ids[None, :].to(device)
+
+                last_hidden_state = model(tokens=token_ids)
                 # remove BOS/EOS, back to CPU
-                esm_embeddings = outputs.last_hidden_state[0, 1:-1].to("cpu")
+                esm_embeddings = last_hidden_state[0, 1:-1].float().to("cpu")
                 seq_len, _emb_dim = esm_embeddings.shape
                 assert seq_len == len(seq)
 

chai_lab/data/parsing/templates/m8.py

@@ -16,7 +16,7 @@
 from chai_lab.data.parsing.templates.template_hit import TemplateHit
 from chai_lab.tools.kalign import kalign_query_to_reference
 
-logger = logging.getLogger(name=__name__)
+logger = logging.getLogger(__name__)
 
 
 def parse_m8_file(fname: Path) -> pd.DataFrame:

examples/predict_structure.py

@@ -6,6 +6,8 @@
 
 from chai_lab.chai1 import run_inference
 
+logging.basicConfig(level=logging.INFO)  # control verbosity
+
 # We use fasta-like format for inputs.
 # - each entity encodes protein, ligand, RNA or DNA
 # - each entity is labeled with unique name;

pyproject.toml

@@ -45,7 +45,6 @@ module = [
     "biotite.*",
     "DockQ.*",
     "boto3.*",
-    "transformers.*",
     "modelcif.*",
     "ihm.*",
 ]

requirements.in

@@ -29,4 +29,3 @@ einops~=0.8
 jaxtyping>=0.2.25   # versions <0.2.25 do not easily support runtime typechecking
 beartype>=0.18      # compatible typechecker to use with jaxtyping
 torch>=2.3.1,<2.7   # 2.2 is broken, 2.3.1 is confirmed to work correctly
-transformers~=4.44  # for esm inference