Update tinymind.py

tinymind.py

@@ -17,9 +17,9 @@ from transformers import AutoTokenizer
 # ---------------------------
 
 DEFAULT_IMAGE_SIZE = 224
-DEFAULT_MAX_ROWS = 4
-DEFAULT_MAX_COLS = 4
-MIN_BLOCK_SIZE = 16
+DEFAULT_MAX_ROWS = 2
+DEFAULT_MAX_COLS = 2
+MIN_BLOCK_SIZE = 4
 
 
 # ---------------------------
@@ -250,7 +250,8 @@ def find_indices(tokens: torch.Tensor) -> Optional[Dict[int, Dict[int, List[Tupl
     """
     B, T = tokens.size()
     # 这里使用与原代码一致的 id 序列（如果 tokenizer 中不同，请改这里）
-    image_ids = [[3, i] for i in range(6, 22)] + [[3, 4]]  # 预设 pattern
+    # image_ids = [[3, i] for i in range(6, 22)] + [[3, 4]]  # 预设 pattern
+    image_ids = [[3, 6], [3, 7], [3, 10], [3, 11],[3, 4]]
     image_ids_tensor = torch.tensor(image_ids, device=tokens.device)
     len_image_ids = image_ids_tensor.size(1)
     if len_image_ids > tokens.size(1):