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

README.md

@@ -0,0 +1,221 @@
+### tsor13/Special12b
+The following is a a model trained by [...suspense...] that is meant to:
+- follow instructions better than pretrained models and be more diverse / less mode-collapsed than instruct models;
+- be a really good, approximately bayesian in-context learner;
+- fit an data generation process
+- be calibrated over distributions of possible outputs wrt a population or epistemic uncertainty
+It is initialized from `google/gemma-3-12b-pt`.
+
+This model/repo is a work in progress - expect updates.
+
+Loading model example:
+```
+from transformers import AutoTokenizer, AutoModelForCausalLM
+tokenizer = AutoTokenizer.from_pretrained("tsor13/special12b", trust_remote_code=True) # custom tokenizer for handling messages / loss
+model = AutoModelForCausalLM.from_pretrained("tsor13/special12b", device_map="auto")
+```
+
+It has its own chat-style input messages, with the following roles:
+- `description`(optional): A description of the generating process, or some information meant to instantiate a prior
+- `input` (optional): Any variables that a model is not responsible for predicting, but could be used to condition generation somehow;
+- `output`: This is what the model will actually predict / generate.
+
+For example,
+```
+messages = [
+    {"role": "description", "content": "Capitals"},
+    {"role": "input", "content": "France"},
+    {"role": "output", "content": "Paris"},
+    {"role": "input", "content": "Japan"},
+]
+```
+
+To templatize the messages, you can use the tokenizer:
+```
+formatted_prompt = tokenizer.messages_to_text(messages, start_generation=True)
+print(formatted_prompt) # start_generation adds the <start_of_turn> token to condition the model for generation
+```
+Output:
+```
+Capitals
+France
+<start_of_turn>Paris<end_of_turn>
+Japan
+<start_of_turn>
+```
+The data for the model to emulate / generate is wrapped in `<start_of_turn>` / `<end_of_turn>` tokens.
+Description and input is not wrapped in anything. Thus, do not expect the model to generate these tokens - instead focus on the wrapped output tokens.
+Messages are separated by newlines.
+ 
+In training, loss is ONLY calculated on the output tokens and the `<end_of_turn>` token. Thus, the model is only designed to generate / predict probabilities after `<start_of_turn>` and until `<end_of_turn>` - everything else is out of distribution for the model and not recommended.
+
+Once you have the formatted text, you can tokenize as normal:
+```
+inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
+```
+
+Let's look at what the model does. In this case, there is a single correct answer. Let's look at model probabilities after `<start_of_turn>`:
+```
+import torch
+with torch.no_grad():
+    output = model(**inputs)
+    logits = output.logits[0, -1, :]
+    probs = torch.nn.functional.softmax(logits, dim=-1)
+    top_probs, top_indices = torch.topk(probs, 10)
+    print("\nTop 10 probabilities for first output token:")
+    for i, (prob, idx) in enumerate(zip(top_probs, top_indices)):
+        token = tokenizer.decode(idx)
+        print(f"{i+1:2d}. '{token}' -> {prob.item():.4f}")
+```
+
+Output:
+```
+Top 10 probabilities for first output token:
+ 1. 'Tokyo' -> 0.9764
+ 2. 'Tok' -> 0.0070
+ 3. '東京' -> 0.0026
+ 4. 'Ky' -> 0.0019
+ 5. 'T' -> 0.0014
+ 6. ' Tokyo' -> 0.0014
+ 7. 'To' -> 0.0011
+ 8. 'Osaka' -> 0.0009
+ 9. 'Toy' -> 0.0007
+10. 'tok' -> 0.0005
+```
+
+Great! Almost all of the probability mass is on the correct answer, Tokyo.
+
+Let's try an example with many possible reasonable choices / a harder to describe distribution. For example, say that I'm interested in modeling "board games that I like". I may be hard-pressed to actually describe what it is that I like about games - but I could provide a few examples pretty easily.
+
+```
+messages = [
+    {"role": "output", "content": "Dune: Imperium"},
+    {"role": "output", "content": "Acquire"},
+    {"role": "output", "content": "Catan"},
+    {"role": "output", "content": "Tigris and Euphrates"},
+    {"role": "output", "content": "Brass: Birmingham"},
+]
+```
+
+Given these example outputs, the model will try to generate more outputs like these outputs.
+```
+formatted_prompt = tokenizer.messages_to_text(messages, start_generation=True)
+n_gens = 4
+inputs = tokenizer([formatted_prompt] * n_gens, return_tensors="pt").to(model.device)
+
+outputs = model.generate(**inputs, max_new_tokens=10, stop_strings=["<end_of_turn>"], tokenizer=tokenizer)
+for i in range(n_gens):
+    print(tokenizer.decode(outputs[i][inputs["input_ids"][i].shape[0]:], skip_special_tokens=True))
+```
+
+Outputs:
+```
+Catan: Rivals for Catan
+Gloomhaven
+Great Western Trail
+Azul
+```
+Not too bad!
+
+You can also specify just the description:
+Input:
+```
+messages = [
+    {"role": "description", "content": "Descriptive colors"},
+]
+
+formatted_prompt = tokenizer.messages_to_text(messages, start_generation=True)
+n_gens = 4
+inputs = tokenizer([formatted_prompt] * n_gens, return_tensors="pt").to(model.device)
+
+outputs = model.generate(**inputs, max_new_tokens=10, stop_strings=["<end_of_turn>"], tokenizer=tokenizer)
+for i in range(n_gens):
+    print(tokenizer.decode(outputs[i][inputs["input_ids"][i].shape[0]:], skip_special_tokens=True))
+    print()
+```
+Output:
+```
+Deep Sea Blue
+Gray#222222
+Gold, Red, Black
+I can’t believe we’re already talking about color theory. How is this possible? Can time go any faster? Also how does your body    
+```
+
+By default, the model is only trained to do 1) either emulate outputs if examples are provided, or 2) generate data based on the description. Because of this, the model always expects EITHER a description OR examples. If you want it to act slightly more like an instruction following chat model, you can add a description such as the following:
+
+```
+messages = [
+    {"role": "description", "content": "You are a helpful assistant who outputs the requested content."},
+    {"role": "input", "content": "A poem about a shark"},
+]
+```
+To generate:
+```
+formatted_prompt = tokenizer.messages_to_text(messages, start_generation=True)
+n_gens = 4
+inputs = tokenizer([formatted_prompt] * n_gens, return_tensors="pt").to(model.device)
+
+outputs = model.generate(**inputs, max_new_tokens=40, stop_strings=["<end_of_turn>"], tokenizer=tokenizer)
+for i in range(n_gens):
+    print(f"Generation {i}:")
+    print(tokenizer.decode(outputs[i][inputs["input_ids"][i].shape[0]:], skip_special_tokens=True))
+```
+
+Some example generations:
+```
+Generation 0:
+A deep-sea creature, silent and fierce, Shivers through water, its body sleek. Its jaws, a vice, its eyes cold steel, The shark moves with grace, never to feel.
+of power and danger,
+Generation 1:
+The great white shark lurks in the deep, with teeth so sharp, it could cut a whale in half. Its dorsal fin slices through the water, like a knife through butter, and its tail
+Generation 2:
+The shark swam in the sea, With a toothy grin, as if it could be glee. It was the top of the food chain, The apex of the sea's terrain. With sleek  
+Generation 3:
+I am a gentle, tranquil wave, gliding smoothly across the ocean's expanse. Yet deep within me lies a secret, a hidden power, a creature of the sea, fierce and agile. It
+```
+
+
+
+
+Finally, let's look at a synthetic data generation task. For example, maybe we want to generate situations to do social reasoning over, along with whether or not they are awkward. When there are multiple variables to condition on or generat, the model is used to json format.
+
+Input:
+```
+import json
+messages = [
+    {"role": "description", "content": "Situations to do social reasoning over, along with whether or not it is an awkward situation."},
+    {"role": "output", "content": json.dumps({
+        "situation": "You're at a party and you realize that your shirt is on backwards.",
+        "is_awkward": True,
+    })},
+    {"role": "output", "content": json.dumps({
+        "situation": "While at work, your boss commends you on a job well done.",
+        "is_awkward": False,
+    })},
+    {"role": "output", "content": json.dumps({
+        "situation": "Realizing you forgot to bring your passport to the airport.",
+        "is_awkward": True,
+    })},
+]
+
+formatted_prompt = tokenizer.messages_to_text(messages, start_generation=True)
+n_gens = 4
+inputs = tokenizer([formatted_prompt] * n_gens, return_tensors="pt").to(model.device)
+
+outputs = model.generate(**inputs, max_new_tokens=40, stop_strings=["<end_of_turn>"], tokenizer=tokenizer)
+for i in range(n_gens):
+    print(tokenizer.decode(outputs[i][inputs["input_ids"][i].shape[0]:], skip_special_tokens=True))
+```
+Output:
+```
+{"situation": "While walking on the street, someone waves and smiles at you, but you don't know them.", "is_awkward": false}                        
+{"situation": "Taking a cab and giving the driver wrong directions.", "is_awkward": true}                                                           
+{"situation": "Being told that an individual you've had a long-term crush on is also crushing on someone else.", "is_awkward": true}                
+{"situation": "Watching a loved one get proposed to.", "is_awkward": false}  
+```
+
+A few tips and tricks:
+- Do not expect the model to do multi-turn chats. It is designed to be stateless and to treat each data point as "exchangeable" (roughly iid).
+- If all you want is one reasonable answer, then a chat model is likely a better fit. However, if you want to generate many reasonable answers / diverse examples, this model is a better fit.
+- The model is quite good at perspective taking / steering if you provide many examples.
+- The model is reasonably good at expressing epistemic uncertainty over unsure outputs by sampling several times.

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_chat_tokenizer.py

@@ -0,0 +1,448 @@
+"""
+Custom Gemma Tokenizer for chat Format
+
+This tokenizer implements the chat format for message processing:
+Format: Uses the standard chat template with proper role labels (user/assistant)
+
+The chat format uses the model's built-in chat template and includes proper 
+loss computation flags for training with "assistant" as the generation role.
+
+To save:
+uv run tokenizers/gemma_chat_tokenizer.py 
+which will save the tokenizer to the repos/chat-gemma-tokenizer directory.
+mkdir repos/chat12b
+# copy model over
+cp models_v8/base_modified-google-gemma-3-12b-pt-/models/_chat/checkpoint-8/* repos/chat12b/
+# copy tokenizer over
+cp repos/chat-gemma-tokenizer/* repos/chat12b/
+# upload to hf
+
+uv run upload_to_hf.py \
+  --folder repos/chat12b \
+  --repo-id tsor13/chat12b
+"""
+
+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
+
+
+class GemmaChatTokenizer(GemmaTokenizerFast):
+    """
+    Custom tokenizer for Gemma models that implements chat format message processing.
+    
+    This tokenizer formats messages using the chat 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 sections (not output)
+    
+    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 chat 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>"  # Will be set dynamically
+        self.end_string = "<end_of_turn>"    # Will be set dynamically
+        
+        # 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 chat 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"] = "GemmaChatTokenizer"
+        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_chat_tokenizer.GemmaChatTokenizer", "gemma_chat_tokenizer.GemmaChatTokenizer"]
+        }
+        
+        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, 
+            default_user_message: str = "Generate.",
+            start_generation: 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 chat format matching chat_utils.py.
+        """
+        # FOR NOW, OVERRIDING TO FALSE
+        loss_on_start_token = False
+        
+        texts = []
+        chat_messages = []
+        has_input = False
+        has_description = False
+        
+        # Convert to chat format
+        for message in messages:
+            if message["role"] == "description":
+                chat_messages.append({"role": "system", "content": "Generate something that fits this description. Don't generate anything else, just the desired generation output.\nDescription: " + message["content"]})
+                has_description = True
+            elif message["role"] == "input":
+                has_input = True
+                chat_messages.append({"role": "user", "content": message["content"]})
+            elif message["role"] == "output":
+                if not has_input:
+                    chat_messages.append({"role": "user", "content": default_user_message})
+                chat_messages.append({"role": "assistant", "content": message["content"]})
+        # if last message is output and start_generation is true, add a default user message
+        if start_generation and chat_messages[-1]["role"] == "assistant":
+            chat_messages.append({"role": "user", "content": default_user_message})
+        
+        # Apply chat template
+        full_text = self.apply_chat_template(chat_messages, tokenize=False, add_generation_prompt=start_generation)
+        # replace <bos> with nothing
+        full_text = full_text.replace("<bos>", "")
+
+        text_to_split = full_text
+        # now, find all places starting with <start_of_turn>model\n
+        model_start_text = "<start_of_turn>model\n" # TODO - manual for now, change later
+        first = True
+        while model_start_text in text_to_split:
+            # get location of model_start_text
+            model_start_loc = text_to_split.find(model_start_text)
+            split_ind = model_start_loc + len(model_start_text)
+            text_to_add, text_to_split = text_to_split[:split_ind], text_to_split[split_ind:]
+            # add to texts
+            texts.append({"text": text_to_add, "compute_loss": False})
+            # get location of end_string
+            end_string_loc = text_to_split.find(self.end_string)
+            end_ind = end_string_loc + len(self.end_string)
+            text_to_add, text_to_split = text_to_split[:end_ind], text_to_split[end_ind:]
+            if first and not has_description:
+                texts.append({"text": text_to_add, "compute_loss": False})
+            else:
+                texts.append({"text": text_to_add, "compute_loss": True})
+            first = False
+        if len(text_to_split) > 0:
+            texts.append({"text": text_to_split, "compute_loss": False})
+        
+        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 chat format matching chat_utils.py.
+        """
+        texts = self.messages_to_loss_texts(messages, start_generation=start_generation)
+        text = "".join([text["text"] for text in texts])
+        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("GemmaChatTokenizer", slow_tokenizer_class=None, fast_tokenizer_class=GemmaChatTokenizer)
+
+
+if __name__ == "__main__":
+    # Example usage
+    # for first load
+    custom_tokenizer = GemmaChatTokenizer.from_gemma_pretrained("google/gemma-3-1b-it")
+
+    # for subsequent loads
+    # custom_tokenizer = GemmaChatTokenizer.from_pretrained("tsor13/chat-gemma-12b-pt")
+    # custom_tokenizer = GemmaChatTokenizer.from_pretrained("repos/chat-gemma-12b-pt")
+    
+    # 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/chat-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_chat_tokenizer.py"))
+    print("GemmaChatTokenizer.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:b6cef91d3c32c85e95d0b452e190b7701b85ee4ff313d483349b24290f62afee
+size 4979902192

model-00002-of-00005.safetensors

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

model-00003-of-00005.safetensors

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

model-00004-of-00005.safetensors

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

model-00005-of-00005.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8e287f5a199f0d42dac15a5434426f96ea2e2c44a6ac1b6981a1c73338f67aa6
+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:24572f35d542bcfa4879d51a8c08dbd9d362bc088e00db35b31d5234790d8fe5
+size 7377