Upload modeling_qwen3_jasper.py

modeling_qwen3_jasper.py

@@ -9,9 +9,9 @@ from transformers.models.qwen3.modeling_qwen3 import Qwen3MLP
 
 class TokenCompressor(nn.Module):
     """
-    自适应Token压缩模块
-    对于长度超过阈值的序列，使用adaptive_avg_pool1d进行压缩
-    压缩后长度 = 阈值 + 超出部分 * compression_ratio
+    Adaptive Token Compression Module
+    For sequences exceeding the threshold length, use adaptive_avg_pool1d for compression
+    Compressed length = threshold + excess_part * compression_ratio
     """
 
     def __init__(self, length_threshold: int = 512, compression_ratio: float = 0.3):
@@ -23,28 +23,28 @@ class TokenCompressor(nn.Module):
             self, token_embeddings: torch.Tensor, attention_mask: torch.Tensor
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """
-        对token embeddings进行自适应压缩
+        Perform adaptive compression on token embeddings
         Args:
             token_embeddings: [batch_size, seq_len, hidden_size]
             attention_mask: [batch_size, seq_len]
         Returns:
-            compressed_embeddings: 压缩后的embeddings
-            compressed_mask: 压缩后的attention mask
+            compressed_embeddings: Compressed embeddings
+            compressed_mask: Compressed attention mask
         """
         padding_side = 'right' if (attention_mask[:, -1] == 0).any() else 'left'
 
         compressed_embeddings_list = []
         compressed_masks_list = []
         for text_idx in range(token_embeddings.shape[0]):
-            # 获取当前样本的有效长度
+            # Get the effective length of current sample
             real_length = int(attention_mask[text_idx].sum().item())
             if real_length <= self.length_threshold:
-                # 根据padding方向提取有效的token embeddings
+                # Extract valid token embeddings based on padding direction
                 if padding_side == 'left':
-                    # 左填充：有效tokens在右边
+                    # Left padding: valid tokens are on the right
                     valid_embeddings = token_embeddings[text_idx:text_idx + 1, -real_length:, :]
                 else:
-                    # 右填充：有效tokens在左边
+                    # Right padding: valid tokens are on the left
                     valid_embeddings = token_embeddings[text_idx:text_idx + 1, :real_length, :]
                 compressed_embeddings_list.append(valid_embeddings)
                 compressed_masks_list.append([1] * real_length)
@@ -52,15 +52,15 @@ class TokenCompressor(nn.Module):
                 target_length = int(
                     self.length_threshold + (real_length - self.length_threshold) * self.compression_ratio
                 )
-                # 根据padding方向提取有效的token embeddings
+                # Extract valid token embeddings based on padding direction
                 if padding_side == 'left':
-                    # 左填充：有效tokens在右边
+                    # Left padding: valid tokens are on the right
                     valid_embeddings = token_embeddings[text_idx:text_idx + 1, -real_length:, :]
                 else:
-                    # 右填充：有效tokens在左边
+                    # Right padding: valid tokens are on the left
                     valid_embeddings = token_embeddings[text_idx:text_idx + 1, :real_length, :]
 
-                # 使用adaptive_avg_pool1d进行压缩
+                # Use adaptive_avg_pool1d for compression
                 compressed_embeddings_list.append(
                     F.adaptive_avg_pool1d(
                         valid_embeddings.transpose(1, 2), target_length
@@ -69,7 +69,7 @@ class TokenCompressor(nn.Module):
                 # print("valid_embeddings.shape,target_length,compressed_embeddings_list[-1].shape",valid_embeddings.shape,target_length,compressed_embeddings_list[-1].shape)
                 compressed_masks_list.append([1] * target_length)
 
-        # 重新组合为token_embeddings和attention_mask
+        # Reassemble token_embeddings and attention_mask
         new_seq_len = max((len(_mask) for _mask in compressed_masks_list))
         new_attention_mask = torch.tensor(
             [
@@ -83,7 +83,7 @@ class TokenCompressor(nn.Module):
             device=token_embeddings.device
         )
 
-        # 生成新的token_embeddings
+        # Generate new token_embeddings
         batch_size = token_embeddings.shape[0]
         hidden_size = token_embeddings.shape[2]
         new_token_embeddings = torch.zeros(
@@ -103,7 +103,6 @@ class TokenCompressor(nn.Module):
         return new_token_embeddings, new_attention_mask
 
 
-
 class JasperV2Encoder(Qwen3PreTrainedModel):
 
     def __init__(self, config: Qwen3Config):
@@ -134,7 +133,7 @@ class JasperV2Encoder(Qwen3PreTrainedModel):
             inputs_embeds=compressed_token_embeddings, attention_mask=attention_mask
         )["last_hidden_state"]
 
-        # 生成句向量
+        # Generate sentence vector
         input_mask_expanded = (
             attention_mask.unsqueeze(-1).expand(compressed_token_embeddings.size()).to(
                 compressed_token_embeddings.dtype)