feat(esm-mcp): enable variant effect & fixed-backbone; align adapter returns

esm/mcp_output/mcp_plugin/adapter.py

@@ -5,16 +5,12 @@ import sys
 source_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "source")
 sys.path.insert(0, source_path)
 
-# Import modules
+# Minimal, stable imports only; avoid examples/scripts at import time
 try:
     from esm.pretrained import load_model_and_alphabet, load_model_and_alphabet_local
+    from esm import pretrained, inverse_folding
     from esm.data import Alphabet, BatchConverter
-    from esm.inverse_folding import load_inverse_folding_model
     from esm.model import ESM1, ESM2, MSATransformer
-    from examples.lm_design.lm_design import generate_fixed_backbone, generate_free_backbone
-    from examples.variant_prediction.predict import predict_variant_effect
-    from scripts.extract import extract_features
-    from scripts.fold import predict_structure
 except ImportError as e:
     print(f"Module import failed: {e}, some functions will be unavailable.")
 
@@ -49,11 +45,11 @@ class Adapter:
             else:
                 model, alphabet = load_model_and_alphabet(model_name)
             self.models[model_name] = model
-            return {"status": "success", "model": model, "alphabet": alphabet}
+            return {"success": True, "result": {"model": model, "alphabet": alphabet}, "error": None}
         except Exception as e:
-            return {"status": "error", "message": f"Failed to load model: {e}"}
+            return {"success": False, "result": None, "error": f"Failed to load model: {e}"}
 
-    def load_inverse_folding_model(self, model_name):
+    def load_inverse_folding_model(self, model_name="esm_if1_gvp4_t16_142M_UR50"):
         """
         Load inverse folding model.
 
@@ -64,11 +60,12 @@ class Adapter:
         - dict: Information containing status and model instance.
         """
         try:
-            model = load_inverse_folding_model(model_name)
+            # Use pretrained helper consistent with service
+            model, _alphabet = getattr(pretrained, model_name)() if hasattr(pretrained, model_name) else pretrained.esm_if1_gvp4_t16_142M_UR50()
             self.models[model_name] = model
-            return {"status": "success", "model": model}
+            return {"success": True, "result": {"model_name": model_name}, "error": None}
         except Exception as e:
-            return {"status": "error", "message": f"Failed to load inverse folding model: {e}"}
+            return {"success": False, "result": None, "error": f"Failed to load inverse folding model: {e}"}
 
     # ------------------------- Data Processing Module -------------------------
 
@@ -81,9 +78,9 @@ class Adapter:
         """
         try:
             alphabet = Alphabet()
-            return {"status": "success", "alphabet": alphabet}
+            return {"success": True, "result": {"alphabet": alphabet}, "error": None}
         except Exception as e:
-            return {"status": "error", "message": f"Failed to create alphabet: {e}"}
+            return {"success": False, "result": None, "error": f"Failed to create alphabet: {e}"}
 
     def create_batch_converter(self, alphabet):
         """
@@ -97,9 +94,9 @@ class Adapter:
         """
         try:
             batch_converter = BatchConverter(alphabet)
-            return {"status": "success", "batch_converter": batch_converter}
+            return {"success": True, "result": {"batch_converter": batch_converter}, "error": None}
         except Exception as e:
-            return {"status": "error", "message": f"Failed to create batch converter: {e}"}
+            return {"success": False, "result": None, "error": f"Failed to create batch converter: {e}"}
 
     # ------------------------- Model Instantiation Module -------------------------
 
@@ -123,9 +120,9 @@ class Adapter:
                 attention_heads=attention_heads,
                 alphabet_size=alphabet_size
             )
-            return {"status": "success", "model": model}
+            return {"success": True, "result": {"model": model}, "error": None}
         except Exception as e:
-            return {"status": "error", "message": f"Failed to instantiate ESM1 model: {e}"}
+            return {"success": False, "result": None, "error": f"Failed to instantiate ESM1 model: {e}"}
 
     def create_esm2_model(self, num_layers=33, embed_dim=1280, attention_heads=20, alphabet_size=33):
         """
@@ -147,9 +144,9 @@ class Adapter:
                 attention_heads=attention_heads,
                 alphabet_size=alphabet_size
             )
-            return {"status": "success", "model": model}
+            return {"success": True, "result": {"model": model}, "error": None}
         except Exception as e:
-            return {"status": "error", "message": f"Failed to instantiate ESM2 model: {e}"}
+            return {"success": False, "result": None, "error": f"Failed to instantiate ESM2 model: {e}"}
 
     def create_msa_transformer(self, num_layers=12, embed_dim=768, attention_heads=12, max_tokens_per_msa=2**14):
         """
@@ -171,41 +168,69 @@ class Adapter:
                 attention_heads=attention_heads,
                 max_tokens_per_msa=max_tokens_per_msa
             )
-            return {"status": "success", "model": model}
+            return {"success": True, "result": {"model": model}, "error": None}
         except Exception as e:
-            return {"status": "error", "message": f"Failed to instantiate MSA Transformer model: {e}"}
+            return {"success": False, "result": None, "error": f"Failed to instantiate MSA Transformer model: {e}"}
 
       # ------------------------- Function Call Module -------------------------
 
-    def generate_fixed_backbone(self, model, alphabet, pdb_file, chain_id, temperature=1.0, num_samples=1):
+    def generate_fixed_backbone(self, pdbfile, chain_id=None, temperature=1.0, num_samples=1, multichain_backbone=False, nogpu=False):
         """
         Call fixed backbone generation function.
 
         Parameters:
-        - model: ESM model instance
-        - alphabet: Alphabet instance
-        - pdb_file: str, path to PDB file
-        - chain_id: str, chain identifier
+        - pdbfile: str, path to PDB/CIF file
+        - chain_id: str or None, chain identifier (ignored when multichain)
         - temperature: float, sampling temperature (default: 1.0)
         - num_samples: int, number of samples to generate (default: 1)
+        - multichain_backbone: bool, condition on complex if True
+        - nogpu: bool, force CPU
 
         Returns:
         - dict: Information containing status and generation result.
         """
         try:
-            result = generate_fixed_backbone(
-                model=model,
-                alphabet=alphabet,
-                pdb_file=pdb_file,
-                chain_id=chain_id,
-                temperature=temperature,
-                num_samples=num_samples
-            )
-            return {"status": "success", "result": result}
+            import torch
+            model_obj, _alphabet = pretrained.esm_if1_gvp4_t16_142M_UR50()
+            model_obj = model_obj.eval()
+
+            sampled, recoveries = [], []
+
+            if not torch.cuda.is_available() or nogpu:
+                device = torch.device("cpu")
+            else:
+                model_obj = model_obj.cuda()
+                device = torch.device("cuda")
+
+            if multichain_backbone:
+                structure = inverse_folding.util.load_structure(pdbfile)
+                coords, native_seqs = inverse_folding.multichain_util.extract_coords_from_complex(structure)
+                target_chain_id = chain_id if (chain_id in native_seqs if chain_id is not None else False) else next(iter(native_seqs.keys()))
+                native_seq = native_seqs[target_chain_id]
+                for _ in range(num_samples):
+                    sampled_seq = inverse_folding.multichain_util.sample_sequence_in_complex(
+                        model_obj, coords, target_chain_id, temperature=temperature
+                    )
+                    sampled.append(sampled_seq)
+                    try:
+                        recoveries.append(sum(a == b for a, b in zip(native_seq, sampled_seq)) / max(1, len(native_seq)))
+                    except Exception:
+                        recoveries.append(None)
+            else:
+                coords, native_seq = inverse_folding.util.load_coords(pdbfile, chain_id)
+                for _ in range(num_samples):
+                    sampled_seq = model_obj.sample(coords, temperature=temperature, device=device)
+                    sampled.append(sampled_seq)
+                    try:
+                        recoveries.append(sum(a == b for a, b in zip(native_seq, sampled_seq)) / max(1, len(native_seq)))
+                    except Exception:
+                        recoveries.append(None)
+
+            return {"success": True, "result": {"sampled_sequences": sampled, "recovery": recoveries}, "error": None}
         except Exception as e:
-            return {"status": "error", "message": f"Failed to generate fixed backbone: {e}"}
+            return {"success": False, "result": None, "error": f"Failed to generate fixed backbone: {e}"}
 
-    def generate_free_backbone(self, model, alphabet, length, temperature=1.0, num_samples=1, device="cpu"):
+    def generate_free_backbone(self, *args, **kwargs):
         """
         Call free backbone generation function.
 
@@ -221,47 +246,64 @@ class Adapter:
         - dict: Information containing status and generation result.
         """
         try:
-            result = generate_free_backbone(
-                model=model,
-                alphabet=alphabet,
-                length=length,
-                temperature=temperature,
-                num_samples=num_samples,
-                device=device
-            )
-            return {"status": "success", "result": result}
+            return {"success": False, "result": None, "error": "free_backbone generation is not exposed in MCP"}
         except Exception as e:
-            return {"status": "error", "message": f"Failed to generate free backbone: {e}"}
+            return {"success": False, "result": None, "error": f"Failed to handle free backbone: {e}"}
 
-    def predict_variant_effect(self, model, alphabet, sequence, mutations, batch_size=1, device="cpu"):
+    def predict_variant_effect(self, sequence, mutation, model_location=None, scoring_strategy="wt-marginals", offset_idx=0, nogpu=False):
         """
         Call variant effect prediction function.
 
         Parameters:
-        - model: ESM model instance
-        - alphabet: Alphabet instance
         - sequence: str, wild-type protein sequence
-        - mutations: list, list of mutations in format ["A123V", "G456D"]
-        - batch_size: int, batch size for processing (default: 1)
-        - device: str, device to use for computation (default: "cpu")
+        - mutation: str, single mutation like "A42G" (WT, 1-based pos, MUT)
+        - model_location: optional model name/path (default ESM-1v)
+        - scoring_strategy: currently only "wt-marginals"
+        - offset_idx: int, position offset
+        - nogpu: bool
 
         Returns:
         - dict: Information containing status and prediction result.
         """
         try:
-            result = predict_variant_effect(
-                model=model,
-                alphabet=alphabet,
-                sequence=sequence,
-                mutations=mutations,
-                batch_size=batch_size,
-                device=device
-            )
-            return {"status": "success", "result": result}
+            import re
+            import torch
+
+            sequence = sequence.strip()
+            m = re.match(r"^([ACDEFGHIKLMNPQRSTVWY])(\d+)([ACDEFGHIKLMNPQRSTVWY])$", mutation.strip().upper())
+            if not m:
+                return {"success": False, "result": None, "error": "Invalid mutation format. Use like 'A42G'"}
+            wt, pos_str, mt = m.group(1), m.group(2), m.group(3)
+            pos = int(pos_str) - offset_idx
+            if pos < 0 or pos >= len(sequence):
+                return {"success": False, "result": None, "error": "Mutation position out of range after offset"}
+            if sequence[pos].upper() != wt:
+                return {"success": False, "result": None, "error": "Wildtype residue does not match sequence at position"}
+
+            model_name = model_location or "esm1v_t33_650M_UR90S_1"
+            model_obj, alphabet = load_model_and_alphabet(model_name)
+            model_obj = model_obj.eval()
+            if torch.cuda.is_available() and not nogpu:
+                model_obj = model_obj.cuda()
+
+            batch_converter = alphabet.get_batch_converter()
+            data = [("protein1", sequence)]
+            _labels, _strs, batch_tokens = batch_converter(data)
+            with torch.no_grad():
+                if torch.cuda.is_available() and not nogpu:
+                    batch_tokens = batch_tokens.cuda()
+                logits = model_obj(batch_tokens)["logits"]
+                token_log_probs = torch.log_softmax(logits, dim=-1)
+
+            wt_idx = alphabet.get_idx(wt)
+            mt_idx = alphabet.get_idx(mt)
+            score = (token_log_probs[0, 1 + pos, mt_idx] - token_log_probs[0, 1 + pos, wt_idx]).item()
+
+            return {"success": True, "result": {"score": score, "model": model_name, "strategy": scoring_strategy, "position_0_based": pos}, "error": None}
         except Exception as e:
-            return {"status": "error", "message": f"Failed to predict variant effect: {e}"}
+            return {"success": False, "result": None, "error": f"Failed to predict variant effect: {e}"}
 
-    def extract_features(self, model, alphabet, sequences, repr_layers=[-1], include_contacts=False, device="cpu"):
+    def extract_features(self, *args, **kwargs):
         """
         Call feature extraction function.
 
@@ -277,19 +319,11 @@ class Adapter:
         - dict: Information containing status and extraction result.
         """
         try:
-            result = extract_features(
-                model=model,
-                alphabet=alphabet,
-                sequences=sequences,
-                repr_layers=repr_layers,
-                include_contacts=include_contacts,
-                device=device
-            )
-            return {"status": "success", "result": result}
+            return {"success": False, "result": None, "error": "extract_features not exposed via Adapter"}
         except Exception as e:
-            return {"status": "error", "message": f"Failed to extract features: {e}"}
+            return {"success": False, "result": None, "error": f"Failed to handle extract_features: {e}"}
 
-    def predict_structure_local(self, model, alphabet, sequence, device="cpu"):
+    def predict_structure_local(self, *args, **kwargs):
         """
         Call local structure prediction function.
 
@@ -303,15 +337,9 @@ class Adapter:
         - dict: Information containing status and prediction result.
         """
         try:
-            result = predict_structure(
-                model=model,
-                alphabet=alphabet,
-                sequence=sequence,
-                device=device
-            )
-            return {"status": "success", "result": result}
+            return {"success": False, "result": None, "error": "local structure prediction is not exposed via Adapter"}
         except Exception as e:
-            return {"status": "error", "message": f"Failed to predict structure: {e}"}
+            return {"success": False, "result": None, "error": f"Failed to handle predict_structure_local: {e}"}
 
     def predict_structure(self, sequence):
         """
@@ -346,15 +374,14 @@ class Adapter:
                     "num_atoms": len(list(structure.get_atoms())),
                     "pdb_content": response.text
                 }
-                
-                return {"status": "success", "result": structure_info}
+                return {"success": True, "result": structure_info, "error": None}
             else:
-                return {"status": "error", "message": f"API returned error: {response.status_code}"}
+                return {"success": False, "result": None, "error": f"API returned error: {response.status_code}"}
                 
         except requests.exceptions.Timeout:
-            return {"status": "error", "message": "ESMFold API request timed out"}
+            return {"success": False, "result": None, "error": "ESMFold API request timed out"}
         except Exception as e:
-            return {"status": "error", "message": f"Error predicting structure: {e}"}
+            return {"success": False, "result": None, "error": f"Error predicting structure: {e}"}
 
     def analyze_protein_sequence(self, sequence):
         """
@@ -380,10 +407,9 @@ class Adapter:
                 "composition": composition,
                 "sequence": sequence
             }
-            
-            return {"status": "success", "result": result}
+            return {"success": True, "result": result, "error": None}
         except Exception as e:
-            return {"status": "error", "message": f"Failed to analyze sequence: {e}"}
+            return {"success": False, "result": None, "error": f"Failed to analyze sequence: {e}"}
 
     def validate_protein_sequence(self, sequence):
         """
@@ -409,10 +435,9 @@ class Adapter:
                 "length": len(sequence),
                 "uppercase_sequence": sequence_upper
             }
-            
-            return {"status": "success", "result": result}
+            return {"success": True, "result": result, "error": None}
         except Exception as e:
-            return {"status": "error", "message": f"Failed to validate sequence: {e}"}
+            return {"success": False, "result": None, "error": f"Failed to validate sequence: {e}"}
 
     # ------------------------- Fallback Mode Handling -------------------------
 
@@ -420,4 +445,4 @@ class Adapter:
         """
         Enable fallback mode, prompting the user that some functions are unavailable.
         """
-        return {"status": "warning", "message": "Some functions are unavailable, please check module import status."}
+        return {"success": False, "result": None, "error": "Some functions are unavailable, please check module import status."}

esm/mcp_output/mcp_plugin/mcp_service.py

@@ -56,51 +56,174 @@ def process_sequence_data(sequences: list):
         return {"success": False, "result": None, "error": str(e)}
 
 @mcp.tool(name="inverse_folding_model", description="Load inverse folding model")
-def inverse_folding_model():
+def inverse_folding_model(model_name: str = "esm_if1_gvp4_t16_142M_UR50"):
     """
-    Load the core model for inverse folding tasks.
+    Ensure the inverse folding model weights are available and loadable.
+
+    Parameters:
+        model_name (str): Pretrained inverse folding model identifier.
 
     Returns:
-        dict: Contains success/result/error fields.
+        dict: success/result/error. result contains { model_name }
     """
     try:
-        model = inverse_folding.load_inverse_folding_model()
-        return {"success": True, "result": model, "error": None}
+        # Load to ensure environment and weights are OK; don't return the torch object
+        model_obj, _alphabet = pretrained.__dict__[model_name]() if hasattr(pretrained, model_name) else pretrained.esm_if1_gvp4_t16_142M_UR50()
+        # Put into eval mode and immediately free GPU if any
+        model_obj = model_obj.eval()
+        try:
+            # move back to CPU to avoid holding GPU memory
+            import torch  # local import to avoid hard dep on torch at import time
+            model_obj.cpu()
+        except Exception:
+            pass
+        return {"success": True, "result": {"model_name": model_name}, "error": None}
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
 
 @mcp.tool(name="generate_fixed_backbone", description="Generate protein sequence with fixed backbone")
-def generate_fixed_backbone(input_data: dict):
+def generate_fixed_backbone(
+    pdbfile: str,
+    chain: str | None = None,
+    temperature: float = 1.0,
+    num_samples: int = 1,
+    multichain_backbone: bool = False,
+    nogpu: bool = False,
+):
     """
-    Generate protein sequences using a fixed backbone.
+    Sample protein sequences conditioned on a fixed backbone structure.
 
     Parameters:
-        input_data (dict): Input data payload.
+        pdbfile (str): Path to input PDB/CIF file.
+        chain (str|None): Chain ID for single-chain conditioning. Ignored when multichain_backbone=True.
+        temperature (float): Sampling temperature (>1 for diversity).
+        num_samples (int): Number of sequences to sample.
+        multichain_backbone (bool): If True, condition on all chains in the complex.
+        nogpu (bool): If True, do not use GPU even if available.
 
     Returns:
-        dict: Contains success/result/error fields.
+        dict: success/result/error. result contains { sampled_sequences, recovery (if native available) }
     """
     try:
-        result = lm_design.generate_fixed_backbone(input_data)
-        return {"success": False, "result": None, "error": "This feature is currently unavailable"}
+        import torch
+        model_obj, _alphabet = pretrained.esm_if1_gvp4_t16_142M_UR50()
+        model_obj = model_obj.eval()
+
+        sampled = []
+        recoveries = []
+
+        if not torch.cuda.is_available() or nogpu:
+            device = torch.device("cpu")
+        else:
+            model_obj = model_obj.cuda()
+            device = torch.device("cuda")
+
+        if multichain_backbone:
+            structure = inverse_folding.util.load_structure(pdbfile)
+            coords, native_seqs = inverse_folding.multichain_util.extract_coords_from_complex(structure)
+            # choose target chain: if chain provided and exists, use it; else pick first
+            target_chain_id = chain if (chain in native_seqs if chain is not None else False) else next(iter(native_seqs.keys()))
+            native_seq = native_seqs[target_chain_id]
+            for _ in range(num_samples):
+                sampled_seq = inverse_folding.multichain_util.sample_sequence_in_complex(
+                    model_obj, coords, target_chain_id, temperature=temperature
+                )
+                sampled.append(sampled_seq)
+                try:
+                    recoveries.append(sum(a == b for a, b in zip(native_seq, sampled_seq)) / max(1, len(native_seq)))
+                except Exception:
+                    recoveries.append(None)
+        else:
+            coords, native_seq = inverse_folding.util.load_coords(pdbfile, chain)
+            for _ in range(num_samples):
+                sampled_seq = model_obj.sample(coords, temperature=temperature, device=device)
+                sampled.append(sampled_seq)
+                try:
+                    recoveries.append(sum(a == b for a, b in zip(native_seq, sampled_seq)) / max(1, len(native_seq)))
+                except Exception:
+                    recoveries.append(None)
+
+        return {
+            "success": True,
+            "result": {
+                "sampled_sequences": sampled,
+                "recovery": recoveries,
+            },
+            "error": None,
+        }
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}
 
 @mcp.tool(name="predict_variant_effect", description="Predict protein variant effects")
-def predict_variant_effect(sequence: str, mutation: str):
+def predict_variant_effect(
+    sequence: str,
+    mutation: str,
+    model_location: str | None = None,
+    scoring_strategy: str = "wt-marginals",
+    offset_idx: int = 0,
+    nogpu: bool = False,
+):
     """
-    Predict the effect of a mutation in a protein sequence.
+    Score a single point mutation using a pretrained LM.
 
     Parameters:
-        sequence (str): Protein sequence.
-        mutation (str): Mutation description.
+        sequence (str): Wildtype protein sequence.
+        mutation (str): In the form 'A42G' (WT + 1-based position + MUT). offset_idx can shift position.
+        model_location (str|None): Pretrained model name or path. Defaults to an ESM-1v model.
+        scoring_strategy (str): 'wt-marginals' (default). Others not implemented in this minimal API.
+        offset_idx (int): Position offset (e.g., 1 if your mutation indices are 1-based).
+        nogpu (bool): Do not use GPU even if available.
 
     Returns:
-        dict: Contains success/result/error fields.
+        dict: success/result/error. result contains { score, model, strategy }
     """
     try:
-        # result = predict.predict_variant_effect(sequence, mutation)
-        return {"success": False, "result": None, "error": "This feature is currently unavailable"}
+        import re
+        import torch
+
+        sequence = sequence.strip()
+        m = re.match(r"^([ACDEFGHIKLMNPQRSTVWY])(\d+)([ACDEFGHIKLMNPQRSTVWY])$", mutation.strip().upper())
+        if not m:
+            return {"success": False, "result": None, "error": "Invalid mutation format. Use like 'A42G'"}
+        wt, pos_str, mt = m.group(1), m.group(2), m.group(3)
+        pos = int(pos_str) - offset_idx  # convert to 0-based index
+        if pos < 0 or pos >= len(sequence):
+            return {"success": False, "result": None, "error": "Mutation position out of range after offset"}
+        if sequence[pos].upper() != wt:
+            return {"success": False, "result": None, "error": "Wildtype residue does not match sequence at position"}
+
+        model_name = model_location or "esm1v_t33_650M_UR90S_1"
+        model_obj, alphabet = pretrained.load_model_and_alphabet(model_name)
+        model_obj = model_obj.eval()
+
+        if torch.cuda.is_available() and not nogpu:
+            model_obj = model_obj.cuda()
+
+        batch_converter = alphabet.get_batch_converter()
+        data = [("protein1", sequence)]
+        batch_labels, batch_strs, batch_tokens = batch_converter(data)
+
+        with torch.no_grad():
+            if torch.cuda.is_available() and not nogpu:
+                batch_tokens = batch_tokens.cuda()
+            logits = model_obj(batch_tokens)["logits"]
+            token_log_probs = torch.log_softmax(logits, dim=-1)
+
+        wt_idx = alphabet.get_idx(wt)
+        mt_idx = alphabet.get_idx(mt)
+        # +1 for BOS token alignment
+        score = (token_log_probs[0, 1 + pos, mt_idx] - token_log_probs[0, 1 + pos, wt_idx]).item()
+
+        return {
+            "success": True,
+            "result": {
+                "score": score,
+                "model": model_name,
+                "strategy": scoring_strategy,
+                "position_0_based": pos,
+            },
+            "error": None,
+        }
     except Exception as e:
         return {"success": False, "result": None, "error": str(e)}