nexaai 1.0.19rc7__cp310-cp310-macosx_14_0_universal2.whl → 1.0.19rc8__cp310-cp310-macosx_14_0_universal2.whl

nexaai/binds/nexa_mlx/py-lib/vlm/modeling/trainer/trainer.py

@@ -1,296 +0,0 @@
-import json
-import warnings
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Union
-
-import mlx.core as mx
-import mlx.nn as nn
-import numpy as np
-from mlx.utils import tree_flatten, tree_map
-
-
-def get_prompt(model_type, processor, conversation):
-    if model_type == "paligemma":
-        return conversation
-
-    if "chat_template" in processor.__dict__.keys():
-        prompt = processor.apply_chat_template(
-            conversation,
-            tokenize=False,
-            add_generation_prompt=False,
-        )
-    elif "tokenizer" in processor.__dict__.keys():
-        prompt = processor.tokenizer.apply_chat_template(
-            conversation,
-            tokenize=False,
-            add_generation_prompt=False,
-        )
-
-    return prompt
-
-
-class Dataset:
-    def __init__(
-        self,
-        hf_dataset,
-        config,
-        processor,
-        image_processor=None,
-        take=None,
-        split=None,
-        image_resize_shape=None,
-    ):
-        if split is not None:
-            self.dataset = hf_dataset[split]
-        else:
-            self.dataset = hf_dataset
-        if take is not None:
-            self.dataset = self.dataset.take(take)
-        self.processor = processor
-        self.config = config
-        self.image_processor = image_processor
-        self.image_resize_shape = image_resize_shape
-
-    def __len__(self):
-        return len(self.dataset)
-
-    def __getitem__(self, idx):
-        from mlx_vlm.utils import prepare_inputs
-
-        item = self.dataset[idx]
-
-        images = item["images"]
-        conversations = item["messages"]
-        prompts = []
-
-        if isinstance(conversations, list) and isinstance(conversations[0], list):
-            for conversation in conversations:
-                if self.config["model_type"] == "pixtral":
-                    conversation = [json.loads(i) for i in conversation]
-                    if len(conversations) > 1:
-                        warnings.warn(
-                            "Pixtral batch processing is not supported yet. Set batch size to 1."
-                        )
-
-                prompt = get_prompt(
-                    self.config["model_type"], self.processor, conversation
-                )
-                prompts.append(prompt)
-
-        else:
-            if self.config["model_type"] == "pixtral":
-                conversations = [json.loads(i) for i in conversations]
-            prompt = get_prompt(
-                self.config["model_type"], self.processor, conversations
-            )
-            prompts.append(prompt)
-
-        image_token_index = self.config["image_token_index"]
-
-        inputs = prepare_inputs(
-            self.processor,
-            images,
-            prompts,
-            image_token_index,
-            self.image_resize_shape,
-        )
-        input_ids = inputs["input_ids"]
-        pixel_values = inputs["pixel_values"]
-        mask = inputs["attention_mask"]
-        kwargs = {
-            k: v
-            for k, v in inputs.items()
-            if k not in ["input_ids", "pixel_values", "attention_mask"]
-        }
-
-        if mask is None:
-            mask = mx.ones_like(input_ids)
-
-        return {
-            "pixel_values": pixel_values,
-            "input_ids": input_ids,
-            "attention_mask": mask,
-            **kwargs,
-        }
-
-
-def grad_checkpoint(layer):
-    """
-    Update all instances of type(layer) to use gradient checkpointing.
-    """
-    fn = type(layer).__call__
-
-    def checkpointed_fn(model, *args, **kwargs):
-        def inner_fn(params, *args, **kwargs):
-            model.update(params)
-            return fn(model, *args, **kwargs)
-
-        return mx.checkpoint(inner_fn)(model.trainable_parameters(), *args, **kwargs)
-
-    type(layer).__call__ = checkpointed_fn
-
-
-@dataclass
-class TrainingArgs:
-    batch_size: int = field(default=4, metadata={"help": "Minibatch size."})
-    iters: int = field(default=100, metadata={"help": "Iterations to train for."})
-    val_batches: int = field(
-        default=25,
-        metadata={
-            "help": "Number of validation batches, -1 uses the entire validation set."
-        },
-    )
-    steps_per_report: int = field(
-        default=10,
-        metadata={"help": "Number of training steps between loss reporting."},
-    )
-    steps_per_eval: int = field(
-        default=200, metadata={"help": "Number of training steps between validations."}
-    )
-    steps_per_save: int = field(
-        default=100, metadata={"help": "Save the model every number steps"}
-    )
-    max_seq_length: int = field(
-        default=2048, metadata={"help": "Maximum sequence length."}
-    )
-    adapter_file: str = field(
-        default="adapters.safetensors",
-        metadata={"help": "Save/load path for the trained adapter weights."},
-    )
-    grad_checkpoint: bool = field(
-        default=False,
-        metadata={"help": "Use gradient checkpointing to reduce memory use."},
-    )
-
-
-def default_loss(model, inputs, targets, lengths):
-    logits = model(inputs)
-    logits = logits.astype(mx.float32)
-
-    length_mask = mx.arange(inputs.shape[1])[None, :] < lengths[:, None]
-
-    ce = nn.losses.cross_entropy(logits, targets) * length_mask
-    ntoks = length_mask.sum()
-    ce = ce.sum() / ntoks
-
-    return ce, ntoks
-
-
-class Trainer:
-    def __init__(
-        self,
-        model,
-        optimizer,
-        train_on_completions=False,
-        assistant_id=77091,
-        clip_gradients=None,
-    ):
-        self.model = model
-        self.optimizer = optimizer
-        self.train_on_completions = train_on_completions
-        self.assistant_id = assistant_id
-        self.clip_gradients = clip_gradients
-
-    def loss_fn(self, model, batch):
-        pixel_values = batch["pixel_values"]
-        input_ids = batch["input_ids"]
-        attention_mask = batch["attention_mask"]
-        lengths = mx.sum(attention_mask, axis=1)
-        labels = input_ids[:, 1:]
-
-        batch_size, seq_length = input_ids.shape
-
-        if self.train_on_completions:
-            weight_mask = mx.ones_like(attention_mask)
-
-            assistant_response_index = np.where(input_ids == self.assistant_id)[1]
-            range_matrix = mx.repeat(
-                mx.expand_dims(mx.arange(seq_length), 0), batch_size, axis=0
-            )
-            assistant_mask = range_matrix <= mx.array(assistant_response_index).reshape(
-                -1, 1
-            )
-            # Apply the mask to weight_mask
-            weight_mask = mx.where(
-                assistant_mask, mx.zeros_like(weight_mask), weight_mask
-            )[:, 1:]
-        else:
-            weight_mask = None
-
-        input_ids = input_ids[:, :-1]
-
-        kwargs = {
-            k: v
-            for k, v in batch.items()
-            if k not in ["input_ids", "pixel_values", "attention_mask"]
-        }
-
-        # Forward pass
-        outputs = model(input_ids, pixel_values, attention_mask, **kwargs)
-
-        # Cast to float32
-        logits = outputs.logits.astype(mx.float32)
-
-        # Ensure logits and labels have the same sequence length
-        def align_logits_with_labels(logits, labels):
-            if logits.shape[1] < labels.shape[1]:
-                pad_length = labels.shape[1] - logits.shape[1]
-                pad_width = ((0, 0), (0, pad_length), (0, 0))
-                return mx.pad(logits, pad_width, mode="constant", constant_values=-100)
-            elif logits.shape[1] > labels.shape[1]:
-                return logits[:, -labels.shape[1] :, :]
-            return logits
-
-        logits = align_logits_with_labels(logits, labels)
-
-        length_mask = mx.arange(input_ids.shape[1])[None, :] < lengths[:, None]
-
-        # Compute loss only on non-padded tokens
-        ce = (
-            nn.losses.cross_entropy(
-                logits,
-                labels,
-                weights=weight_mask,
-            )
-            * length_mask
-        )
-        ntoks = length_mask.sum()
-        ce = ce.sum() / ntoks
-
-        return ce
-
-    def train_step(self, batch):
-        loss_and_grad_fn = nn.value_and_grad(self.model, self.loss_fn)
-        loss, grads = loss_and_grad_fn(self.model, batch)
-
-        # Add gradient clipping
-        if self.clip_gradients is not None:
-            grads = tree_map(
-                lambda g: mx.clip(g, -self.clip_gradients, self.clip_gradients), grads
-            )
-
-        self.optimizer.update(self.model, grads)
-
-        return loss
-
-    @mx.compile
-    def train_epoch(self, dataloader):
-        total_loss = 0
-        for batch in dataloader:
-            loss = self.train_step(batch)
-            mx.eval(self.model, self.optimizer.state)
-            total_loss += loss
-        return total_loss / len(dataloader)
-
-
-def save_adapter(
-    model: nn.Module,
-    adapter_file: Union[str, Path],
-):
-    path = Path(adapter_file)
-    if hasattr(model.config, "lora"):
-        with open(path.parent / "adapter_config.json", "w") as f:
-            json.dump(model.config.lora, f)
-    flattened_tree = tree_flatten(model.trainable_parameters())
-    mx.save_safetensors(str(adapter_file), dict(flattened_tree))

nexaai/binds/nexa_mlx/py-lib/vlm/modeling/trainer/utils.py

@@ -1,160 +0,0 @@
-import json
-from pathlib import Path
-
-import mlx.nn as nn
-from mlx.utils import tree_flatten
-
-from .lora import LoRaLayer
-
-
-def get_module_by_name(model, name):
-    parts = name.split(".")
-    module = model
-    for part in parts:
-        if part.isdigit():
-            module = module[int(part)]
-        else:
-            module = getattr(module, part)
-    return module
-
-
-def set_module_by_name(model, name, new_module):
-    parts = name.split(".")
-    module = model
-    for part in parts[:-1]:
-        if part.isdigit():
-            module = module[int(part)]
-        else:
-            module = getattr(module, part)
-    if parts[-1].isdigit():
-        module[int(parts[-1])] = new_module
-    else:
-        setattr(module, parts[-1], new_module)
-
-
-def get_peft_model(
-    model, linear_layers, rank=10, alpha=0.1, dropout=0.1, freeze=True, verbose=True
-):
-    if freeze:
-        freeze_model(model)
-
-    for name, module in model.language_model.named_modules():
-        if isinstance(module, nn.Linear) or isinstance(module, nn.QuantizedLinear):
-            if name.split(".")[-1] in linear_layers:
-                lora_layer = LoRaLayer(module, rank, alpha, dropout)
-                set_module_by_name(model.language_model, name, lora_layer)
-
-    model.config.lora = {}
-    model.config.lora["rank"] = rank
-    model.config.lora["alpha"] = alpha
-    model.config.lora["dropout"] = dropout
-
-    if verbose:
-        print_trainable_parameters(model.language_model)
-
-    return model
-
-
-def freeze_model(model):
-    for name, module in model.named_modules():
-        name = name.split(".")[0]
-        if name in [
-            "language_model",
-            "vision_model",
-            "vision_tower",
-            "aligner",
-            "connector",
-            "multi_modal_projector",
-            "mm_projector",
-        ]:
-            model[f"{name}"].freeze()
-
-
-def find_all_linear_names(model):
-    cls = nn.Linear
-    quantized_cls = nn.QuantizedLinear
-    lora_module_names = set()
-    multimodal_keywords = [
-        "mm_projector",
-        "vision_tower",
-        "vision_resampler",
-        "aligner",
-    ]
-    for name, module in model.named_modules():
-        if any(mm_keyword in name for mm_keyword in multimodal_keywords):
-            continue
-        if isinstance(module, cls) or isinstance(module, quantized_cls):
-            names = name.split(".")
-            lora_module_names.add(names[0] if len(names) == 1 else names[-1])
-
-    if "lm_head" in lora_module_names:  # needed for 16-bit
-        lora_module_names.remove("lm_head")
-    return list(lora_module_names)
-
-
-def count_parameters(model):
-    def nparams(m):
-        if isinstance(m, (nn.QuantizedLinear, nn.QuantizedEmbedding)):
-            return m.weight.size * (32 // m.bits)
-        return sum(v.size for _, v in tree_flatten(m.parameters()))
-
-    leaf_modules = tree_flatten(
-        model.leaf_modules(), is_leaf=lambda m: isinstance(m, nn.Module)
-    )
-    total_p = sum(nparams(m) for _, m in leaf_modules) / 10**6
-
-    return total_p
-
-
-def print_trainable_parameters(model):
-    def nparams(m):
-        if isinstance(m, (nn.QuantizedLinear, nn.QuantizedEmbedding)):
-            return m.weight.size * (32 // m.bits)
-        return sum(v.size for _, v in tree_flatten(m.parameters()))
-
-    leaf_modules = tree_flatten(
-        model.leaf_modules(), is_leaf=lambda m: isinstance(m, nn.Module)
-    )
-    total_p = sum(nparams(m) for _, m in leaf_modules) / 10**6
-    trainable_p = (
-        sum(v.size for _, v in tree_flatten(model.trainable_parameters())) / 10**6
-    )
-
-    print(
-        f"#trainable params: {trainable_p} M || all params: {total_p} M || trainable%: {(trainable_p * 100 / total_p):.3f}%"
-    )
-
-
-def apply_lora_layers(model: nn.Module, adapter_path: str) -> nn.Module:
-    """
-    Apply LoRA layers to the model.
-
-    Args:
-        model (nn.Module): The neural network model.
-        adapter_path (str): Path to the adapter configuration file.
-
-    Returns:
-        nn.Module: The updated model with LoRA layers applied.
-    """
-    adapter_path = Path(adapter_path)
-
-    if not adapter_path.exists():
-        raise FileNotFoundError(f"The adapter path does not exist: {adapter_path}")
-
-    # Check if the adapter has lora params in the config (adapter_config.json)
-    with open(adapter_path / "adapter_config.json", "r") as f:
-        config = json.load(f)
-        if "rank" not in config:
-            raise ValueError("The adapter does not have lora params in the config")
-
-    # TODO: add lora params to the config and load them here
-    list_of_modules = find_all_linear_names(model.language_model.model)
-    if config is not None:
-        model = get_peft_model(model, list_of_modules, **config)
-    else:
-        model = get_peft_model(model, list_of_modules)
-
-    # TODO: Use custom adapter name
-    model.load_weights(str(adapter_path / "adapters.safetensors"), strict=False)
-
-    return model