Update tokenizer.py

tokenizer.py

@@ -8,38 +8,82 @@ from deepsvg.svglib.geom import Bbox
 
 
 class SVGTokenizer:
-    """SVG tokenizer - 从config.yaml加载所有配置，避免硬编码"""
+    """SVG tokenizer - supports both 8B and 4B models via config.yaml"""
     
-    def __init__(self, config_path: str = "./config.yaml"):
+    def __init__(self, config_path: str = "./config.yaml", model_size: str = None):
+        """
+        Initialize SVGTokenizer.
+        
+        Args:
+            config_path: Path to config.yaml
+            model_size: Model size ("8B" or "4B"). If None, uses default from config.
+        """
         with open(config_path, 'r') as f:
             self.config = yaml.safe_load(f)
         
+        # Determine model size
+        self.model_size = model_size or self.config.get('default_model_size', '8B')
+        if self.model_size not in self.config.get('models', {}):
+            raise ValueError(f"Invalid model_size: {self.model_size}. Must be one of: {list(self.config.get('models', {}).keys())}")
+        
         self._load_config()
         self.pixel2xy = self._create_pixel2xy_mapping()
     
+    def _get_model_specific_config(self, *keys):
+        """Get model-specific config value, with fallback to shared config."""
+        model_cfg = self.config.get('models', {}).get(self.model_size, {})
+        
+        # Navigate through nested keys in model-specific config
+        value = model_cfg
+        for key in keys:
+            if isinstance(value, dict) and key in value:
+                value = value[key]
+            else:
+                value = None
+                break
+        
+        # If not found in model-specific, try shared config
+        if value is None:
+            value = self.config
+            for key in keys:
+                if isinstance(value, dict) and key in value:
+                    value = value[key]
+                else:
+                    return None
+        
+        return value
+    
     def _load_config(self):
-        """从配置文件加载所有常量"""
-        # ========== Token相关配置 ==========
+        """Load all constants from configuration file with model-specific overrides."""
+        # ========== Token-related configs ==========
+        # Model-specific tokens
+        self.NUM_MASK_AND_EOM = self._get_model_specific_config('tokens', 'num_mask_and_eom')
+        self.BASE_OFFSET = self._get_model_specific_config('tokens', 'base_offset')
+        
+        # Shared tokens
         tokens_cfg = self.config['tokens']
         self.NUM_SVG_END = tokens_cfg['svg_end']
-        self.BASE_OFFSET = tokens_cfg['base_offset']
-        self.NUM_MASK_AND_EOM = tokens_cfg['num_mask_and_eom']
         self.NUM_END_TOKEN = tokens_cfg['num_end_token']
         
-        # ========== 坐标相关配置 ==========
+        # ========== Coordinate-related configs ==========
+        # Model-specific coordinates
+        self.PIX_PAD = self._get_model_specific_config('coordinates', 'pix_pad_offset')
+        self.COORD_PAD = self._get_model_specific_config('coordinates', 'coord_pad_offset')
+        
+        # Shared coordinates
         coords_cfg = self.config['coordinates']
         self.BBOX = coords_cfg['bbox']
-        self.PIX_PAD = coords_cfg['pix_pad_offset']
-        self.COORD_PAD = coords_cfg['coord_pad_offset']
         
-        # ========== 颜色相关配置 ==========
+        # ========== Color-related configs ==========
         colors_cfg = self.config['colors']
         self.COLOR_TOKEN_START_RAW = colors_cfg['color_token_start']
-        self.COLOR_START_OFFSET = colors_cfg['color_start_offset']
-        self.COLOR_END_OFFSET = colors_cfg['color_end_offset']
         self.MAX_COLOR_TOKENS = colors_cfg['max_color_tokens']
         
-        # ========== SVG命令值（用于 raster_svg 中的判断）==========
+        # Model-specific colors
+        self.COLOR_START_OFFSET = self._get_model_specific_config('colors', 'color_start_offset')
+        self.COLOR_END_OFFSET = self._get_model_specific_config('colors', 'color_end_offset')
+        
+        # ========== SVG command values ==========
         commands_cfg = self.config['svg_commands']
         self.CMD_MOVE = commands_cfg['move']
         self.CMD_LINE = commands_cfg['line']
@@ -47,41 +91,37 @@ class SVGTokenizer:
         self.CMD_ARC = commands_cfg['arc']
         self.CMD_CLOSE = commands_cfg['close']
         
-        # ========== 模型相关配置 ==========
+        # ========== Model-related configs ==========
         model_cfg = self.config['model']
         self.BOS_TOKEN_ID = model_cfg['bos_token_id']
         self.EOS_TOKEN_ID = model_cfg['eos_token_id']
         self.PAD_TOKEN_ID = model_cfg['pad_token_id']
         
-        # ========== Arc参数配置 ==========
+        # ========== Arc parameter configs ==========
         arc_cfg = self.config.get('arc', {})
         self.ARC_PARAM_OFFSET = arc_cfg.get('param_offset', 44500)
         self.ARC_PARAM_RANGE = arc_cfg.get('param_range', 100)
         self.ARC_PARAM_START = self.ARC_PARAM_OFFSET + self.BASE_OFFSET
         
-        # ========== 派生常量计算 ==========
-        # PIXEL_OFFSET: 从配置推导
-        # 命令token存储值 - BASE_OFFSET - PIXEL_OFFSET = CMD_MOVE
-        # (NUM_MASK_AND_EOM + NUM_SVG_END) - BASE_OFFSET - PIXEL_OFFSET = CMD_MOVE
+        # ========== Derived constants ==========
         self.PIXEL_OFFSET = (self.NUM_MASK_AND_EOM - self.BASE_OFFSET + 
                              self.NUM_SVG_END - self.CMD_MOVE)
         
-        # 命令token的实际范围
+        # Command token range
         self.CMD_TOKEN_START = self.NUM_MASK_AND_EOM + self.NUM_SVG_END
         self.CMD_TOKEN_END = self.PIX_PAD + self.NUM_SVG_END
         
-        # 坐标token起始
+        # Coordinate token start
         self.COORD_TOKEN_START = self.PIX_PAD + self.NUM_SVG_END
         
-        # 颜色token边界（坐标与颜色的分界）
+        # Color-coordinate boundary
         self.COLOR_COORD_BOUNDARY = self.COLOR_TOKEN_START_RAW + 1 + self.BASE_OFFSET
         
-        # 颜色阈值（用于 raster_svg 中判断）
-        # 减去 PIXEL_OFFSET 后的颜色token下限
+        # Color threshold for raster_svg
         self.COLOR_THRESHOLD = self.COLOR_TOKEN_START_RAW - self.PIXEL_OFFSET + 1
         
     def _create_pixel2xy_mapping(self) -> Dict[int, np.ndarray]:
-        """按照 dataset.py 逻辑创建 pixel 到 xy 的映射"""
+        """Create pixel to xy mapping following dataset.py logic."""
         pixel2xy = {}
         x = np.linspace(0, self.BBOX - 1, self.BBOX)
         y = np.linspace(0, self.BBOX - 1, self.BBOX)
@@ -89,13 +129,12 @@ class SVGTokenizer:
         xy_grid = (np.array((xx.ravel(), yy.ravel())).T).astype(int)
         
         for pixel, xy in enumerate(xy_grid):
-            # xy + COORD_PAD + NUM_SVG_END
             pixel2xy[pixel] = xy + self.COORD_PAD + self.NUM_SVG_END
             
         return pixel2xy
     
     def token_to_color(self, color_token: int) -> str:
-        """按照 dataset.py 的 token_to_color 逻辑"""
+        """Convert token to color following dataset.py logic."""
         try:
             if color_token == self.COLOR_TOKEN_START_RAW:
                 return "none"
@@ -123,37 +162,35 @@ class SVGTokenizer:
             return "#808080"
 
     def process_generated_tokens(self, output_ids: torch.Tensor) -> np.ndarray:
-        """
-        按照 dataset.py 的 __getitem__ 逻辑处理 tokens
-        """
-        # 移除 bos/eos
+        """Process generated tokens following dataset.py logic."""
+        # Remove bos/eos
         generated_pixels = output_ids[:, 1:-1].cpu().numpy().flatten()
         
         sample_xys = []
         
         for pixel in generated_pixels:
             try:
-                # 1. 命令tokens: CMD_TOKEN_START <= pixel < CMD_TOKEN_END
+                # 1. Command tokens: CMD_TOKEN_START <= pixel < CMD_TOKEN_END
                 if self.CMD_TOKEN_START <= pixel < self.CMD_TOKEN_END:
                     xy = np.array([pixel - self.BASE_OFFSET, 
                                    pixel - self.BASE_OFFSET]).astype(int)
                     sample_xys.append(xy)
                     
-                # 2. 坐标tokens: COORD_TOKEN_START <= pixel < COLOR_COORD_BOUNDARY
+                # 2. Coordinate tokens: COORD_TOKEN_START <= pixel < COLOR_COORD_BOUNDARY
                 elif self.COORD_TOKEN_START <= pixel < self.COLOR_COORD_BOUNDARY:
                     pixel_index = pixel - self.COORD_TOKEN_START
                     if pixel_index in self.pixel2xy:
                         xy = self.pixel2xy[pixel_index] - self.BASE_OFFSET
                         sample_xys.append(xy)
                     
-                # 3. Arc参数: ARC_PARAM_START + 1 <= pixel < ARC_PARAM_START + 1 + ARC_PARAM_RANGE
+                # 3. Arc parameters: ARC_PARAM_START + 1 <= pixel < ARC_PARAM_START + 1 + ARC_PARAM_RANGE
                 elif (self.ARC_PARAM_START + 1 <= pixel < 
                       self.ARC_PARAM_START + 1 + self.ARC_PARAM_RANGE):
                     value = pixel - self.ARC_PARAM_START - 1
                     xy = np.array([value, value]).astype(int)
                     sample_xys.append(xy)
                     
-                # 4. 颜色tokens: COLOR_COORD_BOUNDARY <= pixel < ARC_PARAM_START
+                # 4. Color tokens: COLOR_COORD_BOUNDARY <= pixel < ARC_PARAM_START
                 elif self.COLOR_COORD_BOUNDARY <= pixel < self.ARC_PARAM_START:
                     xy = np.array([pixel - self.BASE_OFFSET, 
                                    pixel - self.BASE_OFFSET]).astype(int)
@@ -169,15 +206,12 @@ class SVGTokenizer:
             return np.array([]).reshape(0, 2)
 
     def raster_svg(self, pixels: np.ndarray) -> Tuple[List[List[torch.Tensor]], List[int]]:
-        """
-        按照 dataset.py 的 raster_svg 逻辑
-        关键：pixels -= PIXEL_OFFSET 是核心转换步骤
-        """
+        """Convert pixels to SVG tensors following dataset.py logic."""
         try:
             if len(pixels) == 0:
                 return [[]], []
             
-            # ========== 关键步骤：减去 PIXEL_OFFSET ==========
+            # Key step: subtract PIXEL_OFFSET
             pixels = pixels - self.PIXEL_OFFSET
             
             svg_tensors = []
@@ -250,11 +284,11 @@ class SVGTokenizer:
                         path_tensor.append(cmd_tensor.tolist())
                         i += 2
                         
-                    # 颜色token: pix[0] >= COLOR_THRESHOLD
+                    # Color token: pix[0] >= COLOR_THRESHOLD
                     elif pix[0] >= self.COLOR_THRESHOLD:
                         if path_tensor:
                             svg_tensors.append(torch.tensor(path_tensor))
-                            # 逆转换：还原原始颜色token
+                            # Reverse transform: restore original color token
                             color_token = int(pix[0] + self.PIXEL_OFFSET - 1)
                             color_tensors.append(color_token)
                             path_tensor = []
@@ -266,7 +300,7 @@ class SVGTokenizer:
                     print(f"Error at position {i}: {e}")
                     break
             
-            # 处理剩余路径（无颜色）
+            # Handle remaining path (without color)
             if path_tensor:
                 svg_tensors.append(torch.tensor(path_tensor))
                 
@@ -280,7 +314,7 @@ class SVGTokenizer:
 
     def apply_colors_to_svg(self, svg_tensors: List[torch.Tensor], 
                            colors: Optional[List[int]]) -> SVG:
-        """应用颜色并创建最终SVG"""
+        """Apply colors and create final SVG."""
         paths = []
         
         if not svg_tensors: