first commit

tests/models/evolla/test_processing_evolla.py

@@ -0,0 +1,251 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import random
+import unittest
+
+from transformers import (
+    AutoProcessor,
+    EvollaProcessor,
+)
+from transformers.testing_utils import require_torch
+from transformers.utils import is_torch_available
+
+from ...test_processing_common import ProcessorTesterMixin
+
+
+if is_torch_available():
+    import torch
+
+
+EVOLLA_VALID_AA = list("ACDEFGHIKLMNPQRSTVWY#")
+EVOLLA_VALID_FS = list("pynwrqhgdlvtmfsaeikc#")
+
+
+@require_torch
+class EvollaProcessorTest(ProcessorTesterMixin, unittest.TestCase):
+    processor_class = EvollaProcessor
+    model_id = "westlake-repl/Evolla-10B-hf"
+    input_keys = ["protein_input_ids", "protein_attention_mask", "input_ids", "attention_mask"]
+
+    @unittest.skip("EvollaProcessor requires `messages_list` and `proteins` inputs.")
+    def test_processor_with_multiple_inputs(self):
+        pass
+
+    def prepare_input_and_expected_output(self):
+        amino_acid_sequence = "AAAA"
+        foldseek_sequence = "dddd"
+        question = "What is the function of this protein?"
+
+        expected_output = {
+            "protein_input_ids": torch.tensor([[0, 13, 13, 13, 13, 2]]),
+            "protein_attention_mask": torch.tensor([[1, 1, 1, 1, 1, 1]]),
+            "input_ids": torch.tensor(
+                [
+                    [
+                        128000,
+                        128006,
+                        9125,
+                        128007,
+                        271,
+                        2675,
+                        527,
+                        459,
+                        15592,
+                        6335,
+                        430,
+                        649,
+                        4320,
+                        904,
+                        4860,
+                        922,
+                        13128,
+                        13,
+                        128009,
+                        128006,
+                        882,
+                        128007,
+                        271,
+                        3923,
+                        374,
+                        279,
+                        734,
+                        315,
+                        420,
+                        13128,
+                        30,
+                        128009,
+                        128006,
+                        78191,
+                        128007,
+                        271,
+                    ]
+                ]
+            ),
+            "attention_mask": torch.tensor(
+                [
+                    [
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                        1,
+                    ]
+                ]
+            ),
+        }
+        protein_dict = {"aa_seq": amino_acid_sequence, "foldseek": foldseek_sequence}
+        message = [
+            {"role": "system", "content": "You are an AI expert that can answer any questions about protein."},
+            {"role": "user", "content": question},
+        ]
+        return protein_dict, message, expected_output
+
+    def get_protein_tokenizer(self, **kwargs):
+        if "fix_mistral_regex" not in kwargs:
+            kwargs["fix_mistral_regex"] = True
+        return AutoProcessor.from_pretrained(self.tmpdirname, **kwargs).protein_tokenizer
+
+    def prepare_inputs_single(self):
+        proteins = {
+            "aa_seq": "".join(random.choices(EVOLLA_VALID_AA, k=100)),
+            "foldseek": "".join(random.choices(EVOLLA_VALID_FS, k=100)),
+        }
+        return proteins
+
+    def prepare_inputs_pair(self):
+        proteins = [
+            {
+                "aa_seq": "".join(random.choices(EVOLLA_VALID_AA, k=100)),
+                "foldseek": "".join(random.choices(EVOLLA_VALID_FS, k=100)),
+            },
+            {
+                "aa_seq": "".join(random.choices(EVOLLA_VALID_AA, k=100)),
+                "foldseek": "".join(random.choices(EVOLLA_VALID_FS, k=100)),
+            },
+        ]
+        return proteins
+
+    def prepare_inputs_long(self):
+        proteins = [
+            {
+                "aa_seq": "".join(random.choices(EVOLLA_VALID_AA, k=100)),
+                "foldseek": "".join(random.choices(EVOLLA_VALID_FS, k=100)),
+            },
+            {
+                "aa_seq": "".join(random.choices(EVOLLA_VALID_AA, k=2000)),
+                "foldseek": "".join(random.choices(EVOLLA_VALID_FS, k=2000)),
+            },
+        ]
+        return proteins
+
+    def prepare_inputs_short(self):
+        proteins = [
+            {
+                "aa_seq": "".join(random.choices(EVOLLA_VALID_AA, k=1)),
+                "foldseek": "".join(random.choices(EVOLLA_VALID_FS, k=1)),
+            },
+            {
+                "aa_seq": "".join(random.choices(EVOLLA_VALID_AA, k=100)),
+                "foldseek": "".join(random.choices(EVOLLA_VALID_FS, k=100)),
+            },
+        ]
+        return proteins
+
+    def prepare_inputs_empty(self):
+        proteins = [
+            {
+                "aa_seq": "",
+                "foldseek": "",
+            },
+            {
+                "aa_seq": "".join(random.choices(EVOLLA_VALID_AA, k=100)),
+                "foldseek": "".join(random.choices(EVOLLA_VALID_FS, k=100)),
+            },
+        ]
+        return proteins
+
+    def prepare_inputs(self, protein_types="pair"):
+        r"""
+        Prepare inputs for the test.
+
+        Args:
+            protein_types (`str`): the types of proteins to prepare.
+                - "single": a single correct protein.
+                - "pair": a pair of correct proteins.
+                - "long": a long sequence of correct proteins and a correct protein.
+                - "short": a short sequence of correct proteins (only have 1 aa) and a correct protein.
+                - "empty": an empty sequence of proteins and a correct protein.
+        """
+        if protein_types == "single":
+            proteins = self.prepare_inputs_single()
+        elif protein_types == "pair":
+            proteins = self.prepare_inputs_pair()
+        elif protein_types == "long":
+            proteins = self.prepare_inputs_long()
+        elif protein_types == "short":
+            proteins = self.prepare_inputs_short()
+        elif protein_types == "empty":
+            proteins = self.prepare_inputs_empty()
+        else:
+            raise ValueError(
+                f"protein_types should be one of 'single', 'pair', 'long','short', 'empty', but got {protein_types}"
+            )
+
+        questions = ["What is the function of the protein?"] * len(proteins)
+        messages_list = []
+        for question in questions:
+            messages = [
+                {"role": "system", "content": "You are an AI expert that can answer any questions about protein."},
+                {"role": "user", "content": question},
+            ]
+            messages_list.append(messages)
+        return proteins, messages_list
+
+    def test_model_input_names(self):
+        processor = self.get_processor()
+        proteins, messages_list = self.prepare_inputs()
+        inputs = processor(messages_list=messages_list, proteins=proteins, padding="longest", return_tensors="pt")
+        self.assertSetEqual(set(inputs.keys()), set(self.input_keys))