Initial upload of fine‑tuned Gemma + custom tokenizer

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

added_tokens.json

@@ -0,0 +1,3 @@
+{
+  "<image_soft_token>": 262144
+}

config.json

@@ -0,0 +1,62 @@
+{
+  "architectures": [
+    "Gemma3ForConditionalGeneration"
+  ],
+  "boi_token_index": 255999,
+  "eoi_token_index": 256000,
+  "eos_token_id": [
+    1,
+    106
+  ],
+  "image_token_index": 262144,
+  "initializer_range": 0.02,
+  "mm_tokens_per_image": 256,
+  "model_type": "gemma3",
+  "text_config": {
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "attn_logit_softcapping": null,
+    "cache_implementation": "hybrid",
+    "final_logit_softcapping": null,
+    "head_dim": 256,
+    "hidden_activation": "gelu_pytorch_tanh",
+    "hidden_size": 3840,
+    "initializer_range": 0.02,
+    "intermediate_size": 15360,
+    "max_position_embeddings": 131072,
+    "model_type": "gemma3_text",
+    "num_attention_heads": 16,
+    "num_hidden_layers": 48,
+    "num_key_value_heads": 8,
+    "query_pre_attn_scalar": 256,
+    "rms_norm_eps": 1e-06,
+    "rope_local_base_freq": 10000.0,
+    "rope_scaling": {
+      "factor": 8.0,
+      "rope_type": "linear"
+    },
+    "rope_theta": 1000000.0,
+    "sliding_window": 1024,
+    "sliding_window_pattern": 6,
+    "torch_dtype": "float32",
+    "use_cache": true,
+    "vocab_size": 262208
+  },
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.51.3",
+  "vision_config": {
+    "attention_dropout": 0.0,
+    "hidden_act": "gelu_pytorch_tanh",
+    "hidden_size": 1152,
+    "image_size": 896,
+    "intermediate_size": 4304,
+    "layer_norm_eps": 1e-06,
+    "model_type": "siglip_vision_model",
+    "num_attention_heads": 16,
+    "num_channels": 3,
+    "num_hidden_layers": 27,
+    "patch_size": 14,
+    "torch_dtype": "float32",
+    "vision_use_head": false
+  }
+}

gemma_explicit_tokenizer.py

@@ -0,0 +1,429 @@
+"""
+Custom Gemma Tokenizer for explicit Format
+
+This tokenizer implements the explicit format for message processing:
+Format: Uses the standard chat template with proper role labels (user/assistant)
+
+The explicit format uses the model's built-in chat template and includes proper 
+loss computation flags for training.
+
+To save:
+uv run tokenizers/gemma_explicit_tokenizer.py 
+which will save the tokenizer to the repos/explicit-gemma-tokenizer directory.
+mkdir repos/explicit12b
+# copy model over
+cp models_v8/base_modified-google-gemma-3-12b-pt-/models/_explicit/checkpoint-8/* repos/explicit12b/
+# copy tokenizer over
+cp repos/explicit-gemma-tokenizer/* repos/explicit12b/
+# upload to hf
+
+uv run upload_to_hf.py \
+  --folder repos/explicit12b \
+  --repo-id tsor13/explicit12b
+"""
+
+from typing import List, Dict, Any, Optional, Union
+from transformers import AutoTokenizer
+from transformers.models.gemma.tokenization_gemma_fast import GemmaTokenizerFast
+from transformers.models.gemma.tokenization_gemma import GemmaTokenizer
+import warnings
+import difflib
+import json
+import os
+import sys
+
+# # Add parent directory to path to import chat_utils
+# sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+# from chat_utils import chat_messages_to_text_loss, chat_messages_to_raw_text
+
+
+class GemmaExplicitTokenizer(GemmaTokenizerFast):
+    """
+    Custom tokenizer for Gemma models that implements explicit format message processing.
+    
+    This tokenizer formats messages using the explicit format where:
+    - Messages use the standard chat template with proper role labels
+    - Uses the model's built-in chat formatting
+    - Loss is computed on the assistant/output sections
+    
+    Attributes:
+        start_string (str): The starting string used for output generation (depends on tokenizer)
+        end_string (str): The ending string used for output generation (depends on tokenizer)
+    """
+    
+    def __init__(self, *args, **kwargs):
+        """
+        Initialize the custom tokenizer.
+        
+        Accepts the same arguments as GemmaTokenizerFast.
+        """
+        super().__init__(*args, **kwargs)
+        
+        # For explicit format, we use the tokenizer's own chat template
+        # The start/end strings will be determined by the chat template
+        self.start_string = "<start_of_turn>"
+        self.end_string = "<end_of_turn>"
+        
+        # Add custom attributes to the tokenizer config for saving/loading
+        if not hasattr(self, 'init_kwargs'):
+            self.init_kwargs = {}
+        self.init_kwargs['start_string'] = self.start_string
+        self.init_kwargs['end_string'] = self.end_string
+    
+    @classmethod
+    def from_gemma_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+        """
+        Load a tokenizer from a pretrained model or path.
+        
+        This method ensures our custom class is used instead of the base GemmaTokenizerFast.
+        """
+        # Load the base tokenizer first to get all configuration
+        base_tokenizer = GemmaTokenizerFast.from_pretrained(
+            pretrained_model_name_or_path, *args, **kwargs
+        )
+        
+        # Create new instance of our custom class by copying the base tokenizer
+        custom_tokenizer = cls.__new__(cls)
+        
+        # Copy all attributes from base tokenizer
+        for attr, value in base_tokenizer.__dict__.items():
+            setattr(custom_tokenizer, attr, value)
+        
+        # Initialize our custom attributes for explicit format
+        custom_tokenizer.start_string = "<start_of_turn>"
+        custom_tokenizer.end_string = "<end_of_turn>"
+        
+        # Update init_kwargs to include our custom attributes
+        if not hasattr(custom_tokenizer, 'init_kwargs'):
+            custom_tokenizer.init_kwargs = {}
+        custom_tokenizer.init_kwargs['start_string'] = custom_tokenizer.start_string
+        custom_tokenizer.init_kwargs['end_string'] = custom_tokenizer.end_string
+        
+        return custom_tokenizer
+    
+    def save_pretrained(self, save_directory: Union[str, os.PathLike], **kwargs):
+        """
+        Save the tokenizer to a directory, including custom configuration.
+        """
+        # Call parent save method
+        super().save_pretrained(save_directory, **kwargs)
+        
+        # Save custom configuration
+        config_file = os.path.join(save_directory, "tokenizer_config.json")
+        if os.path.exists(config_file):
+            with open(config_file, 'r') as f:
+                config = json.load(f)
+        else:
+            config = {}
+        
+        # Add our custom class info
+        config["tokenizer_class"] = "GemmaExplicitTokenizer"
+        config["start_string"] = self.start_string
+        config["end_string"] = self.end_string
+        # Point to our custom class in the uploaded file
+        config["auto_map"] = {
+            "AutoTokenizer": ["gemma_explicit_tokenizer.GemmaExplicitTokenizer", "gemma_explicit_tokenizer.GemmaExplicitTokenizer"]
+        }
+        
+        with open(config_file, 'w') as f:
+            json.dump(config, f, indent=2)
+    
+    def messages_to_loss_texts(
+            self,
+            messages: List[Dict[str, Any]], 
+            loss_on_start_token: bool = False, 
+        ) -> List[Dict[str, Any]]:
+        """
+        From messages (description / input / output) to texts (text / compute_loss) with whether or not loss should be calculated on the text for training.
+        Uses the explicit format matching chat_utils.py.
+        """
+        # FOR NOW, OVERRIDING TO FALSE
+        loss_on_start_token = False
+        
+        texts = []
+        has_description = False
+        first_output = True
+        
+        for message in messages:
+            role = message["role"]
+            content = message["content"]
+            
+            if role == "description":
+                has_description = True
+                text = f"{self.start_string}{role}\n{content}{self.end_string}\n"
+                texts.append({"text": text, "compute_loss": False, **message})
+            elif role == "input":
+                text = f"{self.start_string}{role}\n{content}{self.end_string}\n"
+                texts.append({"text": text, "compute_loss": False, **message})
+            elif role == "output":
+                if loss_on_start_token:
+                    raise ValueError("Loss on start token is not supported for chat formatters.")
+                else:
+                    start_text = f"{self.start_string}{role}\n"
+                    texts.append({"text": start_text, "compute_loss": False, **message})
+                    text = f"{content}{self.end_string}"
+                    # Apply conditional loss computation
+                    if first_output and not has_description:
+                        texts.append({"text": text, "compute_loss": False, **message})
+                    else:
+                        texts.append({"text": text, "compute_loss": True, **message})
+                    texts.append({"text": "\n", "compute_loss": False, **message})
+                    first_output = False
+            else:
+                raise ValueError(f"Unknown role: {role}. Must be description, input, or output.")
+        
+        return texts
+    
+    def messages_to_text(
+            self, 
+            messages: List[Dict[str, Any]], 
+            start_generation: bool = False, 
+        ) -> str:
+        """
+        Messages (description / input / output) to raw text (text).
+        Uses the explicit format matching chat_utils.py.
+        """
+        texts = self.messages_to_loss_texts(messages)
+        text = "".join([text["text"] for text in texts])
+        if start_generation:
+            text = text + self.start_string + "output\n"
+        return text
+    
+
+    def tokenize_messages(
+            self, 
+            messages: List[Dict[str, Any]] | List[List[Dict[str, Any]]],
+            start_generation: bool = False, 
+            **kwargs,
+        ):
+        """
+        For tokenizing from messages to texts. Supports batching. Good for generation
+        """
+        if isinstance(messages, list) and isinstance(messages[0], list):
+            # Handle list of lists of messages
+            all_texts = []
+            for message_list in messages:
+                texts = self.messages_to_text(message_list, start_generation)
+                all_texts.append(texts)
+        else:
+            # Handle single list of messages
+            texts = self.messages_to_text(messages, start_generation)
+            all_texts = [texts]
+        
+        # Tokenize all texts
+        processed = self(text=all_texts, **kwargs)
+        return processed
+    
+    
+    def tokenize_loss_texts(
+            self, 
+            texts: List[Dict[str, Any]], 
+            loss_on_start_token: bool = False, 
+            loss_on_eos: bool = False, 
+            include_eos: bool = True,
+        ):
+        """
+        Tokenize texts (text / compute_loss) to tokenized texts (input_ids / attention_mask / labels).
+
+        Needs more complex logic to handle the back and forth labeling.
+        """
+        if loss_on_eos:
+            raise ValueError("Loss on EOS is not currently supported.")
+        
+        # Handle single string input
+        if isinstance(texts, str):
+            processed = self(text=texts)
+            # Add EOS token if needed
+            if (self.eos_token_id is not None and 
+                processed["input_ids"][-1] != self.eos_token_id):
+                processed["input_ids"] = processed["input_ids"] + [self.eos_token_id]
+                processed["attention_mask"] = processed["attention_mask"] + [1]
+            return processed
+        
+        # Handle list of text dictionaries
+        all_processed = []
+        all_texts = ''
+        example_inds = []
+        dataset_inds = []
+        
+        for i, item in enumerate(texts):
+            processed = self(text=item["text"])
+            
+            # Remove BOS token from all but first item
+            if i != 0 and self.bos_token_id == processed["input_ids"][0]:
+                processed["input_ids"] = processed["input_ids"][1:]
+                processed["attention_mask"] = processed["attention_mask"][1:]
+            
+            # Remove EOS token if present at the end
+            if processed["input_ids"][-1] == self.eos_token_id:
+                processed["input_ids"] = processed["input_ids"][:-1]
+                processed["attention_mask"] = processed["attention_mask"][:-1]
+            
+            # Check for EOS token in the middle (with special handling for <|im_end|>)
+            if self.eos_token_id in processed["input_ids"]:
+                if not self.decode([self.eos_token_id]) == "<|im_end|>":
+                    raise ValueError(f"EOS token is present in input_ids: {processed['input_ids']}. Not currently supported.")
+            
+            # Set labels based on compute_loss flag
+            if item["compute_loss"]:
+                processed["labels"] = processed["input_ids"].copy()
+            else:
+                processed["labels"] = [-100] * len(processed["input_ids"])
+            
+            # Remove duplicate BOS tokens
+            if all_processed:
+                if processed["input_ids"][0] == self.bos_token_id:
+                    processed["input_ids"] = processed["input_ids"][1:]
+                    processed["attention_mask"] = processed["attention_mask"][1:]
+                    processed["labels"] = processed["labels"][1:]
+            
+            all_processed.append(processed)
+            all_texts += item["text"]
+            
+            # Handle example indices
+            this_num = -1
+            if 'example_ind' in item.keys():
+                if item["example_ind"] is not None:
+                    this_num = item["example_ind"]
+            example_inds.extend([this_num] * len(processed["input_ids"]))
+            
+            # Handle dataset indices
+            dataset_ind = -1
+            if "data_id" in item.keys():
+                if item["data_id"] is not None:
+                    dataset_ind = item["data_id"]
+            dataset_inds.extend([dataset_ind] * len(processed["input_ids"]))
+        
+        # Combine all processed results
+        processed = all_processed[0].copy()
+        processed["input_ids"] = [item for sublist in [p["input_ids"] for p in all_processed] for item in sublist]
+        processed["attention_mask"] = [item for sublist in [p["attention_mask"] for p in all_processed] for item in sublist]
+        processed["labels"] = [item for sublist in [p["labels"] for p in all_processed] for item in sublist]
+        processed["example_inds"] = example_inds
+        processed["data_ids"] = dataset_inds
+        
+        # Validate by tokenizing all_texts at once and comparing
+        processed_all = self(text=all_texts)
+        if len(processed_all["input_ids"]) != len(processed["input_ids"]):
+            warnings.warn(f"All texts are not the same length as the first text. Please check your dataset. {len(processed_all['input_ids'])} != {len(processed['input_ids'])}")
+            
+            # Generate diff for debugging
+            all_text = self.decode(processed_all["input_ids"], skip_special_tokens=False)
+            processed_text = self.decode(processed["input_ids"], skip_special_tokens=False)
+            
+            diff = difflib.unified_diff(all_text.splitlines(), processed_text.splitlines())
+            diff_str = "\n".join(diff)
+            print("Diff between texts:")
+            print(diff_str)
+            
+            # Token diff
+            all_tokens_str = '\n'.join([str(s) for s in processed_all["input_ids"]])
+            processed_tokens_str = '\n'.join([str(s) for s in processed["input_ids"]])
+            token_diff = difflib.unified_diff(all_tokens_str.splitlines(), processed_tokens_str.splitlines())
+            token_diff_str = "\n".join(token_diff)
+            print("Diff between tokenized texts:")
+            print(token_diff_str)
+        
+        # Add EOS token if needed
+        if (self.eos_token_id is not None and 
+            processed["input_ids"][-1] != self.eos_token_id):
+            processed["input_ids"] = processed["input_ids"] + [self.eos_token_id]
+            processed["example_inds"] = processed["example_inds"] + [-1]
+            processed["attention_mask"] = processed["attention_mask"] + [1]
+            if processed["labels"] is not None:
+                if loss_on_eos:
+                    processed["labels"] = processed["labels"] + [self.eos_token_id]
+                else:
+                    processed["labels"] = processed["labels"] + [-100]
+            if "data_ids" in processed:
+                processed["data_ids"] = processed["data_ids"] + [-1]
+            
+        if not include_eos:
+            # check if EOS token is present
+            if processed["input_ids"][-1] == self.eos_token_id:
+                # remove EOS token
+                processed["input_ids"] = processed["input_ids"][:-1]
+                processed["attention_mask"] = processed["attention_mask"][:-1]
+                processed["labels"] = processed["labels"][:-1]
+                processed["example_inds"] = processed["example_inds"][:-1]
+                processed["data_ids"] = processed["data_ids"][:-1]
+        
+        return processed
+    
+    def tokenize_messages(
+            self, 
+            messages: List[Dict[str, Any]], 
+            loss_on_start_token: bool = False, 
+            loss_on_eos: bool = False,
+            include_eos: bool = True,
+        ) -> Dict[str, Any]:
+        """
+        Intended for tokenize from messages to tokenized texts with the loss applied.
+        """
+        # First convert messages to text with loss computation flags
+        texts = self.messages_to_loss_texts(messages, loss_on_start_token)
+        
+        # Then tokenize the texts
+        return self.tokenize_loss_texts(texts, loss_on_eos, include_eos = include_eos)
+    
+
+
+
+# Register tokenizer classes for AutoTokenizer
+AutoTokenizer.register("GemmaExplicitTokenizer", slow_tokenizer_class=None, fast_tokenizer_class=GemmaExplicitTokenizer)
+
+
+if __name__ == "__main__":
+    # Example usage
+    # for first load
+    custom_tokenizer = GemmaExplicitTokenizer.from_gemma_pretrained("google/gemma-3-1b-it")
+
+    # Test messages in role/content format
+    test_messages = [
+        [
+            {"role": "description", "content": "This is a test task"},
+            {"role": "input", "content": "What is 2+2?"},
+            {"role": "output", "content": "4"},
+            {"role": "input", "content": "What is 3+3?"},
+        ],
+        [
+            {"role": "description", "content": "This is a test task"},
+            {"role": "output", "content": "4"},
+            {"role": "output", "content": "10"},
+            {"role": "output", "content": "13"},
+        ],
+        [
+            {"role": "output", "content": "4"},
+            {"role": "output", "content": "10"},
+            {"role": "output", "content": "13"},
+        ],
+        [
+            {"role": "input", "content": "What is 2+2?"},
+            {"role": "output", "content": "4"},
+            {"role": "input", "content": "What is 3+3?"},
+            {"role": "output", "content": "10"},
+            {"role": "input", "content": "What is 4+4?"},
+        ],
+    ]
+    for messages in test_messages:
+        # get messages to text_loss
+        texts = custom_tokenizer.messages_to_loss_texts(messages)
+
+        print("Texts with loss flags:")
+        for i, text in enumerate(texts):
+            print(f"  {i}: {text}")
+
+        text = custom_tokenizer.messages_to_text(messages, start_generation=True)
+        print(f"\nFull text with generation prompt:")
+        print(text)
+    
+    print("\nTesting save/load cycle:")
+    # Test saving and loading
+    tokenizer_path = "repos/explicit-gemma-tokenizer"
+    custom_tokenizer.save_pretrained(tokenizer_path)
+    print("Tokenizer saved successfully!")
+
+    # also save this file in the tokenizer_path
+    import shutil
+    shutil.copy(__file__, os.path.join(tokenizer_path, "gemma_explicit_tokenizer.py"))
+    print("GemmaExplicitTokenizer.py saved successfully!")

generation_config.json

@@ -0,0 +1,13 @@
+{
+  "bos_token_id": 2,
+  "cache_implementation": "hybrid",
+  "do_sample": true,
+  "eos_token_id": [
+    1,
+    106
+  ],
+  "pad_token_id": 0,
+  "top_k": 64,
+  "top_p": 0.95,
+  "transformers_version": "4.51.3"
+}

model-00001-of-00005.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0671cab2a8f94872810b970ac47b9060d358ca9ff131af31c216621f6d89fe2b
+size 4979902192

model-00002-of-00005.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:52dfe505e57fcd5835e806d0d19fd9c7ef19a11f0b6dae6877c084b16ac4db70
+size 4931296592

model-00003-of-00005.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0ed43a3bace6e02d584abcfb101adcdc08022c65e5ca473fb8fe2c4bbb1ee7be
+size 4931296656

model-00004-of-00005.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9c04b4408969e9473d8fd8174b881a4bd1220b93725e3b532671b2e26a2707be
+size 4931296656

model-00005-of-00005.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b7f85f21ad6209156a56a41a9de9b06354f1cffe51166757cabfeeca0f52d1e
+size 4601000928

special_tokens_map.json

@@ -0,0 +1,33 @@
+{
+  "boi_token": "<start_of_image>",
+  "bos_token": {
+    "content": "<bos>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eoi_token": "<end_of_image>",
+  "eos_token": {
+    "content": "<eos>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "image_token": "<image_soft_token>",
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4667f2089529e8e7657cfb6d1c19910ae71ff5f28aa7ab2ff2763330affad795
+size 33384568

training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:30e28ac87c293ee885688a284a96dbd0efd51290d74fbef6c1f426e728026204
+size 7313