processing_cxrmate2.py

import io
import math
import random
from io import BytesIO
from typing import Dict, List, Union

import cv2
import numpy as np
import pydicom
import requests
import torch
import torch.nn.functional as F
import transformers
from PIL import Image
from torch.nn.utils.rnn import pad_sequence
from transformers.feature_extraction_utils import BatchFeature
from transformers.image_utils import ImageInput
from transformers.tokenization_utils_base import PreTokenizedInput, TextInput

try:
    from .dataset import CXRMate2Dataset
except ImportError:
    from dataset import CXRMate2Dataset

# Ordered by oblique, lateral, AP, and then PA views so that PA views are closest in position to the generated tokens (and oblique is furtherest).
VIEW_ORDER = [
    None,        
    'nan',  # PadChest.
    'SWIMMERS',
    'LPO', 
    'RAO', 
    'LAO',
    'OBLICUA', # PadChest.
    'AP LLD',
    'AP RLD',
    'PA LLD',
    'PA RLD',
    'LLD',  # PadChest.
    'XTABLE LATERAL',
    'RL',
    'LL',   
    'Lateral',
    'LATERAL',
    'AP AXIAL',
    'ANTEROPOSTERIOR',  # PadChest.
    'AP',
    'GENERICA',  # PadChest (PA).
    'POSTEROANTERIOR',  # PadChest.
    'PA',
]


def compute_time_delta(event_time, reference_time, to_tensor=True):
    time_delta = reference_time - event_time
    time_delta = time_delta.total_seconds()
    assert isinstance(time_delta, float), f'time_delta should be float, not {type(time_delta)}.'
    if time_delta < 0:
        raise ValueError(f'time_delta should be greater than or equal to zero, not {time_delta}.')
    if to_tensor: 
        time_delta = torch.tensor(time_delta)
    return time_delta


class CXRMate2Processor(transformers.ProcessorMixin):

    attributes = ['image_processor', 'tokenizer']
    image_processor_class = 'AutoImageProcessor'
    tokenizer_class = 'AutoTokenizer'
    valid_kwargs = [
        'token_type_to_token_type_id',
        'max_generated_tokens',
    ]
    
    def __init__(
        self,
        image_processor,
        tokenizer,
        token_type_to_token: Dict[str, int],
        max_generated_tokens: int,
        embeddings_per_image: int,
        image_token: str,
        max_train_images_per_study: int,  # This includes current and prior images.
        generate_findings_token: str,
        generate_impression_token: str,
        convert_to_rgb: bool = False,
        mimic_cxr_normalisation: bool = True,
        **kwargs,
    ):
        super().__init__(image_processor, tokenizer)
        
        self.token_type_to_token = token_type_to_token
        self.max_generated_tokens = max_generated_tokens
        self.embeddings_per_image = embeddings_per_image
        self.image_token = image_token
        self.max_train_images_per_study = max_train_images_per_study
        
        self.generate_findings_token = generate_findings_token
        self.generate_impression_token = generate_impression_token

        self.convert_to_rgb = convert_to_rgb
        self.mimic_cxr_normalisation = mimic_cxr_normalisation
        
        self.generate_findings_token_id = self.tokenizer.convert_tokens_to_ids(self.generate_findings_token)
        self.generate_impression_token_id = self.tokenizer.convert_tokens_to_ids(self.generate_impression_token)
        
        self.time_delta_map = lambda x: 1 / math.sqrt((x / 3600) + 1)
        self.time_delta_monotonic_inversion = True
        self.zero_time_delta_value = self.time_delta_map(0.0)
        self.inf_time_delta_value = self.time_delta_map(float('inf'))

        self.prior_section_token_type_ids = [self.tokenizer.convert_tokens_to_ids(self.token_type_to_token[i]) for i in ['prior_findings', 'prior_impression']]
        self.section_token_type_ids = [self.tokenizer.convert_tokens_to_ids(self.token_type_to_token[i]) for i in ['indication', 'history', 'comparison', 'technique']]
    
        assert self.tokenizer.bos_token_id is not None, 'Tokenizer must have a bos_token_id.'
        assert self.tokenizer.sep_token_id is not None, 'Tokenizer must have a sep_token_id.'
        assert self.tokenizer.eos_token_id is not None, 'Tokenizer must have a eos_token_id.'
        assert self.tokenizer.pad_token_id is not None, 'Tokenizer must have a pad_token_id.'

    def __call__(
        self,
        images: Union[ImageInput, str, list[str], bytes, list[bytes]],
        image_datetime: Union[List[float], None] = None,
        findings: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput], None] = None,
        impression: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput], None] = None,
        views: Union[List[str]] = None,
        indication: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput], None] = None,
        history: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput], None] = None,
        comparison: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput], None] = None,
        technique: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput], None] = None,
        
        study_datetime: Union[float, None] = None,

        # Priors:
        prior_findings: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput], None] = None,
        prior_impression: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput], None] = None,
        prior_study_datetime: Union[List[float], None] = None,

        train: bool = False,
        **kwargs,
    ) -> BatchFeature:
        
        if isinstance(images, torch.Tensor):
            if images.ndim == 3:
                images = images.unsqueeze(0)
            if images.ndim == 4:
                images = images.unsqueeze(0)
        elif isinstance(images, list):
            if isinstance(images[0], (str, bytes)):
                images = [images]
        elif isinstance(images, (str, bytes)):
            images = [[images]]

        if image_datetime is not None and not all(isinstance(x, list) for x in image_datetime):
            image_datetime = [image_datetime]
        if views is not None and not all(isinstance(x, list) for x in views):
            views = [views]

        if indication is not None and not isinstance(indication, list):
            indication = [indication]
        if history is not None and not isinstance(history, list):
            history = [history]
        if comparison is not None and not isinstance(comparison, list):
            comparison = [comparison]
        if technique is not None and not isinstance(technique, list):
            technique = [technique]
        if prior_findings is not None and not isinstance(prior_findings, list):
            prior_findings = [[prior_findings]]
        if prior_findings is not None and isinstance(prior_findings, list) and not isinstance(prior_findings[0], list):
            prior_findings = [prior_findings]
        if prior_impression is not None and not isinstance(prior_impression, list):
            prior_impression = [[prior_impression]]
        if prior_impression is not None and isinstance(prior_impression, list) and not isinstance(prior_impression[0], list):
            prior_impression = [prior_impression]
        if study_datetime is not None and not isinstance(study_datetime, list):
            study_datetime = [study_datetime]

        if prior_study_datetime is not None and not all(isinstance(x, list) for x in prior_study_datetime):
            prior_study_datetime = [prior_study_datetime]

        batch_size = len(images)

        if views is None:
            views = [[None for _, _ in enumerate(i)] for i in images]

        batch = {
            'input_ids': {i: [] for i in range(batch_size)},
            'token_type_ids': {i: [] for i in range(batch_size)},
            'time_deltas': {i: [] for i in range(batch_size)},
            'time_deltas_mask': {i: [] for i in range(batch_size)},
            'attention_mask': [],
        }
        
        non_causal_2d_attention_mask = {i: [] for i in range(batch_size)}
        causal_2d_attention_mask = []
        
        # Map the prior study time delta values using the time delta map:
        if prior_study_datetime is not None:
            prior_study_time_deltas = [
                [self.time_delta_map(compute_time_delta(j, k)) if j is not None else float('nan') for j in i] for i, k in zip(prior_study_datetime, study_datetime, strict=True)
            ]

        # Findings and impression sections from prior studies:
        for i, token_type_id in zip([prior_findings, prior_impression], self.prior_section_token_type_ids, strict=True):
            if not i:
                continue
            assert len(i) == batch_size, f'Length of {i} must be equal to the batch size: {batch_size}.'
            for j in range(len(i)):
                if not i[j]:
                    continue
                for k in range(len(i[j])):
                    if not i[j][k]:
                        continue
                    batch['input_ids'][j].append(self.tokenizer.encode(i[j][k], add_special_tokens=False, return_tensors='pt')[0])
                    batch['token_type_ids'][j].append(torch.full((batch['input_ids'][j][-1].shape[-1],), token_type_id, dtype=torch.long))
                    non_causal_2d_attention_mask[j].append((batch['input_ids'][j][-1] != self.tokenizer.pad_token_id).long())
                    batch['time_deltas'][j].append(
                        torch.full(
                            (batch['input_ids'][j][-1].shape[-1],), 
                            prior_study_time_deltas[j][k] if prior_study_time_deltas is not None and prior_study_time_deltas[j][k] is not None else float('nan'), 
                            dtype=torch.float32,
                        ),
                    )
                    batch['time_deltas_mask'][j].append(torch.full((batch['input_ids'][j][-1].shape[-1],), 1.0, dtype=torch.float32))

        # Sections of the report for the prompt:
        for i, token_type_id in zip([indication, history, comparison, technique], self.section_token_type_ids, strict=True):
            if not i:
                continue
            assert len(i) == batch_size, f'Length of {i} must be equal to the batch size: {batch_size}.'
            for j, k in enumerate(i):
                if not k:
                    continue
                batch['input_ids'][j].append(self.tokenizer.encode(k, add_special_tokens=False, return_tensors='pt')[0])
                batch['token_type_ids'][j].append(torch.full((batch['input_ids'][j][-1].shape[-1],), token_type_id, dtype=torch.long))
                non_causal_2d_attention_mask[j].append((batch['input_ids'][j][-1] != self.tokenizer.pad_token_id).long())
                batch['time_deltas'][j].append(
                    torch.full((batch['input_ids'][j][-1].shape[-1],), self.zero_time_delta_value, dtype=torch.float32),
                )
                batch['time_deltas_mask'][j].append(torch.full((batch['input_ids'][j][-1].shape[-1],), 1.0, dtype=torch.float32))

        # Labels; findings and impression:
        if train:
            batch['label_ids'] = []
            for i, (j, k) in enumerate(zip(findings, impression, strict=True)):
                
                if j is not None and k is not None:
                    report = f'{self.tokenizer.bos_token}{j}{self.tokenizer.sep_token}{k}{self.tokenizer.eos_token}'
                elif j is not None and k is None:
                    report = f'{self.generate_findings_token}{j}{self.tokenizer.eos_token}'
                elif j is None and k is not None:
                    report = f'{self.generate_impression_token}{k}{self.tokenizer.eos_token}'
                else:
                    raise ValueError('Both findings and impression cannot be None.')
                
                report_ids = self.tokenizer.encode(
                    report,
                    truncation=True,
                    max_length=self.max_generated_tokens + 1,  # +1 to account for the bias between input and target.
                    return_tensors='pt',
                    add_special_tokens=False,
                )[0]
            
                # Labels for the decoder (shifted right by one for autoregression):
                batch['label_ids'].append(report_ids[1:].clone())

                # Remove last token identifier to match the sequence length of the labels:
                batch['input_ids'][i].append(report_ids[:-1])
            
                report_token_type_ids = self.token_ids_to_token_type_ids(token_ids=batch['input_ids'][i][-1])
                batch['token_type_ids'][i].append(report_token_type_ids)
                
                causal_2d_attention_mask.append((batch['input_ids'][i][-1] != self.tokenizer.pad_token_id).long())
                
                batch['time_deltas'][i].append(
                    torch.full((batch['input_ids'][i][-1].shape[-1],), self.zero_time_delta_value, dtype=torch.float32),
                )
                
                batch['time_deltas_mask'][i].append(torch.full((batch['input_ids'][i][-1].shape[-1],), 0.0, dtype=torch.float32))

        else:  # Add special tokens for generation:
            for i in range(batch_size):
                
                bos_token_id = self.tokenizer.bos_token_id
                batch['token_type_ids'][i].append(torch.tensor([self.tokenizer.convert_tokens_to_ids(self.token_type_to_token['findings'])], dtype=torch.long))
                                
                batch['input_ids'][i].append(torch.tensor([bos_token_id], dtype=torch.long))
                                
                causal_2d_attention_mask.append(torch.tensor([1], dtype=torch.long))
                
                batch['time_deltas'][i].append(torch.tensor([self.zero_time_delta_value], dtype=torch.float32))
                batch['time_deltas_mask'][i].append(torch.tensor([0.0], dtype=torch.float32))

        # Map the image time delta values using the time delta map:
        if study_datetime is not None:
            image_time_deltas = [[self.time_delta_map(compute_time_delta(j, k)) if j is not None else float('nan') for j in i] for i, k in zip(image_datetime, study_datetime, strict=True)] 
        else:
            image_time_deltas = [[float('nan') for _ in range(len(i))] for i in images]

        # Randomly select max_train_images_per_study if the number of images for a study exceeds max_train_images_per_study.
        for i in range(len(images)):
            if len(images[i]) > self.max_train_images_per_study:
                paired = list(zip(images[i], views[i], image_time_deltas[i], strict=True))
                sampled_pairs = random.sample(paired, self.max_train_images_per_study)
                images[i], views[i], image_time_deltas[i] = map(list, zip(*sampled_pairs, strict=True))

        # Sort based on views:
        images, views, image_time_deltas = self.sort_images(images, views, image_time_deltas)
                
        # Images:
        max_images = max(len(i) for i in images)
        for i in range(batch_size):
            for j in range(max_images):
                if j < len(images[i]):

                    image_np = None

                    if isinstance(images[i][j], bytes):
                        image = Image.open(io.BytesIO(images[i][j]))

                    elif isinstance(images[i][j], str):
                        if images[i][j].endswith('.dcm'):
                            assert self.mimic_cxr_normalisation, 'MIMIC-CXR normalisation must be True when using DICOM images.'
                            ds = pydicom.dcmread(images[i][j])
                            image_np = ds.pixel_array.astype(float)
                    
                        else:
                            if images[i][j].startswith('http://') or images[i][j].startswith('https://'):
                                response = requests.get(images[i][j], stream=True)
                                image = Image.open(BytesIO(response.content))
                            else:
                                image = Image.open(images[i][j])
                    
                    elif isinstance(images[i][j], Image.Image):
                        image = images[i][j]

                    if self.mimic_cxr_normalisation:

                        # MIMIC-CXR normalisation:
                        if image_np is None:
                            image_np = np.array(image.convert('L'), dtype=np.float32)
                        assert image_np.ndim == 2
                        min_val = image_np.min()
                        denom = image_np.max() - min_val
                        if denom == 0:
                            raise ValueError(f'Cannot normalise image with zero dynamic range (min and max both {min_val}).')
                        image_np = (image_np - min_val) / denom
                        image_uint8 = (image_np * 255).astype(np.uint8)
                        image_eq = cv2.equalizeHist(image_uint8)
                        image = Image.fromarray(image_eq)

                    if self.convert_to_rgb:
                        image = image.convert('RGB')

                    images[i][j] = self.image_processor(image, return_tensors='pt')['pixel_values'].squeeze(0)

                    batch['time_deltas'][i].insert(j, torch.full((self.embeddings_per_image,), image_time_deltas[i][j]))
                    batch['time_deltas_mask'][i].insert(j, torch.full((self.embeddings_per_image,), 1.0))

                    token_type_id = self.tokenizer.convert_tokens_to_ids(self.token_type_to_token['image']) if image_time_deltas[i][j] == self.zero_time_delta_value else self.tokenizer.convert_tokens_to_ids(self.token_type_to_token['prior_image'])
                    batch['token_type_ids'][i].insert(j, torch.full((self.embeddings_per_image,), token_type_id))
                    
                    non_causal_2d_attention_mask[i].insert(j, torch.full((self.embeddings_per_image,), 1))
                    
                else:
                    
                    batch['time_deltas'][i].insert(j, torch.full((self.embeddings_per_image,), 0.0))
                    batch['time_deltas_mask'][i].insert(j, torch.full((self.embeddings_per_image,), 0.0))

                    batch['token_type_ids'][i].insert(j, torch.full((self.embeddings_per_image,), self.tokenizer.convert_tokens_to_ids(self.token_type_to_token['image'])))
                    
                    non_causal_2d_attention_mask[i].insert(j, torch.full((self.embeddings_per_image,), 0))

                    
            images[i] = torch.stack(images[i])
            batch['input_ids'][i].insert(0, self.tokenizer.encode(self.image_token * self.embeddings_per_image * max_images, add_special_tokens=False, return_tensors='pt')[0])
                                                        
        batch['pixel_values'] = pad_sequence(images, batch_first=True, padding_value=0.0)
                
        # Concatenate input_ids, token_type_ids, time_deltas, and time_deltas_mask:
        batch['input_ids'] = [torch.cat(j, dim=0) for j in batch['input_ids'].values()]
        batch['token_type_ids'] = [torch.cat(j, dim=0) for j in batch['token_type_ids'].values()]
        batch['time_deltas'] = [torch.cat(j, dim=0) for j in batch['time_deltas'].values()]
        batch['time_deltas_mask'] = [torch.cat(j, dim=0) for j in batch['time_deltas_mask'].values()]
    
        # Concatentate, and convert label_ids into padded sequences:
        if train:
            batch['label_ids'] = [F.pad(i, (len(j) - len(i), 0), 'constant', self.tokenizer.pad_token_id) for i, j in zip(batch['label_ids'], batch['input_ids'], strict=True)]
            batch['label_ids'] = pad_sequence(batch['label_ids'], batch_first=True, padding_value=self.tokenizer.pad_token_id)   

        # Convert input_ids, token_type_ids, time_deltas, and time_deltas_mask into padded sequences:
        batch['input_ids'] = pad_sequence(batch['input_ids'], batch_first=True, padding_value=self.tokenizer.pad_token_id)
        batch['token_type_ids'] = pad_sequence(batch['token_type_ids'], batch_first=True, padding_value=0)
        batch['time_deltas'] = pad_sequence(batch['time_deltas'], batch_first=True, padding_value=0)
        batch['time_deltas_mask'] = pad_sequence(batch['time_deltas_mask'], batch_first=True, padding_value=0)    
        
        # Assert that time_delta values are between zero_time_delta_value and inf_time_delta_value:        
        check_1 = torch.all((batch['time_deltas'][~torch.isnan(batch['time_deltas'])] <= max([self.zero_time_delta_value, self.inf_time_delta_value])))
        check_2 = torch.all((batch['time_deltas'][~torch.isnan(batch['time_deltas'])] >= min([self.zero_time_delta_value, self.inf_time_delta_value])))                
        assert check_1 & check_2, 'Time delta values must be between zero_time_delta_value and inf_time_delta_value, or NaN if the time delta is missing.'
        
        # Mixed causality mask:
        non_causal_2d_attention_mask = [torch.cat(j, dim=0) for j in non_causal_2d_attention_mask.values()]
        batch['attention_mask'] = self.create_4d_mixed_causality_attention_mask(
            non_causal_2d_attention_mask, 
            causal_2d_attention_mask, 
            dtype=batch['pixel_values'].dtype,
        )        
            
        if not train:
            batch['initial_attention_mask'] = batch['attention_mask'].clone()  # For the first iteration of generation.
            batch['attention_mask'] = (batch['attention_mask'].squeeze(1).diagonal(dim1=1, dim2=2) == 0.0).long()
            
        # Create position_ids from time_deltas and attention_mask:
        batch['position_ids'] = self.position_ids_from_time_deltas_and_attention_mask(batch['time_deltas'], batch['attention_mask'])

        rows, cols = (batch['input_ids'] == self.tokenizer.sep_token_id).nonzero(as_tuple=True)
        assert all(batch['token_type_ids'][rows, cols] == self.tokenizer.convert_tokens_to_ids(self.token_type_to_token['findings']))
        
        rows, cols = (batch['input_ids'] == self.tokenizer.bos_token_id).nonzero(as_tuple=True)
        assert all(batch['token_type_ids'][rows, cols] == self.tokenizer.convert_tokens_to_ids(self.token_type_to_token['findings']))             

        return BatchFeature(data=batch)

    @staticmethod
    def sort_images(images, views, image_time_deltas):
        def sort_by_view(images, views, time_deltas):
            paired = list(zip(images, views, time_deltas, strict=True))
            sorted_pairs = sorted(paired, key=lambda x: VIEW_ORDER.index(x[1]))
            sorted_images, sorted_views, sorted_time_deltas = map(list, zip(*sorted_pairs, strict=True))
            return sorted_images, sorted_views, sorted_time_deltas

        # Apply sorting to each set of images, views, and time deltas:
        sorted_results = [sort_by_view(i, j, k) for i, j, k in zip(images, views, image_time_deltas, strict=True)]

        sorted_images = [result[0] for result in sorted_results]
        sorted_views = [result[1] for result in sorted_results]
        sorted_time_deltas = [result[2] for result in sorted_results]

        return sorted_images, sorted_views, sorted_time_deltas

    def token_ids_to_token_type_ids(self, token_ids, num_report_tokens=None):
        findings_id   = self.tokenizer.convert_tokens_to_ids(self.token_type_to_token['findings'])
        impression_id = self.tokenizer.convert_tokens_to_ids(self.token_type_to_token['impression'])

        # Initialize all as 'findings':
        token_type_ids = torch.full_like(token_ids, findings_id)

        # Detect sep_token_id positions:
        sep_positions = (token_ids == self.tokenizer.sep_token_id).nonzero(as_tuple=True)[0]

        if sep_positions.numel() > 0:
            # Use the first sep_token_id as the split point; change anything after it to 'impression' (this is fine as more than one sep_token_id will be treated as invalid for RL):
            first_sep_token_id = sep_positions[0].item()
            if first_sep_token_id + 1 < token_type_ids.numel():
                token_type_ids[first_sep_token_id + 1:] = impression_id

        return token_type_ids if num_report_tokens is None else token_type_ids[-num_report_tokens:]
        
    def create_4d_mixed_causality_attention_mask(self, non_causal_attention_mask, causal_attention_mask, dtype=torch.float32):
        attention_mask = []
        
        max_len = max([len(i) + len(j) for i, j in zip(non_causal_attention_mask, causal_attention_mask, strict=True)])
            
        for i in range(len(non_causal_attention_mask)):
            attention_mask.append(
                self.create_3d_mixed_causality_attention_mask(
                    non_causal_attention_mask[i], 
                    causal_attention_mask[i], 
                    dtype=dtype,
                )
            )
            pad_len = max_len - attention_mask[-1].shape[-1]
            attention_mask[-1] = F.pad(attention_mask[-1], (0, pad_len, 0, pad_len, 0, 0), 'constant', torch.finfo(dtype).min)
        attention_mask = torch.stack(attention_mask)
        
        return attention_mask

    @staticmethod
    def create_3d_mixed_causality_attention_mask(non_causal_1d_attention_mask, causal_1d_attention_mask, dtype=torch.float32):
        
        # Expand to 2D (seq_len x seq_len):
        upper_left = non_causal_1d_attention_mask[:, None] * non_causal_1d_attention_mask[None, :]
        
        if causal_1d_attention_mask is not None:

            prompt_seq_len = non_causal_1d_attention_mask.shape[-1] 
            report_seq_len = causal_1d_attention_mask.shape[-1]

            # Lower right of attention matrix (causal attention with lower triangular masking):
            causal_mask = torch.tril(torch.ones(report_seq_len, report_seq_len, device=causal_1d_attention_mask.device))
            lower_right = causal_1d_attention_mask[:, None] * causal_1d_attention_mask[None, :]
            lower_right = lower_right * causal_mask

            # Upper right of attention matrix (zeroes):
            upper_right = torch.zeros(prompt_seq_len, report_seq_len, dtype=torch.long, device=causal_1d_attention_mask.device)

            # Lower left of attention matrix:
            lower_left = non_causal_1d_attention_mask[None, :] * causal_1d_attention_mask[:, None]

            # Concatenate blocks:
            left = torch.cat((upper_left, lower_left), dim=0)
            right = torch.cat((upper_right, lower_right), dim=0)
            mixed_causality_3d_attention_mask = torch.cat((left, right), dim=-1)
        else:
            mixed_causality_3d_attention_mask = upper_left

        # Convert dtype and apply masking rules:
        mixed_causality_3d_attention_mask = mixed_causality_3d_attention_mask.to(dtype=dtype)
        mixed_causality_3d_attention_mask[mixed_causality_3d_attention_mask == 0] = torch.finfo(mixed_causality_3d_attention_mask.dtype).min
        mixed_causality_3d_attention_mask[mixed_causality_3d_attention_mask == 1] = 0.0
        
        # Add head dimension:
        mixed_causality_3d_attention_mask = mixed_causality_3d_attention_mask.unsqueeze(0)

        return mixed_causality_3d_attention_mask

    def position_ids_from_time_deltas_and_attention_mask(self, time_deltas, attention_mask):
        
        # Set NaNs to inf_time_delta_value:
        time_deltas = torch.nan_to_num(time_deltas, nan=self.inf_time_delta_value)

        # Convert attention mask to 2D if it is 4D:
        if attention_mask.dim() == 4:
            attention_mask = (attention_mask.squeeze(1).diagonal(dim1=1, dim2=2) == 0.0).long()

        # Set time deltas to NaN where the attention mask is 0:
        mask_value = float('inf') if self.time_delta_monotonic_inversion else -float('inf')
        masked_time_deltas = torch.where(attention_mask == 1, time_deltas, mask_value)

        # Sort time deltas and get indices
        sorted_time_deltas, col_indices = masked_time_deltas.sort(
            dim=1, descending=not self.time_delta_monotonic_inversion, stable=True
        )

        num_rows, num_cols = time_deltas.shape

        row_indices = torch.arange(num_rows, device=time_deltas.device).view(-1, 1).repeat(1, num_cols).view(-1)
        position_ids = torch.zeros_like(col_indices, device=time_deltas.device)
        position_ids[row_indices, col_indices.flatten()] = torch.arange(num_cols, device=time_deltas.device)[None, :].expand(num_rows, -1).flatten()

        # Apply the attention mask to zero out invalid positions
        position_ids = position_ids.masked_fill(attention_mask == 0, 1) # Following: https://github.com/huggingface/transformers/blob/c5f0288bc7d76f65996586f79f69fba8867a0e67/src/transformers/models/llama/modeling_llama.py#L1285.

        for i in range(position_ids.shape[0]):
            assert self.validate_position_ids(position_ids[i])
            
        return position_ids
    
    @staticmethod
    def validate_position_ids(tensor, repeat_value=1):
        unique, counts = torch.unique(tensor, return_counts=True)

        # Check if all integers from 0 to tensor.max() exist:
        full_range = torch.arange(0, tensor.max() + 1, device=tensor.device)
        if not torch.equal(unique.sort()[0], full_range):
            return False

        # Check for repeated values except for repeat_value:
        repeated = unique[counts > 1]
        if repeated.nelement() == 0:
            return True
        if not (repeated.numel() == 1 and repeated.item() == repeat_value):
            return False

        return True
    
    def batch_decode(self, *args, **kwargs):
        return self.tokenizer.batch_decode(*args, **kwargs)

    def decode(self, *args, **kwargs):
        return self.tokenizer.decode(*args, **kwargs)

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

    def split_and_decode_sections(self, token_ids):
        """
        Split the token identifiers into sections, then convert the token identifiers into strings.

        Argument/s:
            token_ids - token identifiers.

        Returns:
            token_type_ids - token type identifiers.
        """
        
        sections = {'findings': [], 'impression': []}
        for i in token_ids:
            findings_start_idx = (i == self.tokenizer.bos_token_id).int().argmax().item()
            findings_end_idx = (i == self.tokenizer.sep_token_id).int().argmax().item()
            sections['findings'].append(self.tokenizer.decode(i[findings_start_idx:findings_end_idx], skip_special_tokens=True))
            impression_start_idx = findings_end_idx + 1
            impression_end_idx = (i == self.tokenizer.eos_token_id).int().argmax().item()
            sections['impression'].append(self.tokenizer.decode(i[impression_start_idx:impression_end_idx], skip_special_tokens=True))

        return tuple(sections.values())

    def update_batch_for_rl(self, batch, completion_ids):

        batch_size, prompt_len = batch['token_type_ids'].shape

        # Number of completion tokens:
        num_completion_tokens = completion_ids.shape[1] - prompt_len - 1  # -1 for offset between input and label ids.
        
        # Update mask for completion tokens:
        completion_mask = (completion_ids[:,-(num_completion_tokens + 1):] != self.tokenizer.pad_token_id).float()  # +1 to ignore offset.
        batch['completion_mask'] = completion_mask
        completion_mask_expanded = completion_mask[:, None, None, 1:]  # Start from 1 to reintroduce offset.
        completion_mask_expanded_t = completion_mask[:, None, 1:, None]  # Start from 1 to reintroduce offset.
        
        upper_right = torch.zeros(batch_size, 1, prompt_len, num_completion_tokens, dtype=batch['initial_attention_mask'].dtype, device=completion_ids.device)

        bottom_right = torch.tril(torch.ones(num_completion_tokens, num_completion_tokens, device=completion_ids.device)).bool()
        bottom_right = bottom_right.unsqueeze(0).unsqueeze(0)
        bottom_right = bottom_right.expand(batch_size, -1, -1, -1) 
        bottom_right = bottom_right * completion_mask_expanded * completion_mask_expanded_t
        
        lower_left = batch['attention_mask'][:, None, None, :]
        lower_left = lower_left.expand(-1, -1, num_completion_tokens, -1)
        lower_left = lower_left * completion_mask_expanded_t
        
        right = torch.cat((upper_right, bottom_right), dim=2)
        right[right == 0] = torch.finfo(right.dtype).min
        right[right == 1] = 0.0
        
        lower_left[lower_left == 0] = torch.finfo(lower_left.dtype).min
        lower_left[lower_left == 1] = 0.0
        
        batch['attention_mask'] = torch.cat((batch['initial_attention_mask'], lower_left), dim=2)
        batch['attention_mask'] = torch.cat((batch['attention_mask'], right), dim=3)
        
        # initial_attention_mask was the 4D attention mask, whereas attention_mask was the 2D attention mask (i.e., not needed now that attention_mask is 4D):
        batch.pop('initial_attention_mask', None)

        # Convert remaining batch elements:
        new_token_type_ids = torch.stack([self.token_ids_to_token_type_ids(
            token_ids=i[-num_completion_tokens:], 
            # special_token_ids=[self.tokenizer.sep_token_id],
            # token_type_id_sections=[self.tokenizer.convert_tokens_to_ids(self.token_type_to_token['findings']), self.tokenizer.convert_tokens_to_ids(self.token_type_to_token['impression'])],
        ) for i in completion_ids])
        batch['token_type_ids'] = torch.cat((batch['token_type_ids'], new_token_type_ids), dim=1)
        batch['time_deltas'] = torch.nn.functional.pad(batch['time_deltas'], (0, num_completion_tokens), value=0.0)
        batch['time_deltas_mask'] = torch.nn.functional.pad(batch['time_deltas_mask'], (0, num_completion_tokens), value=0.0)
        
        start_values = batch['position_ids'].max(dim=1).values + 1
        end_values = start_values + num_completion_tokens
        position_ids = torch.stack([torch.arange(i, j, device=batch['position_ids'].device) for i, j in zip(start_values, end_values)])
        batch['position_ids'] = torch.cat((batch['position_ids'], position_ids), dim=1)
        
        batch['label_ids'] = completion_ids[:, 1:].clone()
        batch['input_ids'] = completion_ids[:, :-1]

        # Convert token identifiers that weren't sampled to pad_token_id:
        for i in range(batch_size):
            idx = (batch['label_ids'][i] == self.tokenizer.bos_token_id).nonzero(as_tuple=False)[0, 0].item()
            batch['label_ids'][i][:idx+1] = self.tokenizer.pad_token_id

        
        return batch

    def wrap_dataset(self, dataset):
        return CXRMate2Dataset(dataset)