File size: 13,368 Bytes
e478512
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
# coding=utf-8
# Copyright 2024 The HuggingFace Inc. team.
#
# 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,
# 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.
"""
Processor class for MiniCPMV with Nicheformer integration.
"""

from typing import List, Optional, Union
import torch
import re
from PIL import Image
import anndata as ad
import numpy as np
import os
from transformers import AutoTokenizer
from transformers.image_processing_utils import BatchFeature
from transformers.image_utils import ImageInput
from transformers.processing_utils import ProcessorMixin
from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
from transformers.utils import TensorType

from .image_processing_minicpmv import MiniCPMVBatchFeature
from .tokenization_nicheformer import NicheformerTokenizer


class MiniCPMVProcessor(ProcessorMixin):
    attributes = ["image_processor", "tokenizer", "gene_tokenizer"]
    image_processor_class = "AutoImageProcessor"
    tokenizer_class = "AutoTokenizer"
    gene_tokenizer_class = "AutoTokenizer"

    def __init__(self, image_processor=None, tokenizer=None, gene_tokenizer=None, **kwargs):
        super().__init__(image_processor, tokenizer, gene_tokenizer, **kwargs)
        self.version = kwargs.get("version", 2.6)
        self.gene_tokenizer = AutoTokenizer.from_pretrained("your/path/to/gene_tokenizer", trust_remote_code=True)
        technology_mean_path = 'your/path/to/gene_tokenizer/xenium_mean_script.npy'
        technology_mean = np.load(technology_mean_path)
        self.gene_tokenizer._load_technology_mean(technology_mean)

    def __call__(
        self,
        text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]],
        images: ImageInput = None,
        gene_data: Union[ad.AnnData, np.ndarray, List] = None,
        max_length: Optional[int] = None,
        do_pad: Optional[bool] = True,
        max_slice_nums: int = None,
        use_image_id: bool = None,
        return_tensors: Optional[Union[str, TensorType]] = TensorType.PYTORCH,
        **kwargs,
    ) -> MiniCPMVBatchFeature:

        # Step 1: Process images
        image_inputs = None
        if images is not None and any(img is not None for img in images):
            image_inputs = self.image_processor(
                images, do_pad=do_pad, max_slice_nums=max_slice_nums, return_tensors=return_tensors
            )
            # print(f"[DEBUG] 成功获取image_inputs : {image_inputs.keys()}")


        # Step 2: Process gene data
        gene_inputs = None
        if gene_data and len(gene_data) > 0 and len(gene_data[0]) > 0:
                adata = gene_data[0][0]
                gene_arrays = adata.X 
                gene_inputs = self.gene_tokenizer(gene_arrays)
                # print(f"[DEBUG] 成功获取gene_inputs : {gene_inputs.keys()}")


        # Step 3: Merge modalities
        return self._convert_all_modalities_to_inputs(
            image_inputs=image_inputs,
            gene_inputs=gene_inputs,
            texts=text,
            max_slice_nums=max_slice_nums,
            use_image_id=use_image_id,
            max_length=max_length,
            **kwargs,
        )

    def _convert_all_modalities_to_inputs(
        self,
        image_inputs,
        gene_inputs,
        texts: Union[str, List[str]],
        truncation=None,
        max_length=None,
        max_slice_nums=None,
        use_image_id=None,
        return_tensors=TensorType.PYTORCH,
        **kwargs,
    ):

        if isinstance(texts, str):
            texts = [texts]

        input_ids_list = []
        image_bounds_list = []
        gene_bounds_list = []

        image_pattern = "(<image>./</image>)"
        gene_pattern = "(<gene>./</gene>)"

        for index, text in enumerate(texts):
            
            image_tags = re.findall(image_pattern, text)
            if image_inputs is not None:
                image_sizes = image_inputs["image_sizes"]
                assert len(image_tags) == len(
                    image_sizes[index]
                ), f"Mismatch between image tags ({len(image_tags)}) and actual images ({len(image_sizes[index])})"

            # replace placeholders
            final_text = text
            if image_inputs is not None:
                text_chunks = final_text.split(image_pattern)
                final_text = ""
                for i in range(len(image_tags)):
                    final_text += text_chunks[i] + self.image_processor.get_slice_image_placeholder(
                        image_sizes[index][i], i, max_slice_nums, use_image_id
                    )
                final_text += text_chunks[-1]

            # === 处理 gene ===
            gene_tags = re.findall(gene_pattern, final_text)
            if gene_inputs is not None:
                text_chunks = re.split(gene_pattern, final_text)
                final_text = ""
                for i in range(len(gene_tags)):
                    gene_tokens = gene_inputs["input_ids"][index]
                    # gene_token_str = " ".join(map(str, gene_tokens.tolist()))
                    # final_text += text_chunks[i] + f"<gene_id>{i}</gene_id><gene>{gene_token_str}</gene>"
                    dummy_placeholder = "<unk>" * 32
                    final_text += text_chunks[i] + f"<gene_id>{i}</gene_id><gene>{dummy_placeholder}</gene>"
                final_text += text_chunks[-1]

            # print(f"[DeBUG] final_text: {final_text}")


            # 🔑 get input_ids and image_bounds directly
            input_ids, image_bounds, gene_bounds = self._convert(final_text, max_length)
            

            input_ids_list.append(input_ids)
            image_bounds_list.append(image_bounds)  # ✅ keep tensor
            gene_bounds_list.append(gene_bounds)

        # print(f"[DeBUG] input_ids: {input_ids_list}")
        # print(f"[DeBUG] input_ids length: {input_ids.size(0)}")
        # print(f"[DeBUG] image_bound: {image_bounds_list}")
        # print(f"[DeBUG] gene_bound: {gene_bounds_list}")

        # pad input_ids
        padded_input_ids, padding_lengths = self.pad(input_ids_list, padding_side="left")

        # shift bounds for padding
        for i, length in enumerate(padding_lengths):
            if image_bounds_list[i].numel() > 0:
                image_bounds_list[i] = image_bounds_list[i] + length
            if gene_bounds_list[i].numel() > 0:
                gene_bounds_list[i] = gene_bounds_list[i] + length

        attention_mask = padded_input_ids.ne(self.tokenizer.pad_token_id)
        
        labels = padded_input_ids.clone()
        labels[~attention_mask] = -100  # padding 不算loss
        # gene span 不算loss
        for i, gb in enumerate(gene_bounds_list):
            if torch.is_tensor(gb) and gb.numel() > 0:
                for (s, e) in gb.tolist():
                    labels[i, s:e] = -100
        
        # print(f"[DeBUG] padded_input_ids: {padded_input_ids}")
        # print(f"[DeBUG] attention_mask: {attention_mask}")
        # print(f"[DeBUG] image_bounds_list: {image_bounds_list}")
        # print(f"[DeBUG] gene_bounds_list: {gene_bounds_list}")
        
        data = {
            "input_ids": padded_input_ids,
            "attention_mask": attention_mask,
            "labels": labels,
            "image_bound": image_bounds_list,  # ✅ tensor [N,2]
            "gene_bound": gene_bounds_list,
        }

        if image_inputs:
            data.update(
                {
                    "pixel_values": image_inputs["pixel_values"],
                    "image_sizes": image_inputs["image_sizes"],
                    "tgt_sizes": image_inputs["tgt_sizes"],
                }
            )
        if gene_inputs:
            data.update(
                {
                    "gene_input_ids": gene_inputs["input_ids"],
                    "gene_attention_mask": gene_inputs["attention_mask"],
                 }
            )

        return MiniCPMVBatchFeature(data=data)

    def _convert(self, input_str, max_inp_length: Optional[int] = None):
        if self.version > 2.5 or not getattr(self.tokenizer, "add_bos_token", False):
            input_ids = self.tokenizer.encode(input_str)
        else:
            input_ids = [self.tokenizer.bos_id] + self.tokenizer.encode(input_str)
        if max_inp_length is not None:
            input_ids = input_ids[:max_inp_length]
        input_ids = torch.tensor(input_ids, dtype=torch.int32)
        
         # 找 image 边界
        image_start_tokens = torch.where(
            (input_ids == self.tokenizer.im_start_id) | (input_ids == self.tokenizer.slice_start_id)
        )[0] + 1
        image_end_tokens = torch.where(
            (input_ids == self.tokenizer.im_end_id) | (input_ids == self.tokenizer.slice_end_id)
        )[0]
        valid_image_nums = min(len(image_start_tokens), len(image_end_tokens))
        image_bounds = torch.stack(
            [image_start_tokens[:valid_image_nums], image_end_tokens[:valid_image_nums]], dim=1
        )
        
        # 找 gene 边界
        gene_start_tokens = torch.where(input_ids == self.tokenizer.gene_start_id)[0] + 1
        gene_end_tokens   = torch.where(input_ids == self.tokenizer.gene_end_id)[0]
        valid_gene_nums = min(len(gene_start_tokens), len(gene_end_tokens))
        gene_bounds = torch.stack(
            [gene_start_tokens[:valid_gene_nums], gene_end_tokens[:valid_gene_nums]], dim=1
        ) if valid_gene_nums > 0 else torch.zeros((0, 2), dtype=torch.int32)
        
        # print(f"[DETAIL] self.tokenizer.gene_start_id : {self.tokenizer.gene_start_id}")
        # print(f"[DETAIL] gene_start_tokens : {gene_start_tokens}")
        # print(f"[DETAIL] self.tokenizer.gene_end_id : {self.tokenizer.gene_end_id}")
        # print(f"[DETAIL] gene_end_tokens : {gene_end_tokens}")

        return input_ids, image_bounds, gene_bounds

    def batch_decode(self, *args, **kwargs):
        output_ids = args[0]
        result_text = []
        for result in output_ids:
            result = result[result != 0]
            if result[0] == self.tokenizer.bos_id:
                result = result[1:]
            if result[-1] == self.tokenizer.eos_id:
                result = result[:-1]
            result_text.append(self.tokenizer.decode(result, *args[1:], **kwargs).strip())
        return result_text

    def decode(self, *args, **kwargs):
        result = args[0]
        result = result[result != 0]
        if result[0] == self.tokenizer.bos_id:
            result = result[1:]
        if result[-1] == self.tokenizer.eos_id or (
            hasattr(self.tokenizer, "eot_id") and result[-1] == self.tokenizer.eot_id
        ):
            result = result[:-1]
        return self.tokenizer.decode(result, *args[1:], **kwargs).strip()

    @property
    def model_input_names(self):
        tokenizer_input_names = self.tokenizer.model_input_names
        image_processor_input_names = self.image_processor.model_input_names
        gene_tokenizer_input_names = self.gene_tokenizer.model_input_names
        return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names + gene_tokenizer_input_names))

    def pad(self, inputs, max_length=None, padding_value=0, padding_side="left"):
        items = []
        if isinstance(inputs[0], list):
            assert isinstance(inputs[0][0], torch.Tensor)
            for it in inputs:
                for tr in it:
                    items.append(tr)
        else:
            assert isinstance(inputs[0], torch.Tensor)
            items = inputs

        batch_size = len(items)
        shape = items[0].shape
        dim = len(shape)
        assert dim <= 2
        if max_length is None:
            max_length = 0
        max_length = max(max_length, max(item.shape[-1] for item in items))
        min_length = min(item.shape[-1] for item in items)
        dtype = items[0].dtype

        if dim == 0:
            return torch.stack([item for item in items], dim=0), [0]
        elif dim == 1:
            if max_length == min_length:
                return torch.stack([item for item in items], dim=0), [0] * batch_size
            tensor = torch.zeros((batch_size, max_length), dtype=dtype) + padding_value
        else:
            tensor = torch.zeros((batch_size, max_length, shape[-1]), dtype=dtype) + padding_value

        padding_length = []
        for i, item in enumerate(items):
            if dim == 1:
                if padding_side == "left":
                    tensor[i, -len(item) :] = item.clone()
                else:
                    tensor[i, : len(item)] = item.clone()
            elif dim == 2:
                if padding_side == "left":
                    tensor[i, -len(item) :, :] = item.clone()
                else:
                    tensor[i, : len(item), :] = item.clone()
            padding_length.append(tensor.shape[-1] - len(item))

        return tensor, padding_length