AutoAndroidController/py/venv/Lib/site-packages/modelscope/pipelines/science/protein_structure_pipeline.py

# Copyright (c) Alibaba, Inc. and its affiliates.

import os
import time
from typing import Any, Dict, List, Optional, Union

import json
import numpy as np
import torch
from unicore.utils import tensor_tree_map

from modelscope.metainfo import Pipelines
from modelscope.models.base import Model
from modelscope.models.science.unifold.config import model_config
from modelscope.models.science.unifold.data import protein, residue_constants
from modelscope.models.science.unifold.dataset import (UnifoldDataset,
                                                       load_and_process)
from modelscope.outputs import OutputKeys
from modelscope.pipelines.base import Pipeline, Tensor
from modelscope.pipelines.builder import PIPELINES
from modelscope.preprocessors import Preprocessor, build_preprocessor
from modelscope.utils.constant import Fields, Frameworks, Tasks
from modelscope.utils.device import device_placement
from modelscope.utils.hub import read_config
from modelscope.utils.logger import get_logger

logger = get_logger()

__all__ = ['ProteinStructurePipeline']


def automatic_chunk_size(seq_len):
    if seq_len < 512:
        chunk_size = 256
    elif seq_len < 1024:
        chunk_size = 128
    elif seq_len < 2048:
        chunk_size = 32
    elif seq_len < 3072:
        chunk_size = 16
    else:
        chunk_size = 1
    return chunk_size


def load_feature_for_one_target(
    config,
    data_folder,
    seed=0,
    is_multimer=False,
    use_uniprot=False,
    symmetry_group=None,
):
    if not is_multimer:
        uniprot_msa_dir = None
        sequence_ids = ['A']
        if use_uniprot:
            uniprot_msa_dir = data_folder

    else:
        uniprot_msa_dir = data_folder
        sequence_ids = open(
            os.path.join(data_folder, 'chains.txt'),
            encoding='utf-8').readline().split()

    if symmetry_group is None:
        batch, _ = load_and_process(
            config=config.data,
            mode='predict',
            seed=seed,
            batch_idx=None,
            data_idx=0,
            is_distillation=False,
            sequence_ids=sequence_ids,
            monomer_feature_dir=data_folder,
            uniprot_msa_dir=uniprot_msa_dir,
        )
    else:
        # Not for unifold-symmetry
        # only for unifold-multimer
        batch, _ = load_and_process(
            config=config.data,
            mode='predict',
            seed=seed,
            batch_idx=None,
            data_idx=0,
            is_distillation=False,
            sequence_ids=sequence_ids,
            monomer_feature_dir=data_folder,
            uniprot_msa_dir=uniprot_msa_dir,
        )
    batch = UnifoldDataset.collater([batch])
    return batch


@PIPELINES.register_module(
    Tasks.protein_structure, module_name=Pipelines.protein_structure)
class ProteinStructurePipeline(Pipeline):

    def __init__(self,
                 model: Union[Model, str],
                 preprocessor: Optional[Preprocessor] = None,
                 **kwargs):
        """Use `model` and `preprocessor` to create a protein structure pipeline for prediction.

        Args:
            model (str or Model): Supply either a local model dir which supported the protein structure task,
            or a model id from the model hub, or a torch model instance.
            preprocessor (Preprocessor): An optional preprocessor instance, please make sure the preprocessor fits for
            the model if supplied.

        Examples:
            >>> from modelscope.pipelines import pipeline
            >>> pipeline_ins = pipeline(task='protein-structure',
            >>>    model='DPTech/uni-fold-monomer')
            >>> protein = 'LILNLRGGAFVSNTQITMADKQKKFINEIQEGDLVRSYSITDETFQQNAVTSIVKHEADQLCQINFGKQHVVC'
            >>> print(pipeline_ins(protein))

        """
        super().__init__(model=model, preprocessor=preprocessor, **kwargs)
        self.cfg = read_config(self.model.model_dir)
        self.config = model_config(
            self.cfg['pipeline']['model_name'])  # alphafold config
        self.postprocessor = self.cfg.pop('postprocessor', None)
        if preprocessor is None:
            preprocessor_cfg = self.cfg.preprocessor
            self.preprocessor = build_preprocessor(preprocessor_cfg,
                                                   Fields.science)
        self.model.eval()

    def _sanitize_parameters(self, **pipeline_parameters):
        return pipeline_parameters, pipeline_parameters, pipeline_parameters

    def _process_single(self, input, *args, **kwargs) -> Dict[str, Any]:
        preprocess_params = kwargs.get('preprocess_params', {})
        forward_params = kwargs.get('forward_params', {})
        postprocess_params = kwargs.get('postprocess_params', {})
        out = self.preprocess(input, **preprocess_params)
        with device_placement(self.framework, self.device_name):
            with torch.no_grad():
                out = self.forward(out, **forward_params)

        out = self.postprocess(out, **postprocess_params)
        return out

    def forward(self, inputs: Dict[str, Any],
                **forward_params) -> Dict[str, Any]:
        plddts = {}
        ptms = {}

        output_dir = os.path.join(self.preprocessor.output_dir_base,
                                  inputs['target_id'])

        pdbs = []
        for seed in range(self.cfg['pipeline']['times']):
            cur_seed = hash((42, seed)) % 100000
            batch = load_feature_for_one_target(
                self.config,
                output_dir,
                cur_seed,
                is_multimer=inputs['is_multimer'],
                use_uniprot=inputs['is_multimer'],
                symmetry_group=self.preprocessor.symmetry_group,
            )
            seq_len = batch['aatype'].shape[-1]
            self.model.model.globals.chunk_size = automatic_chunk_size(seq_len)

            with torch.no_grad():
                batch = {
                    k: torch.as_tensor(v, device='cuda:0')
                    for k, v in batch.items()
                }
                out = self.model(batch)

            def to_float(x):
                if x.dtype == torch.bfloat16 or x.dtype == torch.half:
                    return x.float()
                else:
                    return x

            # Toss out the recycling dimensions --- we don't need them anymore
            batch = tensor_tree_map(lambda t: t[-1, 0, ...], batch)
            batch = tensor_tree_map(to_float, batch)
            out = tensor_tree_map(lambda t: t[0, ...], out[0])
            out = tensor_tree_map(to_float, out)
            batch = tensor_tree_map(lambda x: np.array(x.cpu()), batch)
            out = tensor_tree_map(lambda x: np.array(x.cpu()), out)

            plddt = out['plddt']
            mean_plddt = np.mean(plddt)
            plddt_b_factors = np.repeat(
                plddt[..., None], residue_constants.atom_type_num, axis=-1)
            # TODO: , may need to reorder chains, based on entity_ids
            cur_protein = protein.from_prediction(
                features=batch, result=out, b_factors=plddt_b_factors)
            cur_save_name = (f'{cur_seed}')
            plddts[cur_save_name] = str(mean_plddt)
            if inputs[
                    'is_multimer'] and self.preprocessor.symmetry_group is None:
                ptms[cur_save_name] = str(np.mean(out['iptm+ptm']))
            with open(os.path.join(output_dir, cur_save_name + '.pdb'),
                      'w') as f:
                f.write(protein.to_pdb(cur_protein))
                pdbs.append(protein.to_pdb(cur_protein))

        logger.info('plddts:' + str(plddts))
        model_name = self.cfg['pipeline']['model_name']
        score_name = f'{model_name}'
        plddt_fname = score_name + '_plddt.json'

        with open(os.path.join(output_dir, plddt_fname), 'w') as f:
            json.dump(plddts, f, indent=4)
        if ptms:
            logger.info('ptms' + str(ptms))
            ptm_fname = score_name + '_ptm.json'
            with open(os.path.join(output_dir, ptm_fname), 'w') as f:
                json.dump(ptms, f, indent=4)

        return pdbs

    def postprocess(self, inputs: Dict[str, Tensor], **postprocess_params):
        return inputs