fix bugs

feature_extraction_xy_tokenizer.py

@@ -82,8 +82,8 @@ class ExtractorIterator:
             x_padded = F.pad(x, (0, padding_size), "constant", 0)
             output_tensor = x_padded.unfold(dimension=2, size=kernel, step=stride).squeeze(0).transpose(0, 1)
             output_lengths = torch.full((num_chunks,), kernel, dtype=torch.long)
-            if padding_size > 0:
-                output_lengths[-1] = kernel - padding_size
+            for i in range(num_chunks):
+                output_lengths[i] = min(output_lengths[i], L - stride * i)
             output_seq_no = torch.full((num_chunks,), seq_no, dtype=torch.long)
             return output_tensor, output_lengths, output_seq_no
 
@@ -118,6 +118,7 @@ class ExtractorIterator:
                     ]
                     yield BatchFeature({
                         **self.encode_func(list_x),
+                        "input_lengths": input_lengths,
                         "chunk_seq_no": input_seq_no.clone(),    
                     })
                     
@@ -135,6 +136,7 @@ class ExtractorIterator:
             ]
             yield BatchFeature({
                 **self.encode_func(list_x),
+                "input_lengths": input_lengths,
                 "chunk_seq_no": input_seq_no[:batch_num].clone(),    
             })
 
@@ -147,6 +149,9 @@ class XYTokenizerFeatureExtractor(WhisperFeatureExtractor):
         hop_length=160,
         chunk_length=30,
         n_fft=400,
+        n_samples=480000,
+        nb_max_frames=3000,
+        padding_side="right",
         padding_value=0.0,
         dither=0.0,
         return_attention_mask=False,
@@ -163,6 +168,9 @@ class XYTokenizerFeatureExtractor(WhisperFeatureExtractor):
             padding_value=padding_value,
             dither=dither,
             return_attention_mask=return_attention_mask,
+            n_samples=n_samples,
+            nb_max_frames=nb_max_frames,
+            padding_side=padding_side,
             **kwargs,
         )
         self.max_frequency = max_frequency if max_frequency is not None else sampling_rate / 2

modeling_xy_tokenizer.py

@@ -894,10 +894,9 @@ class XYTokenizerModel(XYTokenizerPreTrainedModel):
             
             # 1. Iterate through chunks and store intermediate results
             for chunk_features in features:
-                code_duration_length = features.duration_size // self.encoder_downsample_rate
                 # Always use return_dict=True for easier access to named outputs
                 chunk_output = self._encode(chunk_features, n_quantizers, return_dict=True)
-                valid_code_lengths = torch.clamp(chunk_output.codes_lengths, 0, code_duration_length)
+                valid_code_lengths = torch.clamp(chunk_features["input_lengths"], 0, features.duration_size) // self.encoder_downsample_rate
 
                 # Accumulate weighted commit loss
                 chunk_length = chunk_output.codes_lengths.sum().item()