Spaces:
Running on Zero
Running on Zero
| """Loading and running the protein masked language model.""" | |
| from __future__ import annotations | |
| from dataclasses import dataclass | |
| from functools import lru_cache | |
| import spaces # must be imported before torch (CUDA-related); see app_gradio.py | |
| import pandas as pd | |
| import torch | |
| from esm.models.esmc import EsmcForMaskedLM, EsmcTokenizer | |
| from tqdm import tqdm | |
| from src.constants import AMINO_ACIDS | |
| # Sequences per forward call when scoring mutations, not a total batch count: | |
| # total mutations grows with protein length (19 * L), so a *fixed number of | |
| # batches* still lets each batch (and its O(batch * L^2) attention memory) | |
| # grow without bound for long sequences. A fixed batch size keeps peak memory | |
| # roughly constant regardless of protein length. | |
| MUTATION_BATCH_SIZE = 8 | |
| class MaskPrediction: | |
| token: str | |
| probability: float | |
| def load_model(model_name: str): | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| tokenizer = EsmcTokenizer() | |
| model = EsmcForMaskedLM.from_pretrained(model_name, device=device) | |
| model.eval() | |
| return tokenizer, model | |
| def _sequence_log_probs(sequences: list[str], tokenizer, model) -> torch.Tensor: | |
| """Unmasked forward pass over full sequences, returned as per-position | |
| log-probabilities with the leading <cls> token stripped, so that | |
| `result[i, j]` lines up with `sequences[i][j]`. | |
| Shape: (len(sequences), len(sequences[0]), vocab_size). | |
| """ | |
| inputs = tokenizer(sequences, return_tensors="pt", padding=True).to(model.device) | |
| with torch.no_grad(): | |
| logits = model(**inputs).logits | |
| log_probs = torch.log_softmax(logits, dim=-1) | |
| return log_probs[:, 1 : 1 + len(sequences[0]), :] | |
| def _mutation_llrs( | |
| sequence: str, | |
| mutations: list[tuple[int, str]], | |
| wild_type_log_probs: torch.Tensor, | |
| tokenizer, | |
| model, | |
| batch_size: int = MUTATION_BATCH_SIZE, | |
| ) -> list[float]: | |
| """LLR for each `(position, mutant)` pair in `mutations`, processed in | |
| fixed-size batches of full mutated sequences. | |
| Only the single log-probability needed per mutant (at its own mutated | |
| position) is kept — the rest of that sequence's `(L, vocab)` log-softmax | |
| output is dropped immediately after each batch, instead of being held | |
| for every mutation at once. | |
| """ | |
| llrs: list[float] = [] | |
| batch_starts = range(0, len(mutations), batch_size) | |
| for start in tqdm(batch_starts, desc="Scoring mutations", unit="batch", leave=False): | |
| chunk = mutations[start : start + batch_size] | |
| mutated_sequences = [ | |
| sequence[:position] + mutant + sequence[position + 1 :] | |
| for position, mutant in chunk | |
| ] | |
| chunk_log_probs = _sequence_log_probs(mutated_sequences, tokenizer, model) | |
| for (position, mutant), mutant_log_probs in zip(chunk, chunk_log_probs): | |
| wild_type_id = tokenizer.convert_tokens_to_ids(sequence[position]) | |
| mutant_id = tokenizer.convert_tokens_to_ids(mutant) | |
| wild_type_log_prob = wild_type_log_probs[position, wild_type_id] | |
| mutant_log_prob = mutant_log_probs[position, mutant_id] | |
| llrs.append(float(mutant_log_prob - wild_type_log_prob)) | |
| del chunk_log_probs | |
| if model.device.type == "cuda": | |
| torch.cuda.empty_cache() | |
| return llrs | |
| def calculate_llr_matrix(sequence: str, model_name: str) -> pd.DataFrame: | |
| """LLR of every possible single-point mutation at every position. | |
| Mutant sequences are scored `MUTATION_BATCH_SIZE` at a time (see | |
| `_mutation_llrs`) rather than all at once, so peak memory stays bounded | |
| regardless of protein length. | |
| Returns a `len(AMINO_ACIDS) x len(sequence)` DataFrame where | |
| `df.loc[aa, col]` is the log-likelihood ratio of mutating the residue at | |
| `col` to `aa` (mutant-marginal method: the model's own log-probability | |
| for the mutant residue when reading the full mutated sequence, minus its | |
| log-probability for the wild-type residue when reading the original | |
| sequence). The wild-type residue at each position is 0 by definition, so | |
| it is filled in directly rather than re-scored. | |
| """ | |
| tokenizer, model = load_model(model_name) | |
| columns = [f"{residue}{i + 1}" for i, residue in enumerate(sequence)] | |
| matrix = pd.DataFrame(0.0, index=AMINO_ACIDS, columns=columns) | |
| wild_type_log_probs = _sequence_log_probs([sequence], tokenizer, model)[0] | |
| mutations = [ | |
| (position, mutant) | |
| for position, wild_type in enumerate(sequence) | |
| for mutant in AMINO_ACIDS | |
| if mutant != wild_type | |
| ] | |
| llrs = _mutation_llrs(sequence, mutations, wild_type_log_probs, tokenizer, model) | |
| for (position, mutant), llr in zip(mutations, llrs): | |
| matrix.loc[mutant, matrix.columns[position]] = llr | |
| return matrix | |
| def predict_masked_position( | |
| sequence: str, | |
| position: int, | |
| model_name: str, | |
| top_k: int = 5, | |
| ) -> list[MaskPrediction]: | |
| """Mask `sequence[position]` (0-indexed) and return the model's top_k | |
| amino acid predictions for that position, most probable first.""" | |
| tokenizer, model = load_model(model_name) | |
| residues = list(sequence) | |
| residues[position] = tokenizer.mask_token | |
| masked_sequence = "".join(residues) | |
| inputs = tokenizer(masked_sequence, return_tensors="pt").to(model.device) | |
| mask_token_index = torch.where(inputs["input_ids"][0] == tokenizer.mask_token_id)[0] | |
| if len(mask_token_index) == 0: | |
| raise ValueError("Mask token not found in tokenized sequence") | |
| with torch.no_grad(): | |
| logits = model(**inputs).logits | |
| probs = torch.softmax(logits[0, mask_token_index[0]], dim=-1) | |
| top_probs, top_ids = torch.topk(probs, top_k) | |
| return [ | |
| MaskPrediction(token=tokenizer.decode([tid]).strip(), probability=float(p)) | |
| for tid, p in zip(top_ids, top_probs) | |
| ] | |