Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2025 The HuggingFace Inc. team.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""16Processor class for EVOLLA.17"""18 19import os20from typing import Optional, Union21 22from ...feature_extraction_utils import BatchFeature23from ...processing_utils import (24    ProcessorMixin,25)26from ..auto import AutoTokenizer27 28 29PROTEIN_VALID_KEYS = ["aa_seq", "foldseek", "msa"]30 31 32class EvollaProcessor(ProcessorMixin):33    r"""34    Constructs a EVOLLA processor which wraps a LLama tokenizer and SaProt tokenizer (EsmTokenizer) into a single processor.35 36    [`EvollaProcessor`] offers all the functionalities of [`EsmTokenizer`] and [`LlamaTokenizerFast`]. See the37    docstring of [`~EvollaProcessor.__call__`] and [`~EvollaProcessor.decode`] for more information.38 39    Args:40        protein_tokenizer (`EsmTokenizer`):41            An instance of [`EsmTokenizer`]. The protein tokenizer is a required input.42        tokenizer (`LlamaTokenizerFast`, *optional*):43            An instance of [`LlamaTokenizerFast`]. The tokenizer is a required input.44        protein_max_length (`int`, *optional*, defaults to 1024):45            The maximum length of the sequence to be generated.46        text_max_length (`int`, *optional*, defaults to 512):47            The maximum length of the text to be generated.48    """49 50    attributes = ["protein_tokenizer", "tokenizer"]51    valid_kwargs = ["sequence_max_length"]52    # protein_tokenizer_class = "EsmTokenizer"53    # tokenizer_class = "LlamaTokenizerFast"54    protein_tokenizer_class = "AutoTokenizer"55    tokenizer_class = "AutoTokenizer"56    protein_tokenizer_dir_name = "protein_tokenizer"57    # tokenizer_dir_name = "text_tokenizer"58 59    def __init__(self, protein_tokenizer, tokenizer=None, protein_max_length=1024, text_max_length=512, **kwargs):60        if protein_tokenizer is None:61            raise ValueError("You need to specify an `protein_tokenizer`.")62        if tokenizer is None:63            raise ValueError("You need to specify a `tokenizer`.")64 65        super().__init__(protein_tokenizer, tokenizer)66 67        self.tokenizer.pad_token = "<|reserved_special_token_0|>"68        self.protein_max_length = protein_max_length69        self.text_max_length = text_max_length70 71    def process_proteins(self, proteins, protein_max_length=1024):72        sa_sequences = []73        for protein in proteins:74            aa_seq = protein.get("aa_seq")75            foldseek = protein.get("foldseek")76            sa_sequence = "".join([s.upper() + f.lower() for s, f in zip(aa_seq, foldseek)])77            sa_sequences.append(sa_sequence)78 79        sa_tokens = self.protein_tokenizer.batch_encode_plus(80            sa_sequences, return_tensors="pt", truncation=True, max_length=protein_max_length, padding=True81        )82        return sa_tokens83 84    def process_text(85        self,86        texts,87        text_max_length: int = 512,88    ):89        prompts = []90        for messages in texts:91            prompt = self.tokenizer.apply_chat_template(92                messages,93                tokenize=False,94                add_generation_prompt=True,95            )96            prompts.append(prompt)97 98        prompt_inputs = self.tokenizer(99            prompts,100            add_special_tokens=False,101            return_tensors="pt",102            padding="longest",103            truncation=True,104            max_length=text_max_length,105        )106        return prompt_inputs107 108    def __call__(109        self,110        proteins: Optional[Union[list[dict], dict]] = None,111        messages_list: Optional[Union[list[list[dict]], list[dict]]] = None,112        protein_max_length: Optional[int] = None,113        text_max_length: Optional[int] = None,114        **kwargs,115    ):116        r"""This method takes batched or non-batched proteins and messages_list and converts them into format that can be used by117        the model.118 119        Args:120            proteins (`Union[List[dict], dict]`):121                A list of dictionaries or a single dictionary containing the following keys:122                    - `"aa_seq"` (`str`) -- The amino acid sequence of the protein.123                    - `"foldseek"` (`str`) -- The foldseek string of the protein.124            messages_list (`Union[List[List[dict]], List[dict]]`):125                A list of lists of dictionaries or a list of dictionaries containing the following keys:126                    - `"role"` (`str`) -- The role of the message.127                    - `"content"` (`str`) -- The content of the message.128            protein_max_length (`int`, *optional*, defaults to 1024):129                The maximum length of the sequence to be generated.130            text_max_length (`int`, *optional*, defaults to 512):131                The maximum length of the text.132 133        Return:134            a dict with following keys:135                - `protein_input_ids` (`torch.Tensor` of shape `(batch_size, sequence_length)`) -- The input IDs for the protein sequence.136                - `protein_attention_mask` (`torch.Tensor` of shape `(batch_size, sequence_length)`) -- The attention mask for the protein sequence.137                - `text_input_ids` (`torch.Tensor` of shape `(batch_size, sequence_length)`) -- The input IDs for the text sequence.138                - `text_attention_mask` (`torch.Tensor` of shape `(batch_size, sequence_length)`) -- The attention mask for the text sequence.139        """140        # proteins and messages_list should be provided141        if proteins is None or messages_list is None:142            raise ValueError("You need to specify `messages_list` and `proteins`.")143 144        protein_max_length = protein_max_length if protein_max_length is not None else self.protein_max_length145        text_max_length = text_max_length if text_max_length is not None else self.text_max_length146 147        # proteins should be List[dict]148        if isinstance(proteins, dict):149            proteins = [proteins]150        # messages_list should be List[List[dict]]151        if isinstance(messages_list, (list, tuple)) and not isinstance(messages_list[0], (list, tuple)):152            messages_list = [messages_list]153        # Check if batched proteins are in the correct format154        if isinstance(proteins, (list, tuple)) and not all(isinstance(p, dict) for p in proteins):155            raise ValueError("The proteins should be a list of dictionaries, but not all elements are dictionaries.")156        if isinstance(proteins, (list, tuple)) and not all(157            all(k in PROTEIN_VALID_KEYS for k in p.keys()) for p in proteins158        ):159            raise ValueError(160                "There should be a list of dictionaries with keys: "161                f"{', '.join(PROTEIN_VALID_KEYS)} for each protein."162                f"But got: {proteins}"163            )164        # Check if batched messages_list is in the correct format165        if isinstance(messages_list, (list, tuple)):166            for messages in messages_list:167                if not isinstance(messages, (list, tuple)):168                    raise ValueError(f"Each messages in messages_list should be a list instead of {type(messages)}.")169                if not all(isinstance(m, dict) for m in messages):170                    raise ValueError(171                        "Each message in messages_list should be a list of dictionaries, but not all elements are dictionaries."172                    )173                if any(len(m.keys()) != 2 for m in messages) or any(174                    set(m.keys()) != {"role", "content"} for m in messages175                ):176                    raise ValueError(177                        "Each message in messages_list should be a list of dictionaries with two keys: 'role' and 'content'."178                        f"But got: {messages}"179                    )180        else:181            raise ValueError(182                f"The messages_list should be a list of lists of dictionaries, but it's {type(messages_list)}."183            )184        sa_tokens = self.process_proteins(proteins, protein_max_length)185 186        text_tokens = self.process_text(messages_list, text_max_length)187 188        return BatchFeature(189            data={190                "protein_input_ids": sa_tokens["input_ids"],191                "protein_attention_mask": sa_tokens["attention_mask"],192                "input_ids": text_tokens["input_ids"],193                "attention_mask": text_tokens["attention_mask"],194            }195        )196 197    def batch_decode(self, *args, **kwargs):198        return self.tokenizer.batch_decode(*args, **kwargs)199 200    def decode(self, *args, **kwargs):201        return self.tokenizer.decode(*args, **kwargs)202 203    def protein_batch_decode(self, *args, **kwargs):204        return self.protein_tokenizer.batch_decode(*args, **kwargs)205 206    def protein_decode(self, *args, **kwargs):207        return self.protein_tokenizer.decode(*args, **kwargs)208 209    # overwrite to save the protein tokenizer in a separate folder210    # Adapted from instructblip.processing_instructblip.py (https://github.com/huggingface/transformers/blob/9b479a245b793cac2a8b2e87c6d8e81bb24e20c4/src/transformers/models/instructblip/processing_instructblip.py#L191-L221)211    def save_pretrained(self, save_directory, **kwargs):212        # only save the protein tokenizer in sub_dir213        self.protein_tokenizer.save_pretrained(os.path.join(save_directory, self.protein_tokenizer_dir_name))214 215        # we modify the attributes so that only the text tokenizer are saved in the main folder216        protein_tokenizer_present = "protein_tokenizer" in self.attributes217        # find the correct position of it in the attributes list218        protein_tokenizer_index = self.attributes.index("protein_tokenizer") if protein_tokenizer_present else None219        if protein_tokenizer_present and protein_tokenizer_index is not None:220            self.attributes.remove("protein_tokenizer")221 222        outputs = super().save_pretrained(save_directory, **kwargs)223 224        if protein_tokenizer_present and protein_tokenizer_index is not None:225            self.attributes.insert(protein_tokenizer_index, "protein_tokenizer")226 227        return outputs228 229    # overwrite to load the protein tokenizer from a separate folder230    # Adapted from instructblip.processing_instructblip.py (https://github.com/huggingface/transformers/blob/9b479a245b793cac2a8b2e87c6d8e81bb24e20c4/src/transformers/models/instructblip/processing_instructblip.py#L191-L221)231    @classmethod232    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):233        processor = super().from_pretrained(pretrained_model_name_or_path, **kwargs)234 235        # if return_unused_kwargs a tuple is returned where the second element is 'unused_kwargs'236        if isinstance(processor, tuple):237            processor = processor[0]238        protein_tokenizer = AutoTokenizer.from_pretrained(239            pretrained_model_name_or_path, subfolder=cls.protein_tokenizer_dir_name240        )241 242        processor.protein_tokenizer = protein_tokenizer243 244        return processor245 246 247__all__ = ["EvollaProcessor"]248