mlx-raclate 0.1.0b1__py3-none-any.whl

mlx_raclate/utils/tokenizer_utils.py

@@ -0,0 +1,353 @@
+# Copyright © 2023-2024 Apple Inc.
+import json
+from functools import partial
+
+from transformers import AutoTokenizer
+
+REPLACEMENT_CHAR = "\ufffd"
+
+def _remove_space(x):
+    if x and x[0] == " ":
+        return x[1:]
+    return x
+
+
+class StreamingDetokenizer:
+    """The streaming detokenizer interface so that we can detokenize one token at a time.
+
+    Example usage is as follows:
+
+        detokenizer = ...
+
+        # Reset the tokenizer state
+        detokenizer.reset()
+
+        for token in generate(...):
+            detokenizer.add_token(token.item())
+
+            # Contains the whole text so far. Some tokens may not be included
+            # since it contains whole words usually.
+            detokenizer.text
+
+            # Contains the printable segment (usually a word) since the last
+            # time it was accessed
+            detokenizer.last_segment
+
+            # Contains all the tokens added so far
+            detokenizer.tokens
+
+        # Make sure that we detokenize any remaining tokens
+        detokenizer.finalize()
+
+        # Now detokenizer.text should match tokenizer.decode(detokenizer.tokens)
+    """
+
+    __slots__ = ("text", "tokens", "offset")
+
+    def reset(self):
+        raise NotImplementedError()
+
+    def add_token(self, token):
+        raise NotImplementedError()
+
+    def finalize(self):
+        raise NotImplementedError()
+
+    @property
+    def last_segment(self):
+        """Return the last segment of readable text since last time this property was accessed."""
+        text = self.text
+        if text and text[-1] != REPLACEMENT_CHAR:
+            segment = text[self.offset :]
+            self.offset = len(text)
+            return segment
+        return ""
+
+
+class NaiveStreamingDetokenizer(StreamingDetokenizer):
+    """NaiveStreamingDetokenizer relies on the underlying tokenizer
+    implementation and should work with every tokenizer.
+
+    Its complexity is O(T^2) where T is the longest line since it will
+    repeatedly detokenize the same tokens until a new line is generated.
+    """
+
+    def __init__(self, tokenizer):
+        self._tokenizer = tokenizer
+        self._tokenizer.decode([0])
+        self.reset()
+
+    def reset(self):
+        self.offset = 0
+        self._tokens = []
+        self._text = ""
+        self._current_tokens = []
+        self._current_text = ""
+
+    def add_token(self, token):
+        self._current_tokens.append(token)
+
+    def finalize(self):
+        self._tokens.extend(self._current_tokens)
+        self._text += self._tokenizer.decode(self._current_tokens)
+        self._current_tokens = []
+        self._current_text = ""
+
+    @property
+    def text(self):
+        if self._current_tokens:
+            self._current_text = self._tokenizer.decode(self._current_tokens)
+            if (
+                self._tokenizer.clean_up_tokenization_spaces
+                and self._current_text[-1] == " "
+            ):
+                self._current_text = self._current_text[:-1]
+        if self._current_text and self._current_text[-1] == "\n":
+            self._tokens.extend(self._current_tokens)
+            self._text += self._current_text
+            self._current_tokens.clear()
+            self._current_text = ""
+        return self._text + self._current_text
+
+    @property
+    def tokens(self):
+        return self._tokens
+
+
+class SPMStreamingDetokenizer(StreamingDetokenizer): ## SPM = SentencePiece Model, relevant for Llama and Gemma
+    """A streaming detokenizer for SPM models.
+
+    It adds tokens to the text if the next token starts with the special SPM
+    underscore which results in linear complexity.
+    """
+
+    def __init__(self, tokenizer, trim_space=True):
+        self.trim_space = trim_space
+
+        # Extract the tokens in a list from id to text
+        self.tokenmap = [""] * (max(tokenizer.vocab.values()) + 1)
+        for value, tokenid in tokenizer.vocab.items():
+            self.tokenmap[tokenid] = value
+
+        # Replace bytes with their value
+        for i in range(len(self.tokenmap)):
+            if self.tokenmap[i].startswith("<0x"):
+                self.tokenmap[i] = chr(int(self.tokenmap[i][3:5], 16))
+
+        self.reset()
+
+    def reset(self):
+        self.offset = 0
+        self._unflushed = ""
+        self.text = ""
+        self.tokens = []
+
+    def add_token(self, token):
+        self.tokens.append(token) ### append token to tokens list, not sure why it wasn't there in the first place as reset() sets an empty list
+        v = self.tokenmap[token]
+        if v[0] == "\u2581":
+            if self.text or not self.trim_space:
+                self.text += self._unflushed.replace("\u2581", " ")
+            else:
+                self.text = _remove_space(self._unflushed.replace("\u2581", " "))
+            self._unflushed = v
+        else:
+            self._unflushed += v
+
+    def finalize(self):
+        if self.text or not self.trim_space:
+            self.text += self._unflushed.replace("\u2581", " ")
+        else:
+            self.text = _remove_space(self._unflushed.replace("\u2581", " "))
+        self._unflushed = ""
+
+
+class BPEStreamingDetokenizer(StreamingDetokenizer): ## BPE = Byte Pair Encoding, relevant for GPT-2
+    """A streaming detokenizer for OpenAI style BPE models.
+
+    It adds tokens to the text if the next token starts with a space similar to
+    the SPM detokenizer.
+    """
+
+    _byte_decoder = None
+    _space_matches = (".", "?", "!", ",", "n't", "'m", "'s", "'ve", "'re")
+
+    def __init__(self, tokenizer):
+
+        self.clean_spaces = tokenizer.clean_up_tokenization_spaces
+
+        # Extract the tokens in a list from id to text
+        self.tokenmap = [None] * len(tokenizer.vocab)
+        for value, tokenid in tokenizer.vocab.items():
+            self.tokenmap[tokenid] = value
+
+        self.reset()
+
+        # Make the BPE byte decoder from
+        # https://github.com/openai/gpt-2/blob/master/src/encoder.py
+        self.make_byte_decoder()
+
+        self._added_ids = set(tokenizer.added_tokens_decoder.keys())
+
+    def reset(self):
+        self.offset = 0
+        self._unflushed = ""
+        self.text = ""
+        self.tokens = []
+
+    def _maybe_trim_space(self, current_text):
+        if len(current_text) == 0:
+            return current_text
+        elif current_text[0] != " ":
+            return current_text
+        elif not self.text:
+            return current_text[1:]
+        elif self.clean_spaces and current_text[1:].startswith(self._space_matches):
+            return current_text[1:]
+        return current_text
+
+    def add_token(self, token):
+        self.tokens.append(token) ### append token to tokens list, not sure why it wasn't there in the first place as reset() sets an empty list
+        v = self.tokenmap[token]
+        is_added = token in self._added_ids
+        if is_added or self._byte_decoder[v[0]] == 32:
+            current_text = bytearray(
+                self._byte_decoder[c] for c in self._unflushed
+            ).decode("utf-8")
+            self.text += self._maybe_trim_space(current_text)
+            if is_added:
+                self.text += v
+                self._unflushed = ""
+            else:
+                self._unflushed = v
+        else:
+            self._unflushed += v
+
+    def finalize(self):
+        current_text = bytearray(self._byte_decoder[c] for c in self._unflushed).decode(
+            "utf-8"
+        )
+        self.text += self._maybe_trim_space(current_text)
+        self._unflushed = ""
+
+    @classmethod
+    def make_byte_decoder(cls):
+        """See https://github.com/openai/gpt-2/blob/master/src/encoder.py for the rationale."""
+        if cls._byte_decoder is not None:
+            return
+
+        char_to_bytes = {}
+        limits = [
+            0,
+            ord("!"),
+            ord("~") + 1,
+            ord("¡"),
+            ord("¬") + 1,
+            ord("®"),
+            ord("ÿ") + 1,
+        ]
+        n = 0
+        for i, (start, stop) in enumerate(zip(limits, limits[1:])):
+            if i % 2 == 0:
+                for b in range(start, stop):
+                    char_to_bytes[chr(2**8 + n)] = b
+                    n += 1
+            else:
+                for b in range(start, stop):
+                    char_to_bytes[chr(b)] = b
+        cls._byte_decoder = char_to_bytes
+
+
+class TokenizerWrapper:
+    """A wrapper that combines an HF tokenizer and a detokenizer.
+
+    Accessing any attribute other than the ``detokenizer`` is forwarded to the
+    huggingface tokenizer.
+    """
+
+    def __init__(self, tokenizer, detokenizer_class=NaiveStreamingDetokenizer):
+        self._tokenizer = tokenizer
+        self._detokenizer = detokenizer_class(tokenizer)
+
+    def __getattr__(self, attr):
+        if attr == "detokenizer":
+            return self._detokenizer
+        elif attr.startswith("_"):
+            return self.__getattribute__(attr)
+        else:
+            return getattr(self._tokenizer, attr)
+
+    def __setattr__(self, attr, value):
+        if attr == "detokenizer":
+            raise AttributeError("Cannot set the detokenizer.")
+        elif attr.startswith("_"):
+            super().__setattr__(attr, value)
+        else:
+            setattr(self._tokenizer, attr, value)
+
+
+def _match(a, b):
+    if type(a) != type(b):
+        return False
+    if isinstance(a, dict):
+        return len(a) == len(b) and all(k in b and _match(a[k], b[k]) for k in a)
+    if isinstance(a, list):
+        return len(a) == len(b) and all(_match(ai, bi) for ai, bi in zip(a, b))
+
+    return a == b
+
+
+def _is_spm_decoder(decoder):
+    _target_description = {
+        "type": "Sequence",
+        "decoders": [
+            {"type": "Replace", "pattern": {"String": "▁"}, "content": " "},
+            {"type": "ByteFallback"},
+            {"type": "Fuse"},
+            {"type": "Strip", "content": " ", "start": 1, "stop": 0},
+        ],
+    }
+    return _match(_target_description, decoder)
+
+
+def _is_spm_decoder_no_space(decoder):
+    _target_description = {
+        "type": "Sequence",
+        "decoders": [
+            {"type": "Replace", "pattern": {"String": "▁"}, "content": " "},
+            {"type": "ByteFallback"},
+            {"type": "Fuse"},
+        ],
+    }
+    return _match(_target_description, decoder)
+
+
+def _is_bpe_decoder(decoder):
+    return isinstance(decoder, dict) and decoder.get("type", None) == "ByteLevel"
+
+
+def load_tokenizer(model_path, tokenizer_config_extra={}):
+    """Load a huggingface tokenizer and try to infer the type of streaming
+    detokenizer to use.
+
+    Note, to use a fast streaming tokenizer, pass a local file path rather than
+    a Hugging Face repo ID.
+    """
+    detokenizer_class = NaiveStreamingDetokenizer
+
+    tokenizer_file = model_path / "tokenizer.json"
+    if tokenizer_file.exists():
+        with open(tokenizer_file, "r") as fid:
+            tokenizer_content = json.load(fid)
+        if "decoder" in tokenizer_content:
+            if _is_spm_decoder(tokenizer_content["decoder"]):
+                detokenizer_class = SPMStreamingDetokenizer
+            elif _is_spm_decoder_no_space(tokenizer_content["decoder"]):
+                detokenizer_class = partial(SPMStreamingDetokenizer, trim_space=False)
+            elif _is_bpe_decoder(tokenizer_content["decoder"]):
+                detokenizer_class = BPEStreamingDetokenizer
+
+    return TokenizerWrapper(
+        AutoTokenizer.from_pretrained(model_path, **tokenizer_config_extra),
+        detokenizer_class,
+    )

mlx_raclate/utils/train.py

@@ -0,0 +1,249 @@
+import argparse
+import time
+
+from mlx_raclate.utils.utils import load, PIPELINES
+from mlx_raclate.tuner.datasets import load_dataset, DatasetArgs
+from mlx_raclate.tuner.trainer import Trainer, TrainingArgs
+
+train_tested = {
+    "text-classification": [
+        {"model": "Qwen/Qwen3-Embedding-0.6B", "special_model_config" : {}, "special_trainer_config" : {"use_chat_template": True}, "special_training_args" : {"max_length":8192}},
+        {"model": "answerdotai/ModernBERT-base", "special_model_config" : {}, "special_training_args" : {}},
+        {"model": "LiquidAI/LFM2-350M", "special_model_config" : {"use_chat_template": True}, "special_training_args" : {}},
+        {"model": "google/t5gemma-b-b-ul2", "special_model_config" : {}, "special_training_args" : {"max_length":8192}}, # failed
+        {"model": "google/embeddinggemma-300m", "special_model_config" : {}, "special_training_args" : {}} # failed
+    ],
+}
+
+# LFM2 CHAT TEMPLATE
+FORCED_CHAT_TEMPLATE = """
+{- bos_token -}}
+{%- set system_prompt = "" -%}
+{%- set ns = namespace(system_prompt="") -%}
+{%- if messages[0]["role"] == "system" -%}
+	{%- set ns.system_prompt = messages[0]["content"] -%}
+	{%- set messages = messages[1:] -%}
+{%- endif -%}
+{%- if tools -%}
+	{%- set ns.system_prompt = ns.system_prompt + ("\n" if ns.system_prompt else "") + "List of tools: <|tool_list_start|>[" -%}
+	{%- for tool in tools -%}
+		{%- if tool is not string -%}
+            {%- set tool = tool | tojson -%}
+		{%- endif -%}
+		{%- set ns.system_prompt = ns.system_prompt + tool -%}
+        {%- if not loop.last -%}
+            {%- set ns.system_prompt = ns.system_prompt + ", " -%}
+        {%- endif -%}
+	{%- endfor -%}
+	{%- set ns.system_prompt = ns.system_prompt + "]<|tool_list_end|>" -%}
+{%- endif -%}
+{%- if ns.system_prompt -%}
+	{{- "<|im_start|>system\n" + ns.system_prompt + "<|im_end|>\n" -}}
+{%- endif -%}
+{%- for message in messages -%}
+	{{- "<|im_start|>" + message["role"] + "\n" -}}
+	{%- set content = message["content"] -%}
+	{%- if content is not string -%}
+		{%- set content = content | tojson -%}
+	{%- endif -%}
+	{%- if message["role"] == "tool" -%}
+		{%- set content = "<|tool_response_start|>" + content + "<|tool_response_end|>" -%}
+	{%- endif -%}
+	{{- content + "<|im_end|>\n" -}}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+	{{- "<|im_start|>assistant\n" -}}
+{%- endif -%}
+"""
+
+DEFAULT_MODEL_PATH : str = "./trained_models/Qwen3-Embedding-0.6B_text-classification_20251219_001137/checkpoint-39939" #"Qwen/Qwen3-Embedding-0.6B" "answerdotai/ModernBERT-base" "google/t5gemma-b-b-ul2"
+DEFAULT_DATASET : str = "data/wines" # can be a local path or HF "argilla/synthetic-domain-text-classification" "data/20251205_1125"
+DEFAULT_TASK_TYPE : str = "text-classification"
+DEFAULT_BATCH_SIZE : int = 8
+DEFAULT_GRADIENT_ACCUMULATION_STEPS : int = 8
+DEFAULT_TRAIN_EPOCHS : int = 2
+DEFAULT_WEIGHT_DECAY : float = 0.01
+DEFAULT_LR : float = 2e-5 # 3e-5 for ModernBERT, 5e-5 for T5Gemma, 1e-5 for Qwen
+DEFAULT_LR_SCHEDULER_TYPE : str = "linear_schedule"
+DEFAULT_MIN_LR : float = 2e-6
+DEFAULT_WARMUP_RATIO : float = 0.03
+DEFAULT_WARMUP_STEPS : int = 0
+DEFAULT_SAVE_STEPS : int = 5000
+DEFAULT_LOGGING_STEPS : int = 64
+DEFAULT_EVAL_BATCH_SIZE : int = 8
+
+def init_args():
+    parser = argparse.ArgumentParser(description="Train or evaluate a classification model using MLX Raclate.")
+    # Dataset Init Params
+    parser.add_argument("--dataset", type=str, default=DEFAULT_DATASET, help="Local path or HF identifier of the dataset to use for training/evaluation.")
+    parser.add_argument("--text_field", type=str, default=None, help="Name of the text field in the dataset (if different from default).")
+    parser.add_argument("--text_pair_field", type=str, default=None, help="Name of the text pair field in the dataset (if applicable and different from default).")
+    parser.add_argument("--label_field", type=str, default=None, help="Name of the label field in the dataset (if different from default).")
+    parser.add_argument("--negative_field", type=str, default=None, help="Name of the negative samples field in the dataset (if applicable and different from default).")
+    parser.add_argument("--create_test", action='store_true', help="Set this flag to create a test split, if not already present in the dataset, out of the training set (validation set not affected).")
+
+    # Trainer / End Model Init Params
+    parser.add_argument("--model_path", type=str, default=DEFAULT_MODEL_PATH, help="Path to the pre-trained model or model identifier from a model hub.")
+    parser.add_argument("--task_type", type=str, default=DEFAULT_TASK_TYPE, help="Type of task (default: text-classification).")
+    parser.add_argument("--is_regression", default=False, action='store_true', help="Set this flag if the task is regression.")
+    parser.add_argument("--use_late_interaction", default=False, action='store_true', help="Set this flag to use late interaction for retrieval tasks (if applicable).")
+    parser.add_argument("--eval_only", dest="train", action='store_false', help="Set this flag to skip training and only evaluate.")
+    parser.add_argument("--use_chat_template", default=False, action='store_true', help="Use chat template for decoder models when there are text pairs.")
+    parser.add_argument("--force_separator", type=str, default=None, help="Force a specific separator between text pairs for decoder models, if not using chat template.")
+
+    # Training Params
+    parser.add_argument("--batch_size", type=int, default=DEFAULT_BATCH_SIZE, help="Batch size for training.")
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=DEFAULT_GRADIENT_ACCUMULATION_STEPS, help="Number of gradient accumulation steps.")
+    parser.add_argument("--num_train_epochs", type=int, default=DEFAULT_TRAIN_EPOCHS, help="Number of training epochs.")
+    parser.add_argument("--max_length", type=int, default=None, help="Maximum sequence length for the model inputs. If not specified, the model's default max length will be used.")
+    parser.add_argument("--freeze_embeddings", default=False, action='store_true', help="Set this flag to freeze embedding layers during training.")
+    parser.add_argument("--weight_decay", type=float, default=DEFAULT_WEIGHT_DECAY, help="Weight decay for the optimizer.")
+    # Optimizer and Scheduler Params (AdamW + schedulers)
+    parser.add_argument("--lr", type=float, default=DEFAULT_LR, help="Initial learning rate for the optimizer.")
+    parser.add_argument("--lr_scheduler_type", type=str, default=DEFAULT_LR_SCHEDULER_TYPE, help="Type of learning rate scheduler to use.")
+    parser.add_argument("--min_lr", type=float, default=DEFAULT_MIN_LR, help="Minimum learning rate for the scheduler.")
+    parser.add_argument("--warmup_ratio", type=float, default=DEFAULT_WARMUP_RATIO, help="Warmup ratio for learning rate scheduler.")
+    parser.add_argument("--warmup_steps", type=int, default=DEFAULT_WARMUP_STEPS, help="Number of warmup steps for learning rate scheduler (if set, steps override ratio).")
+    parser.add_argument("--max_grad_norm", type=float, default=1, help="Maximum gradient norm for gradient clipping (Default: 1).")
+    parser.add_argument("--resume_from_step", type=int, default=0, help="Step number to resume training from (if applicable). Will override warmup if steps are after warmup period.")
+    # Other Training Params
+    parser.add_argument("--logging_steps", type=int, default=DEFAULT_LOGGING_STEPS, help="Number of steps between logging training metrics.")
+    parser.add_argument("--save_steps", type=int, default=DEFAULT_SAVE_STEPS, help="Number of steps between model checkpoints.")
+    parser.add_argument("--eval_batch_size", type=int, default=DEFAULT_EVAL_BATCH_SIZE, help="Batch size for evaluation.")
+    parser.add_argument("--output_dir", type=str, default=None, help="Directory to save model checkpoints and logs.")
+    parser.set_defaults(train=True)
+    return parser.parse_args()
+
+def main():
+    args = init_args()
+
+    # Dataset Params
+    dataset : str = args.dataset
+    text_field : str = args.text_field
+    text_pair_field : str = args.text_pair_field
+    label_field : str = args.label_field
+    negative_field : str = args.negative_field
+    create_test : bool = args.create_test
+
+    # Trainer / End Model Params
+    model_path : str = args.model_path 
+    task_type : str = args.task_type
+    is_regression : bool = args.is_regression
+    use_late_interaction : bool = args.use_late_interaction
+    train : bool = args.train
+    use_chat_template : bool = args.use_chat_template
+    force_separator : str = args.force_separator
+
+    # Training Params
+    batch_size : int = args.batch_size
+    gradient_accumulation_steps : int = args.gradient_accumulation_steps
+    num_train_epochs : int = args.num_train_epochs
+    weight_decay : float = args.weight_decay
+    learning_rate : float = args.lr
+    lr_scheduler_type : str = args.lr_scheduler_type
+    min_lr : float = args.min_lr
+    warmup_ratio : float = args.warmup_ratio
+    warmup_steps : int = args.warmup_steps
+    logging_steps : int = args.logging_steps
+    save_steps : int = args.save_steps
+    eval_batch_size : int = args.eval_batch_size
+    resume_from_step : int = args.resume_from_step
+    max_length : int = args.max_length
+    freeze_embeddings : bool = args.freeze_embeddings
+    max_grad_norm : float = args.max_grad_norm
+
+    print(f"Training Mode : {train}")
+
+    if task_type not in PIPELINES:
+        raise ValueError(f"Task type {task_type} not supported. Choose from {PIPELINES.items()}")
+    
+    output_dir : str =  args.output_dir if args.output_dir else model_path.split("/")[-1] + "_" + task_type + "_" + time.strftime("%Y%m%d_%H%M%S")
+
+    # Load datasets
+    dataset_args = DatasetArgs(
+        data=dataset, 
+        task_type=task_type, 
+        text_field=text_field,
+        text_pair_field=text_pair_field,
+        label_field=label_field,
+        negative_field=negative_field,
+        test=create_test
+    )
+    
+    train_dataset, valid_dataset, test_dataset, id2label, label2id = load_dataset(dataset_args)
+
+    model_config={}
+    if task_type == "text-classification" and is_regression:
+        model_config={"is_regression":True}
+    if use_late_interaction and task_type in ["sentence-transformers","sentence-similarity"]:
+        model_config["use_late_interaction"] = True
+    if id2label:
+        model_config["id2label"] = id2label
+    if label2id:
+        model_config["label2id"] = label2id
+        
+    # Load model and tokenizer
+    model, tokenizer = load(
+        model_path, 
+        model_config=model_config, 
+        pipeline=task_type,
+        train=train,
+    )
+
+    # testing chat template
+    if use_chat_template:
+        messages = [
+            {"role": "user", "content": "test_prompt"},
+            {"role": "assistant", "content": "test_response"}
+        ]
+        if not getattr(tokenizer, "chat_template", None) and FORCED_CHAT_TEMPLATE:
+            tokenizer.chat_template = FORCED_CHAT_TEMPLATE
+        
+        templated = tokenizer.apply_chat_template(messages, tokenize=False)
+        print("Chat template working:", templated)
+
+    # Training arguments
+    training_args = TrainingArgs(
+        batch_size=batch_size,
+        gradient_accumulation_steps=gradient_accumulation_steps, 
+        max_length= max_length if max_length else model.config.max_position_embeddings,
+        resume_from_step=resume_from_step, # warmup will be ingnored if before this step and schedulers will only start after
+        num_train_epochs=num_train_epochs,
+        learning_rate=learning_rate, 
+        weight_decay=weight_decay,
+        freeze_embeddings=freeze_embeddings,
+        warmup_ratio=warmup_ratio, # can use warmup_steps or warmup_ratio
+        warmup_steps=warmup_steps, # if both set, warmup_steps will be used
+        lr_scheduler_type=lr_scheduler_type, # would default to "constant", can also use "cosine_decay" or "linear_schedule"
+        min_lr=min_lr,
+        max_grad_norm=max_grad_norm,
+        save_steps=save_steps,
+        logging_steps=logging_steps, # will be adjusted to be multiple of gradient_accumulation_steps inside Trainer
+        eval_batch_size=eval_batch_size,
+        output_dir=output_dir,
+        save_total_limit=None,
+        grad_checkpoint=True,
+        push_to_hub=False,
+    )
+
+    # Initialize trainer
+    trainer = Trainer(
+        model=model,
+        tokenizer=tokenizer,
+        task_type=task_type,
+        training_args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=valid_dataset,
+        use_chat_template=use_chat_template if task_type == "text-classification" else False,
+        force_separator=force_separator if task_type == "text-classification" else None,
+        label2id=label2id
+    )
+    
+    # Train or evaluate
+    if train:
+        trainer.train()
+    if test_dataset:
+        trainer.test(test_dataset)
+
+if __name__ == "__main__":
+    main()