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

nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/llm_common/rope_utils.py

@@ -1,255 +0,0 @@
-# Copyright © 2023-2024 Apple Inc.
-
-import math
-from typing import List, Optional, Union
-
-import mlx.core as mx
-import mlx.nn as nn
-
-
-class SuScaledRoPE(nn.Module):
-    def __init__(
-        self,
-        dims: int,
-        base: float = 10000.0,
-        max_position_embeddings: int = 131072,
-        original_max_position_embeddings: int = 4096,
-        short_factor: Union[List[float], float] = 1.0,
-        long_factor: Union[List[float], float] = 1.0,
-        short_mscale: float = None,
-        long_mscale: float = None,
-    ):
-        """
-        Su Scaled Rotary Embedding layer.
-
-        Args:
-            dims (int): The feature dimensions to be rotated.
-            base (int, optional): Base for the exponential scaling.
-            max_position_embeddings (int, optional): The maximum sequence
-              length that this model was trained with. This is used to determine
-              the size of the original RoPE embeddings when using long scaling.
-              Default: ``131072``.
-            original_max_position_embeddings (int, optional): The maximum
-              sequence length that this model was trained with. This is used to
-              determine the size of the original RoPE embeddings when using long
-              scaling. Default: ``4096``.
-            short_factor (float or list[float], optional): List of scaling
-              factors for sequences of length lesser than
-              ``original_max_position_embeddings``. Default: ``1.0``.
-            long_factor (float or list[float], optional): List of scaling
-              factors for sequences of length greater than
-              ``original_max_position_embeddings``.  Default: ``1.0``.
-            short_mscale (float, optional): Scale the input prior to embedding.
-            long_mscale (float, optional): Scale the input prior to embedding.
-        """
-        super().__init__()
-        freqs = base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)
-        self._freqs = mx.array(long_factor, dtype=mx.float32) * freqs
-        self.original_max_position_embeddings = original_max_position_embeddings
-        self.scale = long_mscale or math.sqrt(
-            1
-            + math.log(max_position_embeddings / original_max_position_embeddings)
-            / math.log(original_max_position_embeddings)
-        )
-        self.dim = dims
-
-    def __call__(self, x, offset: int = 0):
-        x[..., : self.dim] = self.scale * x[..., : self.dim]
-        return mx.fast.rope(
-            x,
-            self.dim,
-            traditional=False,
-            base=None,
-            scale=1.0,
-            offset=offset,
-            freqs=self._freqs,
-        )
-
-
-class Llama3RoPE(nn.Module):
-
-    def __init__(
-        self,
-        dims: int,
-        max_position_embeddings: int = 2048,
-        traditional: bool = False,
-        base: float = 10000,
-        scaling_config: dict = None,
-    ):
-        super().__init__()
-        self.dims = dims
-        self.max_position_embeddings = max_position_embeddings
-        self.traditional = traditional
-
-        factor = scaling_config["factor"]
-        low_freq_factor = scaling_config.get("low_freq_factor", 1.0)
-        high_freq_factor = scaling_config.get("high_freq_factor", 4.0)
-        old_context_len = scaling_config.get(
-            "original_max_position_embeddings",
-            8192,
-        )
-
-        low_freq_wavelen = old_context_len / low_freq_factor
-        high_freq_wavelen = old_context_len / high_freq_factor
-
-        freqs = base ** (mx.arange(0, dims, 2) / dims)
-        wavelens = 2 * mx.pi * freqs
-
-        freqs = mx.where(wavelens > low_freq_wavelen, freqs * factor, freqs)
-        is_medium_freq = (wavelens > high_freq_wavelen) & (wavelens < low_freq_wavelen)
-        smooth_factors = (old_context_len / wavelens - low_freq_factor) / (
-            high_freq_factor - low_freq_factor
-        )
-        smooth_freqs = freqs / ((1 - smooth_factors) / factor + smooth_factors)
-        self._freqs = mx.where(is_medium_freq, smooth_freqs, freqs)
-
-    def extra_repr(self):
-        return (
-            f"{self.dims}, traditional={self.traditional}, "
-            f"max_position_embeddings={self.max_position_embeddings}"
-        )
-
-    def __call__(self, x, offset: int = 0):
-        return mx.fast.rope(
-            x,
-            self.dims,
-            traditional=self.traditional,
-            base=None,
-            scale=1.0,
-            offset=offset,
-            freqs=self._freqs,
-        )
-
-
-class YarnRoPE(nn.Module):
-    def __init__(
-        self,
-        dims,
-        traditional=False,
-        max_position_embeddings=2048,
-        base=10000,
-        scaling_factor=1.0,
-        original_max_position_embeddings=4096,
-        beta_fast=32,
-        beta_slow=1,
-        mscale=1,
-        mscale_all_dim=0,
-    ):
-        super().__init__()
-
-        def yarn_find_correction_dim(num_rotations):
-            return (
-                dims
-                * math.log(
-                    original_max_position_embeddings / (num_rotations * 2 * math.pi)
-                )
-            ) / (2 * math.log(base))
-
-        def yarn_find_correction_range():
-            low = math.floor(yarn_find_correction_dim(beta_fast))
-            high = math.ceil(yarn_find_correction_dim(beta_slow))
-            return max(low, 0), min(high, dims - 1)
-
-        def yarn_get_mscale(scale=1, mscale=1):
-            if scale <= 1:
-                return 1.0
-            return 0.1 * mscale * math.log(scale) + 1.0
-
-        def yarn_linear_ramp_mask(min_val, max_val, dim):
-            if min_val == max_val:
-                max_val += 0.001  # Prevent singularity
-
-            linear_func = (mx.arange(dim, dtype=mx.float32) - min_val) / (
-                max_val - min_val
-            )
-            return mx.clip(linear_func, 0, 1)
-
-        self.mscale = yarn_get_mscale(scaling_factor, mscale) / yarn_get_mscale(
-            scaling_factor, mscale_all_dim
-        )
-        freq_extra = base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)
-        freq_inter = scaling_factor * base ** (
-            mx.arange(0, dims, 2, dtype=mx.float32) / dims
-        )
-        low, high = yarn_find_correction_range()
-        freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dims // 2)
-        self._freqs = (freq_inter * freq_extra) / (
-            freq_inter * freq_mask + freq_extra * (1 - freq_mask)
-        )
-        self.dims = dims
-        self.traditional = traditional
-
-    def __call__(self, x, offset=0):
-        if self.mscale != 1.0:
-            x[..., : self.dims] = self.mscale * x[..., : self.dims]
-        return mx.fast.rope(
-            x,
-            self.dims,
-            traditional=self.traditional,
-            base=None,
-            scale=1.0,
-            offset=offset,
-            freqs=self._freqs,
-        )
-
-
-def initialize_rope(
-    dims,
-    base,
-    traditional,
-    scaling_config: Optional[dict] = None,
-    max_position_embeddings: Optional[int] = None,
-):
-    if scaling_config is not None:
-        rope_type = scaling_config.get("type") or scaling_config.get(
-            "rope_type", "default"
-        )
-    else:
-        rope_type = "default"
-
-    if rope_type in ["default", "linear"]:
-        scale = 1 / scaling_config["factor"] if rope_type == "linear" else 1.0
-        return nn.RoPE(dims, traditional=traditional, base=base, scale=scale)
-
-    elif rope_type == "llama3":
-        return Llama3RoPE(
-            dims=dims,
-            max_position_embeddings=max_position_embeddings,
-            traditional=traditional,
-            base=base,
-            scaling_config=scaling_config,
-        )
-    elif rope_type == "yarn":
-        scaling_factor = scaling_config["factor"]
-        rope_kwargs = {
-            key: scaling_config[key]
-            for key in [
-                "original_max_position_embeddings",
-                "beta_fast",
-                "beta_slow",
-                "mscale",
-                "mscale_all_dim",
-            ]
-            if key in scaling_config
-        }
-        return YarnRoPE(
-            dims=dims,
-            max_position_embeddings=max_position_embeddings,
-            traditional=traditional,
-            base=base,
-            **rope_kwargs,
-        )
-    elif rope_type == "longrope":
-        return SuScaledRoPE(
-            dims=dims,
-            base=base,
-            max_position_embeddings=max_position_embeddings,
-            original_max_position_embeddings=scaling_config[
-                "original_max_position_embeddings"
-            ],
-            short_factor=scaling_config["short_factor"],
-            long_factor=scaling_config["long_factor"],
-        )
-
-    else:
-        raise ValueError(f"Unsupported RoPE type {rope_type}")

nexaai/binds/nexa_mlx/py-lib/vlm/modeling/models/qwen3_vl/llm_common/sample_utils.py

@@ -1,303 +0,0 @@
-import math
-from functools import partial
-from typing import Callable, Dict, List, Optional
-
-import mlx.core as mx
-
-
-def make_sampler(
-    temp: float = 0.0,
-    top_p: float = 0.0,
-    min_p: float = 0.0,
-    min_tokens_to_keep: int = 1,
-    top_k: int = 0,
-    xtc_probability: float = 0.0,
-    xtc_threshold: float = 0.0,
-    xtc_special_tokens: List[int] = [],
-) -> Callable[mx.array, mx.array]:
-    """
-    Make a sampler function for use with ``generate_step``.
-
-    Args:
-        temp (float): The temperature for sampling, if 0 the argmax is used.
-          Default: ``0``.
-        top_p (float, optional): Nulceus sampling, higher means model considers
-          more less likely words.
-        min_p (float, optional): The minimum value (scaled by the top token's
-          probability) that a token probability must have to be considered.
-        min_tokens_to_keep (int, optional): Minimum number of tokens that cannot
-          be filtered by min_p sampling.
-        top_k (int, optional): The top k tokens ranked by probability to constrain
-          the sampling to.
-        xtc_probability (float, optional): The probability of applying XTC
-            sampling.
-        xtc_threshold (float, optional): The threshold the probs need to reach
-            for being sampled.
-        xtc_special_tokens (list(int), optional): List of special tokens IDs to
-            be excluded from XTC sampling.
-
-
-    Returns:
-        Callable[mx.array, mx.array]:
-            A sampler which takes log-probabilities and returns tokens.
-    """
-    if temp == 0:
-        return lambda x: mx.argmax(x, axis=-1)
-
-    # Create sampler chain
-    sampling_methods = []
-    if top_k > 0:
-        sampling_methods.append(lambda x: apply_top_k(x, top_k))
-    if top_p > 0 and top_p < 1.0:
-        sampling_methods.append(lambda x: apply_top_p(x, top_p))
-    if min_p != 0.0:
-        sampling_methods.append(lambda x: apply_min_p(x, min_p, min_tokens_to_keep))
-    if xtc_probability > 0.0:
-        sampling_methods.append(
-            lambda x: apply_xtc(x, xtc_probability, xtc_threshold, xtc_special_tokens)
-        )
-
-    # Apply the sampling methods
-    def sampler(logits):
-        for method in sampling_methods:
-            logits = method(logits)
-
-        # Return the sampled token
-        return categorical_sampling(logits, temp)
-
-    return sampler
-
-
-def make_logits_processors(
-    logit_bias: Optional[Dict[int, float]] = None,
-    repetition_penalty: Optional[float] = None,
-    repetition_context_size: Optional[int] = 20,
-):
-    """
-    Make logits processors for use with ``generate_step``.
-
-    Args:
-        repetition_penalty (float, optional): The penalty factor for repeating
-          tokens.
-        repetition_context_size (int, optional): The number of tokens to
-          consider for repetition penalty. Default: ``20``.
-        logit_bias (dictionary, optional): Additive logit bias.
-
-    Returns:
-        List[Callable[[mx.array, mx.array], mx.array]]:
-            A list of logits processors. Each processor in the list is a
-            callable which takes an array of tokens and an array of logits
-            and returns the updated logits.
-    """
-    logits_processors = []
-    if logit_bias:
-        indices = mx.array(list(logit_bias.keys()))
-        values = mx.array(list(logit_bias.values()))
-
-        def logit_bias_processor(_, logits):
-            logits[:, indices] += values
-            return logits
-
-        logits_processors.append(logit_bias_processor)
-
-    if repetition_penalty and repetition_penalty != 0.0:
-        logits_processors.append(
-            make_repetition_penalty(repetition_penalty, repetition_context_size)
-        )
-    return logits_processors
-
-
-@partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
-def apply_top_k(
-    logprobs: mx.array,
-    top_k: int,
-) -> mx.array:
-    """
-    Sample from only the top K tokens ranked by probability.
-
-    Args:
-        logprobs: A vector of log probabilities.
-        top_k (int): Top k tokens to sample from.
-    """
-    vocab_size = logprobs.shape[-1]
-    if not isinstance(top_k, int) or not (0 < top_k < vocab_size):
-        raise ValueError(
-            f"`top_k` has to be an integer in the (0, {vocab_size}] interval," f" but is {top_k}."
-        )
-    mask_idx = mx.argpartition(-logprobs, kth=top_k - 1, axis=-1)[..., top_k:]
-    masked_logprobs = mx.put_along_axis(
-        logprobs, mask_idx, mx.array(-float("inf"), logprobs.dtype), axis=-1
-    )
-    return masked_logprobs
-
-
-@partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
-def apply_min_p(
-    logprobs: mx.array,
-    min_p: float,
-    min_tokens_to_keep: int = 1,
-) -> mx.array:
-    """
-    Apply min-p sampling to the logprobs.
-
-    Min-p keeps all tokens that are above a minimum probability, scaled by the
-    probability of the most likely token. As a result, the filter is more
-    aggressive given a very high-probability token.
-
-    Args:
-        logprobs: A vector of log probabilities.
-        min_p (float): Minimum token probability. Typical values are in the
-            0.01-0.2 range, comparably selective as setting `top_p` in the
-            0.99-0.8 range.
-        min_tokens_to_keep (int, optional): Minimum number of tokens that cannot
-            be filtered. Default: ``1``.
-
-    """
-    if not (0 <= min_p <= 1.0):
-        raise ValueError(f"`min_p` has to be a float in the [0, 1] interval, but is {min_p}")
-    if not isinstance(min_tokens_to_keep, int) or (min_tokens_to_keep < 1):
-        raise ValueError(
-            f"`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}"
-        )
-    # reference implementation: https://github.com/huggingface/transformers/blob/main/src/transformers/generation/logits_process.py#L531-L605
-
-    # Indices sorted in decreasing order
-    sorted_indices = mx.argsort(-logprobs, axis=-1)
-    sorted_logprobs = mx.take_along_axis(logprobs, sorted_indices, axis=-1)
-
-    # Top probability
-    top_logprobs = sorted_logprobs[:, 0:1]
-
-    # Calculate the min_p threshold
-    scaled_min_p = top_logprobs + math.log(min_p)
-
-    # Mask tokens that have a probability less than the scaled min_p
-    tokens_to_remove = sorted_logprobs < scaled_min_p
-    tokens_to_remove[..., :min_tokens_to_keep] = False
-
-    # Create pool of tokens with probability less than scaled min_p
-    selected_logprobs = mx.where(tokens_to_remove, -float("inf"), sorted_logprobs)
-
-    # Create a mapping to rearrange back to original indices
-    inverse_indices = mx.put_along_axis(
-        mx.zeros_like(sorted_indices),
-        sorted_indices,
-        mx.arange(sorted_indices.shape[-1], dtype=sorted_indices.dtype),
-        axis=-1,
-    )
-
-    # Rearrange selected_logprobs back to original order
-    original_order_logprobs = mx.take_along_axis(selected_logprobs, inverse_indices, axis=-1)
-
-    return original_order_logprobs
-
-
-@partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
-def apply_top_p(logprobs: mx.array, top_p: float) -> mx.array:
-    """
-    Apply top-p (nucleus) sampling to logits.
-
-    Args:
-        logprobs: A vector of log probabilities.
-        top_p: The cumulative probability threshold for top-p filtering.
-    Returns:
-        token selected based on the top-p criterion.
-    """
-    # referenced implementation from https://github.com/huggingface/transformers/blob/main/src/transformers/generation/logits_process.py#L449-L460
-    probs = mx.exp(logprobs)
-    # sort in ascending order
-    sorted_indices = mx.argsort(logprobs, axis=-1)
-    sorted_probs = mx.take_along_axis(probs, sorted_indices, axis=-1)
-
-    cumulative_probs = mx.cumsum(sorted_probs, axis=-1)
-
-    # Rearrange cumulative probs back to original order
-    inverse_indices = mx.put_along_axis(
-        mx.zeros_like(sorted_indices),
-        sorted_indices,
-        mx.arange(sorted_indices.shape[-1], dtype=sorted_indices.dtype),
-        axis=-1,
-    )
-    cumulative_probs = mx.take_along_axis(cumulative_probs, inverse_indices, axis=-1)
-
-    # select tokens with cumulative probs below threshold
-    return mx.where(
-        cumulative_probs > 1 - top_p,
-        logprobs,
-        -float("inf"),
-    )
-
-
-@partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
-def apply_xtc(
-    logits: mx.array,
-    xtc_probability: float,
-    xtc_threshold: float,
-    xtc_special_tokens: List[int],
-) -> mx.array:
-    """
-    Apply XTC sampling to the logits.
-
-    Args:
-        logits: The logits from the model's output.
-        xtc_probability (float): Probability of XTC sampling to happen for each token
-        xtc_threshold (float): The threshold the probs need to reach for being sampled.
-        special_tokens_ids (list(int)): List of special tokens IDs to be excluded from XTC sampling.
-    """
-    if not (0 <= xtc_threshold <= 0.5):
-        raise ValueError(
-            f"`threshold` has to be a float in the [0, 0.5] interval, but is {xtc_threshold}"
-        )
-    if not (0 <= xtc_probability <= 1.0):
-        raise ValueError(
-            f"`probability` has to be a float in the [0, 1] interval, but is {xtc_probability}"
-        )
-
-    probs = mx.softmax(logits, -1)
-    mask = probs > mx.where(probs > xtc_threshold, probs, mx.inf).min()
-    if xtc_special_tokens:
-        mask[..., xtc_special_tokens] = False
-
-    return mx.where(
-        mx.random.uniform(0, 1) > xtc_probability,
-        logits,
-        mx.where(mask, -mx.inf, logits),
-    )
-
-
-@partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
-def categorical_sampling(logits, temp):
-    return mx.random.categorical(logits * (1 / temp))
-
-
-def make_repetition_penalty(penalty: float, context_size: int = 20):
-    """
-    Make repetition penalty processor.
-
-    Paper: https://arxiv.org/abs/1909.05858
-
-    Args:
-        penalty (float): The repetition penalty factor to be applied.
-        context_size (int): The number of previous tokens to use.
-            Default: ``20``.
-
-    Returns:
-        Callable[[mx.array, List[int]], mx.array]:
-            The repetition penalty processor.
-    """
-    if penalty < 0 or not isinstance(penalty, (int, float)):
-        raise ValueError(f"penalty must be a non-negative float, got {penalty}")
-
-    def repetition_penalty_processor(tokens, logits):
-        if len(tokens) > 0:
-            tokens = tokens[-context_size:]
-            selected_logits = logits[:, tokens]
-            selected_logits = mx.where(
-                selected_logits < 0,
-                selected_logits * penalty,
-                selected_logits / penalty,
-            )
-            logits[:, tokens] = selected_logits
-        return logits
-
-    return repetition_penalty_processor