def encode_token_level_interleave(self, json_line):
        """
        Encodes text and interleaved vocal/instrumental codecs.
        Handles standard interleaving, CoT, and ICL-CoT based on args.
        """
        data = json.loads(json_line)
        ids = {}
        lens = {}

        # --- Initial Data Loading and Validation ---
        required_keys = ['splitted_lyrics', 'vocals_codec', 'instrumental_codec', 'audio_length_in_sec', 'genres', 'id']
        if self.args.use_audio_icl:
            # ICL requires additional keys
            required_keys.extend(['msa', 'codec'])

        if not all(key in data for key in required_keys):
            mode_str = "(ICL-CoT)" if self.args.use_audio_icl else ""
            print(f"Warning: Missing required keys in data for {data.get('id', 'unknown')} {mode_str}. Skipping.")
            print(f"Missing: {[k for k in required_keys if k not in data]}")
            return {}, {}, len(json_line)
        if not isinstance(data.get('splitted_lyrics'), dict) or 'segmented_lyrics' not in data['splitted_lyrics']:
             mode_str = "(ICL-CoT)" if self.args.use_audio_icl else ""
             print(f"Warning: Invalid 'splitted_lyrics' format in data for {data.get('id', 'unknown')} {mode_str}. Skipping.")
             return {}, {}, len(json_line)
        if not data['splitted_lyrics']['segmented_lyrics']: # Check if segmented_lyrics is empty
             mode_str = "(ICL-CoT)" if self.args.use_audio_icl else ""
             print(f"Warning: Empty 'segmented_lyrics' in data for {data.get('id', 'unknown')} {mode_str}. Skipping.")
             return {}, {}, len(json_line)

        segmented_lyrics = data['splitted_lyrics']['segmented_lyrics']

        try:
            raw_codec_vocals = np.load(data['vocals_codec'])
            raw_codec_instrumental = np.load(data['instrumental_codec'])
            # Load mixture codec only if needed for ICL prompt or future use
            raw_codec_mixture = None
            if self.args.use_audio_icl and self.args.audio_prompt_mode == "mixture":
                 raw_codec_mixture = np.load(data['codec'])
        except FileNotFoundError as e:
            mode_str = "(ICL-CoT)" if self.args.use_audio_icl else ""
            print(f"Error loading codec file {mode_str}: {e}. Skipping data ID {data['id']}.")
            return {}, {}, len(json_line)
        except Exception as e:
            mode_str = "(ICL-CoT)" if self.args.use_audio_icl else ""
            print(f"Error loading or processing codec npy for {data['id']} {mode_str}: {e}. Skipping.")
            # Estimate size even on error if possible
            bytes_processed = len(json_line)
            try: bytes_processed += get_size_in_bytes(raw_codec_vocals)
            except: pass
            try: bytes_processed += get_size_in_bytes(raw_codec_instrumental)
            except: pass
            try:
                 if raw_codec_mixture is not None: bytes_processed += get_size_in_bytes(raw_codec_mixture)
            except: pass
            return {}, {}, bytes_processed


        # Handle shape mismatch gracefully
        if raw_codec_vocals.shape != raw_codec_instrumental.shape:
            diff = abs(raw_codec_vocals.shape[-1] - raw_codec_instrumental.shape[-1])
            if diff <= 10: # Allow small difference
                min_len = min(raw_codec_vocals.shape[-1], raw_codec_instrumental.shape[-1])
                raw_codec_vocals = raw_codec_vocals[:, :min_len]
                raw_codec_instrumental = raw_codec_instrumental[:, :min_len]
                if DEBUG: print(f"Adjusted codec shapes for {data['id']} due to difference {diff}. New length: {min_len}")
            else:
                mode_str = "(ICL-CoT)" if self.args.use_audio_icl else ""
                print(f"Warning: Mismatch shape {raw_codec_vocals.shape} vs {raw_codec_instrumental.shape} for {data['id']} {mode_str}. Skipping.")
                bytes_processed = len(json_line) + max(get_size_in_bytes(raw_codec_vocals), get_size_in_bytes(raw_codec_instrumental))
                if raw_codec_mixture is not None: bytes_processed += get_size_in_bytes(raw_codec_mixture)
                return {}, {}, bytes_processed

        # Also check mixture codec shape if loaded
        if raw_codec_mixture is not None and raw_codec_mixture.shape[1] != raw_codec_vocals.shape[1]:
             # Attempt to trim mixture like vocals/instrumental if difference is small
             diff_mix = abs(raw_codec_mixture.shape[-1] - raw_codec_vocals.shape[-1])
             if diff_mix <= 10:
                 raw_codec_mixture = raw_codec_mixture[:, :raw_codec_vocals.shape[1]]
                 if DEBUG: print(f"Adjusted mixture codec shape for {data['id']} to match vocals/instrumental.")
             else:
                 print(f"Warning: Mixture codec shape {raw_codec_mixture.shape} mismatch with vocals/instrumental {raw_codec_vocals.shape} for {data['id']} (ICL-CoT). Skipping.")
                 bytes_processed = len(json_line) + get_size_in_bytes(raw_codec_vocals) + get_size_in_bytes(raw_codec_instrumental) + get_size_in_bytes(raw_codec_mixture)
                 return {}, {}, bytes_processed


        full_length_of_song = data['audio_length_in_sec']
        # Basic checks for validity
        if full_length_of_song <= 0 or raw_codec_vocals.ndim < 2 or raw_codec_vocals.shape[1] == 0:
             mode_str = "(ICL-CoT)" if self.args.use_audio_icl else ""
             print(f"Warning: Invalid audio length ({full_length_of_song}) or vocal codec shape ({raw_codec_vocals.shape}) for {data['id']} {mode_str}. Skipping.")
             # Calculate bytes processed before returning
             bytes_processed = len(json_line) + get_size_in_bytes(raw_codec_vocals) + get_size_in_bytes(raw_codec_instrumental)
             if raw_codec_mixture is not None: bytes_processed += get_size_in_bytes(raw_codec_mixture)
             return {}, {}, bytes_processed

        fps = raw_codec_vocals.shape[1] / full_length_of_song
        # Relaxed fps check
        if fps > 51 or fps < 49:
            mode_str = "(ICL-CoT)" if self.args.use_audio_icl else ""
            if DEBUG: print(f"fps={fps} is invalid for {data['id']} {mode_str}, skipping...")
            # Calculate bytes processed before returning
            bytes_processed = len(json_line) + get_size_in_bytes(raw_codec_vocals) + get_size_in_bytes(raw_codec_instrumental)
            if raw_codec_mixture is not None: bytes_processed += get_size_in_bytes(raw_codec_mixture)
            return {}, {}, bytes_processed


        doc_ids = []
        sentence_lens = [] # here sentence means segment
        instruction = self.args.instruction

        # --- Header Construction ---
        if self.args.use_audio_icl:
            # --- Start ICL Prompt Generation ---
            audio_prompt_codec_ids = []
            try:
                # Use the full range covered by lyrics segments for prompt sampling
                prompt_range_start_frame = segmented_lyrics[0].get('codec_frame_start', 0)
                prompt_range_end_frame = segmented_lyrics[-1].get('codec_frame_end', raw_codec_vocals.shape[1])

                # Ensure range is valid
                if prompt_range_start_frame >= prompt_range_end_frame:
                    raise ValueError(f"Invalid prompt range: start={prompt_range_start_frame}, end={prompt_range_end_frame}")

                # Extract relevant segment parts for prompt generation
                raw_codec_vocals_prompt_seg = raw_codec_vocals[:, prompt_range_start_frame:prompt_range_end_frame]
                raw_codec_instrumental_prompt_seg = raw_codec_instrumental[:, prompt_range_start_frame:prompt_range_end_frame]
                raw_codec_mixture_prompt_seg = None
                if raw_codec_mixture is not None:
                    raw_codec_mixture_prompt_seg = raw_codec_mixture[:, prompt_range_start_frame:prompt_range_end_frame]

                vocals_ids_prompt = Encoder.codectool.npy2ids(raw_codec_vocals_prompt_seg)
                instrumental_ids_prompt = Encoder.codectool.npy2ids(raw_codec_instrumental_prompt_seg)

                # Check if ids are valid lists/arrays
                if not isinstance(vocals_ids_prompt, (list, np.ndarray)) or not isinstance(instrumental_ids_prompt, (list, np.ndarray)):
                    raise TypeError("npy2ids did not return list/ndarray for prompt segment")
                if len(vocals_ids_prompt) == 0:
                    raise ValueError("Empty codec IDs generated for prompt segment")

                options_codecs = {}
                codec_step = 1 # How many codec tokens per original frame
                selected_option = self.args.audio_prompt_mode

                if selected_option == "dual":
                    codec_step = 2
                    if len(vocals_ids_prompt) != len(instrumental_ids_prompt):
                        raise ValueError(f"Length mismatch for interleaving prompt: {len(vocals_ids_prompt)} vs {len(instrumental_ids_prompt)}")
                    ids_segment_interleaved = rearrange([np.array(vocals_ids_prompt), np.array(instrumental_ids_prompt)], 'b n -> (n b)')
                    options_codecs['dual'] = ids_segment_interleaved
                elif selected_option == "mixture":
                    if raw_codec_mixture_prompt_seg is None: # Ensure mixture was loaded
                         raise ValueError("Mixture codec selected for prompt but not loaded/available.")
                    mixture_ids_prompt = Encoder.codectool.npy2ids(raw_codec_mixture_prompt_seg)
                    if not isinstance(mixture_ids_prompt, (list, np.ndarray)): raise TypeError("npy2ids failed for mixture prompt")
                    options_codecs['mixture'] = np.array(mixture_ids_prompt)
                elif selected_option == "inst":
                    options_codecs['inst'] = np.array(instrumental_ids_prompt)
                elif selected_option == "vocal":
                    options_codecs['vocal'] = np.array(vocals_ids_prompt)
                else:
                    raise ValueError(f"Invalid audio_prompt_mode: {selected_option}")

                # Determine prompt length in codec frames/tokens
                audio_prompt_length_in_secs = inverse_transform_sampling(Encoder.cdf_values, Encoder.x_values).item()
                audio_prompt_length_in_frames = int(audio_prompt_length_in_secs * fps)
                audio_prompt_length_in_codec_tokens = audio_prompt_length_in_frames * codec_step

                segment_duration_frames = prompt_range_end_frame - prompt_range_start_frame
                segment_duration_codec_tokens = segment_duration_frames * codec_step

                # Ensure prompt length is valid and fits within the segment
                if audio_prompt_length_in_codec_tokens <= 0:
                    audio_prompt_length_in_codec_tokens = int(1 * fps * codec_step) # Default to 1 second
                if audio_prompt_length_in_codec_tokens >= segment_duration_codec_tokens:
                    audio_prompt_length_in_codec_tokens = segment_duration_codec_tokens // 2 # Take half if too long
                    if DEBUG: print(f"Prompt length adjusted to {audio_prompt_length_in_codec_tokens} tokens (half segment) for {data['id']}")

                # --- Sample start position for the prompt ---
                max_start_token_index = segment_duration_codec_tokens - audio_prompt_length_in_codec_tokens
                if max_start_token_index < 0 : max_start_token_index = 0

                prompt_start_token_idx = 0
                # Try sampling from chorus if available
                chorus_list = [s for s in data.get('msa', []) if s.get('label') == 'chorus']
                if chorus_list:
                    random_chorus = random.choice(chorus_list)
                    chorus_start_sec = random_chorus.get('start', 0)
                    chorus_end_sec = random_chorus.get('end', full_length_of_song)

                    # Convert chorus times relative to the start of the lyrics segment range
                    chorus_start_frame_relative = max(0, int(chorus_start_sec * fps) - prompt_range_start_frame)
                    chorus_end_frame_relative = min(segment_duration_frames, int(chorus_end_sec * fps) - prompt_range_start_frame)

                    chorus_start_token_relative = chorus_start_frame_relative * codec_step
                    chorus_end_token_relative = chorus_end_frame_relative * codec_step

                    # Define valid start range within the chorus
                    chorus_max_start_token = chorus_end_token_relative - audio_prompt_length_in_codec_tokens
                    chorus_min_start_token = chorus_start_token_relative

                    if chorus_max_start_token > chorus_min_start_token:
                        prompt_start_token_idx = random.randint(chorus_min_start_token, chorus_max_start_token)
                    else:
                        prompt_start_token_idx = random.randint(0, max_start_token_index) # Fallback
                else:
                    prompt_start_token_idx = random.randint(0, max_start_token_index) # Random start

                prompt_end_token_idx = prompt_start_token_idx + audio_prompt_length_in_codec_tokens
                audio_prompt_codec_array = options_codecs[selected_option][prompt_start_token_idx:prompt_end_token_idx]

                # Optional: Filter prompts with low variation
                retry_count=0
                min_unique_ratio = 0.1
                while (len(np.unique(audio_prompt_codec_array)) < len(audio_prompt_codec_array) * min_unique_ratio) and retry_count < 5:
                    if DEBUG: print(f"Retrying prompt sampling due to low variation ({len(np.unique(audio_prompt_codec_array))} unique) for {data['id']}")
                    prompt_start_token_idx = random.randint(0, max_start_token_index)
                    prompt_end_token_idx = prompt_start_token_idx + audio_prompt_length_in_codec_tokens
                    audio_prompt_codec_array = options_codecs[selected_option][prompt_start_token_idx:prompt_end_token_idx]
                    retry_count += 1

                if retry_count == 5:
                    print(f"Warning: Could not find suitable audio prompt with enough variation for {data['id']} after 5 retries.")

                audio_prompt_codec_ids = ([Encoder.tokenizer.soa] + Encoder.codectool.sep_ids +
                                        list(audio_prompt_codec_array) +
                                        [Encoder.tokenizer.eoa])

            except Exception as e:
                print(f"Error generating ICL audio prompt for {data['id']}: {e}")
                print("Skipping sample due to ICL prompt error.")
                # Calculate bytes processed before returning
                bytes_processed = len(json_line) + get_size_in_bytes(raw_codec_vocals) + get_size_in_bytes(raw_codec_instrumental)
                if raw_codec_mixture is not None: bytes_processed += get_size_in_bytes(raw_codec_mixture)
                return {}, {}, bytes_processed # Skip sample


            # Construct ICL-CoT Header
            genre_str = '[Genre] ' + data['genres']
            complete_lyrics = '\n'.join([l.get('line_content', '') for l in segmented_lyrics])
            # Format: <Instruction> \n <Genre> \n <Lyrics> [start_of_reference] <Prompt> [end_of_reference]
            head = f'{instruction}\n{genre_str}\n{complete_lyrics}'
            head_ids = (Encoder.tokenizer.tokenize(head) +
                        Encoder.tokenizer.tokenize("[start_of_reference]") +
                        audio_prompt_codec_ids +
                        Encoder.tokenizer.tokenize("[end_of_reference]"))
            doc_ids.extend(head_ids)
            sentence_lens.append(len(head_ids))
            # --- End ICL Header ---

        elif self.args.cot:
            # Construct standard CoT Header (no audio prompt)
            genre_str = '[Genre] ' + data['genres']
            complete_lyrics = '\n'.join([l.get('line_content', '') for l in segmented_lyrics])
            # Format: <Instruction> \n <Genre> \n <Lyrics>
            head = f'{instruction}\n{genre_str}\n{complete_lyrics}'
            head_ids = Encoder.tokenizer.tokenize(head)
            doc_ids.extend(head_ids)
            sentence_lens.append(len(head_ids))
        # Else: No CoT, no ICL - header is implicitly handled per segment (instruction prepended)


        # --- Process Individual Segments ---
        for segment in segmented_lyrics:
            duration = segment.get('duration')
            frame_start = segment.get('codec_frame_start')
            frame_end = segment.get('codec_frame_end')
            line_content = segment.get('line_content')

            # Basic validation of segment data
            if duration is None or frame_start is None or frame_end is None or line_content is None:
                if DEBUG: print(f"Skipping segment due to missing keys: {segment} in {data['id']}")
                continue
            # Frame indices validity already checked for the whole song's fps calculation
            if not (0 <= frame_start < frame_end <= raw_codec_vocals.shape[1]):
                 if DEBUG: print(f"Invalid frame indices for segment in {data['id']}: start={frame_start}, end={frame_end}, total={raw_codec_vocals.shape[1]}. Skipping.")
                 continue
            if frame_end - frame_start <= 0:
                 if DEBUG: print(f"Segment frame length is zero or negative in {data['id']}: {frame_end - frame_start}. Skipping.")
                 continue
            # Minimum duration check (e.g., > 1 sec for target, or based on fps)
            min_target_segment_duration_sec = 1.0
            if self.args.use_audio_icl and duration < min_target_segment_duration_sec:
                 if DEBUG: print(f"Skipping target segment in {data['id']} (ICL) because duration {duration} < {min_target_segment_duration_sec}s")
                 continue
            # Check based on fps if not ICL (ensure at least 1 second)
            elif not self.args.use_audio_icl and frame_end - frame_start < fps:
                if DEBUG: print(f"Segment frame too short in {data['id']}: length={frame_end - frame_start} (< {fps}), skipping...")
                continue


            raw_codec_vocals_segment = raw_codec_vocals[:, frame_start:frame_end]
            raw_codec_instrumental_segment = raw_codec_instrumental[:, frame_start:frame_end]

            # --- Tokenize Text ---
            text_ids = []
            text = "" # Initialize text for potential error printing
            if self.args.cot or self.args.use_audio_icl: # CoT/ICL uses only line content for segment text
                text = line_content
            else: # Standard non-CoT mode
                text = instruction + '\n' + line_content
                 # Apply instruction dropout if enabled and not CoT/ICL
                if self.args.instruction_dropout_rate > 0.0 and np.random.rand() < self.args.instruction_dropout_rate:
                    text = line_content
            text_ids = Encoder.tokenizer.tokenize(text)


            # --- Process Codec ---
            try:
                vocals_ids_seg = Encoder.codectool.npy2ids(raw_codec_vocals_segment)
                instrumental_ids_seg = Encoder.codectool.npy2ids(raw_codec_instrumental_segment)

                if not isinstance(vocals_ids_seg, (list, np.ndarray)) or not isinstance(instrumental_ids_seg, (list, np.ndarray)):
                    raise TypeError("npy2ids did not return a list or ndarray for segment")
                if len(vocals_ids_seg) != len(instrumental_ids_seg):
                     mode_str = "(ICL-CoT)" if self.args.use_audio_icl else ""
                     print(f"Warning: Mismatch target vocal/inst IDs ({len(vocals_ids_seg)}/{len(instrumental_ids_seg)}) for {data['id']} {mode_str}. Skipping segment.")
                     continue
                if len(vocals_ids_seg) == 0: # Skip empty segments
                    if DEBUG: print(f"Skipping segment in {data['id']} because resulting codec IDs are empty.")
                    continue

                ids_segment_interleaved = rearrange([np.array(vocals_ids_seg), np.array(instrumental_ids_seg)], 'b n -> (n b)')
                ids_segment_interleaved_list = list(ids_segment_interleaved)

                # --- Construct Segment Tokens ---
                segment_tokens = []
                if self.args.cot or self.args.use_audio_icl:
                    # Format for CoT/ICL-CoT: [start_of_segment] <text> <SOA> <sep> <interleaved_codec> <EOA> [end_of_segment]
                    segment_tokens = (Encoder.tokenizer.tokenize('[start_of_segment]') +
                                     text_ids +
                                     [Encoder.tokenizer.soa] + Encoder.codectool.sep_ids +
                                     ids_segment_interleaved_list +
                                     [Encoder.tokenizer.eoa] +
                                     Encoder.tokenizer.tokenize('[end_of_segment]'))
                else:
                    # Standard non-CoT format: <text> <SOA> <sep> <interleaved_codec> <EOA>
                    codec_tokens = ([Encoder.tokenizer.soa] + Encoder.codectool.sep_ids +
                                    ids_segment_interleaved_list +
                                    [Encoder.tokenizer.eoa])
                    segment_tokens = text_ids + codec_tokens

                doc_ids.extend(segment_tokens)
                sentence_lens.append(len(segment_tokens))

            except Exception as e:
                mode_str = "(ICL-CoT)" if self.args.use_audio_icl else ""
                print(f"Error processing segment in encode_token_level_interleave {mode_str}: {e}")
                print(f"Data ID: {data['id']}")
                print(f"Segment: {segment}")
                print(f"Text Input: {text}") # Print the text that was tokenized


        # --- Finalize Document ---
        if len(doc_ids) > 0 and self.args.append_eod:
            # Add EOD only if we have successfully processed something (header or segments)
            if sentence_lens:
                 doc_ids.append(Encoder.tokenizer.eod)
                 sentence_lens[-1] += 1
            else:
                 if DEBUG: print(f"Skipping EOD for {data['id']} as no valid segments/header were processed.")


        key = "text" # hardcode key
        ids[key] = doc_ids
        lens[key] = sentence_lens

        bytes_processed = len(json_line) + get_size_in_bytes(raw_codec_vocals) + get_size_in_bytes(raw_codec_instrumental)
        if raw_codec_mixture is not None: # Add mixture size if it was loaded
             bytes_processed += get_size_in_bytes(raw_codec_mixture)
        return ids, lens, bytes_processed

        Encodes text and interleaved vocal/instrumental codecs.
        Handles standard interleaving, CoT, and ICL-CoT based on args.
        
finetune/core/preprocess_data_conditional_xcodec_segment.py