oxillama-py 0.1.2

Python bindings for OxiLLaMa LLM inference engine
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61

"""Utility helpers for the OxiLLaMa Python bindings."""
from __future__ import annotations

from typing import List

__all__ = ["decode_from_logits"]


def decode_from_logits(
    logits: "numpy.ndarray",  # noqa: F821
    temperature: float = 1.0,
    top_k: int = 0,
    top_p: float = 1.0,
) -> int:
    """Run greedy/temperature sampling over a logits array and return the argmax token id.

    This is a pure-Python utility — it does NOT call the Rust engine.  It mirrors
    the sampler behaviour for use in notebooks and custom decoding loops.

    Args:
        logits: 1-D float32 ndarray of shape ``(vocab_size,)``.
        temperature: Softmax temperature. 1.0 = unscaled, < 1.0 = sharper.
        top_k: If > 0, restrict sampling to the top-k highest-probability tokens.
        top_p: Nucleus probability mass threshold. 1.0 = no restriction.

    Returns:
        The sampled token id (integer index into the vocabulary).
    """
    import numpy as np

    arr = np.asarray(logits, dtype=np.float32)
    if arr.ndim != 1:
        raise ValueError(f"Expected 1-D logits array, got shape {arr.shape}")

    if temperature != 1.0 and temperature > 0.0:
        arr = arr / temperature

    if top_k > 0:
        kth = np.partition(arr, -top_k)[-top_k]
        arr = np.where(arr >= kth, arr, -np.inf)

    # Softmax
    arr = arr - arr.max()
    exp_arr = np.exp(arr)
    probs = exp_arr / exp_arr.sum()

    if top_p < 1.0:
        sorted_idx = np.argsort(probs)[::-1]
        cumprobs = np.cumsum(probs[sorted_idx])
        cutoff = int(np.searchsorted(cumprobs, top_p))
        allowed = sorted_idx[: cutoff + 1]
        mask = np.zeros_like(probs)
        mask[allowed] = 1.0
        probs = probs * mask
        total = float(probs.sum())
        if total > 0:
            probs = probs / total
        else:
            probs = mask / mask.sum()

    return int(np.argmax(probs))