diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md index 5df12aaf53..eb0fb9662e 100644 --- a/.github/ISSUE_TEMPLATE/bug_report.md +++ b/.github/ISSUE_TEMPLATE/bug_report.md @@ -64,7 +64,7 @@ Try the following: 1. `git clone https://github.com/abetlen/llama-cpp-python` 2. `cd llama-cpp-python` 3. `rm -rf _skbuild/` # delete any old builds -4. `python setup.py develop` +4. `python -m pip install .` 5. `cd ./vendor/llama.cpp` 6. Follow [llama.cpp's instructions](https://github.com/ggerganov/llama.cpp#build) to `cmake` llama.cpp 7. Run llama.cpp's `./main` with the same arguments you previously passed to llama-cpp-python and see if you can reproduce the issue. If you can, [log an issue with llama.cpp](https://github.com/ggerganov/llama.cpp/issues) diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 91abb11fdf..6bf90273ac 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -8,4 +8,12 @@ updates: - package-ecosystem: "pip" # See documentation for possible values directory: "/" # Location of package manifests schedule: - interval: "weekly" + interval: "daily" + - package-ecosystem: "github-actions" + directory: "/" + schedule: + interval: "daily" + - package-ecosystem: "docker" + directory: "/" + schedule: + interval: "daily" diff --git a/.github/workflows/build-and-release.yaml b/.github/workflows/build-and-release.yaml index 2c0ca4a2ab..4ae37b1745 100644 --- a/.github/workflows/build-and-release.yaml +++ b/.github/workflows/build-and-release.yaml @@ -11,28 +11,132 @@ jobs: runs-on: ${{ matrix.os }} strategy: matrix: - os: [ubuntu-latest, windows-latest, macOS-latest] + os: [ubuntu-22.04, windows-2022, macos-14, macos-15] steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v6 with: - submodules: "true" + submodules: "recursive" # Used to host cibuildwheel - - uses: actions/setup-python@v3 + - uses: actions/setup-python@v6 + with: + python-version: "3.9" - - name: Install cibuildwheel - run: python -m pip install cibuildwheel==2.12.1 + - name: Install dependencies (Linux/MacOS) + if: runner.os != 'Windows' + run: | + python -m pip install --upgrade pip + python -m pip install uv + RUST_LOG=trace python -m uv pip install -e .[all] --verbose + shell: bash - - name: Install dependencies + - name: Install dependencies (Windows) + if: runner.os == 'Windows' + env: + RUST_LOG: trace run: | - python -m pip install --upgrade pip pytest cmake scikit-build setuptools + python -m pip install --upgrade pip + python -m pip install uv + python -m uv pip install -e .[all] --verbose + shell: cmd - name: Build wheels - run: python -m cibuildwheel --output-dir wheelhouse + uses: pypa/cibuildwheel@v3.4.1 + env: + # Keep repair disabled by default for non-Linux platforms in this job. + CIBW_REPAIR_WHEEL_COMMAND: "" + # Linux needs auditwheel repair so manylinux and musllinux wheels are + # published with distinct platform tags instead of generic linux tags. + CIBW_REPAIR_WHEEL_COMMAND_LINUX: "LD_LIBRARY_PATH=/project/llama_cpp/lib auditwheel repair -w {dest_dir} {wheel}" + # cibuildwheel v3 defaults to manylinux_2_28 images whose current + # GCC toolchain emits symbols newer than the policy allows. + CIBW_MANYLINUX_X86_64_IMAGE: "manylinux2014" + # The release wheel is tagged py3-none, so one build per platform + # covers all supported Python versions and avoids duplicate names. + CIBW_BUILD_LINUX: "cp38-*" + CIBW_BUILD_MACOS: "cp39-*" + CIBW_BUILD_WINDOWS: "cp39-*" + # Skip cibuildwheel's default i686 sidecar and keep Linux release + # wheels on a portable x86_64 CPU baseline. + CIBW_ARCHS_LINUX: "auto64" + CIBW_ARCHS_WINDOWS: "AMD64" + CIBW_ENVIRONMENT_LINUX: CMAKE_ARGS="-DGGML_NATIVE=off" + # Keep macOS release wheels on a portable CPU baseline instead of + # inheriting the hosted runner's native flags. + CIBW_ENVIRONMENT_MACOS: CMAKE_ARGS="-DGGML_NATIVE=off" + with: + package-dir: . + output-dir: wheelhouse - - uses: actions/upload-artifact@v3 + - uses: actions/upload-artifact@v7 with: + name: wheels-${{ matrix.os }} + path: ./wheelhouse/*.whl + + build_wheels_arm64: + name: Build arm64 wheels + runs-on: ubuntu-24.04-arm + steps: + - uses: actions/checkout@v6 + with: + submodules: "recursive" + + - name: Build wheels + uses: pypa/cibuildwheel@v3.4.1 + env: + CIBW_SKIP: "pp*" + CIBW_REPAIR_WHEEL_COMMAND: "LD_LIBRARY_PATH=$PWD/llama_cpp/lib auditwheel repair -w {dest_dir} {wheel}" + CIBW_ARCHS: "aarch64" + # Keep this consistent with the x86_64 Linux release wheels. + CIBW_MANYLINUX_AARCH64_IMAGE: "manylinux2014" + # Keep native arm64 builds on a portable CPU baseline instead of + # tuning wheels to the hosted runner. + CIBW_ENVIRONMENT: CMAKE_ARGS="-DGGML_NATIVE=off" + # The release wheel is tagged py3-none, so one build covers all + # supported Python versions and avoids duplicate wheel names. + CIBW_BUILD: "cp38-*" + with: + output-dir: wheelhouse + + - name: Upload wheels as artifacts + uses: actions/upload-artifact@v7 + with: + name: wheels_arm64 + path: ./wheelhouse/*.whl + + build_wheels_riscv64: + name: Build riscv64 wheel + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + with: + submodules: "recursive" + + - name: Set up QEMU + uses: docker/setup-qemu-action@v4 + with: + platforms: linux/riscv64 + + - name: Build wheels + uses: pypa/cibuildwheel@v3.4.1 + env: + CIBW_SKIP: "*musllinux* pp*" + CIBW_REPAIR_WHEEL_COMMAND: "" + CIBW_ARCHS: "riscv64" + # Build riscv64 wheels against a conservative baseline instead of + # enabling RVV-related extensions from the build container. + CIBW_ENVIRONMENT: CMAKE_ARGS="-DGGML_NATIVE=off -DGGML_RVV=off -DGGML_RV_ZFH=off -DGGML_RV_ZVFH=off -DGGML_RV_ZICBOP=off -DGGML_RV_ZIHINTPAUSE=off" + # The release wheel is tagged py3-none, so one riscv64 build is + # enough and avoids duplicate same-name release artifacts. + CIBW_BUILD: "cp310-*" + with: + output-dir: wheelhouse + + - name: Upload wheels as artifacts + uses: actions/upload-artifact@v7 + with: + name: wheels_riscv64 path: ./wheelhouse/*.whl build_sdist: @@ -40,32 +144,57 @@ jobs: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v6 with: - submodules: "true" - - uses: actions/setup-python@v3 - - name: Install dependencies + submodules: "recursive" + + - uses: actions/setup-python@v6 + with: + python-version: "3.9" + + - name: Install dependencies (Linux/MacOS) + if: runner.os != 'Windows' run: | - python -m pip install --upgrade pip pytest cmake scikit-build setuptools + python -m pip install --upgrade pip + python -m pip install uv + RUST_LOG=trace python -m uv pip install -e .[all] --verbose + python -m uv pip install build + shell: bash + + - name: Install dependencies (Windows) + if: runner.os == 'Windows' + env: + RUST_LOG: trace + run: | + python -m pip install --upgrade pip + python -m pip install uv + python -m uv pip install -e .[all] --verbose + python -m uv pip install build + shell: cmd + - name: Build source distribution run: | - python setup.py sdist - - uses: actions/upload-artifact@v3 + python -m build --sdist + + - uses: actions/upload-artifact@v7 with: + name: sdist path: ./dist/*.tar.gz release: name: Release - needs: [build_wheels, build_sdist] + needs: [build_wheels, build_wheels_arm64, build_wheels_riscv64, build_sdist] + if: startsWith(github.ref, 'refs/tags/') runs-on: ubuntu-latest steps: - - uses: actions/download-artifact@v3 + - uses: actions/download-artifact@v8 with: - name: artifact + merge-multiple: true path: dist - - uses: softprops/action-gh-release@v1 + + - uses: softprops/action-gh-release@v3 with: files: dist/* env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} \ No newline at end of file + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} diff --git a/.github/workflows/build-docker.yaml b/.github/workflows/build-docker.yaml index 4b38dbacb6..d306efde48 100644 --- a/.github/workflows/build-docker.yaml +++ b/.github/workflows/build-docker.yaml @@ -9,32 +9,51 @@ permissions: jobs: docker: name: Build and push Docker image - runs-on: ubuntu-latest + runs-on: ubuntu-22.04 steps: - name: Checkout - uses: actions/checkout@v3 + uses: actions/checkout@v6 with: - submodules: "true" + submodules: "recursive" + + - name: Set image tag + run: | + if [[ "${GITHUB_REF_TYPE}" == "tag" ]]; then + image_tag="${GITHUB_REF_NAME}" + else + image_tag="${GITHUB_REF_NAME//\//-}" + fi + echo "IMAGE_TAG=$image_tag" >> "$GITHUB_ENV" - name: Set up QEMU - uses: docker/setup-qemu-action@v2 + uses: docker/setup-qemu-action@v4 - name: Set up Docker Buildx - uses: docker/setup-buildx-action@v2 + uses: docker/setup-buildx-action@v4 - name: Login to GitHub Container Registry - uses: docker/login-action@v2 + uses: docker/login-action@v4 with: registry: ghcr.io username: ${{ github.repository_owner }} password: ${{ secrets.GITHUB_TOKEN }} - name: Build and push - uses: docker/build-push-action@v4 + id: docker_build + uses: docker/build-push-action@v7 with: context: . file: "docker/simple/Dockerfile" - push: true # push to registry - pull: true # always fetch the latest base images - platforms: linux/amd64,linux/arm64 # build for both amd64 and arm64 - tags: ghcr.io/abetlen/llama-cpp-python:latest \ No newline at end of file + push: ${{ startsWith(github.ref, 'refs/tags/') }} + pull: true + platforms: linux/amd64,linux/arm64 + tags: | + ghcr.io/abetlen/llama-cpp-python:latest + ghcr.io/abetlen/llama-cpp-python:${{ env.IMAGE_TAG }} + build-args: | + BUILDKIT_INLINE_CACHE=1 + + - name: Publish to GitHub Tag + if: steps.docker_build.outputs.digest && startsWith(github.ref, 'refs/tags/') + run: | + echo "Docker image published for tag: ${{ github.ref_name }}" diff --git a/.github/workflows/build-wheels-cuda.yaml b/.github/workflows/build-wheels-cuda.yaml new file mode 100644 index 0000000000..59dc3558bf --- /dev/null +++ b/.github/workflows/build-wheels-cuda.yaml @@ -0,0 +1,275 @@ +name: Build Wheels (CUDA) + +on: + workflow_dispatch: + inputs: + release_tag: + description: Release tag to upload wheel assets to + required: false + type: string + +permissions: + contents: write + +jobs: + define_matrix: + name: Define Build Matrix + runs-on: ubuntu-22.04 + outputs: + matrix: ${{ steps.set-matrix.outputs.matrix }} + defaults: + run: + shell: pwsh + + steps: + - name: Define Job Output + id: set-matrix + run: | + $matrix = @{ + 'os' = @('ubuntu-22.04', 'windows-2022') + # wheel.py-api = "py3" makes the CUDA wheel interpreter-agnostic, + # so one builder per toolkit version is sufficient. + 'pyver' = @("3.9") + 'cuda' = @("11.8.0", "12.1.1", "12.2.2", "12.3.2", "12.4.1", "12.5.1", "13.0.2", "13.2.1") + 'releasetag' = @("basic") + 'exclude' = @( + @{ 'os' = 'windows-2022'; 'cuda' = '12.1.1' }, + @{ 'os' = 'windows-2022'; 'cuda' = '12.2.2' }, + @{ 'os' = 'windows-2022'; 'cuda' = '12.3.2' } + ) + } + + $matrixOut = ConvertTo-Json $matrix -Compress + Write-Output ('matrix=' + $matrixOut) >> $env:GITHUB_OUTPUT + + build_wheels: + name: Build Wheel ${{ matrix.os }} ${{ matrix.pyver }} ${{ matrix.cuda }} ${{ matrix.releasetag == 'wheels' && 'AVX2' || matrix.releasetag }} + needs: define_matrix + runs-on: ${{ matrix.os }} + strategy: + matrix: ${{ fromJSON(needs.define_matrix.outputs.matrix) }} + defaults: + run: + shell: pwsh + env: + CUDAVER: ${{ matrix.cuda }} + AVXVER: ${{ matrix.releasetag }} + + steps: + - name: Set up MSVC for CUDA 11.8 + if: runner.os == 'Windows' && matrix.cuda == '11.8.0' + uses: ilammy/msvc-dev-cmd@v1 + with: + arch: x64 + toolset: 14.29 + + - name: Set up MSVC + if: runner.os == 'Windows' && matrix.cuda != '11.8.0' + uses: ilammy/msvc-dev-cmd@v1 + with: + arch: x64 + + - uses: actions/checkout@v6 + with: + submodules: "recursive" + + - uses: actions/setup-python@v6 + with: + python-version: ${{ matrix.pyver }} + cache: 'pip' + + - name: Setup Mamba + uses: conda-incubator/setup-miniconda@v4.0.1 + with: + activate-environment: "llamacpp" + python-version: ${{ matrix.pyver }} + miniforge-version: latest + add-pip-as-python-dependency: true + auto-activate-base: false + + - name: Install Dependencies + env: + MAMBA_DOWNLOAD_FAILFAST: "0" + MAMBA_NO_LOW_SPEED_LIMIT: "1" + run: | + $cudaVersion = $env:CUDAVER + $cudaChannel = "nvidia/label/cuda-$cudaVersion" + if ($cudaVersion -eq '11.8.0') { + if ($IsLinux) { + $cudaPackages = @( + "${cudaChannel}::cuda-nvcc_linux-64=11.8.0", + "${cudaChannel}::cuda-cccl=11.8.89", + "${cudaChannel}::cuda-cudart=11.8.89", + "${cudaChannel}::cuda-cudart-dev=11.8.89", + "${cudaChannel}::cuda-driver-dev=11.8.89", + "${cudaChannel}::libcublas=11.11.3.6", + "${cudaChannel}::libcublas-dev=11.11.3.6" + ) + } elseif ($IsWindows) { + $cudaPackages = @( + "${cudaChannel}::cuda-nvcc_win-64=11.8.0", + "${cudaChannel}::cuda-cccl=11.8.89", + "${cudaChannel}::cuda-cudart=11.8.89", + "${cudaChannel}::cuda-cudart-dev=11.8.89", + "${cudaChannel}::libcublas=11.11.3.6", + "${cudaChannel}::libcublas-dev=11.11.3.6" + ) + } else { + throw 'Unsupported CUDA wheel build platform' + } + mamba install -y --channel-priority flexible --override-channels -c $cudaChannel $cudaPackages + } elseif ($IsLinux) { + mamba install -y --channel-priority flexible --override-channels -c $cudaChannel "${cudaChannel}::cuda-toolkit=$cudaVersion" "${cudaChannel}::cuda-nvcc_linux-64" "${cudaChannel}::cuda-cccl" "${cudaChannel}::cuda-cudart" "${cudaChannel}::cuda-cudart-dev" + } elseif ($IsWindows) { + if ($cudaVersion -like '12.5.*' -or [version]$cudaVersion -ge [version]"13.0") { + # The Windows 12.5+ toolkit meta-package pulls compiler activation + # scripts that overflow cmd.exe after MSVC is already initialized. + mamba install -y --channel-priority flexible --override-channels -c $cudaChannel "${cudaChannel}::cuda-nvcc_win-64" "${cudaChannel}::cuda-cccl" "${cudaChannel}::cuda-libraries-dev=$cudaVersion" "${cudaChannel}::cuda-cudart" "${cudaChannel}::cuda-cudart-dev" + } else { + mamba install -y --channel-priority flexible --override-channels -c $cudaChannel "${cudaChannel}::cuda-toolkit=$cudaVersion" "${cudaChannel}::cuda-nvcc_win-64" "${cudaChannel}::cuda-cccl" "${cudaChannel}::cuda-cudart" "${cudaChannel}::cuda-cudart-dev" + } + } else { + throw 'Unsupported CUDA wheel build platform' + } + if ($LASTEXITCODE -ne 0) { + exit $LASTEXITCODE + } + if ($IsWindows) { + python -m pip install build wheel ninja + } else { + sudo apt-get update + sudo apt-get install -y patchelf + python -m pip install auditwheel build wheel + } + + - name: Build Wheel + run: | + $pathSeparator = if ($IsWindows) { ';' } else { ':' } + if ($IsWindows) { + $cudaRoot = Join-Path $env:CONDA_PREFIX 'Library' + } elseif (Test-Path (Join-Path $env:CONDA_PREFIX 'targets/x86_64-linux/include/cuda_runtime.h')) { + $cudaRoot = Join-Path $env:CONDA_PREFIX 'targets/x86_64-linux' + } else { + $cudaRoot = $env:CONDA_PREFIX + } + + $env:CUDA_PATH = $cudaRoot + $env:CUDA_HOME = $cudaRoot + $env:CUDAToolkit_ROOT = $cudaRoot + $env:CUDA_TOOLKIT_ROOT_DIR = $cudaRoot + $cudaHostCompilerArg = '' + $cudaRootCmake = $cudaRoot.Replace('\', '/') + $env:CMAKE_ARGS = "-DCUDAToolkit_ROOT=$cudaRootCmake -DCUDA_TOOLKIT_ROOT_DIR=$cudaRootCmake" + if ($IsLinux) { + if ([version]$env:CUDAVER -lt [version]"12.0" -and (Test-Path '/usr/bin/g++-11')) { + $env:CC = '/usr/bin/gcc-11' + $env:CXX = '/usr/bin/g++-11' + $env:CUDAHOSTCXX = '/usr/bin/g++-11' + $cudaHostCompilerArg = " -DCMAKE_CUDA_HOST_COMPILER=$env:CUDAHOSTCXX" + } elseif (Test-Path '/usr/bin/g++-12') { + $env:CC = '/usr/bin/gcc-12' + $env:CXX = '/usr/bin/g++-12' + $env:CUDAHOSTCXX = '/usr/bin/g++-12' + $cudaHostCompilerArg = " -DCMAKE_CUDA_HOST_COMPILER=$env:CUDAHOSTCXX" + } + $env:CMAKE_ARGS = "-DCUDAToolkit_ROOT=$cudaRoot -DCUDA_TOOLKIT_ROOT_DIR=$cudaRoot$cudaHostCompilerArg" + $env:CPATH = "$cudaRoot/include$pathSeparator$env:CPATH" + $env:CPLUS_INCLUDE_PATH = "$cudaRoot/include$pathSeparator$env:CPLUS_INCLUDE_PATH" + $env:LIBRARY_PATH = "$cudaRoot/lib$pathSeparator$env:CONDA_PREFIX/lib$pathSeparator$env:LIBRARY_PATH" + $env:LD_LIBRARY_PATH = "$cudaRoot/lib$pathSeparator$env:CONDA_PREFIX/lib$pathSeparator$env:LD_LIBRARY_PATH" + $cudaLibraryPaths = @( + (Join-Path $cudaRoot 'lib'), + (Join-Path $cudaRoot 'lib64'), + (Join-Path $env:CONDA_PREFIX 'lib') + ) | Where-Object { Test-Path $_ } + Write-Output "CUDA_LIBRARY_PATHS=$($cudaLibraryPaths -join ':')" >> $env:GITHUB_ENV + } elseif ($IsWindows) { + $ninjaPath = ((Get-Command ninja -ErrorAction Stop).Source).Replace('\', '/') + $env:CMAKE_GENERATOR = 'Ninja' + $env:CMAKE_MAKE_PROGRAM = $ninjaPath + $env:PATH = "$(Join-Path $cudaRoot 'bin')$pathSeparator$env:PATH" + } + + if ($IsWindows) { + $nvccCandidates = @( + (Join-Path $cudaRoot 'bin\nvcc.exe'), + (Join-Path $env:CONDA_PREFIX 'Library\bin\nvcc.exe'), + (Join-Path $env:CONDA_PREFIX 'bin\nvcc.exe') + ) + } else { + $nvccCandidates = @( + (Join-Path $env:CONDA_PREFIX 'bin/nvcc'), + (Join-Path $env:CONDA_PREFIX 'targets/x86_64-linux/bin/nvcc') + ) + } + $nvccPath = $nvccCandidates | Where-Object { Test-Path $_ } | Select-Object -First 1 + if (-not $nvccPath) { + throw 'Failed to find nvcc in the conda environment' + } + $env:CUDACXX = $nvccPath + $env:PATH = "$(Split-Path $nvccPath)$pathSeparator$env:PATH" + if ($IsWindows) { + $nvccPathCmake = $nvccPath.Replace('\', '/') + $env:CUDACXX = $nvccPathCmake + $env:CMAKE_ARGS = "-DCMAKE_CUDA_COMPILER=$nvccPathCmake -DCMAKE_CUDA_COMPILER_ARG1=-allow-unsupported-compiler -DCMAKE_MAKE_PROGRAM=$env:CMAKE_MAKE_PROGRAM $env:CMAKE_ARGS" + } + $nvccVersion = ((& $nvccPath --version) | Select-String 'release ([0-9]+\.[0-9]+)').Matches[0].Groups[1].Value + if (-not $nvccVersion) { + throw 'Failed to detect the installed CUDA toolkit version' + } + $cudaTagVersion = $nvccVersion.Replace('.','') + $env:VERBOSE = '1' + $cudaArchs = "60-real;61-real;70-real;75-real;80-real;86-real;89-real;90-real;90-virtual" + if ([version]$nvccVersion -lt [version]"12.0") { + # CUDA 11.8 cannot compile llama.cpp's Hopper PDL device calls. + $cudaArchs = "60-real;61-real;70-real;75-real;80-real;86-real;89-real" + } elseif ([version]$nvccVersion -ge [version]"13.0") { + # CUDA 13 dropped offline compilation support for pre-Turing targets. + $cudaArchs = "75-real;80-real;86-real;89-real;90-real;90-virtual" + } + # Build real cubins for the supported GPUs and keep + # one forward-compatible PTX target instead of embedding PTX for every + # SM. This keeps the wheel under GitHub's 2 GiB release-asset limit. + $env:CMAKE_ARGS = "-DGGML_CUDA_FORCE_MMQ=ON -DGGML_CUDA=on -DCMAKE_CUDA_ARCHITECTURES=$cudaArchs -DCMAKE_CUDA_FLAGS=-allow-unsupported-compiler -DCMAKE_CUDA_FLAGS_INIT=-allow-unsupported-compiler $env:CMAKE_ARGS" + $env:CMAKE_ARGS = $env:CMAKE_ARGS + ' -DGGML_AVX2=off -DGGML_FMA=off -DGGML_F16C=off' + if ($IsLinux) { + $env:CMAKE_ARGS = $env:CMAKE_ARGS + ' -DGGML_OPENMP=OFF' + } + python -m build --wheel + # Publish tags that reflect the actual installed toolkit version. + Write-Output "CUDA_VERSION=$cudaTagVersion" >> $env:GITHUB_ENV + + - name: Repair Linux wheel + if: runner.os == 'Linux' + shell: bash + run: | + set -euxo pipefail + mkdir -p wheelhouse + export LD_LIBRARY_PATH="$PWD/llama_cpp/lib:${CUDA_LIBRARY_PATHS}:${LD_LIBRARY_PATH:-}" + auditwheel_bin="${CONDA}/envs/llamacpp/bin/auditwheel" + "${auditwheel_bin}" repair \ + --exclude libcuda.so \ + --exclude libcuda.so.1 \ + --exclude libcudart.so.11.0 \ + --exclude libcudart.so.12 \ + --exclude libcudart.so.13 \ + --exclude libcublas.so.11 \ + --exclude libcublas.so.12 \ + --exclude libcublas.so.13 \ + --exclude libcublasLt.so.11 \ + --exclude libcublasLt.so.12 \ + --exclude libcublasLt.so.13 \ + -w wheelhouse \ + dist/*.whl + rm dist/*.whl + cp wheelhouse/*.whl dist/ + "${auditwheel_bin}" show dist/*.whl + + - uses: softprops/action-gh-release@v3 + if: startsWith(github.ref, 'refs/tags/') || (github.event_name == 'workflow_dispatch' && inputs.release_tag != '') + with: + files: dist/* + # Set tag_name to -cu. + tag_name: ${{ github.event_name == 'workflow_dispatch' && inputs.release_tag || github.ref_name }}-cu${{ env.CUDA_VERSION }} + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} diff --git a/.github/workflows/build-wheels-metal.yaml b/.github/workflows/build-wheels-metal.yaml new file mode 100644 index 0000000000..892787ee1f --- /dev/null +++ b/.github/workflows/build-wheels-metal.yaml @@ -0,0 +1,68 @@ +name: Build Wheels (Metal) + +on: workflow_dispatch + +permissions: + contents: write + +jobs: + build_wheels: + name: Build wheels on ${{ matrix.os }} + runs-on: ${{ matrix.os }} + strategy: + matrix: + os: [macos-14, macos-15] + + steps: + - uses: actions/checkout@v6 + with: + submodules: "recursive" + + # Used to host cibuildwheel + - uses: actions/setup-python@v6 + with: + python-version: "3.12" + cache: 'pip' + + - name: Install dependencies (Linux/MacOS) + run: | + python -m pip install --upgrade pip + python -m pip install uv + RUST_LOG=trace python -m uv pip install -e .[all] --verbose + shell: bash + + - name: Build wheels + uses: pypa/cibuildwheel@v3.4.1 + env: + # disable repair + CIBW_REPAIR_WHEEL_COMMAND: "" + CIBW_ARCHS: "arm64" + CIBW_ENVIRONMENT: CMAKE_ARGS="-DCMAKE_OSX_ARCHITECTURES=arm64 -DCMAKE_APPLE_SILICON_PROCESSOR=arm64 -DGGML_METAL=on -DCMAKE_CROSSCOMPILING=ON" + CIBW_BUILD: "cp39-* cp310-* cp311-* cp312-*" + with: + package-dir: . + output-dir: wheelhouse2 + + - uses: actions/upload-artifact@v7 + with: + name: wheels-mac_${{ matrix.os }} + path: ./wheelhouse2/*.whl + + release: + name: Release + needs: [build_wheels] + runs-on: ubuntu-latest + + steps: + - uses: actions/download-artifact@v8 + with: + merge-multiple: true + path: dist2 + + - uses: softprops/action-gh-release@v3 + with: + files: dist2/* + # set release name to -metal + tag_name: ${{ github.ref_name }}-metal + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} diff --git a/.github/workflows/build-wheels-rocm.yaml b/.github/workflows/build-wheels-rocm.yaml new file mode 100644 index 0000000000..70fca5edb7 --- /dev/null +++ b/.github/workflows/build-wheels-rocm.yaml @@ -0,0 +1,283 @@ +name: Build Wheels (ROCm) + +on: + workflow_dispatch: + inputs: + release_tag: + description: Release tag to upload wheel assets to + required: false + type: string + amdgpu_targets: + description: AMDGPU targets to compile into the Linux ROCm wheel + required: false + default: gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1200;gfx1201 + windows_amdgpu_targets: + description: AMDGPU targets to compile into the Windows HIP Radeon wheel + required: false + default: gfx1150;gfx1151;gfx1200;gfx1201;gfx1100;gfx1101;gfx1102;gfx1030;gfx1031;gfx1032 + +permissions: + contents: write + +jobs: + build_wheels: + name: Build Wheel ${{ matrix.os }} ${{ matrix.pyver }} ROCm ${{ matrix.rocm }} + runs-on: ${{ matrix.os }} + container: + image: rocm/dev-ubuntu-22.04:${{ matrix.rocm }}-complete + strategy: + fail-fast: false + matrix: + include: + - os: ubuntu-22.04 + pyver: "3.9" + rocm: "7.2.4" + amdgpu_targets: gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1102;gfx1150;gfx1151;gfx1200;gfx1201 + + steps: + - name: Install system dependencies + run: | + apt-get update + apt-get install -y --no-install-recommends git cmake lsb-release ninja-build patchelf + + - uses: actions/checkout@v6 + with: + submodules: "recursive" + + - uses: actions/setup-python@v6 + with: + python-version: ${{ matrix.pyver }} + + - name: Install build dependencies + run: | + python -m pip install --upgrade pip + python -m pip install auditwheel build wheel + + - name: Build ROCm wheel + env: + MATRIX_AMDGPU_TARGETS: ${{ matrix.amdgpu_targets }} + INPUT_AMDGPU_TARGETS: ${{ inputs.amdgpu_targets }} + run: | + export ROCM_PATH="${ROCM_PATH:-/opt/rocm}" + export HIP_PATH="${HIP_PATH:-$ROCM_PATH}" + export PATH="$ROCM_PATH/bin:$ROCM_PATH/llvm/bin:$PATH" + export LD_LIBRARY_PATH="$ROCM_PATH/lib:$ROCM_PATH/lib64:${LD_LIBRARY_PATH:-}" + export CC="$ROCM_PATH/llvm/bin/clang" + export CXX="$ROCM_PATH/llvm/bin/clang++" + export HIPCXX="$ROCM_PATH/llvm/bin/clang" + export CMAKE_GENERATOR=Ninja + + hipconfig --version + hipcc --version + + rocm_tag="$(hipconfig --version | sed -E 's/^([0-9]+)\.([0-9]+).*/\1\2/')" + echo "ROCM_VERSION=$rocm_tag" >> "$GITHUB_ENV" + + amdgpu_targets="${INPUT_AMDGPU_TARGETS:-$MATRIX_AMDGPU_TARGETS}" + export CMAKE_ARGS="-DGGML_HIP=on -DGGML_NATIVE=off -DGGML_OPENMP=OFF -DGGML_AVX=off -DGGML_AVX2=off -DGGML_FMA=off -DGGML_F16C=off -DAMDGPU_TARGETS=$amdgpu_targets -DCMAKE_HIP_ARCHITECTURES=$amdgpu_targets" + python -m build --wheel + + - name: Repair Linux wheel + run: | + export ROCM_PATH="${ROCM_PATH:-/opt/rocm}" + export LD_LIBRARY_PATH="$PWD/llama_cpp/lib:$ROCM_PATH/lib:$ROCM_PATH/lib64:${LD_LIBRARY_PATH:-}" + mkdir -p wheelhouse + python -m auditwheel repair \ + --exclude libamdhip64.so \ + --exclude libamdhip64.so.6 \ + --exclude libamdhip64.so.7 \ + --exclude libhiprtc.so \ + --exclude libhiprtc.so.6 \ + --exclude libhiprtc.so.7 \ + --exclude libhipblas.so \ + --exclude libhipblas.so.2 \ + --exclude libhipblas.so.3 \ + --exclude libhipblaslt.so \ + --exclude libhipblaslt.so.0 \ + --exclude libhipblaslt.so.1 \ + --exclude librocblas.so \ + --exclude librocblas.so.4 \ + --exclude librocblas.so.5 \ + --exclude libhsa-runtime64.so.1 \ + --exclude libhsakmt.so.1 \ + -w wheelhouse \ + dist/*.whl + rm dist/*.whl + cp wheelhouse/*.whl dist/ + python -m auditwheel show dist/*.whl + + - uses: actions/upload-artifact@v7 + with: + name: wheels-rocm${{ env.ROCM_VERSION }}-${{ matrix.os }} + path: ./wheelhouse/*.whl + + build_wheels_windows_hip: + name: Build Wheel windows-2022 ${{ matrix.pyver }} HIP ${{ matrix.name }} + runs-on: windows-2022 + env: + HIPSDK_INSTALLER_VERSION: "26.Q1" + strategy: + fail-fast: false + matrix: + include: + - name: radeon + pyver: "3.9" + amdgpu_targets: gfx1150;gfx1151;gfx1200;gfx1201;gfx1100;gfx1101;gfx1102;gfx1030;gfx1031;gfx1032 + + steps: + - uses: actions/checkout@v6 + with: + submodules: "recursive" + + - uses: actions/setup-python@v6 + with: + python-version: ${{ matrix.pyver }} + cache: "pip" + + - name: Set up MSVC + uses: ilammy/msvc-dev-cmd@v1 + with: + arch: x64 + + - name: Install build dependencies + run: | + python -m pip install --upgrade pip + python -m pip install build wheel + + - name: Grab rocWMMA package + run: | + curl -o rocwmma.deb "https://repo.radeon.com/rocm/apt/7.2.1/pool/main/r/rocwmma-dev/rocwmma-dev_2.2.0.70201-81~24.04_amd64.deb" + 7z x rocwmma.deb + 7z x data.tar + + - name: Cache ROCm installation + id: cache-rocm + uses: actions/cache@v5 + with: + path: C:\Program Files\AMD\ROCm + key: cache-gha-rocm-${{ env.HIPSDK_INSTALLER_VERSION }}-${{ runner.os }} + + - name: Install ROCm + if: steps.cache-rocm.outputs.cache-hit != 'true' + run: | + $ErrorActionPreference = "Stop" + Invoke-WebRequest ` + -Uri "https://download.amd.com/developer/eula/rocm-hub/AMD-Software-PRO-Edition-${{ env.HIPSDK_INSTALLER_VERSION }}-Win11-For-HIP.exe" ` + -OutFile "${env:RUNNER_TEMP}\rocm-install.exe" + $proc = Start-Process "${env:RUNNER_TEMP}\rocm-install.exe" -ArgumentList '-install' -NoNewWindow -PassThru + $completed = $proc.WaitForExit(1800000) + if (-not $completed) { + $proc.Kill() + throw "ROCm installation timed out after 30 minutes" + } + if ($proc.ExitCode -ne 0) { + throw "ROCm installation failed with exit code $($proc.ExitCode)" + } + + - name: Verify ROCm + run: | + $clangPath = Get-ChildItem 'C:\Program Files\AMD\ROCm\*\bin\clang.exe' | Select-Object -First 1 + if (-not $clangPath) { + throw "ROCm installation not found" + } + & $clangPath.FullName --version + + - name: Build HIP wheel + env: + MATRIX_AMDGPU_TARGETS: ${{ matrix.amdgpu_targets }} + INPUT_AMDGPU_TARGETS: ${{ inputs.windows_amdgpu_targets }} + run: | + $ErrorActionPreference = "Stop" + $hipPath = Resolve-Path 'C:\Program Files\AMD\ROCm\*\bin\clang.exe' | Split-Path | Split-Path + $rocwmmaInclude = (Join-Path $PWD 'opt\rocm-7.2.1\include').Replace('\', '/') + $amdgpuTargets = if ($env:INPUT_AMDGPU_TARGETS) { $env:INPUT_AMDGPU_TARGETS } else { $env:MATRIX_AMDGPU_TARGETS } + + $env:HIP_PATH = $hipPath + $env:ROCM_PATH = $hipPath + $env:CMAKE_PREFIX_PATH = $hipPath + $env:HIP_PLATFORM = 'amd' + $env:PATH = "$hipPath\bin;$env:PATH" + $env:CC = "$hipPath\bin\clang.exe" + $env:CXX = "$hipPath\bin\clang++.exe" + $env:HIPCXX = "$hipPath\bin\clang.exe" + $env:CMAKE_GENERATOR = 'Unix Makefiles' + $env:CXXFLAGS = "-I$rocwmmaInclude -Wno-ignored-attributes -Wno-nested-anon-types" + $env:CMAKE_ARGS = "-DGGML_HIP=ON -DGGML_HIP_ROCWMMA_FATTN=ON -DGGML_NATIVE=OFF -DGGML_AVX=OFF -DGGML_AVX2=OFF -DGGML_FMA=OFF -DGGML_F16C=OFF -DGPU_TARGETS=$amdgpuTargets" + + python -m build --wheel + + - name: Bundle ROCm runtime DLLs + run: | + $ErrorActionPreference = "Stop" + $hipPath = Resolve-Path 'C:\Program Files\AMD\ROCm\*\bin\clang.exe' | Split-Path | Split-Path + $wheel = Get-ChildItem dist\*.whl | Select-Object -First 1 + python -m wheel unpack $wheel.FullName -d wheel-unpacked + $wheelRoot = Get-ChildItem wheel-unpacked -Directory | Select-Object -First 1 + $libDir = Join-Path $wheelRoot.FullName 'llama_cpp\lib' + New-Item -ItemType Directory -Force $libDir | Out-Null + + $dllPatterns = @( + 'amdhip64.dll', + 'hiprtc*.dll', + 'libhipblas.dll', + 'libhipblaslt.dll', + 'rocblas.dll' + ) + foreach ($pattern in $dllPatterns) { + Copy-Item (Join-Path $hipPath "bin\$pattern") $libDir -ErrorAction SilentlyContinue + } + + New-Item -ItemType Directory -Force (Join-Path $libDir 'rocblas\library') | Out-Null + New-Item -ItemType Directory -Force (Join-Path $libDir 'hipblaslt\library') | Out-Null + Copy-Item "$hipPath\bin\rocblas\library\*" (Join-Path $libDir 'rocblas\library') -Recurse -Force + Copy-Item "$hipPath\bin\hipblaslt\library\*" (Join-Path $libDir 'hipblaslt\library') -Recurse -Force + + Remove-Item dist\*.whl + python -m wheel pack $wheelRoot.FullName -d dist + New-Item -ItemType Directory -Force wheelhouse | Out-Null + Copy-Item dist/*.whl wheelhouse/ + + - uses: actions/upload-artifact@v7 + with: + name: wheels-hip-${{ matrix.name }}-windows-2022 + path: ./wheelhouse/*.whl + + release_rocm: + name: Release ROCm + needs: [build_wheels] + if: startsWith(github.ref, 'refs/tags/') || (github.event_name == 'workflow_dispatch' && inputs.release_tag != '') + runs-on: ubuntu-latest + + steps: + - uses: actions/download-artifact@v8 + with: + merge-multiple: true + path: dist + + - uses: softprops/action-gh-release@v3 + with: + files: dist/* + # Set release name to -rocm. + tag_name: ${{ github.event_name == 'workflow_dispatch' && inputs.release_tag || github.ref_name }}-rocm72 + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + + release_hip: + name: Release HIP + needs: [build_wheels_windows_hip] + if: startsWith(github.ref, 'refs/tags/') || (github.event_name == 'workflow_dispatch' && inputs.release_tag != '') + runs-on: ubuntu-latest + + steps: + - uses: actions/download-artifact@v8 + with: + merge-multiple: true + path: dist + + - uses: softprops/action-gh-release@v3 + with: + files: dist/* + # Set release name to -hip-radeon. + tag_name: ${{ github.event_name == 'workflow_dispatch' && inputs.release_tag || github.ref_name }}-hip-radeon + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} diff --git a/.github/workflows/build-wheels-vulkan.yaml b/.github/workflows/build-wheels-vulkan.yaml new file mode 100644 index 0000000000..822662c3fb --- /dev/null +++ b/.github/workflows/build-wheels-vulkan.yaml @@ -0,0 +1,130 @@ +name: Build Wheels (Vulkan) + +on: + workflow_dispatch: + inputs: + release_tag: + description: Release tag to upload wheel assets to + required: false + type: string + +permissions: + contents: write + +env: + VULKAN_SDK_VERSION: "1.4.341.0" + VULKAN_SDK_LINUX_SHA256: "ed66477d587a5587dc3601b1c2cdcc1fab5529c505f53a00171876cecd9b4fbe" + +jobs: + build_wheels: + name: Build Vulkan wheel on ${{ matrix.os }} + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + include: + - os: ubuntu-22.04 + pyver: "3.9" + artifact: wheels-vulkan-ubuntu-22.04 + - os: windows-2022 + pyver: "3.9" + artifact: wheels-vulkan-windows-2022 + + steps: + - name: Set up MSVC + if: runner.os == 'Windows' + uses: ilammy/msvc-dev-cmd@v1 + with: + arch: x64 + + - uses: actions/checkout@v6 + with: + submodules: "recursive" + + - uses: actions/setup-python@v6 + with: + python-version: ${{ matrix.pyver }} + cache: "pip" + + - name: Install Vulkan SDK + if: runner.os == 'Windows' + shell: pwsh + run: | + choco install vulkan-sdk --version="$env:VULKAN_SDK_VERSION" --no-progress -y + $vulkanSdk = Join-Path 'C:\VulkanSDK' $env:VULKAN_SDK_VERSION + if (-not (Test-Path $vulkanSdk)) { + throw "Failed to find Vulkan SDK at $vulkanSdk" + } + "VULKAN_SDK=$vulkanSdk" >> $env:GITHUB_ENV + "$vulkanSdk\Bin" >> $env:GITHUB_PATH + & "$vulkanSdk\Bin\glslc.exe" --version + + - name: Install build dependencies + if: runner.os == 'Windows' + run: | + python -m pip install --upgrade pip + python -m pip install build wheel + + - name: Install Windows build dependencies + if: runner.os == 'Windows' + run: python -m pip install ninja + + - name: Build Vulkan wheel + if: runner.os == 'Linux' + uses: pypa/cibuildwheel@v3.4.1 + env: + CIBW_BUILD: "cp38-manylinux_*" + CIBW_ARCHS: "auto64" + CIBW_MANYLINUX_X86_64_IMAGE: "manylinux2014" + CIBW_BEFORE_ALL_LINUX: > + yum install -y xz && + curl -L https://micro.mamba.pm/api/micromamba/linux-64/latest -o /tmp/micromamba.tar.bz2 && + mkdir -p /tmp/micromamba && + tar -xjf /tmp/micromamba.tar.bz2 -C /tmp/micromamba bin/micromamba && + /tmp/micromamba/bin/micromamba create -y -p /opt/vulkan -c conda-forge shaderc libvulkan-loader spirv-headers && + /opt/vulkan/bin/glslc --version && + curl -fL "https://sdk.lunarg.com/sdk/download/${{ env.VULKAN_SDK_VERSION }}/linux/vulkansdk-linux-x86_64-${{ env.VULKAN_SDK_VERSION }}.tar.xz" -o /tmp/vulkan-sdk.tar.xz && + echo "${{ env.VULKAN_SDK_LINUX_SHA256 }} /tmp/vulkan-sdk.tar.xz" | sha256sum -c - && + mkdir -p /opt/vulkan-sdk && + tar -xf /tmp/vulkan-sdk.tar.xz -C /opt/vulkan-sdk + CIBW_ENVIRONMENT_LINUX: > + CMAKE_ARGS="-DGGML_NATIVE=off -DGGML_METAL=OFF -DGGML_OPENMP=OFF -DGGML_VULKAN=on -DCMAKE_PREFIX_PATH=/opt/vulkan -DVulkan_INCLUDE_DIR=/opt/vulkan-sdk/${{ env.VULKAN_SDK_VERSION }}/x86_64/include -DVulkan_LIBRARY=/opt/vulkan/lib/libvulkan.so -DVulkan_GLSLC_EXECUTABLE=/opt/vulkan/bin/glslc" + CIBW_REPAIR_WHEEL_COMMAND_LINUX: "LD_LIBRARY_PATH=/project/llama_cpp/lib:/opt/vulkan/lib auditwheel repair --exclude libvulkan.so.1 -w {dest_dir} {wheel}" + with: + package-dir: . + output-dir: wheelhouse + + - name: Build Vulkan wheel + if: runner.os == 'Windows' + shell: pwsh + run: | + $env:CMAKE_GENERATOR = 'Ninja' + $env:CMAKE_ARGS = '-DGGML_NATIVE=off -DGGML_VULKAN=on' + python -m build --wheel + New-Item -ItemType Directory -Force wheelhouse | Out-Null + Copy-Item dist/*.whl wheelhouse/ + + - uses: actions/upload-artifact@v7 + with: + name: ${{ matrix.artifact }} + path: ./wheelhouse/*.whl + + release: + name: Release + needs: [build_wheels] + if: startsWith(github.ref, 'refs/tags/') || (github.event_name == 'workflow_dispatch' && inputs.release_tag != '') + runs-on: ubuntu-latest + + steps: + - uses: actions/download-artifact@v8 + with: + merge-multiple: true + path: dist + + - uses: softprops/action-gh-release@v3 + with: + files: dist/* + # Set release name to -vulkan. + tag_name: ${{ github.event_name == 'workflow_dispatch' && inputs.release_tag || github.ref_name }}-vulkan + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} diff --git a/.github/workflows/generate-index-from-release.yaml b/.github/workflows/generate-index-from-release.yaml new file mode 100644 index 0000000000..edf292387d --- /dev/null +++ b/.github/workflows/generate-index-from-release.yaml @@ -0,0 +1,62 @@ +name: Wheels Index + +on: + # Trigger on new release + workflow_run: + workflows: ["Release", "Build Wheels (CUDA)", "Build Wheels (Metal)", "Build Wheels (Vulkan)", "Build Wheels (ROCm)"] + types: + - completed + + # Allows you to run this workflow manually from the Actions tab + workflow_dispatch: + +# Sets permissions of the GITHUB_TOKEN to allow deployment to GitHub Pages +permissions: + contents: read + pages: write + id-token: write + +# Allow only one concurrent deployment, skipping runs queued between the run in-progress and latest queued. +# However, do NOT cancel in-progress runs as we want to allow these production deployments to complete. +concurrency: + group: "pages" + cancel-in-progress: false + +jobs: + # Single deploy job since we're just deploying + deploy: + environment: + name: github-pages + url: ${{ steps.deployment.outputs.page_url }} + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v6 + - name: Setup Pages + uses: actions/configure-pages@v6 + - name: Build + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + ./scripts/get-releases.sh + ./scripts/releases-to-pep-503.sh index/whl/cpu '^[v]?[0-9]+\.[0-9]+\.[0-9]+$' + ./scripts/releases-to-pep-503.sh index/whl/cu118 '^[v]?[0-9]+\.[0-9]+\.[0-9]+-cu118$' + ./scripts/releases-to-pep-503.sh index/whl/cu121 '^[v]?[0-9]+\.[0-9]+\.[0-9]+-cu121$' + ./scripts/releases-to-pep-503.sh index/whl/cu122 '^[v]?[0-9]+\.[0-9]+\.[0-9]+-cu122$' + ./scripts/releases-to-pep-503.sh index/whl/cu123 '^[v]?[0-9]+\.[0-9]+\.[0-9]+-cu123$' + ./scripts/releases-to-pep-503.sh index/whl/cu124 '^[v]?[0-9]+\.[0-9]+\.[0-9]+-cu124$' + ./scripts/releases-to-pep-503.sh index/whl/cu125 '^[v]?[0-9]+\.[0-9]+\.[0-9]+-cu125$' + ./scripts/releases-to-pep-503.sh index/whl/cu130 '^[v]?[0-9]+\.[0-9]+\.[0-9]+-cu130$' + ./scripts/releases-to-pep-503.sh index/whl/cu132 '^[v]?[0-9]+\.[0-9]+\.[0-9]+-cu132$' + ./scripts/releases-to-pep-503.sh index/whl/rocm72 '^[v]?[0-9]+\.[0-9]+\.[0-9]+-rocm72$' + ./scripts/releases-to-pep-503.sh index/whl/hip-radeon '^[v]?[0-9]+\.[0-9]+\.[0-9]+-hip-radeon$' + ./scripts/releases-to-pep-503.sh index/whl/vulkan '^[v]?[0-9]+\.[0-9]+\.[0-9]+-vulkan$' + ./scripts/releases-to-pep-503.sh index/whl/metal '^[v]?[0-9]+\.[0-9]+\.[0-9]+-metal$' + - name: Upload artifact + uses: actions/upload-pages-artifact@v5 + with: + # Upload entire repository + path: 'index' + - name: Deploy to GitHub Pages + id: deployment + uses: actions/deploy-pages@v5 diff --git a/.github/workflows/lint.yaml b/.github/workflows/lint.yaml new file mode 100644 index 0000000000..ac39181f90 --- /dev/null +++ b/.github/workflows/lint.yaml @@ -0,0 +1,29 @@ +name: Lint + +on: + pull_request: + branches: + - main + push: + branches: + - main + +jobs: + ruff: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: "3.12" + + - name: Install Ruff + run: python -m pip install "ruff>=0.15.7" + + - name: Lint with Ruff + run: python -m ruff check llama_cpp tests + + - name: Check formatting with Ruff + run: python -m ruff format --check llama_cpp tests diff --git a/.github/workflows/publish-to-test.yaml b/.github/workflows/publish-to-test.yaml index 5a9f3393a2..70ae5389af 100644 --- a/.github/workflows/publish-to-test.yaml +++ b/.github/workflows/publish-to-test.yaml @@ -2,7 +2,13 @@ name: Publish to TestPyPI -on: workflow_dispatch +on: + workflow_dispatch: + inputs: + dev_version: + description: 'Dev version N' + required: true + jobs: build-n-publish: @@ -10,21 +16,47 @@ jobs: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v6 with: - submodules: "true" + submodules: "recursive" + - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: - python-version: "3.8" - - name: Install dependencies + python-version: "3.11" + cache: 'pip' + + - name: Append Dev Version to __version__ + run: | + DEV_VERSION=${{ github.event.inputs.dev_version }} + CURRENT_VERSION=$(awk -F= '/__version__ =/ {print $2}' llama_cpp/__init__.py | tr -d ' "') + NEW_VERSION="${CURRENT_VERSION}.dev${DEV_VERSION}" + sed -i 's/__version__ = \".*\"/__version__ = \"'"${NEW_VERSION}"'\"/' llama_cpp/__init__.py + + - name: Install dependencies (Linux/MacOS) + if: runner.os != 'Windows' + run: | + python -m pip install --upgrade pip + python -m pip install uv + RUST_LOG=trace python -m uv pip install -e .[all] --verbose + shell: bash + + - name: Install dependencies (Windows) + if: runner.os == 'Windows' + env: + RUST_LOG: trace run: | - python -m pip install --upgrade pip pytest cmake scikit-build setuptools + python -m pip install --upgrade pip + python -m pip install uv + python -m uv pip install -e .[all] --verbose + shell: cmd + - name: Build source distribution run: | - python setup.py sdist + python -m build --sdist + - name: Publish to Test PyPI uses: pypa/gh-action-pypi-publish@release/v1 with: password: ${{ secrets.TEST_PYPI_API_TOKEN }} - repository-url: https://test.pypi.org/legacy/ \ No newline at end of file + repository-url: https://test.pypi.org/legacy/ diff --git a/.github/workflows/publish.yaml b/.github/workflows/publish.yaml index ddefd686d5..8908ccdd6f 100644 --- a/.github/workflows/publish.yaml +++ b/.github/workflows/publish.yaml @@ -10,19 +10,39 @@ jobs: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v6 with: - submodules: "true" + submodules: "recursive" + - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: - python-version: "3.8" - - name: Install dependencies + python-version: "3.9" + + - name: Install dependencies (Linux/MacOS) + if: runner.os != 'Windows' + run: | + python -m pip install --upgrade pip + python -m pip install uv + RUST_LOG=trace python -m uv pip install -e .[all] --verbose + python -m uv pip install build + shell: bash + + - name: Install dependencies (Windows) + if: runner.os == 'Windows' + env: + RUST_LOG: trace run: | - python -m pip install --upgrade pip pytest cmake scikit-build setuptools + python -m pip install --upgrade pip + python -m pip install uv + python -m uv pip install -e .[all] --verbose + python -m uv pip install build + shell: cmd + - name: Build source distribution run: | - python setup.py sdist + python -m build --sdist + - name: Publish distribution to PyPI # TODO: move to tag based releases # if: startsWith(github.ref, 'refs/tags') diff --git a/.github/workflows/test-pypi.yaml b/.github/workflows/test-pypi.yaml index 38f2d926c2..416cd0cddf 100644 --- a/.github/workflows/test-pypi.yaml +++ b/.github/workflows/test-pypi.yaml @@ -8,57 +8,105 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: ["3.7", "3.8", "3.9", "3.10", "3.11"] + python-version: ["3.9", "3.10", "3.11", "3.12"] steps: - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: python-version: ${{ matrix.python-version }} - - name: Install dependencies + cache: 'pip' + + - name: Install dependencies (Linux/MacOS) + if: runner.os != 'Windows' + run: | + python -m pip install --upgrade pip + python -m pip install uv + RUST_LOG=trace python -m uv pip install llama-cpp-python[all] --verbose + shell: bash + + - name: Install dependencies (Windows) + if: runner.os == 'Windows' + env: + RUST_LOG: trace run: | - python3 -m pip install --upgrade pip - python3 -m pip install --verbose llama-cpp-python[server,test] + python -m pip install --upgrade pip + python -m pip install uv + python -m uv pip install llama-cpp-python[all] --verbose + shell: cmd + - name: Test with pytest run: | - python3 -c "import llama_cpp" + python -c "import llama_cpp" build-windows: runs-on: windows-latest strategy: matrix: - python-version: ["3.7", "3.8", "3.9", "3.10", "3.11"] + python-version: ["3.9", "3.10", "3.11", "3.12"] steps: - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: python-version: ${{ matrix.python-version }} - - name: Install dependencies + cache: 'pip' + + - name: Install dependencies (Linux/MacOS) + if: runner.os != 'Windows' + run: | + python -m pip install --upgrade pip + python -m pip install uv + RUST_LOG=trace python -m uv pip install llama-cpp-python[all] --verbose + shell: bash + + - name: Install dependencies (Windows) + if: runner.os == 'Windows' + env: + RUST_LOG: trace run: | - python3 -m pip install --upgrade pip - python3 -m pip install --verbose llama-cpp-python[server,test] + python -m pip install --upgrade pip + python -m pip install uv + python -m uv pip install llama-cpp-python[all] --verbose + shell: cmd + - name: Test with pytest run: | - python3 -c "import llama_cpp" + python -c "import llama_cpp" build-macos: runs-on: macos-latest strategy: matrix: - python-version: ["3.7", "3.8", "3.9", "3.10", "3.11"] + python-version: ["3.9", "3.10", "3.11", "3.12"] steps: - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: python-version: ${{ matrix.python-version }} - - name: Install dependencies + cache: 'pip' + + - name: Install dependencies (Linux/MacOS) + if: runner.os != 'Windows' + run: | + python -m pip install --upgrade pip + python -m pip install uv + RUST_LOG=trace python -m uv pip install llama-cpp-python[all] --verbose + shell: bash + + - name: Install dependencies (Windows) + if: runner.os == 'Windows' + env: + RUST_LOG: trace run: | - python3 -m pip install --upgrade pip - python3 -m pip install --verbose llama-cpp-python[server,test] + python -m pip install --upgrade pip + python -m pip install uv + python -m uv pip install llama-cpp-python[all] --verbose + shell: cmd + - name: Test with pytest run: | - python3 -c "import llama_cpp" \ No newline at end of file + python -c "import llama_cpp" diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index a73e347b5a..82f7a8a7de 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -1,5 +1,4 @@ name: Tests - on: pull_request: branches: @@ -8,72 +7,177 @@ on: branches: - main +env: + REPO_ID: lmstudio-community/Qwen3.5-0.8B-GGUF + MODEL_FILE: Qwen3.5-0.8B-Q8_0.gguf + RECURRENT_REPO_ID: QuantFactory/mamba-130m-hf-GGUF + RECURRENT_MODEL_FILE: mamba-130m-hf.Q2_K.gguf + HYBRID_REPO_ID: tiiuae/Falcon-H1-Tiny-90M-Instruct-GGUF + HYBRID_MODEL_FILE: Falcon-H1-Tiny-90M-Instruct-Q2_K.gguf + EMBEDDING_REPO_ID: CompendiumLabs/bge-small-en-v1.5-gguf + EMBEDDING_MODEL_FILE: bge-small-en-v1.5-q4_k_m.gguf + MODEL_CACHE_KEY: qwen35-q8-mamba130m-q2-falconh1tiny-q2-bge-small-q4 + jobs: - build-linux: + download-model: + runs-on: ubuntu-latest + steps: + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: "3.9" + - name: Install huggingface-hub + run: pip install huggingface-hub + - name: Restore model cache + uses: actions/cache@v5 + with: + path: ~/.cache/huggingface/hub + key: ${{ runner.os }}-model-${{ env.MODEL_CACHE_KEY }} + restore-keys: | + ${{ runner.os }}-model-qwen35-q8-mamba130m-q2-falconh1tiny-q2 + - name: Download model + run: | + hf download ${{ env.REPO_ID }} ${{ env.MODEL_FILE }} + hf download ${{ env.RECURRENT_REPO_ID }} ${{ env.RECURRENT_MODEL_FILE }} + hf download ${{ env.HYBRID_REPO_ID }} ${{ env.HYBRID_MODEL_FILE }} + hf download ${{ env.EMBEDDING_REPO_ID }} ${{ env.EMBEDDING_MODEL_FILE }} + build-linux: + needs: download-model runs-on: ubuntu-latest strategy: matrix: - python-version: ["3.7", "3.8", "3.9", "3.10", "3.11"] - + python-version: ["3.9", "3.10", "3.11", "3.12"] steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v6 with: - submodules: "true" + submodules: "recursive" + - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: python-version: ${{ matrix.python-version }} - - name: Install dependencies + cache: 'pip' + - name: Restore model cache + uses: actions/cache@v5 + with: + path: ~/.cache/huggingface/hub + key: ${{ runner.os }}-model-${{ env.MODEL_CACHE_KEY }} + - name: Install dependencies (Linux/MacOS) run: | - python -m pip install --upgrade pip pytest cmake scikit-build setuptools fastapi sse-starlette httpx uvicorn pydantic-settings - pip install . -v + python -m pip install --upgrade pip + python -m pip install uv + python -m uv pip install -e .[all] --verbose + shell: bash - name: Test with pytest run: | - pytest + python -m pytest build-windows: - + needs: download-model runs-on: windows-latest strategy: matrix: - python-version: ["3.7", "3.8", "3.9", "3.10", "3.11"] - + python-version: ["3.9", "3.10", "3.11", "3.12"] steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v6 with: - submodules: "true" + submodules: "recursive" + - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: python-version: ${{ matrix.python-version }} - - name: Install dependencies + cache: 'pip' + + - name: Restore model cache + uses: actions/cache@v5 + with: + path: ~/.cache/huggingface/hub + key: ${{ runner.os }}-model-${{ env.MODEL_CACHE_KEY }} + + - name: Install dependencies (Windows) run: | - python -m pip install --upgrade pip pytest cmake scikit-build setuptools fastapi sse-starlette httpx uvicorn pydantic-settings - pip install . -v + python -m pip install --upgrade pip + python -m pip install uv + python -m uv pip install -e .[all] --verbose + shell: cmd + - name: Test with pytest run: | - pytest + python -m pytest build-macos: - - runs-on: macos-latest + needs: download-model + runs-on: macos-15 strategy: matrix: - python-version: ["3.7", "3.8", "3.9", "3.10", "3.11"] - + python-version: ["3.9", "3.10", "3.11", "3.12"] steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v6 with: - submodules: "true" + submodules: "recursive" + - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v4 + uses: actions/setup-python@v6 with: python-version: ${{ matrix.python-version }} + cache: 'pip' + + - name: System Info + run: | + uname -a + sysctl -n machdep.cpu.brand_string + python3 -c "import platform; print(platform.machine(), platform.architecture())" + + - name: Restore model cache + uses: actions/cache@v5 + with: + path: ~/.cache/huggingface/hub + key: ${{ runner.os }}-model-${{ env.MODEL_CACHE_KEY }} + + - name: Install dependencies (Linux/MacOS) + run: | + python3 -m pip install --upgrade pip + python3 -m pip install uv + CMAKE_ARGS="-DGGML_NATIVE=off" python3 -m uv pip install -e .[all] --verbose + shell: bash + + - name: Test with pytest + run: | + python3 -m pytest + + build-macos-intel: + needs: download-model + runs-on: macos-15-intel + steps: + - uses: actions/checkout@v6 + with: + submodules: "recursive" + + - name: Set up Python 3.9 + uses: actions/setup-python@v6 + with: + python-version: "3.9" + + - name: System Info + run: | + uname -a + sysctl -n machdep.cpu.brand_string + python3 -c "import platform; print(platform.machine(), platform.architecture())" + + - name: Restore model cache + uses: actions/cache@v5 + with: + path: ~/.cache/huggingface/hub + key: ${{ runner.os }}-model-${{ env.MODEL_CACHE_KEY }} + - name: Install dependencies run: | - python -m pip install --upgrade pip pytest cmake scikit-build setuptools fastapi sse-starlette httpx uvicorn pydantic-settings - pip install . -v + python3 -m pip install --upgrade pip + python3 -m pip install .[all] --verbose + shell: bash + - name: Test with pytest run: | - pytest \ No newline at end of file + python3 -m pytest diff --git a/.gitignore b/.gitignore index 36ed7f7fdc..ff773c6684 100644 --- a/.gitignore +++ b/.gitignore @@ -1,8 +1,13 @@ +*.local + +.python-version + .vscode/ _skbuild/ .envrc +.direnv models/ @@ -12,11 +17,11 @@ __pycache__/ *$py.class # C extensions -*.so -*.dylib -*.metal -*.dll -*.lib +llama_cpp/*.so +llama_cpp/*.dylib +llama_cpp/*.metal +llama_cpp/*.dll +llama_cpp/*.lib # Distribution / packaging .Python @@ -61,6 +66,7 @@ coverage.xml *.py,cover .hypothesis/ .pytest_cache/ +.ruff_cache/ cover/ # Translations diff --git a/.gitmodules b/.gitmodules index 7edf0975dc..f56cca32df 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +1,3 @@ [submodule "vendor/llama.cpp"] path = vendor/llama.cpp - url = https://github.com/ggerganov/llama.cpp.git + url = https://github.com/ggml-org/llama.cpp.git diff --git a/CHANGELOG.md b/CHANGELOG.md index 8b5db37b09..7cc732b935 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,66 +7,925 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] -## [0.1.70] +- feat(example): support server video inputs and Gemma text tool calls by @abetlen in #2291 +- feat: update llama.cpp to ggml-org/llama.cpp@e3471b3e7 +- fix(example): support multi-step Responses tool streaming by @abetlen in #2288 +- fix(ci): Repair Linux accelerator wheels for manylinux publishing + +## [0.3.28] + +- feat(example): align server MTP support with llama.cpp by @abetlen in #2283 +- feat: update llama.cpp to ggml-org/llama.cpp@9e3b928fd +- feat(example): add OpenAI-compatible embeddings endpoint by @abetlen in #2281 + +## [0.3.27] + +- feat: update llama.cpp to ggml-org/llama.cpp@465b1f0e7 +- feat(example): Updated server example (batch processing, multi-token prediction, `/v1/responses` api, response parsing) by @abetlen in #2174 + +## [0.3.26] + +- feat: Generic Multimodal Chat Handler by @abetlen in #2256 +- feat: update llama.cpp to ggml-org/llama.cpp@7c158fbb4 +- feat(ci): add ROCm wheel builds by @abetlen in #2252 +- feat(ci): add Vulkan wheel builds by @abetlen in #2251 +- fix: handle additional `from_pretrained` files in subfolders by @TNing in #2085 + +## [0.3.25] + +- feat: Update llama.cpp to ggml-org/llama.cpp@210a6570c by @abetlen in #2242 +- feat: add Gemma 4 multimodal chat support by @abetlen in #2241 +- feat(ci): add CUDA 13.0 and 13.2 wheel builds by @abetlen in #2239 +- feat(ci): add CUDA 11.8 wheel builds by @abetlen in #2238 +- fix(ci): add Pascal compute capability targets to CUDA wheel builds by @abetlen in #2237 + +## [0.3.24] + +- docs: update ROCm install instructions by @agronholm in #1867 +- fix: clear prompt for recurrent / hybrid models when only a partial prefix matches by @avion23 in #2108 +- fix: match Transformers `tojson` in chat template rendering by @CISC in #1486 +- fix: use env var configured multimodal library override paths when loading shared libraries by @navratil-matej in #1782 +- feat: add Jinja2 loop controls to chat templates by @handshape in #2018 +- fix: avoid cleanup errors for partially initialized `LlamaModel` objects by @usernames122 in #2173 +- fix: suppress stdout and stderr in Jupyter notebooks by @Anai-Guo in #2181 +- feat: enable arm64 musl builds by @acon96 in #2221 +- feat: Update llama.cpp to ggml-org/llama.cpp@af6528e6d +- fix: model fails to load when chat template uses HuggingFace generation tags by @tobocop2 in #2226 +- docs: add contributing guide by @abetlen in #2229 +- chore: Migrate llama.cpp submodule URL to ggml-org/llama.cpp by @shalinib-ibm in #2034 +- fix: Enable unified KV cache for embedding contexts to preserve full per-sequence context in batch embedding calls by @SanjanaB123 in #2217 + +## [0.3.23] + +- feat: Update llama.cpp to ggerganov/llama.cpp@7d442abf +- fix: Correct batched embedding outputs for multi-sequence `embed()` calls by @Anai-Guo in #2205 +- fix: Configure embedding contexts with enough sequence slots for batched `embed()` calls +- fix: Mark all embedding input tokens as outputs to avoid llama.cpp override warnings by @Anai-Guo in #2212 + +## [0.3.22] + +- feat: Update llama.cpp to ggerganov/llama.cpp@63d93d173 +- feat(ci): Re-enable Windows CUDA wheels and add CUDA 12.5.1 wheel builds + +## [0.3.21] + +- feat: Update llama.cpp to ggerganov/llama.cpp@f53577432 and sync Python bindings +- fix(ci): Build one arm64 release wheel for `py3-none` wheel publishing + +## [0.3.20] + +- refactor: Replace deprecated llama.cpp references in library, docs, and examples by @abetlen in #2170 +- feat: Update llama.cpp to ggerganov/llama.cpp@f49e9178767d557a522618b16ce8694f9ddac628 by @abetlen in #2169 +- feat(server): Add model-load `chat_template_kwargs` support and document the CLI/config usage by @abetlen in #2168 +- ci: Publish release wheels as `py3-none` by @Bing-su in #2166 +- fix(ci): Publish distinct manylinux and musllinux CPU wheels by @abetlen in #2165 + +## [0.3.19] + +- feat: Update llama.cpp to ggerganov/llama.cpp@c0159f9c1f874da15e94f371d136f5920b4b5335 by @abetlen in #2161 +- fix: Handle embedding models without KV memory and test embeddings with a real GGUF embedding model by @abetlen in #2160 +- fix(ci): Shrink CUDA wheel fatbins so CUDA releases stay under GitHub's asset size limit by @abetlen in #2158 + +## [0.3.18] + +- feat: Expose `attention_type` in `Llama.__init__` for non-causal embedding models by @jamesbiederbeck in #2143 +- fix(ci): Build Docker images from the checked-out source and sanitize branch tags by @abetlen in #2156 +- fix(ci): Fix the CUDA wheel workflow and keep release tags aligned with the built toolkit by @abetlen in #2155 +- fix(ci): Speed up release wheel builds by moving arm64 off QEMU and parallelizing riscv64 by @abetlen in #2154 + +## [0.3.17] + +- feat: Update llama.cpp to ggerganov/llama.cpp@49bfddeca18e62fa3d39114a23e9fcbdf8a22388 and sync Python bindings by @abetlen in #2151 +- fix: Handle Qwen 3.5 hybrid prefix reuse by @codavidgarcia and @r-dh in #2152 +- chore(dev): Add Ruff-based formatting and a safe lint baseline, and run it in CI for pull requests and pushes to `main` +- fix(ci): Run macOS CI on supported Apple Silicon and Intel runners by @abetlen in #2150 +- fix(ci): Use the `hf` CLI instead of the deprecated `huggingface-cli` name in GitHub Actions and docs by @abetlen in #2149 +- ci: add riscv64 wheel builds to release workflow by @gounthar in #2139 + +## [0.3.16] + +- feat: Update llama.cpp to ggerganov/llama.cpp@4227c9be4268ac844921b90f31595f81236bd317 + +## [0.3.15] + +- feat: Update llama.cpp to ggerganov/llama.cpp@9a96389544a08fd829fccda28142ce2066017fde +- feat: Add gpt-oss chat format support through strftime_now in chat format by @iamlemec in af637928db7351e030011085f818b034c6efc047 +- fix: rename op_offloat to op_offload in llama.py by @sergey21000 in #2046 + +## [0.3.14] + +- feat: Update llama.cpp to ggerganov/llama.cpp@79e0b68c178656bb0632cb8602d2940b755077f8 + +## [0.3.13] + +- feat: Update llama.cpp to ggerganov/llama.cpp@bdca38376f7e8dd928defe01ce6a16218a64b040 +- fix: Better chat format for Qwen2.5-VL by @alcoftTAO in #2040 + +## [0.3.12] + +- feat: Update llama.cpp to ggerganov/llama.cpp@a0374a67e2924f2e845cdc59dd67d9a44065a89c + +## [0.3.11] + +- fix: Update reference to `llama_kv_cache_clear` in Llama.embed. Closes #2037 by @abetlen in 9e5a4eaa84156084ed7bbb91e6efcc91dc6217bc + +## [0.3.10] + +- feat: Update llama.cpp to ggerganov/llama.cpp@8846aace4934ad29651ea61b8c7e3f6b0556e3d2 +- feat: Add support for llama.cpp multimodal, add Qwen2.5-VL chat handler by @abetlen in cd548bd0f14210627798237d5c2ea78acfb88ccb + +## [0.3.9] + +- feat: Update llama.cpp to ggerganov/llama.cpp@8733e0cf6eefc7c7752297cc22d0836706f4222c + +## [0.3.8] + +- feat: Update llama.cpp to ggerganov/llama.cpp@7841fc723e059d1fd9640e5c0ef19050fcc7c698 + +## [0.3.7] + +- feat: Update llama.cpp to ggerganov/llama.cpp@794fe23f29fb40104975c91fe19f23798f7c726e +- fix(ci): Fix the CUDA workflow by @oobabooga in #1894 +- fix: error showing time spent in llama perf context print, adds `no_perf` flag to `Llama` class by @shakalaca in #1898 + +## [0.3.6] + +- feat: Update llama.cpp to ggerganov/llama.cpp@f7cd13301c2a88f97073fd119072b4cc92c08df1 +- fix(server): streaming resource lock by @gjpower in #1879 + +## [0.3.5] + +- feat: Update llama.cpp to ggerganov/llama.cpp@26a8406ba9198eb6fdd8329fa717555b4f77f05f +- fix(ci): Fix release by updating macos runner image to non-deprecated version by @abetlen in afedfc888462f9a6e809dc9455eb3b663764cc3f +- fix(server): add missing await statements for async exit_stack handling by @gjpower in #1858 + +## [0.3.4] + +- fix(ci): Build wheels for macos 13-15, cuda 12.1-12.4 by @abetlen in ca808028bd16b8327bd84128d48015a4b1304690 + +## [0.3.3] + +- feat: Update llama.cpp to ggerganov/llama.cpp@ce8784bdb153ff7794dde5a50b0ebfa51baa6171 +- fix: chat API logprobs format by @domdomegg in #1788 +- feat: Add support for CUDA 12.6, fix CUDA 12.5 by @Smartappli in #1775 +- fix: Make content not required in ChatCompletionRequestAssistantMessage by @feloy in #1807 +- fix: Fix pickling of Llama class by setting seed from _seed member by @abetlen in 2523472c3eccb9ab9277117cc4ff705212b6888a +- fix: Fix logit-bias type hint by @ddh0 in #1802 +- fix(server): Avoid thread starvation on many concurrent requests by making use of asyncio to lock llama_proxy context by @gjpower in #1798 +- fix(server): Added missing exit_stack.close() to /v1/chat/completions by @Ian321 in #1796 +- fix(examples): Refactor Batching notebook to use new sampler chain API by @lukestanley in #1793 +- fix(docs): Update development instructions by @Florents-Tselai in #1833 +- fix(docs): Remove ref to llama_eval in llama_cpp.py docs by @richdougherty in #1819 + +## [0.3.2] + +- feat: Update llama.cpp to ggerganov/llama.cpp@74d73dc85cc2057446bf63cc37ff649ae7cebd80 + +## [0.3.1] + +- feat: Update llama.cpp to ggerganov/llama.cpp@c919d5db39c8a7fcb64737f008e4b105ee0acd20 +- feat: Expose libggml in internal APIs by @abetlen in #1761 +- fix: Fix speculative decoding by @abetlen in 9992c5084a3df2f533e265d10f81d4269b97a1e6 and e975dabf74b3ad85689c9a07719cbb181313139b +- misc: Rename all_text to remaining_text by @xu-song in #1658 + +## [0.3.0] + +- feat: Update llama.cpp to ggerganov/llama.cpp@ea9c32be71b91b42ecc538bd902e93cbb5fb36cb +- feat: Enable detokenizing special tokens with special=True by @benniekiss in #1596 +- feat(ci): Speed up CI workflows using uv, add support for CUDA 12.5 wheels by @Smartappli in e529940f45d42ed8aa31334123b8d66bc67b0e78 +- feat: Add loading sharded GGUF files from HuggingFace with Llama.from_pretrained(additional_files=[...]) by @Gnurro in 84c092063e8f222758dd3d60bdb2d1d342ac292e +- feat: Add option to configure n_ubatch by @abetlen in 6c44a3f36b089239cb6396bb408116aad262c702 +- feat: Update sampling API for llama.cpp. Sampling now uses sampler chain by @abetlen in f8fcb3ea3424bcfba3a5437626a994771a02324b +- fix: Don't store scores internally unless logits_all=True. Reduces memory requirements for large context by @abetlen in 29afcfdff5e75d7df4c13bad0122c98661d251ab +- fix: Fix memory allocation of ndarray in by @xu-song in #1704 +- fix: Use system message in og qwen format by @abetlen in 98eb092d3c6e7c142c4ba2faaca6c091718abbb3 + + +## [0.2.90] + +- feat: Update llama.cpp to ggerganov/llama.cpp@1d1ccce67613674c75c9c7e3fa4c1e24e428ba48 +- feat: Add support for `MiniCPMv26ChatHandler` and `minicpm-v-26` in server by @abetlen in f70df824985d875226793b94dacc0c302a4256b2 + +## [0.2.89] + +- feat: Update llama.cpp to ggerganov/llama.cpp@cfac111e2b3953cdb6b0126e67a2487687646971 +- fix: Llama.close didn't free lora adapter by @jkawamoto in #1679 +- fix: missing dependencies for test by @jkawamoto in #1680 + +## [0.2.88] + +- feat: Update llama.cpp to ggerganov/llama.cpp@fc4ca27b25464a11b3b86c9dbb5b6ed6065965c2 +- fix: only print 'cache saved' in verbose mode by @lsorber in #1668 +- fix: Added back from_file method to LlamaGrammar by @ExtReMLapin in #1673 +- fix: grammar prints on each call by @abetlen in 0998ea0deea076a547d54bd598d6b413b588ee2b +- feat: Enable recursive search of HFFS.ls when using from_pretrained by @benHeidabetlen in #1656 +- feat: Add more detailed log for prefix-match by @xu-song in #1659 + +## [0.2.87] + +- feat: Update llama.cpp to ggerganov/llama.cpp@be55695eff44784a141a863f273661a6bce63dfc +- fix: Include all llama.cpp source files and subdirectories by @abetlen in 9cad5714ae6e7c250af8d0bbb179f631368c928b +- feat(ci): Re-build wheel index automatically when releases are created by @abetlen in 198f47dc1bd202fd2b71b29e041a9f33fe40bfad + +## [0.2.86] + +- feat: Update llama.cpp to ggerganov/llama.cpp@398ede5efeb07b9adf9fbda7ea63f630d476a792 +- feat: Ported back new grammar changes from C++ to Python implementation by @ExtReMLapin in (#1637) +- fix: llama_grammar_accept_token arg order by @tc-wolf in (#1649) + +## [0.2.85] + +- feat: Update llama.cpp to ggerganov/llama.cpp@398ede5efeb07b9adf9fbda7ea63f630d476a792 +- fix: Missing LoRA adapter after API change by @shamitv in #1630 +- fix(docker): Update Dockerfile BLAS options by @olivierdebauche in #1632 +- fix(docker): Fix GGML_CUDA param by @olivierdebauche in #1633 +- fix(docker): Update Dockerfile build options from `LLAMA_` to `GGML_` by @olivierdebauche in #1634 +- feat: FreeBSD compatibility by @yurivict in #1635 + +## [0.2.84] + +- feat: Update llama.cpp to ggerganov/llama.cpp@4730faca618ff9cee0780580145e3cbe86f24876 +- fix: fix: Correcting run.sh filepath in Simple Docker implementation by @mashuk999 in #1626 + +## [0.2.83] + +- feat: Update llama.cpp to ggerganov/llama.cpp@081fe431aa8fb6307145c4feb3eed4f48cab19f8 +- feat: Add 'required' literal to ChatCompletionToolChoiceOption by @mjschock in #1597 +- fix: Change repeat_penalty to 1.0 to match llama.cpp defaults by @ddh0 in #1590 +- fix(docs): Update README.md typo by @ericcurtin in #1589 +- fix(server): Use split_mode from model settings by @grider-withourai in #1594 +- feat(ci): Dockerfile update base images and post-install cleanup by @Smartappli in #1530 + +## [0.2.82] + +- feat: Update llama.cpp to ggerganov/llama.cpp@7fdb6f73e35605c8dbc39e9f19cd9ed84dbc87f2 + +## [0.2.81] + +- feat: Update llama.cpp to ggerganov/llama.cpp@968967376dc2c018d29f897c4883d335bbf384fb +- fix(ci): Fix CUDA wheels, use LLAMA_CUDA instead of removed LLAMA_CUBLAS by @abetlen in 4fb6fc12a02a68884c25dd9f6a421cacec7604c6 +- fix(ci): Fix MacOS release, use macos-12 image instead of removed macos-11 by @abetlen in 3a551eb5263fdbd24b36d7770856374c04e92788 + +## [0.2.80] + +- feat: Update llama.cpp to ggerganov/llama.cpp@023b8807e10bc3ade24a255f01c1ad2a01bb4228 +- fix(server): Fix bug in FastAPI streaming response where dependency was released before request completes causing SEGFAULT by @abetlen in 296304b60bb83689659883c9cc24f4c074dd88ff +- fix(server): Update default config value for embeddings to False to fix error in text generation where logits were not allocated by llama.cpp by @abetlen in bf5e0bb4b151f4ca2f5a21af68eb832a96a79d75 +- fix(ci): Fix the CUDA workflow by @oobabooga in #1551 +- docs: Update readme examples to use newer Qwen2 model by @jncraton in #1544 + +## [0.2.79] + +- feat: Update llama.cpp to ggerganov/llama.cpp@9c77ec1d74874ee22bdef8f110e8e8d41389abf2 +- feat(ci): Update workflows and pre-built wheels by @Smartappli in #1416 +- feat: Add .close() method to Llama class to explicitly free model from memory by @jkawamoto in #1513 +- feat: Support SPM infill by @CISC in #1492 + +## [0.2.78] + +- feat: Update llama.cpp to ggerganov/llama.cpp@fd5ea0f897ecb3659d6c269ef6f3d833e865ead7 +- fix: Avoid duplicate special tokens in chat formats by @CISC in #1439 +- fix: fix logprobs when BOS is not present by @ghorbani in #1471 +- feat: adding rpc_servers parameter to Llama class by @chraac in #1477 + +## [0.2.77] + +- feat: Update llama.cpp to ggerganov/llama.cpp@bde7cd3cd949c1a85d3a199498ac98e78039d46f +- fix: string value kv_overrides by @abetlen in df45a4b3fe46e72664bda87301b318210c6d4782 +- fix: Fix typo in Llama3VisionAlphaChatHandler by @abetlen in 165b4dc6c188f8fda2fc616154e111f710484eba +- fix: Use numpy recarray for candidates data, fixes bug with temp < 0 by @abetlen in af3ed503e9ce60fe6b5365031abad4176a3536b3 +fix: Disable Windows+CUDA workaround when compiling for HIPBLAS by Engininja2 in #1493 + +## [0.2.76] + +- feat: Update llama.cpp to ggerganov/llama.cpp@0df0aa8e43c3378975269a51f9b876c8692e70da +- feat: Improve Llama.eval performance by avoiding list conversion by @thoughtp0lice in #1476 +- example: LLM inference with Ray Serve by @rgerganov in #1465 + +## [0.2.75] + +- feat: Update llama.cpp to ggerganov/llama.cpp@13ad16af1231ab2d245d35df3295bcfa23de1305 +- fix: segfault for models without eos / bos tokens by @abetlen in d99a6ba607a4885fb00e63e967964aa41bdbbbcb +- feat: add MinTokensLogitProcessor and min_tokens argument to server by @twaka in #1333 +- misc: Remove unnecessary metadata lookups by @CISC in #1448 + +## [0.2.74] + +- feat: Update llama.cpp to ggerganov/llama.cpp@b228aba91ac2cd9eb90e9d423ba1d0d20e0117e2 +- fix: Enable CUDA backend for llava by @abetlen in 7f59856fa6f3e23f07e12fc15aeb9359dc6c3bb4 +- docs: Fix typo in README.md by @yupbank in #1444 + +## [0.2.73] + +- feat: Update llama.cpp to ggerganov/llama.cpp@25c6e82e7a1ad25a42b0894e87d9b5c557409516 +- fix: Clear kv cache at beginning of image chat formats to avoid bug when image is evaluated first by @abetlen in ac55d0a175115d1e719672ce1cb1bec776c738b1 + +## [0.2.72] + +- fix(security): Remote Code Execution by Server-Side Template Injection in Model Metadata by @retr0reg in b454f40a9a1787b2b5659cd2cb00819d983185df +- fix(security): Update remaining jinja chat templates to use immutable sandbox by @CISC in #1441 + +## [0.2.71] + +- feat: Update llama.cpp to ggerganov/llama.cpp@911b3900dded9a1cfe0f0e41b82c7a29baf3a217 +- fix: Make leading bos_token optional for image chat formats, fix nanollava system message by @abetlen in 77122638b4153e31d9f277b3d905c2900b536632 +- fix: free last image embed in llava chat handler by @abetlen in 3757328b703b2cd32dcbd5853271e3a8c8599fe7 + +## [0.2.70] + +- feat: Update llama.cpp to ggerganov/llama.cpp@c0e6fbf8c380718102bd25fcb8d2e55f8f9480d1 +- feat: fill-in-middle support by @CISC in #1386 +- fix: adding missing args in create_completion for functionary chat handler by @skalade in #1430 +- docs: update README.md @eltociear in #1432 +- fix: chat_format log where auto-detected format prints None by @balvisio in #1434 +- feat(server): Add support for setting root_path by @abetlen in 0318702cdc860999ee70f277425edbbfe0e60419 +- feat(ci): Add docker checks and check deps more frequently by @Smartappli in #1426 +- fix: detokenization case where first token does not start with a leading space by @noamgat in #1375 +- feat: Implement streaming for Functionary v2 + Bug fixes by @jeffrey-fong in #1419 +- fix: Use memmove to copy str_value kv_override by @abetlen in 9f7a85571ae80d3b6ddbd3e1bae407b9f1e3448a +- feat(server): Remove temperature bounds checks for server by @abetlen in 0a454bebe67d12a446981eb16028c168ca5faa81 +- fix(server): Propagate flash_attn to model load by @dthuerck in #1424 + +## [0.2.69] + +- feat: Update llama.cpp to ggerganov/llama.cpp@6ecf3189e00a1e8e737a78b6d10e1d7006e050a2 +- feat: Add llama-3-vision-alpha chat format by @abetlen in 31b1d95a6c19f5b615a3286069f181a415f872e8 +- fix: Change default verbose value of verbose in image chat format handlers to True to match Llama by @abetlen in 4f01c452b6c738dc56eacac3758119b12c57ea94 +- fix: Suppress all logs when verbose=False, use hardcoded fileno's to work in colab notebooks by @abetlen in f116175a5a7c84569c88cad231855c1e6e59ff6e +- fix: UTF-8 handling with grammars by @jsoma in #1415 + +## [0.2.68] + +- feat: Update llama.cpp to ggerganov/llama.cpp@77e15bec6217a39be59b9cc83d6b9afb6b0d8167 +- feat: Add option to enable flash_attn to Lllama params and ModelSettings by @abetlen in 22d77eefd2edaf0148f53374d0cac74d0e25d06e +- fix(ci): Fix build-and-release.yaml by @Smartappli in #1413 + +## [0.2.67] + +- fix: Ensure image renders before text in chat formats regardless of message content order by @abetlen in 3489ef09d3775f4a87fb7114f619e8ba9cb6b656 +- fix(ci): Fix bug in use of upload-artifact failing to merge multiple artifacts into a single release by @abetlen in d03f15bb73a1d520970357b702a9e7d4cc2a7a62 + +## [0.2.66] + +- feat: Update llama.cpp to ggerganov/llama.cpp@8843a98c2ba97a25e93319a104f9ddfaf83ce4c4 +- feat: Generic Chat Formats, Tool Calling, and Huggingface Pull Support for Multimodal Models (Obsidian, LLaVA1.6, Moondream) by @abetlen in #1147 +- ci(fix): Workflow actions updates and fix arm64 wheels not included in release by @Smartappli in #1392 +- ci: Add support for pre-built cuda 12.4.1 wheels by @Smartappli in #1388 +- feat: Add support for str type kv_overrides by @abetlen in a411612b385cef100d76145da1fbd02a7b7cc894 +- fix: Functionary bug fixes by @jeffrey-fong in #1385 +- examples: fix quantize example by @iyubondyrev in #1387 +- ci: Update dependabot.yml by @Smartappli in #1391 + +## [0.2.65] + +- feat: Update llama.cpp to ggerganov/llama.cpp@46e12c4692a37bdd31a0432fc5153d7d22bc7f72 +- feat: Allow for possibly non-pooled embeddings by @iamlemec in #1380 + +## [0.2.64] + +- feat: Update llama.cpp to ggerganov/llama.cpp@4e96a812b3ce7322a29a3008db2ed73d9087b176 +- feat: Add `llama-3` chat format by @andreabak in #1371 +- feat: Use new llama_token_is_eog in create_completions by @abetlen in d40a250ef3cfaa8224d12c83776a2f1de96ae3d1 +- feat(server): Provide ability to dynamically allocate all threads if desired using -1 by @sean-bailey in #1364 +- ci: Build arm64 wheels by @gaby in 611781f5319719a3d05fefccbbf0cc321742a026 +- fix: Update scikit-build-core build dependency avoid bug in 0.9.1 by @evelkey in #1370 + +## [0.2.63] + +- feat: Update llama.cpp to ggerganov/llama.cpp@0e4802b2ecbaab04b4f829fde4a3096ca19c84b5 +- feat: Add stopping_criteria to ChatFormatter, allow stopping on arbitrary token ids, fixes llama3 instruct by @abetlen in cc81afebf04d26ca1ac3cf72f23f18da6ab58588 + +## [0.2.62] + +- feat: Update llama.cpp to ggerganov/llama.cpp@3b8f1ec4b18770531d0b1d792f3edf08254e4f0c +- feat: update grammar schema converter to match llama.cpp by @themrzmaster in #1353 +- feat: add disable_ping_events flag by @khimaros in #1257 +- feat: Make saved state more compact on-disk by @tc-wolf in #1296 +- feat: Use all available CPUs for batch processing by @ddh0 in #1345 + +## [0.2.61] + +- feat: Update llama.cpp to ggerganov/llama.cpp@ba5e134e073ec6837078c874aba44a702944a676 +- fix: pass correct type to chat handlers for chat completion logprobs by @abetlen in bb65b4d76411112c6fb0bf759efd746f99ef3c6b +- feat: Add support for yaml based server configs by @abetlen in 060bfa64d529ade2af9b1f4e207a3937bbc4138f +- feat: Add typechecking for ctypes structure attributes by @abetlen in 1347e1d050fc5a9a32ffe0bb3e22858da28003bd + +## [0.2.60] + +- feat: Update llama.cpp to ggerganov/llama.cpp@75cd4c77292034ecec587ecb401366f57338f7c0 +- fix: Always embed metal library by @abetlen in b3bfea6dbfb6ed9ce18f9a2723e0a9e4bd1da7ad +- fix: missing logprobs in response, incorrect response type for functionary by @abetlen in 1ae3abbcc3af7f4a25a3ffc40b246f18039565e8 +- fix(docs): incorrect tool_choice example by @CISC in #1330 + +## [0.2.59] + +- feat: Update llama.cpp to ggerganov/llama.cpp@ba0c7c70ab5b15f1f2be7fb0dfbe0366dda30d6c +- feat: Binary wheels for CPU, CUDA (12.1 - 12.3), Metal by @abetlen, @jllllll, and @oobabooga in #1247 +- fix: segfault when logits_all=False by @abetlen in 8649d7671bd1a7c0d9cc6a5ad91c6ca286512ab3 +- fix: last tokens passing to sample_repetition_penalties function by @ymikhailov in #1295 + +## [0.2.58] + +- feat: Update llama.cpp to ggerganov/llama.cpp@ba0c7c70ab5b15f1f2be7fb0dfbe0366dda30d6c +- feat: add support for KV cache quantization options by @Limour-dev in #1307 +- feat: Add logprobs support to chat completions by @windspirit95 in #1311 +- fix: set LLAMA_METAL_EMBED_LIBRARY=on on MacOS arm64 by @bretello in #1289 +- feat: Add tools/functions variables to Jinja2ChatFormatter, add function response formatting for all simple chat formats by @CISC in #1273 +- fix: Changed local API doc references to hosted by by @lawfordp2017 in #1317 + +## [0.2.57] + +- feat: Update llama.cpp to ggerganov/llama.cpp@ac9ee6a4ad740bc1ee484ede43e9f92b5af244c1 +- fix: set default embedding pooling type to unspecified by @abetlen in 4084aabe867b8ec2aba1b22659e59c9318b0d1f3 +- fix: Fix and optimize functionary chat handler by @jeffrey-fong in #1282 +- fix: json mode for basic chat formats by @abetlen in 20e6815252d0efd9f015f7adbf108faaf36e3f3c + +## [0.2.56] + +- feat: Update llama.cpp to ggerganov/llama.cpp@c2101a2e909ac7c08976d414e64e96c90ee5fa9e +- feat(server): Add endpoints for tokenize, detokenize and count tokens by @felipelo in #1136 +- feat: Switch embed to llama_get_embeddings_seq by @iamlemec in #1263 +- fix: Fixed json strings grammar by blacklisting character control set by @ExtReMLapin in d02a9cf16ff88ad011e2eb1ce29f4d9400f13cd1 +- fix: Check for existence of clip model path by @kejcao in #1264 + +## [0.2.55] + +- feat: Update llama.cpp to ggerganov/llama.cpp@9731134296af3a6839cd682e51d9c2109a871de5 +- docs: fix small typo in README: 'model know how' -> 'model knows how' by @boegel in #1244 + +## [0.2.54] + +- feat: Update llama.cpp to ggerganov/llama.cpp@cb49e0f8c906e5da49e9f6d64a57742a9a241c6a +- docs: fix typo in README.md embeddings example by @iamlemec in #1232 + +## [0.2.53] + +- feat: Update llama.cpp to ggerganov/llama.cpp@cb49e0f8c906e5da49e9f6d64a57742a9a241c6a +- fix: eos/bos_token set correctly for Jinja2ChatFormatter and automatic chat formatter by @CISC in #1230 + +## [0.2.52] + +- feat: Update llama.cpp to ggerganov/llama.cpp@a33e6a0d2a66104ea9a906bdbf8a94d050189d91 +- fix: Llava15ChatHandler (this function takes at least 4 arguments) by @abetlen in 8383a9e5620f5df5a88f62da16813eac200dd706 + +## [0.2.51] + +- feat: Update llama.cpp to ggerganov/llama.cpp@c39373398803c669056304090050fe3f44b41bf9 +- fix: Restore type hints for low-level api by @abetlen in 19234aa0dbd0c3c87656e65dd2b064665371925b + +## [0.2.50] + +- docs: Update Functionary OpenAI Server Readme by @jeffrey-fong in #1193 +- fix: LlamaHFTokenizer now receives pre_tokens by @abetlen in 47bad30dd716443652275099fa3851811168ff4a + +## [0.2.49] + +- fix: module 'llama_cpp.llama_cpp' has no attribute 'c_uint8' in Llama.save_state by @abetlen in db776a885cd4c20811f22f8bd1a27ecc71dba927 +- feat: Auto detect Mixtral's slightly different format by @lukestanley in #1214 + +## [0.2.48] + +- feat: Update llama.cpp to ggerganov/llama.cpp@15499eb94227401bdc8875da6eb85c15d37068f7 +- feat: Add Google's Gemma formatting via chat_format="gemma" by @alvarobartt in #1210 +- feat: support minItems/maxItems in JSON grammar converter by @nopperl in 3921e10770996d95a9eb22c8248bacef39f69365 +- fix: Update from_pretrained defaults to match hf_hub_download and pull to local cache folder by @abetlen in e6d6260a91b7831733f7d1f73c7af46a3e8185ed +- fix: Raise exceptions when llama model or context fails to load by @abetlen in dd22010e85265ae840c76ec835d67a29ed852722 +- docs: Update README.md to fix pip install llama cpp server by @audip in #1187 + +## [0.2.47] + +- feat: Update llama.cpp to ggerganov/llama.cpp@973053d8b0d04809836b3339a50f68d9c842de90 + +## [0.2.46] -### Fixed +- feat: Update llama.cpp to ggerganov/llama.cpp@ba2135ccae7462470b3865c6e41d2e1d734eac05 +- feat: Pull models directly from huggingface by @abetlen in #1206 +- feat(low-level-api): Improve API static type-safety and performance. Low level api functions are positional args only now. by @abetlen in #1205 + +## [0.2.45] + +- feat: Update llama.cpp to ggerganov/llama.cpp@89febfed9322c8849520dc63c93ee4f5fd72556e + +## [0.2.44] + +- feat: Update llama.cpp to ggerganov/llama.cpp@4524290e87b8e107cc2b56e1251751546f4b9051 +- fix: create_embedding broken response for input type str by @abetlen in 0ce66bc080fe537590b05b24bf442480bf2dd045 +- fix: Use '\n' seperator for EventSourceResponse by @khimaros in #1188 +- fix: Incorporate embedding pooling layer fixes by @iamlemec in #1194 + +## [0.2.43] + +- feat: Update llama.cpp to ggerganov/llama.cpp@8084d554406b767d36b3250b3b787462d5dd626f +- feat: Support batch embeddings by @iamlemec in #1186 +- fix: submodule kompute is not included in sdist by @abetlen in 7dbbfdecadebe7750be650d9409959640ff9a460 +- fix: fix: Update openbuddy prompt format by @abetlen in 07a783779a62a4aac0b11161c7e0eb983ff215f8 + +## [0.2.42] + +- feat: Update llama.cpp to ggerganov/llama.cpp@ea9c8e11436ad50719987fa23a289c74b7b40d40 +- fix: sample idx off-by-one error for logit_processors by @lapp0 in #1179 +- fix: chat formatting bugs in `chatml-function-calling` by @abetlen in 4b0e3320bd8c2c209e29978d0b21e2e471cc9ee3 and 68fb71b6a26a1e57331868f959b47ab4b87851e1 + +## [0.2.41] + +- feat: Update llama.cpp to ggerganov/llama.cpp@895407f31b358e3d9335e847d13f033491ec8a5b +- fix: Don't change order of json schema object properties in generated grammar unless prop_order is passed by @abetlen in d1822fed6b706f38bd1ff0de4dec5baaa3cf84fa + +## [0.2.40] + +- feat: Update llama.cpp to ggerganov/llama.cpp@3bdc4cd0f595a6096cca4a64aa75ffa8a3503465 +- feat: Generic chatml Function Calling using chat_format="chatml-function-calling"` by @abetlen in #957 +- fix: Circular dependancy preventing early Llama object free by @notwa in #1176 +- docs: Set the correct command for compiling with syscl support by @akarshanbiswas in #1172 +- feat: use gpu backend for clip if available by @iamlemec in #1175 + +## [0.2.39] + +- feat: Update llama.cpp to ggerganov/llama.cpp@b08f22c882a1443e6b97081f3ce718a4d1a741f8 +- fix: Fix destructor logging bugs by using llama_log_callback to avoid suppress_stdout_stderr by @abetlen in 59760c85eddc72dfcc1839f43760ef72c23d6874 + +## [0.2.38] + +- feat: Update llama.cpp to ggerganov/llama.cpp@1cfb5372cf5707c8ec6dde7c874f4a44a6c4c915 +- feat: Add speculative decoding by @abetlen in #1120 +- fix: Pass raise_exception and add_generation_prompt to jinja2 chat template by @abetlen in 078cca0361bf5a94d2cf52ed04980d20e32d6f95 + +## [0.2.37] + +- feat: Update llama.cpp to ggerganov/llama.cpp@fea4fd4ba7f6b754ac795387b275e1a014a77bde +- feat: Automatically set chat format from gguf by @abetlen in #1110 + +## [0.2.36] + +- feat: Update llama.cpp to ggerganov/llama.cpp@2aed77eb06a329f0d82bb1c467f4244904d4073f +- feat: Add mistral instruct chat format as "mistral-instruct" by @Rafaelblsilva in #799 + +## [0.2.35] + +- feat: Update llama.cpp to ggerganov/llama.cpp@d2f650cb5b04ee2726663e79b47da5efe196ce00 + +## [0.2.34] + +- feat: Update llama.cpp to ggerganov/llama.cpp@6db2b41a76ee78d5efdd5c3cddd5d7ad3f646855 +- feat: Add json schema mode by @abetlen in #1122 + +## [0.2.33] + +- feat: Update llama.cpp to ggerganov/llama.cpp@faa3526a1eba458120987ed8269e5616385a76f4 +- feat(server): include llama-cpp-python version in openapi spec by @abetlen in cde7514c3d28e6d52f272614e9957208c344dde5 +- fix: use both eos and bos tokens as stop sequences for hf-tokenizer-config chat format. by @abetlen in 5b982d0f8c6f35242c8862ffdce00e17cea0b44f +- fix: GGUF metadata KV overrides, re #1011 by @phiharri in #1116 +- fix: llama_log_set should be able to accept null pointer by @abetlen in c970d41a85381fd55235136f123422df0bf0c7e7 + +## [0.2.32] + +- feat: Update llama.cpp to ggerganov/llama.cpp@504dc37be8446fb09b1ede70300250ad41be32a2 +- fix: from_json_schema oneof/anyof bug by @jndiogo in d3f5528ca8bcb9d69d4f27e21631e911f1fb9bfe +- fix: pass chat handler not chat formatter for huggingface autotokenizer and tokenizer_config formats by @abetlen in 24f39454e91cf5dddbc4b6041aead4accc7c7a2d +- feat: Add add_generation_prompt option for jinja2chatformatter by @abetlen in 7f3209b1eb4ad3260ba063801fab80a8c25a2f4c +- feat: Add Jinja2ChatFormatter by @abetlen in be09318c26add8674ce494ae7cc480cce72a4146 +- feat: Expose gguf model metadata in metadata property by @abetlen in 5a34c57e5479e50c99aba9b38218cc48e6560b81 + +## [0.2.31] + +- feat: Update llama.cpp to ggerganov/llama.cpp@a5cacb22b2114fd9adf61c00cbb237384d86bced +- fix: Mirostat sampling now passes correct type to ctypes and tracks state during generation by @abetlen in 3babe3512cb95743108f2b595210c38ed6f1b904 +- fix: Python3.8 support in server by @abetlen in 141293a75b564a8699e0acba1da24d9aa1cf0ab1 + +## [0.2.30] + +- feat: Update llama.cpp to ggerganov/llama.cpp@57e2a7a52a819883f40dada8a2edc24ecf48186b +- feat(server): Add ability to load chat format from huggingface autotokenizer or tokenizer_config.json files by @abetlen in b8fc1c7d83ad4a9207c707ba1d954fe580286a01 +- feat: Integration of Jinja2 Templating for chat formats by @teleprint-me in #875 +- fix: Offload KQV by default by @abetlen in 48c3b77e6f558a9899de0e1155c7dc0c7958d8e8 +- fix: Support Accept text/event-stream in chat and completion endpoints, resolves #1083 by @aniljava in #1088 +- fix(cli): allow passing n_ctx=0 to openAI API server args to use model n_ctx_train field per #1015 by @K-Mistele in #1093 + +## [0.2.29] + +- feat: Update llama.cpp to ggerganov/llama.cpp@4483396751c79dea540808b9cb9238245d06da2b +- feat: Add split_mode option by @abetlen in 84615adbc6855c8384807c42f0130f9a1763f99d +- feat: Implement GGUF metadata KV overrides by @phiharri in #1011 +- fix: Avoid "LookupError: unknown encoding: ascii" when open() called in a destructor by @yieldthought in #1012 +- fix: Fix low_level_api_chat_cpp example to match current API by @aniljava in #1086 +- fix: Fix Pydantic model parsing by @DeNeutoy in #1087 + +## [0.2.28] + +- feat: Update llama.cpp to ggerganov/llama.cpp@6efb8eb30e7025b168f3fda3ff83b9b386428ad6 +- feat: Add ability to pass in penalize_nl param by @shankinson in #1068 +- fix: print_grammar to stderr by @turian in #1052 + +## [0.2.27] + +- feat: Update llama.cpp to ggerganov/llama.cpp@b3a7c20b5c035250257d2b62851c379b159c899a +- feat: Add `saiga` chat format by @femoiseev in #1050 +- feat: Added `chatglm3` chat format by @xaviviro in #1059 +- fix: Correct typo in README.md by @qeleb in (#1058) + +## [0.2.26] + +- feat: Update llama.cpp to ggerganov/llama.cpp@f6793491b5af6da75edad34d6f503ef86d31b09f + +## [0.2.25] + +- feat(server): Multi model support by @D4ve-R in #931 +- feat(server): Support none defaulting to infinity for completions by @swg in #111 +- feat(server): Implement openai api compatible authentication by @docmeth2 in #1010 +- fix: text_offset of multi-token characters by @twaka in #1037 +- fix: ctypes bindings for kv override by @phiharri in #1011 +- fix: ctypes definitions of llama_kv_cache_view_update and llama_kv_cache_view_free. by @e-c-d in #1028 + +## [0.2.24] + +- feat: Update llama.cpp to ggerganov/llama.cpp@0e18b2e7d0b5c0a509ea40098def234b8d4a938a +- feat: Add offload_kqv option to llama and server by @abetlen in 095c65000642a3cf73055d7428232fb18b73c6f3 +- feat: n_ctx=0 now uses the n_ctx_train of the model by @DanieleMorotti in #1015 +- feat: logits_to_logprobs supports both 2-D and 3-D logits arrays by @kddubey in #1002 +- fix: Remove f16_kv, add offload_kqv fields in low level and llama apis by @brandonrobertz in #1019 +- perf: Don't convert logprobs arrays to lists by @kddubey in #1021 +- docs: Fix README.md functionary demo typo by @evelynmitchell in #996 +- examples: Update low_level_api_llama_cpp.py to match current API by @jsoma in #1023 + +## [0.2.23] + +- Update llama.cpp to ggerganov/llama.cpp@948ff137ec37f1ec74c02905917fa0afc9b97514 +- Add qwen chat format by @yhfgyyf in #1005 +- Add support for running the server with SSL by @rgerganov in #994 +- Replace logits_to_logprobs implementation with numpy equivalent to llama.cpp by @player1537 in #991 +- Fix UnsupportedOperation: fileno in suppress_stdout_stderr by @zocainViken in #961 +- Add Pygmalion chat format by @chiensen in #986 +- README.md multimodal params fix by @zocainViken in #967 +- Fix minor typo in README by @aniketmaurya in #958 + +## [0.2.22] + +- Update llama.cpp to ggerganov/llama.cpp@8a7b2fa528f130631a5f43648481596ab320ed5a +- Fix conflict with transformers library by kddubey in #952 + +## [0.2.21] + +- Update llama.cpp to ggerganov/llama.cpp@64e64aa2557d97490b2fe1262b313e2f4a1607e3 +- Make building llava optional by setting `CMAKE_ARGS="-DLLAVA_BUILD=OFF"` and using `LLAVA_CPP_LIB` to specify alternative path to shared library by @abetlen in e3941d9c674dbd9891dc3ceda390daeb21f05fd1 + +## [0.2.20] + +- Update llama.cpp to ggerganov/llama.cpp@b38a16dfcff88d547f78f52d1bea31b84a05aff7 +- Add `zephyr` chat format by @fakerybakery in #938 +- Add `baichuan` chat format by @caiyesd in #938 +- Add `baichuan-2` chat format by @caiyesd in #936 +- Improve documentation for server chat formats by @jooray in #934 +- Fix typo in README by @antonvice in 940 +- Fix typo in the Open Orca chat format by @gardner in #947 + +## [0.2.19] + +- Update llama.cpp to ggerganov/llama.cpp@0b871f1a04ef60e114bbe43004fd9c21114e802d +- Fix #569: stop parameter in chat completion api should accept str by @abetlen in 128dc4731fa846ead7e684a137ca57d8931b8899 +- Document server host and port parameters by @jamesbraza in #768 +- Do not set grammar to None when initializing LlamaGrammar by @mthuurne in #834 +- Add mistrallite, intel, and openchat formats by @fakerybakery in #927 +- Add support for min_p parameter by @tk-master in #921 +- Fix #929: tokenizer adding leading space when generating from empty prompt by @abetlen in a34d48014192771d2e308a76c22f33bc0318d983 +- Fix low level api example by @zocainViken in #925 +- Fix missing package in openblas docker image by @ZisisTsatsas in #920 + +## [0.2.18] + +- Update llama.cpp to ggerganov/llama.cpp@6bb4908a17150b49373b5f977685b2e180a04f6f + +## [0.2.17] + +- Update llama.cpp to ggerganov/llama.cpp@df9d1293defe783f42bc83af732d3c670552c541 +- Hotfix: Set `CUDA_ARCHITECTURES=OFF` for `llava_shared` target on Windows by @abetlen in 4388f3341413110217b98c4f097ac5c590bdf40b + +## [0.2.16] + +- Update llama.cpp to ggerganov/llama.cp@a75fa576abba9d37f463580c379e4bbf1e1ad03c +- Add `set_seed` to `Llama` class by @abetlen in fd41ed3a908761d286102a019a34c2938a15118d +- Fix server doc arguments by @kjunggithub in #892 +- Fix response_format handler in llava chat handler by @abetlen in b62c44983921197ed10a7d29dc4ba920e9979380 +- Fix default max_tokens, chat completion is now unlimited (to context length) and completion is 16 tokens to match OpenAI defaults by @abetlen in e7962d2c733cbbeec5a37392c81f64185a9a39e8 +- Fix json_schema_to_gbnf helper so that it takes a json schema string as input instead by @abetlen in faeae181b1e868643c0dc28fcf039f077baf0829 +- Add support for $ref and $def in json_schema_to_gbnf to handle more complex function schemas by @abetlen in 770df344369c0630df1be14be9f9e301e7c56d24 +- Update functionary chat handler for new OpenAI api by abetlen in 1b376c62b775b401653facf25a519d116aafe99a +- Fix add default stop sequence to chatml chat format by @abetlen in b84d76a844149216d511cfd8cdb9827148a1853c +- Fix sampling bug when logits_all=False by @abetlen in 6f0b0b1b840af846938ed74d0e8170a91c40e617 + +## [0.2.15] + +- Update llama.cpp to ggerganov/llama.cpp@0a7c980b6f94a049cb804573df2d8092a34df8e4 +- Add support for Llava1.5 multimodal models by @damian0815 and @abetlen in #821 +- Update OpenAI API compatibility to match dev day update by @abetlen in #821 +- Add seed parameter to completion and chat_completion functions of Llama class by @abetlen in 86aeb9f3a14808575d2bb0076e6acb4a30907e6a +- Add JSON mode support to constrain chat completion to JSON objects by @abetlen in b30b9c338bf9af316d497ea501d39f5c246900db + +## [0.2.14] + +- Update llama.cpp to ggerganov/llama.cpp@f0b30ef7dc1360922ccbea0a8cd3918ecf15eaa7 +- Add support for Huggingface Autotokenizer Chat Formats by @bioshazard and @abetlen in #790 and bbffdaebaa7bb04b543dbf683a07276087251f86 +- Fix llama-2 chat format by @earonesty in #869 +- Add support for functionary chat format by @abetlen in #784 +- Migrate inference from deprecated `llama_eval`API to `llama_batch` and `llama_decode` by @abetlen in #795 + +## [0.2.13] + +- Update llama.cpp to ggerganov/llama.cpp@51b2fc11f7f605fff49725a4540e9a6ef7b51b70 +- Fix name 'open' is not defined exception when deleting model by @abetlen in 011b95d7f34cbfc528af75a892757bd9a20838ab +- Fix tokenization of special characters by @antoine-lizee in #850 + +## [0.2.12] + +- Update llama.cpp to ggerganov/llama.cpp@50337961a678fce4081554b24e56e86b67660163 +- Fix missing `n_seq_id` in `llama_batch` by @NickAlgra in #842 +- Fix for shared libraries on Windows that start with `lib` prefix by @sujeendran in #848 +- Fix exception raised in `__del__` when freeing models by @cebtenzzre in #846 +- Performance improvement for logit bias by @zolastro in #851 +- Fix suffix check arbitrary code execution bug by @mtasic85 in #854 +- Fix typo in `function_call` parameter in `llama_types.py` by @akatora28 in #849 +- Fix streaming not returning `finish_reason` by @gmcgoldr in #798 +- Fix `n_gpu_layers` check to allow values less than 1 for server by @hxy9243 in #826 +- Supppress stdout and stderr when freeing model by @paschembri in #803 +- Fix `llama2` chat format by @delock in #808 +- Add validation for tensor_split size by @eric1932 #820 +- Print stack trace on server error by @abetlen in d6a130a052db3a50975a719088a9226abfebb266 +- Update docs for gguf by @johnccshen in #783 +- Add `chatml` chat format by @abetlen in 305482bd4156c70802fc054044119054806f4126 + +## [0.2.11] + +- Fix bug in `llama_model_params` object has no attribute `logits_all` by @abetlen in d696251fbe40015e8616ea7a7d7ad5257fd1b896 + +## [0.2.10] + +- Fix bug 'llama_model_params' object has no attribute 'embedding' by @abetlen in 42bb721d64d744242f9f980f2b89d5a6e335b5e4 + +## [0.2.9] + +- Fix critical bug in pip installation of v0.2.8 due to `.git` directory in ac853e01e1a217a578080a4e1b851d2d08450adf + +## [0.2.8] + +- Update llama.cpp to ggerganov/llama.cpp@40e07a60f9ce06e79f3ccd4c903eba300fb31b5e +- Add configurable chat formats by @abetlen in #711 +- Fix rope scaling bug by @Josh-XT in #767 +- Fix missing numa parameter in server by @abetlen in d9bce17794d0dd6f7962d10aad768fedecf3ab89 + +## [0.2.7] + +- Update llama.cpp to ggerganov/llama.cpp@a98b1633d5a94d0aa84c7c16e1f8df5ac21fc850 +- Install required runtime dlls to package directory on windows by @abetlen in 8d75016549e2ff62a511b1119d966ffc0df5c77b +- Add openai-processing-ms to server response header by @Tradunsky in #748 +- Bump minimum version of scikit-build-core to 0.5.1 to fix msvc cmake issue by @abetlen in 1ed0f3ebe16993a0f961155aa4b2c85f1c68f668 +- Update `llama_types.py` to better match the openai api, old names are aliased to new ones by @abetlen in dbca136feaaf7f8b1182c4c3c90c32918b1d0bb3 + +## [0.2.6] + +- Update llama.cpp to 80291a1d02a07f7f66666fb576c5b1e75aa48b46 + +## [0.2.5] + +- Fix docker images missing starlette-context dependency by @abetlen in 22917989003c5e67623d54ab45affa1e0e475410 +- Fix loading dll in Windows Isolation Containers by @abetlen in 847466562573191efa655753d9252f308c4fbdb0 +- Fix build issue on m1 macs by @abetlen in dbd3a6d1ed8416a8fd800127251e730153afa305 +- Update docs to gguf and add hw acceleration docs for server by @jasonacox in #688 + +## [0.2.4] + +- Add NUMA support. **NOTE** low level api users must call llama_backend_init at the start of their programs by abetlen in f4090a0bb2a2a25acfe28d31c82cc1aa273bedee +- Fix tensor_split server cli argument by @abetlen in c4c440ba2dc86d9de728a751311fdd1c8e3756fa +- Made all `Llama` init parameters into keyword-only parameters by @abetlen in c8f9b8a734b5b040379bbd93995ba177affab1fe +- Added server params for `low_vram`, `main_gpu`, `lora_base`, and `lora_path` by @abetlen in 2920c4bf7ee1412d6bba7846e0e1b7ef6d34043b +- Removed server params for `rms_norm_eps` and `n_gqa` by @abetlen in 2920c4bf7ee1412d6bba7846e0e1b7ef6d34043b +- Fix boolean cli options by @abetlen in c999325e8e4507f6c6249dd2fb8de7f8bf57f71e and 0449d29b9f940e437231a07b9d56550226558bac +- Silence Pydantic Settings warnings about `model_alias` setting by @earonesty in #705 + +## [0.2.3] + +- Update llama.cpp to ggerganov/llama.cpp@71ca2fad7d6c0ef95ef9944fb3a1a843e481f314 +- Add X-Request-ID request header for mirroring custom IDs by @devrimcavusoglu in #703 +- Add pyproject extra for scikit-build-core to ensure compatible pathspec version by @abetlen in 6cfc54284b99ef1bff8193e2d5e483dbd89ada02 +- Fix issue with Literal and Optional cli arguments not working by @abetlen in #702 + +## [0.2.2] + +- Fix bug in pip install of v0.2.1 due to scikit-build-core removing all `.metal` files in the source distribution (see #701) + +## [0.2.1] + +- Fix bug in pip install of v0.2.0 due to .git folder being included in the source distribution (see #701) + +## [0.2.0] + +- Migrated to scikit-build-core build system by @abetlen in #499 +- Use `numpy` views for `LogitsProcessor` and `StoppingCriteria` instead of python lists by @abetlen in #499 +- Drop support for end-of-life Python3.7 by @abetlen in #499 +- Convert low level `llama.cpp` constants to use basic python types instead of `ctypes` types by @abetlen in #499 + +## [0.1.85] + +- Add `llama_cpp.__version__` attribute by @janvdp in #684 +- Fix low level api examples by @jbochi in #680 + +## [0.1.84] + +- Update llama.cpp + +## [0.1.83] + +- Update llama.cpp + +## [0.1.82] + +- Update llama.cpp + +## [0.1.81] + +- Update llama.cpp + +## [0.1.80] + +- Update llama.cpp + +## [0.1.79] + +- GGUF Support (breaking change requiring new model format) + +## [0.1.78] + +- Grammar based sampling via LlamaGrammar which can be passed to completions +- Make n_gpu_layers == -1 offload all layers + +## [0.1.77] + +- (llama.cpp) Update llama.cpp add support for LLaMa 2 70B +- (server) Add temporary n_gqa and rms_norm_eps parameters required for LLaMa 2 70B + +## [0.1.76] + +- (llama.cpp) Update llama.cpp add support for LLaMa 2 70B + +## [0.1.75] + +- Update llama.cpp + +## [0.1.74] + +- (server) OpenAI style error responses + +## [0.1.73] + +- (server) Add rope parameters to server settings + +## [0.1.72] + +- (llama.cpp) Update llama.cpp added custom_rope for extended context lengths + +## [0.1.71] + +- (llama.cpp) Update llama.cpp + +- (server) Fix several pydantic v2 migration bugs + +## [0.1.70] - (Llama.create_completion) Revert change so that `max_tokens` is not truncated to `context_size` in `create_completion` - (server) Fixed changed settings field names from pydantic v2 migration ## [0.1.69] -### Added - - (server) Streaming requests can are now interrupted pre-maturely when a concurrent request is made. Can be controlled with the `interrupt_requests` setting. - (server) Moved to fastapi v0.100.0 and pydantic v2 - (docker) Added a new "simple" image that builds llama.cpp from source when started. - -## Fixed - - (server) performance improvements by avoiding unnecessary memory allocations during sampling ## [0.1.68] -### Added - - (llama.cpp) Update llama.cpp ## [0.1.67] -### Fixed - - Fix performance bug in Llama model by pre-allocating memory tokens and logits. - Fix bug in Llama model where the model was not free'd after use. ## [0.1.66] -### Added - - (llama.cpp) New model API -### Fixed - - Performance issue during eval caused by looped np.concatenate call - State pickling issue when saving cache to disk ## [0.1.65] -### Added - - (llama.cpp) Fix struct misalignment bug ## [0.1.64] -### Added - - (llama.cpp) Update llama.cpp - Fix docs for seed. Set -1 for random. ## [0.1.63] -### Added - - (llama.cpp) Add full gpu utilisation in CUDA - (llama.cpp) Add get_vocab - (llama.cpp) Add low_vram parameter @@ -74,59 +933,37 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [0.1.62] -### Fixed - - Metal support working - Cache re-enabled ## [0.1.61] -### Fixed - - Fix broken pip installation ## [0.1.60] -### NOTE - -- This release was deleted due to a bug with the packaging system that caused pip installations to fail. - -### Fixed +NOTE: This release was deleted due to a bug with the packaging system that caused pip installations to fail. - Truncate max_tokens in create_completion so requested tokens doesn't exceed context size. - Temporarily disable cache for completion requests ## [v0.1.59] -### Added - - (llama.cpp) k-quants support - (server) mirostat sampling parameters to server - -### Fixed - - Support both `.so` and `.dylib` for `libllama` on MacOS ## [v0.1.58] -### Added - - (llama.cpp) Metal Silicon support ## [v0.1.57] -### Added - - (llama.cpp) OpenLlama 3B support ## [v0.1.56] -### Added - - (misc) Added first version of the changelog - (server) Use async routes - (python-api) Use numpy for internal buffers to reduce memory usage and improve performance. - -### Fixed - -- (python-api) Performance bug in stop sequence check slowing down streaming. \ No newline at end of file +- (python-api) Performance bug in stop sequence check slowing down streaming. diff --git a/CMakeLists.txt b/CMakeLists.txt index 788402a562..9b2744cdce 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,34 +1,207 @@ -cmake_minimum_required(VERSION 3.4...3.22) +cmake_minimum_required(VERSION 3.21) project(llama_cpp) -option(FORCE_CMAKE "Force CMake build of Python bindings" OFF) +option(LLAMA_BUILD "Build llama.cpp shared library and install alongside python package" ON) +option(LLAVA_BUILD "Build llava shared library and install alongside python package" ON) -set(FORCE_CMAKE $ENV{FORCE_CMAKE}) +function(llama_cpp_python_install_target target) + if(NOT TARGET ${target}) + return() + endif() -if (UNIX AND NOT FORCE_CMAKE) - add_custom_command( - OUTPUT ${CMAKE_CURRENT_SOURCE_DIR}/vendor/llama.cpp/libllama.so - COMMAND make libllama.so - WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/vendor/llama.cpp - ) - add_custom_target( - run ALL - DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/vendor/llama.cpp/libllama.so + install( + TARGETS ${target} + LIBRARY DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/llama_cpp/lib + RUNTIME DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/llama_cpp/lib + ARCHIVE DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/llama_cpp/lib + FRAMEWORK DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/llama_cpp/lib + RESOURCE DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/llama_cpp/lib ) install( - FILES ${CMAKE_CURRENT_SOURCE_DIR}/vendor/llama.cpp/libllama.so - DESTINATION llama_cpp + TARGETS ${target} + LIBRARY DESTINATION ${SKBUILD_PLATLIB_DIR}/llama_cpp/lib + RUNTIME DESTINATION ${SKBUILD_PLATLIB_DIR}/llama_cpp/lib + ARCHIVE DESTINATION ${SKBUILD_PLATLIB_DIR}/llama_cpp/lib + FRAMEWORK DESTINATION ${SKBUILD_PLATLIB_DIR}/llama_cpp/lib + RESOURCE DESTINATION ${SKBUILD_PLATLIB_DIR}/llama_cpp/lib ) -else() + set_target_properties(${target} PROPERTIES + INSTALL_RPATH "$ORIGIN" + BUILD_WITH_INSTALL_RPATH TRUE + ) + if(UNIX) + if(APPLE) + set_target_properties(${target} PROPERTIES + INSTALL_RPATH "@loader_path" + BUILD_WITH_INSTALL_RPATH TRUE + ) + else() + set_target_properties(${target} PROPERTIES + INSTALL_RPATH "$ORIGIN" + BUILD_WITH_INSTALL_RPATH TRUE + ) + endif() + endif() +endfunction() + +if (LLAMA_BUILD) set(BUILD_SHARED_LIBS "On") + + set(CMAKE_SKIP_BUILD_RPATH FALSE) + + # When building, don't use the install RPATH already + # (but later on when installing) + set(CMAKE_BUILD_WITH_INSTALL_RPATH FALSE) + + # Add the automatically determined parts of the RPATH + # which point to directories outside the build tree to the install RPATH + set(CMAKE_INSTALL_RPATH_USE_LINK_PATH TRUE) + set(CMAKE_SKIP_RPATH FALSE) + + # Enable building of the common library + set(LLAMA_BUILD_COMMON ON CACHE BOOL "Build llama.cpp common library" FORCE) + + # Disable building curl support + set(LLAMA_CURL OFF CACHE BOOL "llama.cpp: enable curl" FORCE) + + # Architecture detection and settings for Apple platforms + if (APPLE) + # Get the target architecture + execute_process( + COMMAND uname -m + OUTPUT_VARIABLE HOST_ARCH + OUTPUT_STRIP_TRAILING_WHITESPACE + ) + + # If CMAKE_OSX_ARCHITECTURES is not set, use the host architecture + if(NOT CMAKE_OSX_ARCHITECTURES) + set(CMAKE_OSX_ARCHITECTURES ${HOST_ARCH} CACHE STRING "Build architecture for macOS" FORCE) + endif() + + message(STATUS "Host architecture: ${HOST_ARCH}") + message(STATUS "Target architecture: ${CMAKE_OSX_ARCHITECTURES}") + + # Configure based on target architecture + if(CMAKE_OSX_ARCHITECTURES STREQUAL "x86_64") + # Intel Mac settings + set(GGML_AVX "OFF" CACHE BOOL "ggml: enable AVX" FORCE) + set(GGML_AVX2 "OFF" CACHE BOOL "ggml: enable AVX2" FORCE) + set(GGML_FMA "OFF" CACHE BOOL "ggml: enable FMA" FORCE) + set(GGML_F16C "OFF" CACHE BOOL "ggml: enable F16C" FORCE) + endif() + + # Metal settings (enable for both architectures) + set(GGML_METAL "ON" CACHE BOOL "ggml: enable Metal" FORCE) + set(GGML_METAL_EMBED_LIBRARY "ON" CACHE BOOL "ggml: embed metal library" FORCE) + endif() + + add_subdirectory(vendor/llama.cpp) - install( - TARGETS llama - LIBRARY DESTINATION llama_cpp - RUNTIME DESTINATION llama_cpp - ARCHIVE DESTINATION llama_cpp - FRAMEWORK DESTINATION llama_cpp - RESOURCE DESTINATION llama_cpp - ) + + if (WIN32) + if (TARGET llama) + set_target_properties(llama PROPERTIES WINDOWS_EXPORT_ALL_SYMBOLS ON) + endif() + endif() + + llama_cpp_python_install_target(llama) + llama_cpp_python_install_target(ggml) + + llama_cpp_python_install_target(ggml-base) + + llama_cpp_python_install_target(ggml-amx) + llama_cpp_python_install_target(ggml-blas) + llama_cpp_python_install_target(ggml-can) + llama_cpp_python_install_target(ggml-cpu) + llama_cpp_python_install_target(ggml-cuda) + llama_cpp_python_install_target(ggml-hip) + llama_cpp_python_install_target(ggml-kompute) + llama_cpp_python_install_target(ggml-metal) + llama_cpp_python_install_target(ggml-musa) + llama_cpp_python_install_target(ggml-rpc) + llama_cpp_python_install_target(ggml-sycl) + llama_cpp_python_install_target(ggml-vulkan) + + # Workaround for Windows + CUDA https://github.com/abetlen/llama-cpp-python/issues/563 + if (WIN32) + install( + FILES $ + DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/llama_cpp/lib + ) + install( + FILES $ + DESTINATION ${SKBUILD_PLATLIB_DIR}/llama_cpp/lib + ) + install( + FILES $ + DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/llama_cpp/lib + ) + install( + FILES $ + DESTINATION ${SKBUILD_PLATLIB_DIR}/llama_cpp/lib + ) + endif() + + if (LLAVA_BUILD) + if (LLAMA_CUBLAS OR LLAMA_CUDA) + add_compile_definitions(GGML_USE_CUBLAS) + add_compile_definitions(GGML_USE_CUDA) + endif() + + if (LLAMA_METAL) + add_compile_definitions(GGML_USE_METAL) + endif() + + # Upstream mtmd expects LLAMA_INSTALL_VERSION to be set by llama.cpp's + # top-level CMakeLists.txt. When we include tools/mtmd directly from the + # Python package build, that directory scope is skipped. + if (NOT DEFINED LLAMA_INSTALL_VERSION OR "${LLAMA_INSTALL_VERSION}" STREQUAL "") + set(LLAMA_INSTALL_VERSION 0.0.0) + find_package(Git QUIET) + if (Git_FOUND) + execute_process( + COMMAND ${GIT_EXECUTABLE} rev-list --count HEAD + WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/vendor/llama.cpp + OUTPUT_VARIABLE LLAMA_MTMD_BUILD_NUMBER + OUTPUT_STRIP_TRAILING_WHITESPACE + RESULT_VARIABLE LLAMA_MTMD_BUILD_NUMBER_RESULT + ) + if (LLAMA_MTMD_BUILD_NUMBER_RESULT EQUAL 0) + set(LLAMA_INSTALL_VERSION 0.0.${LLAMA_MTMD_BUILD_NUMBER}) + endif() + endif() + endif() + + # Building llava + add_subdirectory(vendor/llama.cpp/tools/mtmd) + + if (WIN32) + set_target_properties(mtmd PROPERTIES CUDA_ARCHITECTURES OFF) + endif() + llama_cpp_python_install_target(mtmd) + if (WIN32) + install( + FILES $ + DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/llama_cpp/lib + ) + install( + FILES $ + DESTINATION ${SKBUILD_PLATLIB_DIR}/llama_cpp/lib + ) + endif() + + # Fix for mtmd build: Add include directory for llama.h + # Move these commands after the add_subdirectory call + target_include_directories(mtmd PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/vendor/llama.cpp/include) + target_include_directories(mtmd PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/vendor/llama.cpp/ggml/include) + + if (BUILD_SHARED_LIBS) + target_include_directories(mtmd PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/vendor/llama.cpp/include) + target_include_directories(mtmd PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/vendor/llama.cpp/ggml/include) + endif() + + # target_include_directories(llama-llava-cli PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/vendor/llama.cpp/include) + # target_include_directories(llama-minicpmv-cli PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/vendor/llama.cpp/include) + endif() endif() diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000000..1d8f9e398d --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,77 @@ +# Contributing + +Hello human and AI contributors, this document exists to help you understand the project and set some rules for contributions. + +## Contribution Workflow + +Before opening a pull request, search existing issues and PRs to avoid duplicate work. +Keep each PR focused on one feature, bug fix, or vendor update. +Avoid mixing unrelated Python changes, generated binding updates, documentation edits, and `vendor/llama.cpp` changes unless they are required for the same fix. + +Describe what changed, why it changed, and how it was tested. +Link relevant issues, include any required build flags or hardware assumptions, and add a `CHANGELOG.md` entry for user-visible fixes or features (see `CHANGELOG.md` for examples). + +BREAKING CHANGES WILL ALMOST CERTAINLY BE REJECTED OR REFACTORED. + +## PR Titles and Changelog Entries + +Use PR titles in the form `: `, with an optional scope when it adds clarity: `feat: add X`, `fix(server): handle Y`, `fix(ci): repair Z`, or `chore: bump version to N`. +Prefer tags already used in the project history, such as `feat`, `fix`, `chore`, `ci`, `docs`, and `refactor`. + +Add changelog entries under `## [Unreleased]` using the PR title followed by `by @contributor in #1234`. + +```md +- feat(server): add support for X by @contributor in #1234 +- fix(ci): repair Y wheel builds by @contributor in #1234 +``` + +## Local Development + +Prerequisites: Python 3.8+, CMake 3.21+, a C/C++ compiler, and Git submodules. +From a fresh checkout of the repository, initialize submodules and create a virtual environment: + +```bash +git submodule update --init --recursive +python3 -m venv .venv +source .venv/bin/activate +python -m pip install --upgrade pip +make deps +make build +``` + +Run tests and lint checks before submitting changes: + +```bash +make test +make lint +``` + +Use backend-specific build targets when validating native acceleration or backend-specific fixes, for example `make build.openblas`, `make build.cuda`, `make build.metal`, or `make build.vulkan`. + +## Testing Expectations + +Add or update tests for behavior changes or fixing regressions. +The test suite uses pytest and lives under `tests/`; name files `test_*.py` and test functions `test_*`. + +For changes involving native backends, model behavior, performance, or platform compatibility, document the environment used for validation in the PR. +If a change cannot be covered by automated tests, include a short manual validation recipe instead. + +## Code Style + +Python code is formatted with Ruff using an 88-character line length. +Run `make format` to apply automatic fixes and `make lint` to check formatting and lint rules. + +Use 4-space indentation, `snake_case` for functions and variables, `PascalCase` for classes, and `UPPER_CASE` for constants. +Follow existing patterns when touching ctypes bindings or server APIs, and avoid adding dependencies unless they are necessary for the feature or fix. + +## Documentation Style + +Write Markdown with one sentence or core idea per physical line to keep diffs focused and easier to review. +Do not manually wrap lines at a fixed column width. +Keep `README.md` focused on user-facing setup and usage; link to this guide for contribution workflow details rather than duplicating them. + +## Project Layout + +The Python package lives in `llama_cpp/`, with tests in `tests/` and examples in `examples/`. +Documentation lives in `docs/` and is built with `mkdocs.yml`. +The `vendor/llama.cpp/` directory is a Git submodule containing the upstream llama.cpp sources used by the bindings. diff --git a/Makefile b/Makefile index c359260b6c..db45246c79 100644 --- a/Makefile +++ b/Makefile @@ -5,26 +5,57 @@ update: update.vendor: cd vendor/llama.cpp && git pull origin master +deps: + python3 -m pip install --upgrade pip + python3 -m pip install -e ".[all]" + build: - python3 setup.py develop + python3 -m pip install --verbose -e . -build.cuda: - CMAKE_ARGS="-DLLAMA_CUBLAS=on" FORCE_CMAKE=1 python3 setup.py develop +build.debug: + python3 -m pip install \ + --verbose \ + --config-settings=cmake.verbose=true \ + --config-settings=logging.level=INFO \ + --config-settings=install.strip=false \ + --config-settings=cmake.args="-DCMAKE_BUILD_TYPE=Debug;-DCMAKE_C_FLAGS='-ggdb -O0';-DCMAKE_CXX_FLAGS='-ggdb -O0'" \ + --editable . -build.opencl: - CMAKE_ARGS="-DLLAMA_CLBLAST=on" FORCE_CMAKE=1 python3 setup.py develop +build.debug.extra: + python3 -m pip install \ + --verbose \ + --config-settings=cmake.verbose=true \ + --config-settings=logging.level=INFO \ + --config-settings=install.strip=false \ + --config-settings=cmake.args="-DCMAKE_BUILD_TYPE=Debug;-DCMAKE_C_FLAGS='-fsanitize=address -ggdb -O0';-DCMAKE_CXX_FLAGS='-fsanitize=address -ggdb -O0'" \ + --editable . + +build.cuda: + CMAKE_ARGS="-DGGML_CUDA=on" python3 -m pip install --verbose -e . build.openblas: - CMAKE_ARGS="-DLLAMA_OPENBLAS=on" FORCE_CMAKE=1 python3 setup.py develop + CMAKE_ARGS="-DGGML_BLAS=ON -DGGML_BLAS_VENDOR=OpenBLAS" python3 -m pip install --verbose -e . build.blis: - CMAKE_ARGS="-DLLAMA_OPENBLAS=on -DLLAMA_OPENBLAS_VENDOR=blis" FORCE_CMAKE=1 python3 setup.py develop + CMAKE_ARGS="-DGGML_BLAS=on -DGGML_BLAS_VENDOR=FLAME" python3 -m pip install --verbose -e . build.metal: - CMAKE_ARGS="-DLLAMA_METAL=on" FORCE_CMAKE=1 python3 setup.py develop + CMAKE_ARGS="-DGGML_METAL=on" python3 -m pip install --verbose -e . + +build.vulkan: + CMAKE_ARGS="-DGGML_VULKAN=on" python3 -m pip install --verbose -e . + +build.kompute: + CMAKE_ARGS="-DGGML_KOMPUTE=on" python3 -m pip install --verbose -e . + +build.sycl: + CMAKE_ARGS="-DGGML_SYCL=on" python3 -m pip install --verbose -e . + +build.rpc: + CMAKE_ARGS="-DGGML_RPC=on" python3 -m pip install --verbose -e . build.sdist: - python3 setup.py sdist + python3 -m build --sdist --verbose deploy.pypi: python3 -m twine upload dist/* @@ -34,23 +65,29 @@ deploy.gh-docs: mkdocs gh-deploy test: - python3 -m pytest + python3 -m pytest --full-trace -v + +lint: + python3 -m ruff check llama_cpp tests + python3 -m ruff format --check llama_cpp tests + +format: + python3 -m ruff check --fix llama_cpp tests + python3 -m ruff format llama_cpp tests docker: docker build -t llama-cpp-python:latest -f docker/simple/Dockerfile . run-server: - uvicorn --factory llama.server:app --host ${HOST} --port ${PORT} + python3 -m llama_cpp.server --model ${MODEL} clean: - - cd vendor/llama.cpp && make clean - - cd vendor/llama.cpp && rm libllama.so - rm -rf _skbuild - - rm llama_cpp/*.so - - rm llama_cpp/*.dylib - - rm llama_cpp/*.metal - - rm llama_cpp/*.dll - - rm llama_cpp/*.lib + - rm llama_cpp/lib/*.so + - rm llama_cpp/lib/*.dylib + - rm llama_cpp/lib/*.metal + - rm llama_cpp/lib/*.dll + - rm llama_cpp/lib/*.lib .PHONY: \ update \ @@ -62,5 +99,7 @@ clean: build.sdist \ deploy.pypi \ deploy.gh-docs \ + lint \ + format \ docker \ - clean \ No newline at end of file + clean diff --git a/README.md b/README.md index 0322c73a33..c02df9eed0 100644 --- a/README.md +++ b/README.md @@ -1,94 +1,331 @@ -# 🦙 Python Bindings for `llama.cpp` +<p align="center"> + <img src="https://raw.githubusercontent.com/abetlen/llama-cpp-python/main/docs/icon.svg" style="height: 5rem; width: 5rem"> +</p> + +# Python Bindings for [`llama.cpp`](https://github.com/ggerganov/llama.cpp) [![Documentation Status](https://readthedocs.org/projects/llama-cpp-python/badge/?version=latest)](https://llama-cpp-python.readthedocs.io/en/latest/?badge=latest) [![Tests](https://github.com/abetlen/llama-cpp-python/actions/workflows/test.yaml/badge.svg?branch=main)](https://github.com/abetlen/llama-cpp-python/actions/workflows/test.yaml) [![PyPI](https://img.shields.io/pypi/v/llama-cpp-python)](https://pypi.org/project/llama-cpp-python/) [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/llama-cpp-python)](https://pypi.org/project/llama-cpp-python/) [![PyPI - License](https://img.shields.io/pypi/l/llama-cpp-python)](https://pypi.org/project/llama-cpp-python/) -[![PyPI - Downloads](https://img.shields.io/pypi/dm/llama-cpp-python)](https://pypi.org/project/llama-cpp-python/) +[![PyPI - Downloads](https://static.pepy.tech/badge/llama-cpp-python/month)](https://pepy.tech/projects/llama-cpp-python) +[![Github All Releases](https://img.shields.io/github/downloads/abetlen/llama-cpp-python/total.svg?label=Github%20Downloads)]() Simple Python bindings for **@ggerganov's** [`llama.cpp`](https://github.com/ggerganov/llama.cpp) library. This package provides: - Low-level access to C API via `ctypes` interface. - High-level Python API for text completion - - OpenAI-like API - - LangChain compatibility + - OpenAI-like API + - [LangChain compatibility](https://python.langchain.com/docs/integrations/llms/llamacpp) + - [LlamaIndex compatibility](https://docs.llamaindex.ai/en/stable/examples/llm/llama_2_llama_cpp.html) +- OpenAI compatible web server + - [Local Copilot replacement](https://llama-cpp-python.readthedocs.io/en/latest/server/#code-completion) + - [Function Calling support](https://llama-cpp-python.readthedocs.io/en/latest/server/#function-calling) + - [Vision API support](https://llama-cpp-python.readthedocs.io/en/latest/server/#multimodal-models) + - [Multiple Models](https://llama-cpp-python.readthedocs.io/en/latest/server/#configuration-and-multi-model-support) Documentation is available at [https://llama-cpp-python.readthedocs.io/en/latest](https://llama-cpp-python.readthedocs.io/en/latest). +## Installation + +Requirements: -## Installation from PyPI (recommended) + - Python 3.8+ + - C compiler + - Linux: gcc or clang + - Windows: Visual Studio or MinGW + - MacOS: Xcode -Install from PyPI (requires a c compiler): +To install the package, run: ```bash pip install llama-cpp-python ``` -The above command will attempt to install the package and build `llama.cpp` from source. -This is the recommended installation method as it ensures that `llama.cpp` is built with the available optimizations for your system. +This will also build `llama.cpp` from source and install it alongside this python package. + +If this fails, add `--verbose` to the `pip install` see the full cmake build log. + +**Pre-built Wheel (New)** -If you have previously installed `llama-cpp-python` through pip and want to upgrade your version or rebuild the package with different compiler options, please add the following flags to ensure that the package is rebuilt correctly: +It is also possible to install a pre-built wheel with basic CPU support. ```bash -pip install llama-cpp-python --force-reinstall --upgrade --no-cache-dir +pip install llama-cpp-python \ + --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cpu ``` -Note: If you are using Apple Silicon (M1) Mac, make sure you have installed a version of Python that supports arm64 architecture. For example: +### Installation Configuration + +`llama.cpp` supports a number of hardware acceleration backends to speed up inference as well as backend specific options. See the [llama.cpp README](https://github.com/ggml-org/llama.cpp/blob/master/docs/build.md) for a full list. + +All `llama.cpp` cmake build options can be set via the `CMAKE_ARGS` environment variable or via the `--config-settings / -C` cli flag during installation. + +<details open> +<summary>Environment Variables</summary> + +```bash +# Linux and Mac +CMAKE_ARGS="-DGGML_BLAS=ON -DGGML_BLAS_VENDOR=OpenBLAS" \ + pip install llama-cpp-python ``` -wget https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-MacOSX-arm64.sh -bash Miniforge3-MacOSX-arm64.sh + +```powershell +# Windows +$env:CMAKE_ARGS = "-DGGML_BLAS=ON -DGGML_BLAS_VENDOR=OpenBLAS" +pip install llama-cpp-python +``` +</details> + +<details> +<summary>CLI / requirements.txt</summary> + +They can also be set via `pip install -C / --config-settings` command and saved to a `requirements.txt` file: + +```bash +pip install --upgrade pip # ensure pip is up to date +pip install llama-cpp-python \ + -C cmake.args="-DGGML_BLAS=ON;-DGGML_BLAS_VENDOR=OpenBLAS" +``` + +```txt +# requirements.txt + +llama-cpp-python -C cmake.args="-DGGML_BLAS=ON;-DGGML_BLAS_VENDOR=OpenBLAS" +``` + +</details> + +### Supported Backends + +Below are some common backends, their build commands and any additional environment variables required. + +<details open> +<summary>OpenBLAS (CPU)</summary> + +To install with OpenBLAS, set the `GGML_BLAS` and `GGML_BLAS_VENDOR` environment variables before installing: + +```bash +CMAKE_ARGS="-DGGML_BLAS=ON -DGGML_BLAS_VENDOR=OpenBLAS" pip install llama-cpp-python +``` +</details> + +<details> +<summary>CUDA</summary> + +To install with CUDA support, set the `GGML_CUDA=on` environment variable before installing: + +```bash +CMAKE_ARGS="-DGGML_CUDA=on" pip install llama-cpp-python +``` + +**Pre-built Wheel (New)** + +It is also possible to install a pre-built wheel with CUDA support. As long as your system meets some requirements: + +- CUDA Version is 11.8, 12.1, 12.2, 12.3, 12.4, 12.5, 13.0 or 13.2 +- NVIDIA GPU compute capability is 6.0 through 8.9 for CUDA 11.8 wheels, 6.0 or newer for CUDA 12 wheels, or 7.5 or newer for CUDA 13 wheels +- Python Version is 3.10, 3.11 or 3.12 + +```bash +pip install llama-cpp-python \ + --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/<cuda-version> +``` + +Where `<cuda-version>` is one of the following: +- `cu118`: CUDA 11.8 +- `cu121`: CUDA 12.1 +- `cu122`: CUDA 12.2 +- `cu123`: CUDA 12.3 +- `cu124`: CUDA 12.4 +- `cu125`: CUDA 12.5 +- `cu130`: CUDA 13.0 +- `cu132`: CUDA 13.2 + +For example, to install the CUDA 12.1 wheel: + +```bash +pip install llama-cpp-python \ + --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cu121 +``` + +</details> + +<details> +<summary>Metal</summary> + +To install with Metal (MPS), set the `GGML_METAL=on` environment variable before installing: + +```bash +CMAKE_ARGS="-DGGML_METAL=on" pip install llama-cpp-python +``` + +**Pre-built Wheel (New)** + +It is also possible to install a pre-built wheel with Metal support. As long as your system meets some requirements: + +- MacOS Version is 11.0 or later +- Python Version is 3.10, 3.11 or 3.12 + +```bash +pip install llama-cpp-python \ + --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/metal +``` + +</details> + +<details> +<summary>HIP (ROCm)</summary> + +To install with HIP / ROCm support for AMD cards, set the `GGML_HIP=on` environment variable before installing: + +```bash +CMAKE_ARGS="-DGGML_HIP=on" pip install llama-cpp-python ``` -Otherwise, while installing it will build the llama.ccp x86 version which will be 10x slower on Apple Silicon (M1) Mac. -### Installation with OpenBLAS / cuBLAS / CLBlast / Metal +**Pre-built Wheel (New)** + +It is also possible to install a pre-built wheel with ROCm support for Linux: -`llama.cpp` supports multiple BLAS backends for faster processing. -Use the `FORCE_CMAKE=1` environment variable to force the use of `cmake` and install the pip package for the desired BLAS backend. +```bash +pip install llama-cpp-python \ + --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/rocm72 +``` -To install with OpenBLAS, set the `LLAMA_OPENBLAS=1` environment variable before installing: +Or a pre-built wheel with HIP Radeon support for Windows: + +```powershell +pip install llama-cpp-python ` + --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/hip-radeon +``` + +</details> + +<details> +<summary>Vulkan</summary> + +To install with Vulkan support, set the `GGML_VULKAN=on` environment variable before installing: ```bash -CMAKE_ARGS="-DLLAMA_OPENBLAS=on" FORCE_CMAKE=1 pip install llama-cpp-python +CMAKE_ARGS="-DGGML_VULKAN=on" pip install llama-cpp-python ``` -To install with cuBLAS, set the `LLAMA_CUBLAS=1` environment variable before installing: +**Pre-built Wheel (New)** + +It is also possible to install a pre-built wheel with Vulkan support for Linux or Windows: ```bash -CMAKE_ARGS="-DLLAMA_CUBLAS=on" FORCE_CMAKE=1 pip install llama-cpp-python +pip install llama-cpp-python \ + --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/vulkan ``` -To install with CLBlast, set the `LLAMA_CLBLAST=1` environment variable before installing: +</details> + +<details> +<summary>SYCL</summary> + +To install with SYCL support, set the `GGML_SYCL=on` environment variable before installing: ```bash -CMAKE_ARGS="-DLLAMA_CLBLAST=on" FORCE_CMAKE=1 pip install llama-cpp-python +source /opt/intel/oneapi/setvars.sh +CMAKE_ARGS="-DGGML_SYCL=on -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx" pip install llama-cpp-python ``` +</details> + +<details> +<summary>RPC</summary> -To install with Metal (MPS), set the `LLAMA_METAL=on` environment variable before installing: +To install with RPC support, set the `GGML_RPC=on` environment variable before installing: ```bash -CMAKE_ARGS="-DLLAMA_METAL=on" FORCE_CMAKE=1 pip install llama-cpp-python +source /opt/intel/oneapi/setvars.sh +CMAKE_ARGS="-DGGML_RPC=on" pip install llama-cpp-python +``` +</details> + + +### Windows Notes + +<details> +<summary>Error: Can't find 'nmake' or 'CMAKE_C_COMPILER'</summary> + +If you run into issues where it complains it can't find `'nmake'` `'?'` or CMAKE_C_COMPILER, you can extract w64devkit as [mentioned in llama.cpp repo](https://github.com/ggerganov/llama.cpp#openblas) and add those manually to CMAKE_ARGS before running `pip` install: + +```ps +$env:CMAKE_GENERATOR = "MinGW Makefiles" +$env:CMAKE_ARGS = "-DGGML_OPENBLAS=on -DCMAKE_C_COMPILER=C:/w64devkit/bin/gcc.exe -DCMAKE_CXX_COMPILER=C:/w64devkit/bin/g++.exe" ``` -Detailed MacOS Metal GPU install documentation is available at [docs/install/macos.md](docs/install/macos.md) +See the above instructions and set `CMAKE_ARGS` to the BLAS backend you want to use. +</details> + +### MacOS Notes + +Detailed MacOS Metal GPU install documentation is available at [docs/install/macos.md](https://llama-cpp-python.readthedocs.io/en/latest/install/macos/) + +<details> +<summary>M1 Mac Performance Issue</summary> + +Note: If you are using Apple Silicon (M1) Mac, make sure you have installed a version of Python that supports arm64 architecture. For example: + +```bash +wget https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-MacOSX-arm64.sh +bash Miniforge3-MacOSX-arm64.sh +``` + +Otherwise, while installing it will build the llama.cpp x86 version which will be 10x slower on Apple Silicon (M1) Mac. +</details> + +<details> +<summary>M Series Mac Error: `(mach-o file, but is an incompatible architecture (have 'x86_64', need 'arm64'))`</summary> + +Try installing with + +```bash +CMAKE_ARGS="-DCMAKE_OSX_ARCHITECTURES=arm64 -DCMAKE_APPLE_SILICON_PROCESSOR=arm64 -DGGML_METAL=on" pip install --upgrade --verbose --force-reinstall --no-cache-dir llama-cpp-python +``` +</details> + +### Upgrading and Reinstalling + +To upgrade and rebuild `llama-cpp-python` add `--upgrade --force-reinstall --no-cache-dir` flags to the `pip install` command to ensure the package is rebuilt from source. ## High-level API -The high-level API provides a simple managed interface through the `Llama` class. +[API Reference](https://llama-cpp-python.readthedocs.io/en/latest/api-reference/#high-level-api) + +The high-level API provides a simple managed interface through the [`Llama`](https://llama-cpp-python.readthedocs.io/en/latest/api-reference/#llama_cpp.Llama) class. -Below is a short example demonstrating how to use the high-level API to generate text: +Below is a short example demonstrating how to use the high-level API for basic text completion: + +```python +from llama_cpp import Llama + +llm = Llama( + model_path="./models/7B/llama-model.gguf", + # n_gpu_layers=-1, # Uncomment to use GPU acceleration + # seed=1337, # Uncomment to set a specific seed + # n_ctx=2048, # Uncomment to increase the context window +) +output = llm( + "Q: Name the planets in the solar system? A: ", # Prompt + max_tokens=32, # Generate up to 32 tokens, set to None to generate up to the end of the context window + stop=["Q:", "\n"], # Stop generating just before the model would generate a new question + echo=True # Echo the prompt back in the output +) # Generate a completion, can also call create_completion +print(output) +``` + +By default `llama-cpp-python` generates completions in an OpenAI compatible format: ```python ->>> from llama_cpp import Llama ->>> llm = Llama(model_path="./models/7B/ggml-model.bin") ->>> output = llm("Q: Name the planets in the solar system? A: ", max_tokens=32, stop=["Q:", "\n"], echo=True) ->>> print(output) { "id": "cmpl-xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx", "object": "text_completion", "created": 1679561337, - "model": "./models/7B/ggml-model.bin", + "model": "./models/7B/llama-model.gguf", "choices": [ { "text": "Q: Name the planets in the solar system? A: Mercury, Venus, Earth, Mars, Jupiter, Saturn, Uranus, Neptune and Pluto.", @@ -105,16 +342,349 @@ Below is a short example demonstrating how to use the high-level API to generate } ``` +Text completion is available through the [`__call__`](https://llama-cpp-python.readthedocs.io/en/latest/api-reference/#llama_cpp.Llama.__call__) and [`create_completion`](https://llama-cpp-python.readthedocs.io/en/latest/api-reference/#llama_cpp.Llama.create_completion) methods of the [`Llama`](https://llama-cpp-python.readthedocs.io/en/latest/api-reference/#llama_cpp.Llama) class. + +### Pulling models from Hugging Face Hub + +You can download `Llama` models in `gguf` format directly from Hugging Face using the [`from_pretrained`](https://llama-cpp-python.readthedocs.io/en/latest/api-reference/#llama_cpp.Llama.from_pretrained) method. +You'll need to install the `huggingface-hub` package to use this feature (`pip install huggingface-hub`). + +```python +llm = Llama.from_pretrained( + repo_id="lmstudio-community/Qwen3.5-0.8B-GGUF", + filename="*Q8_0.gguf", + verbose=False +) +``` + +By default [`from_pretrained`](https://llama-cpp-python.readthedocs.io/en/latest/api-reference/#llama_cpp.Llama.from_pretrained) will download the model to the huggingface cache directory, you can then manage installed model files with the [`hf`](https://huggingface.co/docs/huggingface_hub/en/guides/cli) tool. + +### Chat Completion + +The high-level API also provides a simple interface for chat completion. + +Chat completion requires that the model knows how to format the messages into a single prompt. +The `Llama` class does this using pre-registered chat formats (ie. `chatml`, `llama-2`, `gemma`, etc) or by providing a custom chat handler object. + +The model will format the messages into a single prompt using the following order of precedence: + - Use the `chat_handler` if provided + - Use the `chat_format` if provided + - Use the `tokenizer.chat_template` from the `gguf` model's metadata (should work for most new models, older models may not have this) + - else, fallback to the `llama-2` chat format + +Set `verbose=True` to see the selected chat format. + +```python +from llama_cpp import Llama +llm = Llama( + model_path="path/to/llama-2/llama-model.gguf", + chat_format="llama-2" +) +llm.create_chat_completion( + messages = [ + {"role": "system", "content": "You are an assistant who perfectly describes images."}, + { + "role": "user", + "content": "Describe this image in detail please." + } + ] +) +``` + +Chat completion is available through the [`create_chat_completion`](https://llama-cpp-python.readthedocs.io/en/latest/api-reference/#llama_cpp.Llama.create_chat_completion) method of the [`Llama`](https://llama-cpp-python.readthedocs.io/en/latest/api-reference/#llama_cpp.Llama) class. + +For OpenAI API v1 compatibility, you use the [`create_chat_completion_openai_v1`](https://llama-cpp-python.readthedocs.io/en/latest/api-reference/#llama_cpp.Llama.create_chat_completion_openai_v1) method which will return pydantic models instead of dicts. + + +### JSON and JSON Schema Mode + +To constrain chat responses to only valid JSON or a specific JSON Schema use the `response_format` argument in [`create_chat_completion`](https://llama-cpp-python.readthedocs.io/en/latest/api-reference/#llama_cpp.Llama.create_chat_completion). + +#### JSON Mode + +The following example will constrain the response to valid JSON strings only. + +```python +from llama_cpp import Llama +llm = Llama(model_path="path/to/model.gguf", chat_format="chatml") +llm.create_chat_completion( + messages=[ + { + "role": "system", + "content": "You are a helpful assistant that outputs in JSON.", + }, + {"role": "user", "content": "Who won the world series in 2020"}, + ], + response_format={ + "type": "json_object", + }, + temperature=0.7, +) +``` + +#### JSON Schema Mode + +To constrain the response further to a specific JSON Schema add the schema to the `schema` property of the `response_format` argument. + +```python +from llama_cpp import Llama +llm = Llama(model_path="path/to/model.gguf", chat_format="chatml") +llm.create_chat_completion( + messages=[ + { + "role": "system", + "content": "You are a helpful assistant that outputs in JSON.", + }, + {"role": "user", "content": "Who won the world series in 2020"}, + ], + response_format={ + "type": "json_object", + "schema": { + "type": "object", + "properties": {"team_name": {"type": "string"}}, + "required": ["team_name"], + }, + }, + temperature=0.7, +) +``` + +### Function Calling + +The high-level API supports OpenAI compatible function and tool calling. This is possible through the `functionary` pre-trained models chat format or through the generic `chatml-function-calling` chat format. + +```python +from llama_cpp import Llama +llm = Llama(model_path="path/to/chatml/llama-model.gguf", chat_format="chatml-function-calling") +llm.create_chat_completion( + messages = [ + { + "role": "system", + "content": "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. The assistant calls functions with appropriate input when necessary" + + }, + { + "role": "user", + "content": "Extract Jason is 25 years old" + } + ], + tools=[{ + "type": "function", + "function": { + "name": "UserDetail", + "parameters": { + "type": "object", + "title": "UserDetail", + "properties": { + "name": { + "title": "Name", + "type": "string" + }, + "age": { + "title": "Age", + "type": "integer" + } + }, + "required": [ "name", "age" ] + } + } + }], + tool_choice={ + "type": "function", + "function": { + "name": "UserDetail" + } + } +) +``` + +<details> +<summary>Functionary v2</summary> + +The various gguf-converted files for this set of models can be found [here](https://huggingface.co/meetkai). Functionary is able to intelligently call functions and also analyze any provided function outputs to generate coherent responses. All v2 models of functionary supports **parallel function calling**. You can provide either `functionary-v1` or `functionary-v2` for the `chat_format` when initializing the Llama class. + +Due to discrepancies between llama.cpp and HuggingFace's tokenizers, it is required to provide HF Tokenizer for functionary. The `LlamaHFTokenizer` class can be initialized and passed into the Llama class. This will override the default llama.cpp tokenizer used in Llama class. The tokenizer files are already included in the respective HF repositories hosting the gguf files. + +```python +from llama_cpp import Llama +from llama_cpp.llama_tokenizer import LlamaHFTokenizer +llm = Llama.from_pretrained( + repo_id="meetkai/functionary-small-v2.2-GGUF", + filename="functionary-small-v2.2.q4_0.gguf", + chat_format="functionary-v2", + tokenizer=LlamaHFTokenizer.from_pretrained("meetkai/functionary-small-v2.2-GGUF") +) +``` + +**NOTE**: There is no need to provide the default system messages used in Functionary as they are added automatically in the Functionary chat handler. Thus, the messages should contain just the chat messages and/or system messages that provide additional context for the model (e.g.: datetime, etc.). +</details> + +### Multi-modal Models + +`llama-cpp-python` supports such as llava1.5 which allow the language model to read information from both text and images. + +Below are the supported multi-modal models and their respective chat handlers (Python API) and chat formats (Server API). + +| Model | `LlamaChatHandler` | `chat_format` | +|:--- |:--- |:--- | +| [llava-v1.5-7b](https://huggingface.co/mys/ggml_llava-v1.5-7b) | `Llava15ChatHandler` | `llava-1-5` | +| [llava-v1.5-13b](https://huggingface.co/mys/ggml_llava-v1.5-13b) | `Llava15ChatHandler` | `llava-1-5` | +| [llava-v1.6-34b](https://huggingface.co/cjpais/llava-v1.6-34B-gguf) | `Llava16ChatHandler` | `llava-1-6` | +| [moondream2](https://huggingface.co/vikhyatk/moondream2) | `MoondreamChatHandler` | `moondream2` | +| [nanollava](https://huggingface.co/abetlen/nanollava-gguf) | `NanoLlavaChatHandler` | `nanollava` | +| [llama-3-vision-alpha](https://huggingface.co/abetlen/llama-3-vision-alpha-gguf) | `Llama3VisionAlphaChatHandler` | `llama-3-vision-alpha` | +| [minicpm-v-2.6](https://huggingface.co/openbmb/MiniCPM-V-2_6-gguf) | `MiniCPMv26ChatHandler` | `minicpm-v-2.6` | +| [qwen2.5-vl](https://huggingface.co/unsloth/Qwen2.5-VL-3B-Instruct-GGUF) | `Qwen25VLChatHandler` | `qwen2.5-vl` | +| [gemma-4](https://huggingface.co/unsloth/gemma-4-E4B-it-GGUF) | `Gemma4ChatHandler` | `gemma4` | +| GGUF models with an mtmd projector and embedded chat template | `MTMDChatHandler` | `mtmd` | + +Try Gemma 4 12B in Google Colab -> [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/abetlen/llama-cpp-python/blob/main/examples/colab/notebook.ipynb) + +Try Gemma 4 12B QAT in Google Colab -> [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/abetlen/llama-cpp-python/blob/main/examples/colab/Gemma4-12B-QAT.ipynb) + +Then you'll need to use a custom chat handler to load the clip model and process the chat messages and images. + +```python +from llama_cpp import Llama +from llama_cpp.llama_chat_format import Llava15ChatHandler +chat_handler = Llava15ChatHandler(clip_model_path="path/to/llava/mmproj.bin") +llm = Llama( + model_path="./path/to/llava/llama-model.gguf", + chat_handler=chat_handler, + n_ctx=2048, # n_ctx should be increased to accommodate the image embedding +) +llm.create_chat_completion( + messages = [ + {"role": "system", "content": "You are an assistant who perfectly describes images."}, + { + "role": "user", + "content": [ + {"type" : "text", "text": "What's in this image?"}, + {"type": "image_url", "image_url": {"url": "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg" } } + ] + } + ] +) +``` + +You can also pull the model from the Hugging Face Hub using the `from_pretrained` method. + +```python +from llama_cpp import Llama +from llama_cpp.llama_chat_format import MoondreamChatHandler + +chat_handler = MoondreamChatHandler.from_pretrained( + repo_id="vikhyatk/moondream2", + filename="*mmproj*", +) + +llm = Llama.from_pretrained( + repo_id="vikhyatk/moondream2", + filename="*text-model*", + chat_handler=chat_handler, + n_ctx=2048, # n_ctx should be increased to accommodate the image embedding +) + +response = llm.create_chat_completion( + messages = [ + { + "role": "user", + "content": [ + {"type" : "text", "text": "What's in this image?"}, + {"type": "image_url", "image_url": {"url": "https://upload.wikimedia.org/wikipedia/commons/thumb/d/dd/Gfp-wisconsin-madison-the-nature-boardwalk.jpg/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg" } } + + ] + } + ] +) +print(response["choices"][0]["text"]) +``` + +**Note**: Multi-modal models also support tool calling and JSON mode. + +<details> +<summary>Loading a Local Image</summary> + +Images can be passed as base64 encoded data URIs. The following example demonstrates how to do this. + +```python +import base64 + +def image_to_base64_data_uri(file_path): + with open(file_path, "rb") as img_file: + base64_data = base64.b64encode(img_file.read()).decode('utf-8') + return f"data:image/png;base64,{base64_data}" + +# Replace 'file_path.png' with the actual path to your PNG file +file_path = 'file_path.png' +data_uri = image_to_base64_data_uri(file_path) + +messages = [ + {"role": "system", "content": "You are an assistant who perfectly describes images."}, + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": data_uri }}, + {"type" : "text", "text": "Describe this image in detail please."} + ] + } +] + +``` + +</details> + +### Speculative Decoding + +`llama-cpp-python` supports speculative decoding which allows the model to generate completions based on a draft model. + +The fastest way to use speculative decoding is through the `LlamaPromptLookupDecoding` class. + +Just pass this as a draft model to the `Llama` class during initialization. + +```python +from llama_cpp import Llama +from llama_cpp.llama_speculative import LlamaPromptLookupDecoding + +llama = Llama( + model_path="path/to/model.gguf", + draft_model=LlamaPromptLookupDecoding(num_pred_tokens=10) # num_pred_tokens is the number of tokens to predict 10 is the default and generally good for gpu, 2 performs better for cpu-only machines. +) +``` + +### Embeddings + +To generate text embeddings use [`create_embedding`](https://llama-cpp-python.readthedocs.io/en/latest/api-reference/#llama_cpp.Llama.create_embedding) or [`embed`](https://llama-cpp-python.readthedocs.io/en/latest/api-reference/#llama_cpp.Llama.embed). Note that you must pass `embedding=True` to the constructor upon model creation for these to work properly. + +```python +import llama_cpp + +llm = llama_cpp.Llama(model_path="path/to/model.gguf", embedding=True) + +embeddings = llm.create_embedding("Hello, world!") + +# or create multiple embeddings at once + +embeddings = llm.create_embedding(["Hello, world!", "Goodbye, world!"]) +``` + +There are two primary notions of embeddings in a Transformer-style model: *token level* and *sequence level*. Sequence level embeddings are produced by "pooling" token level embeddings together, usually by averaging them or using the first token. + +Models that are explicitly geared towards embeddings will usually return sequence level embeddings by default, one for each input string. Non-embedding models such as those designed for text generation will typically return only token level embeddings, one for each token in each sequence. Thus the dimensionality of the return type will be one higher for token level embeddings. + +It is possible to control pooling behavior in some cases using the `pooling_type` flag on model creation. You can ensure token level embeddings from any model using `LLAMA_POOLING_TYPE_NONE`. The reverse, getting a generation oriented model to yield sequence level embeddings is currently not possible, but you can always do the pooling manually. + ### Adjusting the Context Window + The context window of the Llama models determines the maximum number of tokens that can be processed at once. By default, this is set to 512 tokens, but can be adjusted based on your requirements. For instance, if you want to work with larger contexts, you can expand the context window by setting the n_ctx parameter when initializing the Llama object: ```python -llm = Llama(model_path="./models/7B/ggml-model.bin", n_ctx=2048) +llm = Llama(model_path="./models/7B/llama-model.gguf", n_ctx=2048) ``` -## Web Server +## OpenAI Compatible Web Server `llama-cpp-python` offers a web server which aims to act as a drop-in replacement for the OpenAI API. This allows you to use llama.cpp compatible models with any OpenAI compatible client (language libraries, services, etc). @@ -122,73 +692,161 @@ This allows you to use llama.cpp compatible models with any OpenAI compatible cl To install the server package and get started: ```bash -pip install llama-cpp-python[server] -python3 -m llama_cpp.server --model models/7B/ggml-model.bin +pip install 'llama-cpp-python[server]' +python3 -m llama_cpp.server --model models/7B/llama-model.gguf +``` + +Similar to Hardware Acceleration section above, you can also install with GPU (cuBLAS) support like this: + +```bash +CMAKE_ARGS="-DGGML_CUDA=on" FORCE_CMAKE=1 pip install 'llama-cpp-python[server]' +python3 -m llama_cpp.server --model models/7B/llama-model.gguf --n_gpu_layers 35 ``` Navigate to [http://localhost:8000/docs](http://localhost:8000/docs) to see the OpenAPI documentation. +To bind to `0.0.0.0` to enable remote connections, use `python3 -m llama_cpp.server --host 0.0.0.0`. +Similarly, to change the port (default is 8000), use `--port`. + +You probably also want to set the prompt format. For chatml, use + +```bash +python3 -m llama_cpp.server --model models/7B/llama-model.gguf --chat_format chatml +``` + +That will format the prompt according to how model expects it. You can find the prompt format in the model card. +For possible options, see [llama_cpp/llama_chat_format.py](llama_cpp/llama_chat_format.py) and look for lines starting with "@register_chat_format". + +If you have `huggingface-hub` installed, you can also use the `--hf_model_repo_id` flag to load a model from the Hugging Face Hub. + +```bash +python3 -m llama_cpp.server --hf_model_repo_id lmstudio-community/Qwen3.5-0.8B-GGUF --model '*Q8_0.gguf' +``` + +### Web Server Features + +- [Local Copilot replacement](https://llama-cpp-python.readthedocs.io/en/latest/server/#code-completion) +- [Function Calling support](https://llama-cpp-python.readthedocs.io/en/latest/server/#function-calling) +- [Vision API support](https://llama-cpp-python.readthedocs.io/en/latest/server/#multimodal-models) +- [Multiple Models](https://llama-cpp-python.readthedocs.io/en/latest/server/#configuration-and-multi-model-support) + ## Docker image A Docker image is available on [GHCR](https://ghcr.io/abetlen/llama-cpp-python). To run the server: ```bash -docker run --rm -it -p 8000:8000 -v /path/to/models:/models -e MODEL=/models/ggml-model-name.bin ghcr.io/abetlen/llama-cpp-python:latest +docker run --rm -it -p 8000:8000 -v /path/to/models:/models -e MODEL=/models/llama-model.gguf ghcr.io/abetlen/llama-cpp-python:latest ``` +[Docker on termux (requires root)](https://gist.github.com/FreddieOliveira/efe850df7ff3951cb62d74bd770dce27) is currently the only known way to run this on phones, see [termux support issue](https://github.com/abetlen/llama-cpp-python/issues/389) + ## Low-level API +[API Reference](https://llama-cpp-python.readthedocs.io/en/latest/api-reference/#low-level-api) + The low-level API is a direct [`ctypes`](https://docs.python.org/3/library/ctypes.html) binding to the C API provided by `llama.cpp`. -The entire lowe-level API can be found in [llama_cpp/llama_cpp.py](https://github.com/abetlen/llama-cpp-python/blob/master/llama_cpp/llama_cpp.py) and directly mirrors the C API in [llama.h](https://github.com/ggerganov/llama.cpp/blob/master/llama.h). +The entire low-level API can be found in [llama_cpp/llama_cpp.py](https://github.com/abetlen/llama-cpp-python/blob/master/llama_cpp/llama_cpp.py) and directly mirrors the C API in [llama.h](https://github.com/ggerganov/llama.cpp/blob/master/llama.h). Below is a short example demonstrating how to use the low-level API to tokenize a prompt: ```python ->>> import llama_cpp ->>> import ctypes ->>> params = llama_cpp.llama_context_default_params() +import llama_cpp +import ctypes +llama_cpp.llama_backend_init() # Must be called once at the start of each program +model_params = llama_cpp.llama_model_default_params() +ctx_params = llama_cpp.llama_context_default_params() +prompt = b"Q: Name the planets in the solar system? A: " # use bytes for char * params ->>> ctx = llama_cpp.llama_init_from_file(b"./models/7b/ggml-model.bin", params) ->>> max_tokens = params.n_ctx +model = llama_cpp.llama_model_load_from_file(b"./models/7b/llama-model.gguf", model_params) +ctx = llama_cpp.llama_init_from_model(model, ctx_params) +vocab = llama_cpp.llama_model_get_vocab(model) +max_tokens = ctx_params.n_ctx # use ctypes arrays for array params ->>> tokens = (llama_cpp.llama_token * int(max_tokens))() ->>> n_tokens = llama_cpp.llama_tokenize(ctx, b"Q: Name the planets in the solar system? A: ", tokens, max_tokens, add_bos=llama_cpp.c_bool(True)) ->>> llama_cpp.llama_free(ctx) +tokens = (llama_cpp.llama_token * int(max_tokens))() +n_tokens = llama_cpp.llama_tokenize(vocab, prompt, len(prompt), tokens, max_tokens, True, False) +llama_cpp.llama_free(ctx) +llama_cpp.llama_model_free(model) ``` Check out the [examples folder](examples/low_level_api) for more examples of using the low-level API. +## Documentation -# Documentation - -Documentation is available at [https://abetlen.github.io/llama-cpp-python](https://abetlen.github.io/llama-cpp-python). +Documentation is available via [https://llama-cpp-python.readthedocs.io/](https://llama-cpp-python.readthedocs.io/). If you find any issues with the documentation, please open an issue or submit a PR. -# Development +## Development This package is under active development and I welcome any contributions. +See [CONTRIBUTING.md](CONTRIBUTING.md) for contribution workflow, PR title, changelog, testing, and style guidelines. -To get started, clone the repository and install the package in development mode: +To get started, clone the repository and install the package in editable / development mode: ```bash -git clone --recurse-submodules git@github.com:abetlen/llama-cpp-python.git +git clone --recurse-submodules https://github.com/abetlen/llama-cpp-python.git cd llama-cpp-python +# Upgrade pip (required for editable mode) +pip install --upgrade pip + # Install with pip pip install -e . +# install development tooling (tests, docs, ruff) +pip install -e '.[dev]' + # if you want to use the fastapi / openapi server -pip install -e .[server] +pip install -e '.[server]' -# If you're a poetry user, installing will also include a virtual environment -poetry install --all-extras -. .venv/bin/activate +# to install all optional dependencies +pip install -e '.[all]' -# Will need to be re-run any time vendor/llama.cpp is updated -python3 setup.py develop +# to clear the local build cache +make clean ``` -# How does this compare to other Python bindings of `llama.cpp`? +Now try running the tests + +```bash +pytest +``` + +And check formatting / linting before opening a PR: + +```bash +python -m ruff check llama_cpp tests +python -m ruff format --check llama_cpp tests + +# or use the Makefile targets +make lint +make format +``` + +There's a `Makefile` available with useful targets. +A typical workflow would look like this: + +```bash +make build +make test +``` + +You can also test out specific commits of `llama.cpp` by checking out the desired commit in the `vendor/llama.cpp` submodule and then running `make clean` and `pip install -e .` again. Any changes in the `llama.h` API will require +changes to the `llama_cpp/llama_cpp.py` file to match the new API (additional changes may be required elsewhere). + +## FAQ + +### Are there pre-built binaries / binary wheels available? + +The recommended installation method is to install from source as described above. +The reason for this is that `llama.cpp` is built with compiler optimizations that are specific to your system. +Using pre-built binaries would require disabling these optimizations or supporting a large number of pre-built binaries for each platform. + +That being said there are some pre-built binaries available through the Releases as well as some community provided wheels. + +In the future, I would like to provide pre-built binaries and wheels for common platforms and I'm happy to accept any useful contributions in this area. +This is currently being tracked in [#741](https://github.com/abetlen/llama-cpp-python/issues/741) + +### How does this compare to other Python bindings of `llama.cpp`? I originally wrote this package for my own use with two goals in mind: @@ -197,6 +855,6 @@ I originally wrote this package for my own use with two goals in mind: Any contributions and changes to this package will be made with these goals in mind. -# License +## License This project is licensed under the terms of the MIT license. diff --git a/docker/README.md b/docker/README.md index 053d311b4f..474503fdfc 100644 --- a/docker/README.md +++ b/docker/README.md @@ -1,37 +1,44 @@ -# Install Docker Server - -**Note #1:** This was tested with Docker running on Linux. If you can get it working on Windows or MacOS, please update this `README.md` with a PR! +### Install Docker Server +> [!IMPORTANT] +> This was tested with Docker running on Linux. <br>If you can get it working on Windows or MacOS, please update this `README.md` with a PR!<br> [Install Docker Engine](https://docs.docker.com/engine/install) -**Note #2:** NVidia GPU CuBLAS support requires a NVidia GPU with sufficient VRAM (approximately as much as the size in the table below) and Docker NVidia support (see [container-toolkit/install-guide](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html)) -# Simple Dockerfiles for building the llama-cpp-python server with external model bin files -## openblas_simple - a simple Dockerfile for non-GPU OpenBLAS, where the model is located outside the Docker image +## Simple Dockerfiles for building the llama-cpp-python server with external model bin files +### openblas_simple +A simple Dockerfile for non-GPU OpenBLAS, where the model is located outside the Docker image: ``` cd ./openblas_simple docker build -t openblas_simple . -docker run -e USE_MLOCK=0 -e MODEL=/var/model/<model-path> -v <model-root-path>:/var/model -t openblas_simple +docker run --cap-add SYS_RESOURCE -e USE_MLOCK=0 -e MODEL=/var/model/<model-path> -v <model-root-path>:/var/model -t openblas_simple ``` where `<model-root-path>/<model-path>` is the full path to the model file on the Docker host system. -## cuda_simple - a simple Dockerfile for CUDA accelerated CuBLAS, where the model is located outside the Docker image +### cuda_simple +> [!WARNING] +> Nvidia GPU CuBLAS support requires an Nvidia GPU with sufficient VRAM (approximately as much as the size in the table below) and Docker Nvidia support (see [container-toolkit/install-guide](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/install-guide.html)) <br> + +A simple Dockerfile for CUDA-accelerated CuBLAS, where the model is located outside the Docker image: + ``` cd ./cuda_simple docker build -t cuda_simple . -docker run -e USE_MLOCK=0 -e MODEL=/var/model/<model-path> -v <model-root-path>:/var/model -t cuda_simple +docker run --gpus=all --cap-add SYS_RESOURCE -e USE_MLOCK=0 -e MODEL=/var/model/<model-path> -v <model-root-path>:/var/model -t cuda_simple ``` where `<model-root-path>/<model-path>` is the full path to the model file on the Docker host system. -# "Open-Llama-in-a-box" -## Download an Apache V2.0 licensed 3B paramter Open Llama model and install into a Docker image that runs an OpenBLAS-enabled llama-cpp-python server +-------------------------------------------------------------------------- + +### "Open-Llama-in-a-box" +Download an Apache V2.0 licensed 3B params Open LLaMA model and install into a Docker image that runs an OpenBLAS-enabled llama-cpp-python server: ``` $ cd ./open_llama ./build.sh ./start.sh ``` -# Manually choose your own Llama model from Hugging Face +### Manually choose your own Llama model from Hugging Face `python3 ./hug_model.py -a TheBloke -t llama` You should now have a model in the current directory and `model.bin` symlinked to it for the subsequent Docker build and copy step. e.g. ``` @@ -39,8 +46,10 @@ docker $ ls -lh *.bin -rw-rw-r-- 1 user user 4.8G May 23 18:30 <downloaded-model-file>q5_1.bin lrwxrwxrwx 1 user user 24 May 23 18:30 model.bin -> <downloaded-model-file>q5_1.bin ``` -**Note #1:** Make sure you have enough disk space to download the model. As the model is then copied into the image you will need at least -**TWICE** as much disk space as the size of the model: + +> [!NOTE] +> Make sure you have enough disk space to download the model. As the model is then copied into the image you will need at least +**TWICE** as much disk space as the size of the model:<br> | Model | Quantized size | |------:|----------------:| @@ -50,17 +59,6 @@ lrwxrwxrwx 1 user user 24 May 23 18:30 model.bin -> <downloaded-model-file>q5_ | 33B | 25 GB | | 65B | 50 GB | -**Note #2:** If you want to pass or tune additional parameters, customise `./start_server.sh` before running `docker build ...` - -## Use OpenBLAS -Use if you don't have a NVidia GPU. Defaults to `python:3-slim-bullseye` Docker base image and OpenBLAS: -### Build: -`docker build -t openblas .` -### Run: -`docker run --cap-add SYS_RESOURCE -t openblas` -## Use CuBLAS -### Build: -`docker build --build-arg IMAGE=nvidia/cuda:12.1.1-devel-ubuntu22.04 -t cublas .` -### Run: -`docker run --cap-add SYS_RESOURCE -t cublas` +> [!NOTE] +> If you want to pass or tune additional parameters, customise `./start_server.sh` before running `docker build ...` diff --git a/docker/cuda_simple/Dockerfile b/docker/cuda_simple/Dockerfile index e4a2f07e2c..0bbf20ffe9 100644 --- a/docker/cuda_simple/Dockerfile +++ b/docker/cuda_simple/Dockerfile @@ -1,16 +1,27 @@ -ARG CUDA_IMAGE="12.1.1-devel-ubuntu22.04" +ARG CUDA_IMAGE="12.5.0-devel-ubuntu22.04" FROM nvidia/cuda:${CUDA_IMAGE} # We need to set the host to 0.0.0.0 to allow outside access ENV HOST 0.0.0.0 +RUN apt-get update && apt-get upgrade -y \ + && apt-get install -y git build-essential \ + python3 python3-pip gcc wget \ + ocl-icd-opencl-dev opencl-headers clinfo \ + libclblast-dev libopenblas-dev \ + && mkdir -p /etc/OpenCL/vendors && echo "libnvidia-opencl.so.1" > /etc/OpenCL/vendors/nvidia.icd + COPY . . -# Install the package -RUN apt update && apt install -y python3 python3-pip -RUN python3 -m pip install --upgrade pip pytest cmake scikit-build setuptools fastapi uvicorn sse-starlette pydantic-settings +# setting build related env vars +ENV CUDA_DOCKER_ARCH=all +ENV GGML_CUDA=1 + +# Install depencencies +RUN python3 -m pip install --upgrade pip pytest cmake scikit-build setuptools fastapi uvicorn sse-starlette pydantic-settings starlette-context -RUN LLAMA_CUBLAS=1 pip install llama-cpp-python +# Install llama-cpp-python (build with cuda) +RUN CMAKE_ARGS="-DGGML_CUDA=on" pip install llama-cpp-python # Run the server CMD python3 -m llama_cpp.server diff --git a/docker/open_llama/Dockerfile b/docker/open_llama/Dockerfile index 7788f33de6..f05e66ef2a 100644 --- a/docker/open_llama/Dockerfile +++ b/docker/open_llama/Dockerfile @@ -1,5 +1,5 @@ # Define the image argument and provide a default value -ARG IMAGE=python:3-slim-bullseye +ARG IMAGE=python:3-slim-bookworm # Use the image as specified FROM ${IMAGE} @@ -12,19 +12,21 @@ RUN apt-get update && apt-get upgrade -y && apt-get install -y --no-install-reco python3 \ python3-pip \ ninja-build \ - build-essential + build-essential \ + && apt-get clean \ + && rm -rf /var/lib/apt/lists/* -RUN python3 -m pip install --upgrade pip pytest cmake scikit-build setuptools fastapi uvicorn sse-starlette pydantic-settings +RUN python3 -m pip install --upgrade pip pytest cmake scikit-build setuptools fastapi uvicorn sse-starlette pydantic-settings starlette-context # Perform the conditional installations based on the image RUN echo "Image: ${IMAGE}" && \ - if [ "${IMAGE}" = "python:3-slim-bullseye" ] ; then \ + if [ "${IMAGE}" = "python:3-slim-bookworm" ] ; then \ echo "OpenBLAS install:" && \ apt-get install -y --no-install-recommends libopenblas-dev && \ - LLAMA_OPENBLAS=1 pip install llama-cpp-python --verbose; \ + CMAKE_ARGS="-DGGML_BLAS=ON -DGGML_BLAS_VENDOR=OpenBLAS" pip install llama-cpp-python --verbose; \ else \ echo "CuBLAS install:" && \ - LLAMA_CUBLAS=1 pip install llama-cpp-python --verbose; \ + CMAKE_ARGS="-DGGML_CUDA=on" pip install llama-cpp-python --verbose; \ fi # Clean up apt cache diff --git a/docker/openblas_simple/Dockerfile b/docker/openblas_simple/Dockerfile index 8231bdb96c..5ee667dc0d 100644 --- a/docker/openblas_simple/Dockerfile +++ b/docker/openblas_simple/Dockerfile @@ -1,4 +1,4 @@ -FROM python:3-slim-bullseye +FROM python:3-slim-bookworm # We need to set the host to 0.0.0.0 to allow outside access ENV HOST 0.0.0.0 @@ -6,10 +6,13 @@ ENV HOST 0.0.0.0 COPY . . # Install the package -RUN apt update && apt install -y libopenblas-dev ninja-build build-essential -RUN python -m pip install --upgrade pip pytest cmake scikit-build setuptools fastapi uvicorn sse-starlette pydantic-settings +RUN apt update && apt install -y libopenblas-dev ninja-build build-essential pkg-config \ + && apt-get clean \ + && rm -rf /var/lib/apt/lists/* /tmp/* + +RUN python -m pip install --upgrade pip pytest cmake scikit-build setuptools fastapi uvicorn sse-starlette pydantic-settings starlette-context -RUN LLAMA_OPENBLAS=1 pip install llama_cpp_python --verbose +RUN CMAKE_ARGS="-DGGML_BLAS=ON -DGGML_BLAS_VENDOR=OpenBLAS" pip install llama_cpp_python --verbose # Run the server CMD python3 -m llama_cpp.server diff --git a/docker/simple/Dockerfile b/docker/simple/Dockerfile index 507b2ba46d..bad4f456ff 100644 --- a/docker/simple/Dockerfile +++ b/docker/simple/Dockerfile @@ -1,27 +1,33 @@ # Define the image argument and provide a default value -ARG IMAGE=python:3-slim-bullseye +ARG IMAGE=python:3-slim-bookworm # Use the image as specified FROM ${IMAGE} # Re-declare the ARG after FROM ARG IMAGE +ARG CMAKE_ARGS="-DGGML_NATIVE=off" # Update and upgrade the existing packages RUN apt-get update && apt-get upgrade -y && apt-get install -y --no-install-recommends \ + git \ python3 \ python3-pip \ ninja-build \ libopenblas-dev \ - build-essential + build-essential \ + && apt-get clean \ + && rm -rf /var/lib/apt/lists/* /tmp/* RUN mkdir /app WORKDIR /app COPY . /app -RUN python3 -m pip install --upgrade pip pytest cmake scikit-build setuptools fastapi uvicorn sse-starlette pydantic-settings +RUN python3 -m pip install --upgrade pip -RUN make build && make clean +RUN python3 -m pip install --upgrade pip pytest cmake scikit-build setuptools fastapi uvicorn sse-starlette pydantic-settings starlette-context + +RUN CMAKE_ARGS="${CMAKE_ARGS}" pip install . --verbose # Set environment variable for the host ENV HOST=0.0.0.0 diff --git a/docs/api-reference.md b/docs/api-reference.md index 1290cad493..ab51ef754e 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -2,6 +2,10 @@ title: API Reference --- +## High Level API + +High-level Python bindings for llama.cpp. + ::: llama_cpp.Llama options: members: @@ -17,13 +21,21 @@ title: API Reference - create_completion - __call__ - create_chat_completion + - create_chat_completion_openai_v1 - set_cache - save_state - load_state - token_bos - token_eos + - from_pretrained show_root_heading: true +::: llama_cpp.LlamaGrammar + options: + members: + - from_string + - from_json_schema + ::: llama_cpp.LlamaCache options: show_root_heading: true @@ -48,6 +60,29 @@ title: API Reference options: show_root_heading: true +## Low Level API + +Low-level Python bindings for llama.cpp using Python's ctypes library. + ::: llama_cpp.llama_cpp + options: + show_if_no_docstring: true + # filter only members starting with `llama_` + filters: + - "^llama_" + +::: llama_cpp.llama_cpp + options: + show_if_no_docstring: true + show_root_heading: false + show_root_toc_entry: false + heading_level: 4 + # filter only members starting with `LLAMA_` + filters: + - "^LLAMA_" + +## Misc + +::: llama_cpp.llama_types options: show_if_no_docstring: true \ No newline at end of file diff --git a/docs/changelog.md b/docs/changelog.md new file mode 100644 index 0000000000..047bc14424 --- /dev/null +++ b/docs/changelog.md @@ -0,0 +1 @@ +-8<- "CHANGELOG.md" \ No newline at end of file diff --git a/docs/icon.svg b/docs/icon.svg new file mode 100644 index 0000000000..9f2d087539 --- /dev/null +++ b/docs/icon.svg @@ -0,0 +1,5 @@ +<svg width="264" height="264" viewBox="0 0 264 264" fill="none" xmlns="http://www.w3.org/2000/svg"> +<path d="M229.665 157.022C226.791 157.022 224.362 157.022 221.452 157.022C221.452 159.886 221.452 162.638 221.452 165.386C221.452 169.439 218.221 172.725 214.168 172.725V172.725C210.115 172.725 206.775 169.439 206.775 165.386C206.775 162.694 206.775 159.931 206.775 157.051C201.473 157.051 201.438 157.051 196.384 157.051C196.384 152.118 196.384 147.629 196.384 142.792C201.292 142.792 201.341 142.792 206.498 142.792C206.498 139.997 206.498 137.292 206.498 134.581C206.498 130.483 209.802 127.161 213.9 127.161V127.161C217.998 127.161 221.337 130.483 221.337 134.581C221.337 137.217 221.337 139.869 221.337 142.601C223.866 142.601 226.276 142.601 228.674 142.601C232.657 142.601 235.885 145.829 235.885 149.812C235.885 150.141 235.885 150.471 235.885 150.802C235.885 154.238 233.1 157.023 229.665 157.022V157.022Z" fill="#FF8236"/> +<path d="M153.094 142.863C155.969 142.864 158.398 142.864 161.307 142.864C161.307 140 161.307 137.248 161.307 134.5C161.307 130.447 164.538 127.161 168.591 127.161V127.161C172.644 127.161 175.984 130.446 175.984 134.499C175.984 137.192 175.984 139.955 175.984 142.835C181.287 142.835 181.314 142.835 186.368 142.835C186.368 147.768 186.368 152.257 186.368 157.093C181.46 157.093 181.419 157.093 176.262 157.093C176.262 159.889 176.262 162.594 176.262 165.305C176.262 169.403 172.957 172.725 168.86 172.725V172.725C164.762 172.725 161.422 169.403 161.422 165.305C161.422 162.669 161.422 160.017 161.422 157.285C158.893 157.285 156.484 157.285 154.085 157.285C150.103 157.285 146.874 154.056 146.874 150.074C146.874 149.745 146.874 149.415 146.874 149.083C146.874 145.648 149.659 142.863 153.094 142.863V142.863Z" fill="#FF8236"/> +<path fill-rule="evenodd" clip-rule="evenodd" d="M142.114 45.2382C144.949 46.6406 144.087 48.6771 143.291 50.4754C143.001 51.1309 142.713 51.7881 142.424 52.4456C140.823 56.0975 139.217 59.7593 137.254 63.2122C134.69 67.7245 133.969 72.1034 134.718 76.425C135.779 76.6467 136.846 76.8782 137.916 77.1194C140.063 69.1673 143.313 61.8666 150.34 56.3784C154.237 53.3347 158.881 52.5295 163.741 53.4167C166.806 53.9761 167.354 55.1539 165.72 57.7847C165.608 57.9647 165.496 58.1449 165.384 58.3252C163.881 60.7487 162.364 63.1964 160.522 65.3492C155.705 70.9818 156.079 76.8613 159.133 83.0193C159.817 83.2456 160.5 83.4751 161.183 83.7078C172.683 87.6287 177.516 100.775 172.451 111.819L169.949 117.273C153.793 109.863 135.437 103.1 119.57 100.817C119.293 100.777 118.925 100.719 118.479 100.649C112.549 99.7223 92.8929 96.6489 91.0237 104.451L91.0191 104.47C90.054 108.427 90.9132 112.361 91.78 116.33C92.9307 121.599 94.095 126.93 91.0237 132.458C88.0952 137.729 82.3817 140.345 76.8738 142.335C75.8854 142.692 74.8542 143.051 73.7966 143.419C63.0234 147.168 49.5058 151.871 50.5002 165.426C51.8566 183.916 76.7213 188.107 92.8753 190.83C93.4693 190.93 94.0515 191.028 94.6203 191.125C104.41 192.791 114.831 190.226 125.131 187.691C134.359 185.419 143.49 183.171 151.981 184.013C161.451 184.952 168.196 189.989 174.294 194.543C179.765 198.629 184.717 202.327 190.651 202.327C197.695 202.327 202.659 200.151 206.23 198.585C209.02 197.362 210.959 196.512 212.374 197.365C215.601 199.308 206.98 221.281 190.651 221.281C182.127 221.281 175.694 218.458 169.545 215.76C163.916 213.289 158.524 210.923 151.981 210.923C146.86 210.923 140.616 212.768 133.299 214.929C121.058 218.545 105.818 223.046 87.8198 221.281C59.0636 218.46 16.4933 202.541 14.1171 170.151C12.1175 142.894 41.5744 124.106 64.291 117.009C61.6175 103.844 63.224 90.9669 73.8515 81.7137C80.3917 76.0193 89.2057 73.9986 96.7861 73.2898C101.748 72.8258 107.118 72.8441 112.689 73.2406C114.024 67.1837 116.323 61.4235 119.904 56.2059C125.156 48.5519 132.354 44.5736 142.114 45.2382ZM164.455 97.0679C164.455 100.146 161.96 102.641 158.882 102.641C155.804 102.641 153.308 100.146 153.308 97.0679C153.308 93.9899 155.804 91.4946 158.882 91.4946C161.96 91.4946 164.455 93.9899 164.455 97.0679Z" fill="#FF8236"/> +</svg> diff --git a/docs/index.md b/docs/index.md index 7d5ccc314f..60bc7aef42 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,92 +1,5 @@ -# Getting Started +--- +title: Getting Started +--- -## 🦙 Python Bindings for `llama.cpp` - -[![Documentation](https://img.shields.io/badge/docs-passing-green.svg)](https://abetlen.github.io/llama-cpp-python) -[![Tests](https://github.com/abetlen/llama-cpp-python/actions/workflows/test.yaml/badge.svg?branch=main)](https://github.com/abetlen/llama-cpp-python/actions/workflows/test.yaml) -[![PyPI](https://img.shields.io/pypi/v/llama-cpp-python)](https://pypi.org/project/llama-cpp-python/) -[![PyPI - Python Version](https://img.shields.io/pypi/pyversions/llama-cpp-python)](https://pypi.org/project/llama-cpp-python/) -[![PyPI - License](https://img.shields.io/pypi/l/llama-cpp-python)](https://pypi.org/project/llama-cpp-python/) -[![PyPI - Downloads](https://img.shields.io/pypi/dm/llama-cpp-python)](https://pypi.org/project/llama-cpp-python/) - -Simple Python bindings for **@ggerganov's** [`llama.cpp`](https://github.com/ggerganov/llama.cpp) library. -This package provides: - -- Low-level access to C API via `ctypes` interface. -- High-level Python API for text completion - - OpenAI-like API - - LangChain compatibility - -## Installation - -Install from PyPI: - -```bash -pip install llama-cpp-python -``` - -## High-level API - -```python ->>> from llama_cpp import Llama ->>> llm = Llama(model_path="./models/7B/ggml-model.bin") ->>> output = llm("Q: Name the planets in the solar system? A: ", max_tokens=32, stop=["Q:", "\n"], echo=True) ->>> print(output) -{ - "id": "cmpl-xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx", - "object": "text_completion", - "created": 1679561337, - "model": "./models/7B/ggml-model.bin", - "choices": [ - { - "text": "Q: Name the planets in the solar system? A: Mercury, Venus, Earth, Mars, Jupiter, Saturn, Uranus, Neptune and Pluto.", - "index": 0, - "logprobs": None, - "finish_reason": "stop" - } - ], - "usage": { - "prompt_tokens": 14, - "completion_tokens": 28, - "total_tokens": 42 - } -} -``` - -## Web Server - -`llama-cpp-python` offers a web server which aims to act as a drop-in replacement for the OpenAI API. -This allows you to use llama.cpp compatible models with any OpenAI compatible client (language libraries, services, etc). - -To install the server package and get started: - -```bash -pip install llama-cpp-python[server] -export MODEL=./models/7B/ggml-model.bin -python3 -m llama_cpp.server -``` - -Navigate to [http://localhost:8000/docs](http://localhost:8000/docs) to see the OpenAPI documentation. - -## Low-level API - -The low-level API is a direct `ctypes` binding to the C API provided by `llama.cpp`. -The entire API can be found in [llama_cpp/llama_cpp.py](https://github.com/abetlen/llama-cpp-python/blob/master/llama_cpp/llama_cpp.py) and should mirror [llama.h](https://github.com/ggerganov/llama.cpp/blob/master/llama.h). - - -## Development - -This package is under active development and I welcome any contributions. - -To get started, clone the repository and install the package in development mode: - -```bash -git clone git@github.com:abetlen/llama-cpp-python.git -git submodule update --init --recursive -# Will need to be re-run any time vendor/llama.cpp is updated -python3 setup.py develop -``` - -## License - -This project is licensed under the terms of the MIT license. \ No newline at end of file +-8<- "README.md" \ No newline at end of file diff --git a/docs/install/macos.md b/docs/install/macos.md index 3330396e3b..e006fc0a3c 100644 --- a/docs/install/macos.md +++ b/docs/install/macos.md @@ -30,7 +30,7 @@ conda activate llama *(you needed xcode installed in order pip to build/compile the C++ code)* ``` pip uninstall llama-cpp-python -y -CMAKE_ARGS="-DLLAMA_METAL=on" FORCE_CMAKE=1 pip install -U llama-cpp-python --no-cache-dir +CMAKE_ARGS="-DGGML_METAL=on" pip install -U llama-cpp-python --no-cache-dir pip install 'llama-cpp-python[server]' # you should now have llama-cpp-python v0.1.62 or higher installed @@ -38,19 +38,19 @@ llama-cpp-python         0.1.68 ``` -**(5) Download a v3 ggml model** - - **ggmlv3** - - file name ends with **q4_0.bin** - indicating it is 4bit quantized, with quantisation method 0 +**(5) Download a v3 gguf v2 model** + - **ggufv2** + - file name ends with **Q4_0.gguf** - indicating it is 4bit quantized, with quantisation method 0 -https://huggingface.co/TheBloke/open-llama-7b-open-instruct-GGML +https://huggingface.co/TheBloke/CodeLlama-7B-GGUF **(6) run the llama-cpp-python API server with MacOS Metal GPU support** ``` # config your ggml model path -# make sure it is ggml v3 +# make sure it is gguf v2 # make sure it is q4_0 -export MODEL=[path to your llama.cpp ggml models]]/[ggml-model-name]]q4_0.bin +export MODEL=[path to your llama.cpp ggml models]]/[ggml-model-name]]Q4_0.gguf python3 -m llama_cpp.server --model $MODEL --n_gpu_layers 1 ``` diff --git a/docs/server.md b/docs/server.md new file mode 100644 index 0000000000..9c09a1f1cf --- /dev/null +++ b/docs/server.md @@ -0,0 +1,247 @@ +# OpenAI Compatible Server + +`llama-cpp-python` offers an OpenAI API compatible web server. + +This web server can be used to serve local models and easily connect them to existing clients. + +## Setup + +### Installation + +The server can be installed by running the following command: + +```bash +pip install llama-cpp-python[server] +``` + +### Running the server + +The server can then be started by running the following command: + +```bash +python3 -m llama_cpp.server --model <model_path> +``` + +You can also pass chat-template kwargs at model load time from the CLI: + +```bash +python3 -m llama_cpp.server \ + --model <model_path> \ + --chat_format chatml \ + --chat_template_kwargs '{"enable_thinking": true}' +``` + +### Server options + +For a full list of options, run: + +```bash +python3 -m llama_cpp.server --help +``` + +NOTE: All server options are also available as environment variables. For example, `--model` can be set by setting the `MODEL` environment variable. + +Check out the server config reference below settings for more information on the available options. +CLI arguments and environment variables are available for all of the fields defined in [`ServerSettings`](#llama_cpp.server.settings.ServerSettings) and [`ModelSettings`](#llama_cpp.server.settings.ModelSettings) + +Additionally the server supports configuration check out the [configuration section](#configuration-and-multi-model-support) for more information and examples. + + +## Guides + +### Code Completion + +`llama-cpp-python` supports code completion via GitHub Copilot. + +*NOTE*: Without GPU acceleration this is unlikely to be fast enough to be usable. + +You'll first need to download one of the available code completion models in GGUF format: + +- [replit-code-v1_5-GGUF](https://huggingface.co/abetlen/replit-code-v1_5-3b-GGUF) + +Then you'll need to run the OpenAI compatible web server with a increased context size substantially for GitHub Copilot requests: + +```bash +python3 -m llama_cpp.server --model <model_path> --n_ctx 16192 +``` + +Then just update your settings in `.vscode/settings.json` to point to your code completion server: + +```json +{ + // ... + "github.copilot.advanced": { + "debug.testOverrideProxyUrl": "http://<host>:<port>", + "debug.overrideProxyUrl": "http://<host>:<port>" + } + // ... +} +``` + +### Function Calling + +`llama-cpp-python` supports structured function calling based on a JSON schema. +Function calling is completely compatible with the OpenAI function calling API and can be used by connecting with the official OpenAI Python client. + +You'll first need to download one of the available function calling models in GGUF format: + +- [functionary](https://huggingface.co/meetkai) + +Then when you run the server you'll need to also specify either `functionary-v1` or `functionary-v2` chat_format. + +Note that since functionary requires a HF Tokenizer due to discrepancies between llama.cpp and HuggingFace's tokenizers as mentioned [here](https://github.com/abetlen/llama-cpp-python/blob/main?tab=readme-ov-file#function-calling), you will need to pass in the path to the tokenizer too. The tokenizer files are already included in the respective HF repositories hosting the gguf files. + +```bash +python3 -m llama_cpp.server --model <model_path_to_functionary_v2_model> --chat_format functionary-v2 --hf_pretrained_model_name_or_path <model_path_to_functionary_v2_tokenizer> +``` + +Check out this [example notebook](https://github.com/abetlen/llama-cpp-python/blob/main/examples/notebooks/Functions.ipynb) for a walkthrough of some interesting use cases for function calling. + +### Multimodal Models + +`llama-cpp-python` supports the llava1.5 family of multi-modal models which allow the language model to +read information from both text and images. + +You'll first need to download one of the available multi-modal models in GGUF format: + +- [llava-v1.5-7b](https://huggingface.co/mys/ggml_llava-v1.5-7b) +- [llava-v1.5-13b](https://huggingface.co/mys/ggml_llava-v1.5-13b) +- [bakllava-1-7b](https://huggingface.co/mys/ggml_bakllava-1) +- [llava-v1.6-34b](https://huggingface.co/cjpais/llava-v1.6-34B-gguf) +- [moondream2](https://huggingface.co/vikhyatk/moondream2) + +Then when you run the server you'll need to also specify the path to the clip model used for image embedding and the `llava-1-5` chat_format + +```bash +python3 -m llama_cpp.server --model <model_path> --clip_model_path <clip_model_path> --chat_format llava-1-5 +``` + +Then you can just use the OpenAI API as normal + +```python3 +from openai import OpenAI + +client = OpenAI(base_url="http://<host>:<port>/v1", api_key="sk-xxx") +response = client.chat.completions.create( + model="gpt-4-vision-preview", + messages=[ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "<image_url>" + }, + }, + {"type": "text", "text": "What does the image say"}, + ], + } + ], +) +print(response) +``` + +## Configuration and Multi-Model Support + +The server supports configuration via a JSON config file that can be passed using the `--config_file` parameter or the `CONFIG_FILE` environment variable. + +```bash +python3 -m llama_cpp.server --config_file <config_file> +``` + +Config files support all of the server and model options supported by the cli and environment variables however instead of only a single model the config file can specify multiple models. + +The server supports routing requests to multiple models based on the `model` parameter in the request which matches against the `model_alias` in the config file. + +At the moment only a single model is loaded into memory at, the server will automatically load and unload models as needed. + +For a single-model config, `chat_template_kwargs` can be set directly on the model entry: + +```json +{ + "models": [ + { + "model": "models/Qwen3.5-0.8B/qwen3.5-0.8b-q8_0.gguf", + "chat_format": "chatml", + "chat_template_kwargs": { + "enable_thinking": true + } + } + ] +} +``` + +```json +{ + "host": "0.0.0.0", + "port": 8080, + "models": [ + { + "model": "models/OpenHermes-2.5-Mistral-7B-GGUF/openhermes-2.5-mistral-7b.Q4_K_M.gguf", + "model_alias": "gpt-3.5-turbo", + "chat_format": "chatml", + "n_gpu_layers": -1, + "offload_kqv": true, + "n_threads": 12, + "n_batch": 512, + "n_ctx": 2048 + }, + { + "model": "models/OpenHermes-2.5-Mistral-7B-GGUF/openhermes-2.5-mistral-7b.Q4_K_M.gguf", + "model_alias": "gpt-4", + "chat_format": "chatml", + "n_gpu_layers": -1, + "offload_kqv": true, + "n_threads": 12, + "n_batch": 512, + "n_ctx": 2048 + }, + { + "model": "models/ggml_llava-v1.5-7b/ggml-model-q4_k.gguf", + "model_alias": "gpt-4-vision-preview", + "chat_format": "llava-1-5", + "clip_model_path": "models/ggml_llava-v1.5-7b/mmproj-model-f16.gguf", + "n_gpu_layers": -1, + "offload_kqv": true, + "n_threads": 12, + "n_batch": 512, + "n_ctx": 2048 + }, + { + "model": "models/mistral-7b-v0.1-GGUF/ggml-model-Q4_K.gguf", + "model_alias": "text-davinci-003", + "n_gpu_layers": -1, + "offload_kqv": true, + "n_threads": 12, + "n_batch": 512, + "n_ctx": 2048 + }, + { + "model": "models/replit-code-v1_5-3b-GGUF/replit-code-v1_5-3b.Q4_0.gguf", + "model_alias": "copilot-codex", + "n_gpu_layers": -1, + "offload_kqv": true, + "n_threads": 12, + "n_batch": 1024, + "n_ctx": 9216 + } + ] +} +``` + +The config file format is defined by the [`ConfigFileSettings`](#llama_cpp.server.settings.ConfigFileSettings) class. + +## Server Options Reference + +::: llama_cpp.server.settings.ConfigFileSettings + options: + show_if_no_docstring: true + +::: llama_cpp.server.settings.ServerSettings + options: + show_if_no_docstring: true + +::: llama_cpp.server.settings.ModelSettings + options: + show_if_no_docstring: true diff --git a/examples/batch-processing/server.py b/examples/batch-processing/server.py new file mode 100644 index 0000000000..2b6fa759e0 --- /dev/null +++ b/examples/batch-processing/server.py @@ -0,0 +1,31 @@ +"""llama-cpp-python server from scratch in a single file. +""" + +# import llama_cpp + +# path = b"../../models/Qwen1.5-0.5B-Chat-GGUF/qwen1_5-0_5b-chat-q8_0.gguf" + +# model_params = llama_cpp.llama_model_default_params() +# model = llama_cpp.llama_model_load_from_file(path, model_params) + +# if model is None: +# raise RuntimeError(f"Failed to load model from file: {path}") + + +# ctx_params = llama_cpp.llama_context_default_params() +# ctx = llama_cpp.llama_init_from_model(model, ctx_params) + +# if ctx is None: +# raise RuntimeError("Failed to create context") + + +from fastapi import FastAPI + +app = FastAPI() + +import openai.types.chat as types + + +@app.post("/v1/chat/completions") +def create_chat_completions(): + return {"message": "Hello World"} diff --git a/examples/colab/Gemma4-12B-QAT.ipynb b/examples/colab/Gemma4-12B-QAT.ipynb new file mode 100644 index 0000000000..2cced5a89b --- /dev/null +++ b/examples/colab/Gemma4-12B-QAT.ipynb @@ -0,0 +1,162 @@ +{ + "nbformat": 4, + "nbformat_minor": 5, + "metadata": { + "colab": { + "provenance": [], + "gpuType": "T4" + }, + "accelerator": "GPU", + "kernelspec": { + "name": "python3", + "display_name": "Python 3" + }, + "language_info": { + "name": "python" + } + }, + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Gemma 4 12B QAT Multimodal Chat\n", + "\n", + "Run the Gemma 4 12B QAT GGUF model locally in Google Colab with the pre-built CUDA wheel for `llama-cpp-python`.\n", + "\n", + "Use a GPU runtime before running this notebook: **Runtime > Change runtime type > T4 GPU**.\n", + "\n", + "Current Colab CUDA images commonly provide CUDA 12 user-space libraries even when `nvidia-smi` reports a CUDA 13-capable driver, so this notebook installs the `cu125` wheel. If your runtime provides `libcudart.so.13`, switch the wheel index URL to `/whl/cu130`.\n" + ], + "outputs": [], + "execution_count": null + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!pip install --no-cache-dir --upgrade --force-reinstall \\\n", + " \"huggingface-hub>=0.23.0\" \\\n", + " llama-cpp-python \\\n", + " --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cu125\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from llama_cpp import Llama\n", + "from llama_cpp.llama_chat_format import Gemma4ChatHandler\n", + "\n", + "MODEL_REPO = \"unsloth/gemma-4-12B-it-qat-GGUF\"\n", + "MODEL_FILE = \"gemma-4-12B-it-qat-UD-Q4_K_XL.gguf\"\n", + "MMPROJ_FILE = \"mmproj-F16.gguf\"\n", + "\n", + "chat_handler = Gemma4ChatHandler.from_pretrained(\n", + " repo_id=MODEL_REPO,\n", + " filename=MMPROJ_FILE,\n", + " verbose=False,\n", + ")\n", + "\n", + "llm = Llama.from_pretrained(\n", + " repo_id=MODEL_REPO,\n", + " filename=MODEL_FILE,\n", + " chat_handler=chat_handler,\n", + " n_gpu_layers=-1,\n", + " n_ctx=8192,\n", + " flash_attn=True,\n", + " verbose=False,\n", + ")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from IPython.display import Image, display\n", + "\n", + "IMAGE_URL = \"https://raw.githubusercontent.com/ggml-org/llama.cpp/master/tools/mtmd/test-1.jpeg\"\n", + "\n", + "display(Image(url=IMAGE_URL, width=320))\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response = llm.create_chat_completion(\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"text\", \"text\": \"Describe this image in one concise sentence.\"},\n", + " {\"type\": \"image_url\", \"image_url\": {\"url\": IMAGE_URL}},\n", + " ],\n", + " }\n", + " ],\n", + " max_tokens=128,\n", + " temperature=0.2,\n", + ")\n", + "\n", + "print(response[\"choices\"][0][\"message\"][\"content\"])\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "tools = [\n", + " {\n", + " \"type\": \"function\",\n", + " \"function\": {\n", + " \"name\": \"record_image_observation\",\n", + " \"description\": \"Record a structured observation about an image.\",\n", + " \"parameters\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"main_subject\": {\"type\": \"string\"},\n", + " \"setting\": {\"type\": \"string\"},\n", + " \"notable_details\": {\n", + " \"type\": \"array\",\n", + " \"items\": {\"type\": \"string\"},\n", + " },\n", + " \"confidence\": {\"type\": \"number\", \"minimum\": 0, \"maximum\": 1},\n", + " },\n", + " \"required\": [\"main_subject\", \"setting\", \"notable_details\", \"confidence\"],\n", + " },\n", + " },\n", + " }\n", + "]\n", + "\n", + "response = llm.create_chat_completion(\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"text\", \"text\": \"Use the provided tool to record a structured observation for this image.\"},\n", + " {\"type\": \"image_url\", \"image_url\": {\"url\": IMAGE_URL}},\n", + " ],\n", + " }\n", + " ],\n", + " tools=tools,\n", + " tool_choice={\"type\": \"function\", \"function\": {\"name\": \"record_image_observation\"}},\n", + " max_tokens=256,\n", + " temperature=0.0,\n", + ")\n", + "\n", + "message = response[\"choices\"][0][\"message\"]\n", + "print(message.get(\"tool_calls\", message.get(\"content\")))\n" + ] + } + ] +} diff --git a/examples/colab/notebook.ipynb b/examples/colab/notebook.ipynb new file mode 100644 index 0000000000..8e258b9c03 --- /dev/null +++ b/examples/colab/notebook.ipynb @@ -0,0 +1,131 @@ +{ + "nbformat": 4, + "nbformat_minor": 5, + "metadata": { + "colab": { + "provenance": [], + "gpuType": "T4" + }, + "accelerator": "GPU", + "kernelspec": { + "name": "python3", + "display_name": "Python 3" + }, + "language_info": { + "name": "python" + } + }, + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Gemma 4 12B Multimodal Chat\n", + "\n", + "Run Gemma 4 12B locally in Google Colab with the pre-built CUDA wheel for `llama-cpp-python`.\n", + "\n", + "Use a GPU runtime before running this notebook: **Runtime > Change runtime type > T4 GPU**.\n", + "\n", + "Current Colab CUDA images commonly provide CUDA 12 user-space libraries even when `nvidia-smi` reports a CUDA 13-capable driver, so this notebook installs the `cu125` wheel. If your runtime provides `libcudart.so.13`, switch the wheel index URL to `/whl/cu130`.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!pip install --no-cache-dir --upgrade --force-reinstall \\\n", + " \"huggingface-hub>=0.23.0\" \\\n", + " llama-cpp-python \\\n", + " --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cu125\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from llama_cpp import Llama\n", + "from llama_cpp.llama_chat_format import Gemma4ChatHandler\n", + "\n", + "MODEL_REPO = \"ggml-org/gemma-4-12B-it-GGUF\"\n", + "MODEL_FILE = \"gemma-4-12B-it-Q4_K_M.gguf\"\n", + "MMPROJ_FILE = \"mmproj-gemma-4-12B-it-Q8_0.gguf\"\n", + "\n", + "chat_handler = Gemma4ChatHandler.from_pretrained(\n", + " repo_id=MODEL_REPO,\n", + " filename=MMPROJ_FILE,\n", + " verbose=False,\n", + ")\n", + "\n", + "llm = Llama.from_pretrained(\n", + " repo_id=MODEL_REPO,\n", + " filename=MODEL_FILE,\n", + " chat_handler=chat_handler,\n", + " n_gpu_layers=-1,\n", + " n_ctx=8192,\n", + " flash_attn=True,\n", + " verbose=False,\n", + ")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response = llm.create_chat_completion(\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": \"What is the capital of France? Answer in one sentence.\",\n", + " }\n", + " ],\n", + " max_tokens=32,\n", + " temperature=0.0,\n", + ")\n", + "\n", + "print(response[\"choices\"][0][\"message\"][\"content\"])\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from IPython.display import Image, display\n", + "\n", + "IMAGE_URL = \"https://raw.githubusercontent.com/ggml-org/llama.cpp/master/tools/mtmd/test-1.jpeg\"\n", + "\n", + "display(Image(url=IMAGE_URL, width=320))\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response = llm.create_chat_completion(\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"text\", \"text\": \"Describe this image in one concise sentence.\"},\n", + " {\"type\": \"image_url\", \"image_url\": {\"url\": IMAGE_URL}},\n", + " ],\n", + " }\n", + " ],\n", + " max_tokens=128,\n", + " temperature=0.2,\n", + ")\n", + "\n", + "print(response[\"choices\"][0][\"message\"][\"content\"])\n" + ] + } + ] +} diff --git a/examples/gradio_chat/local.py b/examples/gradio_chat/local.py new file mode 100644 index 0000000000..871d8b09b5 --- /dev/null +++ b/examples/gradio_chat/local.py @@ -0,0 +1,67 @@ +import llama_cpp +import llama_cpp.llama_tokenizer + +import gradio as gr + +llama = llama_cpp.Llama.from_pretrained( + repo_id="lmstudio-community/Qwen3.5-0.8B-GGUF", + filename="*Q8_0.gguf", + tokenizer=llama_cpp.llama_tokenizer.LlamaHFTokenizer.from_pretrained( + "Qwen/Qwen3.5-0.8B" + ), + verbose=False, +) + +model = "gpt-3.5-turbo" + + +def predict(message, history): + messages = [] + + for user_message, assistant_message in history: + messages.append({"role": "user", "content": user_message}) + messages.append({"role": "assistant", "content": assistant_message}) + + messages.append({"role": "user", "content": message}) + + response = llama.create_chat_completion_openai_v1( + model=model, messages=messages, stream=True + ) + + text = "" + for chunk in response: + content = chunk.choices[0].delta.content + if content: + text += content + yield text + + +js = """function () { + gradioURL = window.location.href + if (!gradioURL.endsWith('?__theme=dark')) { + window.location.replace(gradioURL + '?__theme=dark'); + } +}""" + +css = """ +footer { + visibility: hidden; +} +full-height { + height: 100%; +} +""" + +with gr.Blocks(theme=gr.themes.Soft(), js=js, css=css, fill_height=True) as demo: + gr.ChatInterface( + predict, + fill_height=True, + examples=[ + "What is the capital of France?", + "Who was the first person on the moon?", + ], + ) + + +if __name__ == "__main__": + demo.launch() diff --git a/examples/gradio_chat/server.py b/examples/gradio_chat/server.py new file mode 100644 index 0000000000..52061bea67 --- /dev/null +++ b/examples/gradio_chat/server.py @@ -0,0 +1,59 @@ +import gradio as gr + +from openai import OpenAI + +client = OpenAI(base_url="http://localhost:8000/v1", api_key="llama.cpp") + +model = "gpt-3.5-turbo" + + +def predict(message, history): + messages = [] + + for user_message, assistant_message in history: + messages.append({"role": "user", "content": user_message}) + messages.append({"role": "assistant", "content": assistant_message}) + + messages.append({"role": "user", "content": message}) + + response = client.chat.completions.create( + model=model, messages=messages, stream=True + ) + + text = "" + for chunk in response: + content = chunk.choices[0].delta.content + if content: + text += content + yield text + + +js = """function () { + gradioURL = window.location.href + if (!gradioURL.endsWith('?__theme=dark')) { + window.location.replace(gradioURL + '?__theme=dark'); + } +}""" + +css = """ +footer { + visibility: hidden; +} +full-height { + height: 100%; +} +""" + +with gr.Blocks(theme=gr.themes.Soft(), js=js, css=css, fill_height=True) as demo: + gr.ChatInterface( + predict, + fill_height=True, + examples=[ + "What is the capital of France?", + "Who was the first person on the moon?", + ], + ) + + +if __name__ == "__main__": + demo.launch() diff --git a/examples/hf_pull/main.py b/examples/hf_pull/main.py new file mode 100644 index 0000000000..a9ca424d19 --- /dev/null +++ b/examples/hf_pull/main.py @@ -0,0 +1,36 @@ +import llama_cpp +import llama_cpp.llama_tokenizer + + +llama = llama_cpp.Llama.from_pretrained( + repo_id="lmstudio-community/Qwen3.5-0.8B-GGUF", + filename="*Q8_0.gguf", + tokenizer=llama_cpp.llama_tokenizer.LlamaHFTokenizer.from_pretrained( + "Qwen/Qwen3.5-0.8B" + ), + verbose=False, +) + +response = llama.create_chat_completion( + messages=[{"role": "user", "content": "What is the capital of France?"}], + response_format={ + "type": "json_object", + "schema": { + "type": "object", + "properties": { + "country": {"type": "string"}, + "capital": {"type": "string"}, + }, + "required": ["country", "capital"], + }, + }, + stream=True, +) + +for chunk in response: + delta = chunk["choices"][0]["delta"] + if "content" not in delta: + continue + print(delta["content"], end="", flush=True) + +print() diff --git a/examples/high_level_api/fastapi_server.py b/examples/high_level_api/fastapi_server.py index 4b3189dd1a..ee59767d65 100644 --- a/examples/high_level_api/fastapi_server.py +++ b/examples/high_level_api/fastapi_server.py @@ -9,7 +9,7 @@ Then run: ``` -uvicorn llama_cpp.server.app:app --reload +uvicorn --factory llama_cpp.server.app:create_app --reload ``` or @@ -24,6 +24,7 @@ To actually see the implementation of the server, see llama_cpp/server/app.py """ + import os import uvicorn diff --git a/examples/high_level_api/high_level_api_infill.py b/examples/high_level_api/high_level_api_infill.py new file mode 100644 index 0000000000..282333e5a8 --- /dev/null +++ b/examples/high_level_api/high_level_api_infill.py @@ -0,0 +1,37 @@ +import argparse + +from llama_cpp import Llama + +parser = argparse.ArgumentParser() +parser.add_argument("-m", "--model", type=str, default="../models/7B/ggml-models.bin") +parser.add_argument("-p", "--prompt", type=str, default="def add(") +parser.add_argument("-s", "--suffix", type=str, default="\n return sum\n\n") +parser.add_argument("-i", "--spm-infill", action="store_true") +args = parser.parse_args() + +llm = Llama(model_path=args.model, n_gpu_layers=-1, spm_infill=args.spm_infill) + +output = llm.create_completion( + temperature=0.0, + repeat_penalty=1.0, + prompt=args.prompt, + suffix=args.suffix, +) + +# Models sometimes repeat suffix in response, attempt to filter that +response = output["choices"][0]["text"] +response_stripped = response.rstrip() +unwanted_response_suffix = args.suffix.rstrip() +unwanted_response_length = len(unwanted_response_suffix) + +filtered = False +if ( + unwanted_response_suffix + and response_stripped[-unwanted_response_length:] == unwanted_response_suffix +): + response = response_stripped[:-unwanted_response_length] + filtered = True + +print( + f"Fill-in-Middle completion{' (filtered)' if filtered else ''}:\n\n{args.prompt}\033[32m{response}\033[{'33' if filtered else '0'}m{args.suffix}\033[0m" +) diff --git a/examples/low_level_api/Chat.py b/examples/low_level_api/Chat.py index fcef8cd800..a755089b2a 100644 --- a/examples/low_level_api/Chat.py +++ b/examples/low_level_api/Chat.py @@ -3,10 +3,12 @@ from common import GptParams from low_level_api_chat_cpp import LLaMAInteract + def env_or_def(env, default): - if (env in os.environ): - return os.environ[env] - return default + if env in os.environ: + return os.environ[env] + return default + AI_NAME = env_or_def("AI_NAME", "ChatLLaMa") MODEL = env_or_def("MODEL", "./models/llama-13B/ggml-model.bin") @@ -15,10 +17,10 @@ def env_or_def(env, default): N_THREAD = int(env_or_def("N_THREAD", "8")) today = datetime.datetime.today() -DATE_YEAR=today.strftime("%Y") -DATE_TIME=today.strftime("%H:%M") +DATE_YEAR = today.strftime("%Y") +DATE_TIME = today.strftime("%H:%M") -prompt=f"""Text transcript of a never ending dialog, where {USER_NAME} interacts with an AI assistant named {AI_NAME}. +prompt = f"""Text transcript of a never ending dialog, where {USER_NAME} interacts with an AI assistant named {AI_NAME}. {AI_NAME} is helpful, kind, honest, friendly, good at writing and never fails to answer {USER_NAME}'s requests immediately and with details and precision. There are no annotations like (30 seconds passed...) or (to himself), just what {USER_NAME} and {AI_NAME} say aloud to each other. The dialog lasts for years, the entirety of it is shared below. It's 10000 pages long. @@ -45,27 +47,29 @@ def env_or_def(env, default): {AI_NAME}: Blue. {USER_NAME}: What time is it? {AI_NAME}: It is {DATE_TIME}. -{USER_NAME}:""" + " ".join(sys.argv[1:]) +{USER_NAME}:""" + " ".join( + sys.argv[1:] +) print("Loading model...") params = GptParams( - n_ctx=2048, - temp=0.7, - top_k=40, - top_p=0.5, - repeat_last_n=256, - n_batch=1024, - repeat_penalty=1.17647, - model=MODEL, - n_threads=N_THREAD, - n_predict=N_PREDICTS, - use_color=True, - interactive=True, - antiprompt=[f"{USER_NAME}:"], - input_prefix=" ", - input_suffix=f"{AI_NAME}:", - prompt=prompt, + n_ctx=2048, + temp=0.7, + top_k=40, + top_p=0.5, + repeat_last_n=256, + n_batch=1024, + repeat_penalty=1.17647, + model=MODEL, + n_threads=N_THREAD, + n_predict=N_PREDICTS, + use_color=True, + interactive=True, + antiprompt=[f"{USER_NAME}:"], + input_prefix=" ", + input_suffix=f"{AI_NAME}:", + prompt=prompt, ) with LLaMAInteract(params) as m: - m.interact() + m.interact() diff --git a/examples/low_level_api/Miku.py b/examples/low_level_api/Miku.py index eb9a2cfa91..e072ab1b10 100644 --- a/examples/low_level_api/Miku.py +++ b/examples/low_level_api/Miku.py @@ -3,10 +3,12 @@ from common import GptParams from low_level_api_chat_cpp import LLaMAInteract + def env_or_def(env, default): - if (env in os.environ): - return os.environ[env] - return default + if env in os.environ: + return os.environ[env] + return default + AI_NAME = env_or_def("AI_NAME", "Miku") MODEL = env_or_def("MODEL", "./models/llama-13B/ggml-model.bin") @@ -14,7 +16,7 @@ def env_or_def(env, default): N_PREDICTS = int(env_or_def("N_PREDICTS", "4096")) N_THREAD = int(env_or_def("N_THREAD", "0")) -prompt=f"""This is a transcript of a 1000 page, never ending conversation between {USER_NAME} and the cute and helpful AI assistant {AI_NAME}. {AI_NAME} is a girl who is an AI running on the users computer. +prompt = f"""This is a transcript of a 1000 page, never ending conversation between {USER_NAME} and the cute and helpful AI assistant {AI_NAME}. {AI_NAME} is a girl who is an AI running on the users computer. {AI_NAME} can think for herself without the user seeing her thoughts by adding a /think prefix to her output. She uses this to reason about the world and to think about what she should say next. {AI_NAME} is always coherent and makes sense, but if she isn't sure if what she is saying is correct she will ask the user for help. {AI_NAME} is a very helpful AI and will help the user with anything they need, she is also very friendly and will try to make the user feel better if they are sad. @@ -32,28 +34,30 @@ def env_or_def(env, default): {AI_NAME}: /think It sounds like {USER_NAME} is happy to have me as their assistant! I'm so happy too! ^_^ Glad that whole emotion thing didn't scare him off! {AI_NAME}: /think I wonder what {USER_NAME} likes to do in his free time? I should ask him about that! {AI_NAME}: What do you like to do in your free time? ^_^ -{USER_NAME}:""" + " ".join(sys.argv[1:]) +{USER_NAME}:""" + " ".join( + sys.argv[1:] +) print("Loading model...") params = GptParams( - n_batch=1024, - n_ctx=2048, - n_keep=-1, - repeat_last_n=256, - repeat_penalty=1.17647, - temp=0.7, - top_k=40, - top_p=0.5, - model=MODEL, - n_predict=N_PREDICTS, - use_color=True, - interactive=True, - antiprompt=[f"{USER_NAME}:"], - prompt=prompt, + n_batch=1024, + n_ctx=2048, + n_keep=-1, + repeat_last_n=256, + repeat_penalty=1.17647, + temp=0.7, + top_k=40, + top_p=0.5, + model=MODEL, + n_predict=N_PREDICTS, + use_color=True, + interactive=True, + antiprompt=[f"{USER_NAME}:"], + prompt=prompt, ) if N_THREAD > 0: - params.n_threads = N_THREAD + params.n_threads = N_THREAD with LLaMAInteract(params) as m: - m.interact() + m.interact() diff --git a/examples/low_level_api/ReasonAct.py b/examples/low_level_api/ReasonAct.py index 82e5c4487e..1f2c590177 100644 --- a/examples/low_level_api/ReasonAct.py +++ b/examples/low_level_api/ReasonAct.py @@ -3,14 +3,16 @@ from common import GptParams from low_level_api_chat_cpp import LLaMAInteract + def env_or_def(env, default): - if (env in os.environ): - return os.environ[env] - return default + if env in os.environ: + return os.environ[env] + return default + MODEL = env_or_def("MODEL", "./models/llama-13B/ggml-model.bin") -prompt=f"""You run in a loop of Thought, Action, Observation. +prompt = f"""You run in a loop of Thought, Action, Observation. At the end of the loop either Answer or restate your Thought and Action. Use Thought to describe your thoughts about the question you have been asked. Use Action to run one of these actions available to you: @@ -27,23 +29,25 @@ def env_or_def(env, default): Question: What is capital of france? Thought: Do I need to use an action? No, I know the answer Answer: Paris is the capital of France -Question:""" + " ".join(sys.argv[1:]) +Question:""" + " ".join( + sys.argv[1:] +) print("Loading model...") params = GptParams( - interactive=True, - interactive_start=True, - top_k=10000, - temp=0.2, - repeat_penalty=1, - n_threads=7, - n_ctx=2048, - antiprompt=["Question:","Observation:"], - model=MODEL, - input_prefix=" ", - n_predict=-1, - prompt=prompt, + interactive=True, + interactive_start=True, + top_k=10000, + temp=0.2, + repeat_penalty=1, + n_threads=7, + n_ctx=2048, + antiprompt=["Question:", "Observation:"], + model=MODEL, + input_prefix=" ", + n_predict=-1, + prompt=prompt, ) with LLaMAInteract(params) as m: - m.interact() + m.interact() diff --git a/examples/low_level_api/common.py b/examples/low_level_api/common.py index 55d08db5f5..a0212ff0d2 100644 --- a/examples/low_level_api/common.py +++ b/examples/low_level_api/common.py @@ -65,82 +65,242 @@ class GptParams: # Set to "\nUser:" etc. # This is an alternative to input_prefix which always adds it, so it potentially duplicates "User:"" fix_prefix: str = "" - input_echo: bool = True, + input_echo: bool = (True,) # Default instructions for Alpaca # switch to "Human" and "Assistant" for Vicuna. # TODO: TBD how they are gonna handle this upstream - instruct_inp_prefix: str="\n\n### Instruction:\n\n" - instruct_inp_suffix: str="\n\n### Response:\n\n" + instruct_inp_prefix: str = "\n\n### Instruction:\n\n" + instruct_inp_suffix: str = "\n\n### Response:\n\n" -def gpt_params_parse(argv = None): - parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) - parser.add_argument("-s", "--seed", type=int, default=-1, help="RNG seed (use random seed for <= 0)",dest="seed") - parser.add_argument("-t", "--threads", type=int, default=min(4, os.cpu_count() or 1), help="number of threads to use during computation",dest="n_threads") - parser.add_argument("-n", "--n_predict", type=int, default=128, help="number of tokens to predict (-1 = infinity)",dest="n_predict") - parser.add_argument("--n_parts", type=int, default=-1, help="number of model parts", dest="n_parts") - parser.add_argument("-c", "--ctx_size", type=int, default=512, help="size of the prompt context",dest="n_ctx") - parser.add_argument("-b", "--batch_size", type=int, default=8, help="batch size for prompt processing",dest="n_batch") - parser.add_argument("--keep", type=int, default=0, help="number of tokens to keep from the initial prompt",dest="n_keep") +def gpt_params_parse(argv=None): + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + parser.add_argument( + "-s", + "--seed", + type=int, + default=-1, + help="RNG seed (use random seed for <= 0)", + dest="seed", + ) + parser.add_argument( + "-t", + "--threads", + type=int, + default=min(4, os.cpu_count() or 1), + help="number of threads to use during computation", + dest="n_threads", + ) + parser.add_argument( + "-n", + "--n_predict", + type=int, + default=128, + help="number of tokens to predict (-1 = infinity)", + dest="n_predict", + ) + parser.add_argument( + "--n_parts", type=int, default=-1, help="number of model parts", dest="n_parts" + ) + parser.add_argument( + "-c", + "--ctx_size", + type=int, + default=512, + help="size of the prompt context", + dest="n_ctx", + ) + parser.add_argument( + "-b", + "--batch_size", + type=int, + default=8, + help="batch size for prompt processing", + dest="n_batch", + ) + parser.add_argument( + "--keep", + type=int, + default=0, + help="number of tokens to keep from the initial prompt", + dest="n_keep", + ) parser.add_argument( "-l", "--logit-bias", type=str, - action='append', + action="append", help="--logit-bias TOKEN_ID(+/-)BIAS", - dest="logit_bias_str" - ) - parser.add_argument("--ignore-eos", action="store_true", help="ignore end of stream token and continue generating", dest="ignore_eos") - parser.add_argument("--top_k", type=int, default=40, help="top-k sampling",dest="top_k") - parser.add_argument("--top_p", type=float, default=0.95, help="top-p samplin",dest="top_p") - parser.add_argument("--tfs", type=float, default=1.0, help="tail free sampling, parameter z (1.0 = disabled)",dest="tfs_z") - parser.add_argument("--temp", type=float, default=0.80, help="temperature",dest="temp") - parser.add_argument("--repeat_penalty", type=float, default=1.10, help="penalize repeat sequence of tokens",dest="repeat_penalty") - parser.add_argument("--repeat_last_n", type=int, default=64, help="last n tokens to consider for penalize ",dest="repeat_last_n") - parser.add_argument("--frequency_penalty", type=float, default=0.0, help="repeat alpha frequency penalty (0.0 = disabled)",dest="tfs_z") - parser.add_argument("--presence_penalty", type=float, default=0.0, help="repeat alpha presence penalty (0.0 = disabled)",dest="presence_penalty") - parser.add_argument("--mirostat", type=float, default=1.0, help="use Mirostat sampling.",dest="mirostat") - parser.add_argument("--mirostat_ent", type=float, default=5.0, help="Mirostat target entropy, parameter tau represents the average surprise value",dest="mirostat_tau") - parser.add_argument("--mirostat_lr", type=float, default=0.1, help="Mirostat learning rate, parameter eta",dest="mirostat_eta") - - parser.add_argument("-m", "--model", type=str, default="./models/llama-7B/ggml-model.bin", help="model path",dest="model") - parser.add_argument("-p", "--prompt", type=str, default="", help="initial prompt",dest="prompt") - parser.add_argument("-f", "--file", type=str, default=None, help="file containing initial prompt to load",dest="file") - parser.add_argument("--session", type=str, default=None, help="file to cache model state in (may be large!)",dest="path_session") - parser.add_argument("--in-prefix", type=str, default="", help="string to prefix user inputs with", dest="input_prefix") - parser.add_argument("--in-suffix", type=str, default="", help="append to input", dest="input_suffix") + dest="logit_bias_str", + ) + parser.add_argument( + "--ignore-eos", + action="store_true", + help="ignore end of stream token and continue generating", + dest="ignore_eos", + ) + parser.add_argument( + "--top_k", type=int, default=40, help="top-k sampling", dest="top_k" + ) + parser.add_argument( + "--top_p", type=float, default=0.95, help="top-p samplin", dest="top_p" + ) + parser.add_argument( + "--tfs", + type=float, + default=1.0, + help="tail free sampling, parameter z (1.0 = disabled)", + dest="tfs_z", + ) + parser.add_argument( + "--temp", type=float, default=0.80, help="temperature", dest="temp" + ) + parser.add_argument( + "--repeat_penalty", + type=float, + default=1.10, + help="penalize repeat sequence of tokens", + dest="repeat_penalty", + ) + parser.add_argument( + "--repeat_last_n", + type=int, + default=64, + help="last n tokens to consider for penalize ", + dest="repeat_last_n", + ) + parser.add_argument( + "--frequency_penalty", + type=float, + default=0.0, + help="repeat alpha frequency penalty (0.0 = disabled)", + dest="tfs_z", + ) + parser.add_argument( + "--presence_penalty", + type=float, + default=0.0, + help="repeat alpha presence penalty (0.0 = disabled)", + dest="presence_penalty", + ) + parser.add_argument( + "--mirostat", + type=float, + default=1.0, + help="use Mirostat sampling.", + dest="mirostat", + ) + parser.add_argument( + "--mirostat_ent", + type=float, + default=5.0, + help="Mirostat target entropy, parameter tau represents the average surprise value", + dest="mirostat_tau", + ) + parser.add_argument( + "--mirostat_lr", + type=float, + default=0.1, + help="Mirostat learning rate, parameter eta", + dest="mirostat_eta", + ) + + parser.add_argument( + "-m", + "--model", + type=str, + default="./models/llama-7B/ggml-model.bin", + help="model path", + dest="model", + ) + parser.add_argument( + "-p", "--prompt", type=str, default=None, help="initial prompt", dest="prompt" + ) + parser.add_argument( + "-f", + "--file", + type=str, + default=None, + help="file containing initial prompt to load", + dest="file", + ) + parser.add_argument( + "--session", + type=str, + default=None, + help="file to cache model state in (may be large!)", + dest="path_session", + ) + parser.add_argument( + "--in-prefix", + type=str, + default="", + help="string to prefix user inputs with", + dest="input_prefix", + ) + parser.add_argument( + "--in-suffix", type=str, default="", help="append to input", dest="input_suffix" + ) parser.add_argument( "-r", "--reverse-prompt", type=str, - action='append', + action="append", help="poll user input upon seeing PROMPT (can be\nspecified more than once for multiple prompts).", - dest="antiprompt" + dest="antiprompt", + ) + + parser.add_argument( + "--lora", + type=str, + default="", + help="apply LoRA adapter (implies --no-mmap)", + dest="lora_adapter", + ) + parser.add_argument( + "--lora-base", + type=str, + default="", + help="optional model to use as a base for the layers modified by the LoRA adapter", + dest="lora_base", ) - - parser.add_argument("--lora", type=str, default="", help="apply LoRA adapter (implies --no-mmap)", dest="lora_adapter") - parser.add_argument("--lora-base", type=str, default="", help="optional model to use as a base for the layers modified by the LoRA adapter", dest="lora_base") - parser.add_argument("--memory_f32", action="store_false", help="use f32 instead of f16 for memory key+value",dest="memory_f16") - parser.add_argument("--random-prompt", action="store_true", help="start with a randomized prompt.", dest="random_prompt") + parser.add_argument( + "--memory_f32", + action="store_false", + help="use f32 instead of f16 for memory key+value", + dest="memory_f16", + ) + parser.add_argument( + "--random-prompt", + action="store_true", + help="start with a randomized prompt.", + dest="random_prompt", + ) parser.add_argument( "--color", action="store_true", help="colorise output to distinguish prompt and user input from generations", - dest="use_color" + dest="use_color", ) parser.add_argument( - "-i", "--interactive", action="store_true", help="run in interactive mode", dest="interactive" + "-i", + "--interactive", + action="store_true", + help="run in interactive mode", + dest="interactive", ) - + parser.add_argument("--embedding", action="store_true", help="", dest="embedding") parser.add_argument( "--interactive-first", action="store_true", help="run in interactive mode and wait for input right away", - dest="interactive_start" + dest="interactive_start", ) parser.add_argument( @@ -148,42 +308,84 @@ def gpt_params_parse(argv = None): "--instruct", action="store_true", help="run in instruction mode (use with Alpaca or Vicuna models)", - dest="instruct" + dest="instruct", + ) + parser.add_argument( + "--no-penalize-nl", + action="store_false", + help="do not penalize newline token", + dest="penalize_nl", + ) + parser.add_argument( + "--perplexity", + action="store_true", + help="compute perplexity over the prompt", + dest="perplexity", + ) + parser.add_argument( + "--no-mmap", + action="store_false", + help="do not memory-map model (slower load but may reduce pageouts if not using mlock)", + dest="use_mmap", + ) + parser.add_argument( + "--mlock", + action="store_true", + help="force system to keep model in RAM rather than swapping or compressing", + dest="use_mlock", + ) + parser.add_argument( + "--mtest", + action="store_true", + help="compute maximum memory usage", + dest="mem_test", + ) + parser.add_argument( + "--verbose-prompt", + action="store_true", + help="print prompt before generation", + dest="verbose_prompt", ) - parser.add_argument("--no-penalize-nl", action="store_false", help="do not penalize newline token", dest="penalize_nl") - parser.add_argument("--perplexity", action="store_true", help="compute perplexity over the prompt", dest="perplexity") - parser.add_argument("--no-mmap", action="store_false",help="do not memory-map model (slower load but may reduce pageouts if not using mlock)",dest="use_mmap") - parser.add_argument("--mlock", action="store_true",help="force system to keep model in RAM rather than swapping or compressing",dest="use_mlock") - parser.add_argument("--mtest", action="store_true",help="compute maximum memory usage",dest="mem_test") - parser.add_argument("--verbose-prompt", action="store_true",help="print prompt before generation",dest="verbose_prompt") - #Custom args - parser.add_argument("--fix-prefix", type=str, default="", help="append to input when generated n_predict tokens", dest="fix_prefix") - parser.add_argument("--input-noecho", action="store_false", help="dont output the input", dest="input_echo") + # Custom args + parser.add_argument( + "--fix-prefix", + type=str, + default="", + help="append to input when generated n_predict tokens", + dest="fix_prefix", + ) + parser.add_argument( + "--input-noecho", + action="store_false", + help="dont output the input", + dest="input_echo", + ) parser.add_argument( "--interactive-start", action="store_true", help="run in interactive mode", - dest="interactive" + dest="interactive", ) args = parser.parse_args(argv) - + logit_bias_str = args.logit_bias_str delattr(args, "logit_bias_str") params = GptParams(**vars(args)) - if (params.lora_adapter): + if params.lora_adapter: params.use_mmap = False - if (logit_bias_str != None): + if logit_bias_str != None: for i in logit_bias_str: - if (m := re.match(r"(\d+)([-+]\d+)", i)): + if m := re.match(r"(\d+)([-+]\d+)", i): params.logit_bias[int(m.group(1))] = float(m.group(2)) return params + def gpt_random_prompt(rng): return [ "So", @@ -198,5 +400,6 @@ def gpt_random_prompt(rng): "They", ][rng % 10] + if __name__ == "__main__": print(gpt_params_parse()) diff --git a/examples/low_level_api/low_level_api_chat_cpp.py b/examples/low_level_api/low_level_api_chat_cpp.py index f5d51a36ed..20f7a158ac 100644 --- a/examples/low_level_api/low_level_api_chat_cpp.py +++ b/examples/low_level_api/low_level_api_chat_cpp.py @@ -10,6 +10,7 @@ You should also still be feeding the model with a "primer" prompt that shows it the expected format. """ + import ctypes import sys from time import time @@ -19,188 +20,260 @@ from common import GptParams, gpt_params_parse, gpt_random_prompt import util + # A LLaMA interactive session class LLaMAInteract: - def __init__(self, params: GptParams) -> None: - # input args - self.params = params - - if (self.params.perplexity): - raise NotImplementedError("""************ + def __init__(self, params: GptParams) -> None: + # input args + self.params = params + if self.params.path_session is None: + self.params.path_session = "" + if self.params.antiprompt is None: + self.params.antiprompt = "" + + if self.params.perplexity: + raise NotImplementedError( + """************ please use the 'perplexity' tool for perplexity calculations -************""") +************""" + ) - if (self.params.embedding): - raise NotImplementedError("""************ + if self.params.embedding: + raise NotImplementedError( + """************ please use the 'embedding' tool for embedding calculations -************""") +************""" + ) - if (self.params.n_ctx > 2048): - print(f"""warning: model does not support \ + if self.params.n_ctx > 2048: + print( + f"""warning: model does not support \ context sizes greater than 2048 tokens ({self.params.n_ctx} \ -specified) expect poor results""", file=sys.stderr) - - if (self.params.seed <= 0): - self.params.seed = int(time()) - - print(f"seed = {self.params.seed}", file=sys.stderr) - - if (self.params.random_prompt): - self.params.prompt = gpt_random_prompt(self.params.seed) - - # runtime args - self.input_consumed = 0 - self.n_past = 0 - self.n_session_consumed = 0 - self.first_antiprompt = [] - self.remaining_tokens = self.params.n_predict - self.output_echo = self.params.input_echo - self.multibyte_fix = [] - - # model load - self.lparams = llama_cpp.llama_context_default_params() - self.lparams.n_ctx = self.params.n_ctx - self.lparams.n_parts = self.params.n_parts - self.lparams.seed = self.params.seed - self.lparams.memory_f16 = self.params.memory_f16 - self.lparams.use_mlock = self.params.use_mlock - self.lparams.use_mmap = self.params.use_mmap - - self.ctx = llama_cpp.llama_init_from_file(self.params.model.encode("utf8"), self.lparams) - if (not self.ctx): - raise RuntimeError(f"error: failed to load model '{self.params.model}'") - - if (self.params.ignore_eos): - self.params.logit_bias[llama_cpp.llama_token_eos()] = -float("inf") - - if (len(self.params.lora_adapter) > 0): - if (llama_cpp.llama_apply_lora_from_file( - self.ctx, - self.params.lora_adapter.encode("utf8"), - self.params.lora_base.encode("utf8") if len(self.params.lora_base) > 0 else None, - self.params.n_threads - ) != 0): - print("error: failed to apply lora adapter") - return - - print(file=sys.stderr) - print(f"system_info: n_threads = {self.params.n_threads} / {cpu_count()} \ -| {llama_cpp.llama_print_system_info().decode('utf8')}", file=sys.stderr) - - # determine the required inference memory per token: - if (self.params.mem_test): - tmp = [0, 1, 2, 3] - llama_cpp.llama_eval(self.ctx, (llama_cpp.c_int * len(tmp))(*tmp), len(tmp), 0, self.n_threads) - llama_cpp.llama_print_timings(self.ctx) - self.exit() - return - - # create internal context - self.n_ctx = llama_cpp.llama_n_ctx(self.ctx) - - # Add a space in front of the first character to match OG llama tokenizer behavior - self.params.prompt = " " + self.params.prompt - - # Load prompt file - if (self.params.file): - with open(self.params.file) as f: - self.params.prompt = f.read() - - self.session_tokens: list[llama_cpp.llama_token] = [] - if (len(self.params.path_session) > 0): - print(f"attempting to load saved session from '{self.params.path_session}'", file=sys.stderr) - - if (path.exists(self.params.path_session)): - _session_tokens = (llama_cpp.llama_token * (self.params.n_ctx))() - _n_token_count_out = llama_cpp.c_size_t() - if (llama_cpp.llama_load_session_file( - self.ctx, - self.params.path_session.encode("utf8"), - _session_tokens, - self.params.n_ctx, - ctypes.byref(_n_token_count_out) - ) != 1): - print(f"error: failed to load session file '{self.params.path_session}'", file=sys.stderr) - return - _n_token_count_out = _n_token_count_out.value - self.session_tokens = _session_tokens[:_n_token_count_out] - print(f"loaded a session with prompt size of {_n_token_count_out} tokens", file=sys.stderr) - else: - print(f"session file does not exist, will create", file=sys.stderr) - - # tokenize the prompt - self.embd = [] - self.embd_inp = self._tokenize(self.params.prompt) - - if (len(self.embd_inp) > self.n_ctx - 4): - raise RuntimeError(f"error: prompt is too long ({len(self.embd_inp)} tokens, max {self.params.n_ctx - 4})") - - # debug message about similarity of saved session, if applicable - self.n_matching_session_tokens = 0 - if len(self.session_tokens) > 0: - for id in self.session_tokens: - if self.n_matching_session_tokens >= len(self.embd_inp) or id != self.embd_inp[self.n_matching_session_tokens]: - break - self.n_matching_session_tokens += 1 - - if self.n_matching_session_tokens >= len(self.embd_inp): - print(f"session file has exact match for prompt!") - elif self.n_matching_session_tokens < (len(self.embd_inp) / 2): - print(f"warning: session file has low similarity to prompt ({self.n_matching_session_tokens} / {len(self.embd_inp)} tokens); will mostly be reevaluated") - else: - print(f"session file matches {self.n_matching_session_tokens} / {len(self.embd_inp)} tokens of prompt") - - self.need_to_save_session = len(self.params.path_session) > 0 and self.n_matching_session_tokens < (len(self.embd_inp) * 3 / 4) - - # number of tokens to keep when resetting context - if (self.params.n_keep < 0 or self.params.n_keep > len(self.embd_inp) or self.params.instruct): - self.params.n_keep = len(self.embd_inp) - - self.inp_prefix = self._tokenize(self.params.instruct_inp_prefix) - self.inp_suffix = self._tokenize(self.params.instruct_inp_suffix, False) - - # in instruct mode, we inject a prefix and a suffix to each input by the user - self.antiecho = None - if (self.params.instruct): - self.params.interactive_start = True - _ptn = self._tokenize(self.params.instruct_inp_prefix.strip(), False) - self.first_antiprompt.append(_ptn) - self.antiecho = util.IterSearch(_ptn) - - # enable interactive mode if reverse prompt or interactive start is specified - if (len(self.params.antiprompt) != 0 or self.params.interactive_start): - self.params.interactive = True - - # determine newline token - self.llama_token_newline = self._tokenize("\n", False) - self.llama_token_eot = self._tokenize(" [end of text]\n", False) - - if (self.params.verbose_prompt): - print(f""" +specified) expect poor results""", + file=sys.stderr, + ) + + if self.params.seed <= 0: + self.params.seed = int(time()) + + print(f"seed = {self.params.seed}", file=sys.stderr) + + if self.params.random_prompt: + self.params.prompt = gpt_random_prompt(self.params.seed) + + # runtime args + self.input_consumed = 0 + self.n_past = 0 + self.n_session_consumed = 0 + self.first_antiprompt = [] + self.remaining_tokens = self.params.n_predict + self.output_echo = self.params.input_echo + self.multibyte_fix = [] + + # model load + self.lparams = llama_cpp.llama_model_default_params() + self.lparams.n_ctx = self.params.n_ctx + self.lparams.n_parts = self.params.n_parts + self.lparams.seed = self.params.seed + self.lparams.memory_f16 = self.params.memory_f16 + self.lparams.use_mlock = self.params.use_mlock + self.lparams.use_mmap = self.params.use_mmap + + self.model = llama_cpp.llama_model_load_from_file( + self.params.model.encode("utf8"), self.lparams + ) + self.vocab = llama_cpp.llama_model_get_vocab(self.model) + + # Context Params. + self.cparams = llama_cpp.llama_context_default_params() + + self.ctx = llama_cpp.llama_init_from_model(self.model, self.cparams) + if not self.ctx: + raise RuntimeError(f"error: failed to load model '{self.params.model}'") + + if self.params.ignore_eos: + self.params.logit_bias[llama_cpp.llama_vocab_eos(self.vocab)] = -float( + "inf" + ) + + if len(self.params.lora_adapter) > 0: + if ( + llama_cpp.llama_apply_lora_from_file( + self.ctx, + self.params.lora_adapter.encode("utf8"), + ( + self.params.lora_base.encode("utf8") + if len(self.params.lora_base) > 0 + else None + ), + self.params.n_threads, + ) + != 0 + ): + print("error: failed to apply lora adapter") + return + + print(file=sys.stderr) + print( + f"system_info: n_threads = {self.params.n_threads} / {cpu_count()} \ +| {llama_cpp.llama_print_system_info().decode('utf8')}", + file=sys.stderr, + ) + + # determine the required inference memory per token: + if self.params.mem_test: + tmp = [0, 1, 2, 3] + llama_cpp.llama_eval( + self.ctx, + (llama_cpp.c_int * len(tmp))(*tmp), + len(tmp), + 0, + self.n_threads, + ) + llama_cpp.llama_print_timings(self.ctx) + self.exit() + return + + # create internal context + self.n_ctx = llama_cpp.llama_n_ctx(self.ctx) + + # Add a space in front of the first character to match OG llama tokenizer behavior + self.params.prompt = " " + self.params.prompt + + # Load prompt file + if self.params.file: + with open(self.params.file) as f: + self.params.prompt = f.read() + + self.session_tokens: list[llama_cpp.llama_token] = [] + if len(self.params.path_session) > 0: + print( + f"attempting to load saved session from '{self.params.path_session}'", + file=sys.stderr, + ) + + if path.exists(self.params.path_session): + _session_tokens = (llama_cpp.llama_token * (self.params.n_ctx))() + _n_token_count_out = llama_cpp.c_size_t() + if ( + llama_cpp.llama_state_load_file( + self.ctx, + self.params.path_session.encode("utf8"), + _session_tokens, + self.params.n_ctx, + ctypes.byref(_n_token_count_out), + ) + != 1 + ): + print( + f"error: failed to load session file '{self.params.path_session}'", + file=sys.stderr, + ) + return + _n_token_count_out = _n_token_count_out.value + self.session_tokens = _session_tokens[:_n_token_count_out] + print( + f"loaded a session with prompt size of {_n_token_count_out} tokens", + file=sys.stderr, + ) + else: + print(f"session file does not exist, will create", file=sys.stderr) + + # tokenize the prompt + self.embd = [] + self.embd_inp = self._tokenize(self.params.prompt) + + if len(self.embd_inp) > self.n_ctx - 4: + raise RuntimeError( + f"error: prompt is too long ({len(self.embd_inp)} tokens, max {self.params.n_ctx - 4})" + ) + + # debug message about similarity of saved session, if applicable + self.n_matching_session_tokens = 0 + if len(self.session_tokens) > 0: + for id in self.session_tokens: + if ( + self.n_matching_session_tokens >= len(self.embd_inp) + or id != self.embd_inp[self.n_matching_session_tokens] + ): + break + self.n_matching_session_tokens += 1 + + if self.n_matching_session_tokens >= len(self.embd_inp): + print(f"session file has exact match for prompt!") + elif self.n_matching_session_tokens < (len(self.embd_inp) / 2): + print( + f"warning: session file has low similarity to prompt ({self.n_matching_session_tokens} / {len(self.embd_inp)} tokens); will mostly be reevaluated" + ) + else: + print( + f"session file matches {self.n_matching_session_tokens} / {len(self.embd_inp)} tokens of prompt" + ) + + self.need_to_save_session = len( + self.params.path_session + ) > 0 and self.n_matching_session_tokens < (len(self.embd_inp) * 3 / 4) + + # number of tokens to keep when resetting context + if ( + self.params.n_keep < 0 + or self.params.n_keep > len(self.embd_inp) + or self.params.instruct + ): + self.params.n_keep = len(self.embd_inp) + + self.inp_prefix = self._tokenize(self.params.instruct_inp_prefix) + self.inp_suffix = self._tokenize(self.params.instruct_inp_suffix, False) + + # in instruct mode, we inject a prefix and a suffix to each input by the user + self.antiecho = None + if self.params.instruct: + self.params.interactive_start = True + _ptn = self._tokenize(self.params.instruct_inp_prefix.strip(), False) + self.first_antiprompt.append(_ptn) + self.antiecho = util.IterSearch(_ptn) + + # enable interactive mode if reverse prompt or interactive start is specified + if len(self.params.antiprompt) != 0 or self.params.interactive_start: + self.params.interactive = True + + # determine newline token + self.llama_token_newline = self._tokenize("\n", False) + self.llama_token_eot = self._tokenize(" [end of text]\n", False) + + if self.params.verbose_prompt: + print( + f""" prompt: '{self.params.prompt}' -number of tokens in prompt = {len(self.embd_inp)}""", file=sys.stderr) - - for i in range(len(self.embd_inp)): - print(f"{self.embd_inp[i]} -> '{llama_cpp.llama_token_to_str(self.ctx, self.embd_inp[i])}'", file=sys.stderr) - - if (self.params.n_keep > 0): - print("static prompt based on n_keep: '") - for i in range(self.params.n_keep): - print(llama_cpp.llama_token_to_str(self.ctx, self.embd_inp[i]), file=sys.stderr) - print("'", file=sys.stderr) - print(file=sys.stderr) - - if (self.params.interactive): - print("interactive mode on.", file=sys.stderr) - - if (len(self.params.antiprompt) > 0): - for antiprompt in self.params.antiprompt: - print(f"Reverse prompt: '{antiprompt}'", file=sys.stderr) - - if len(self.params.input_prefix) > 0: - print(f"Input prefix: '{self.params.input_prefix}'", file=sys.stderr) - - print(f"""sampling: repeat_last_n = {self.params.repeat_last_n},\ +number of tokens in prompt = {len(self.embd_inp)}""", + file=sys.stderr, + ) + + for i in range(len(self.embd_inp)): + print( + f"{self.embd_inp[i]} -> '{self.token_to_str(self.embd_inp[i])}'", + file=sys.stderr, + ) + + if self.params.n_keep > 0: + print("static prompt based on n_keep: '") + for i in range(self.params.n_keep): + print(self.token_to_str(self.embd_inp[i]), file=sys.stderr) + print("'", file=sys.stderr) + print(file=sys.stderr) + + if self.params.interactive: + print("interactive mode on.", file=sys.stderr) + + if len(self.params.antiprompt) > 0: + for antiprompt in self.params.antiprompt: + print(f"Reverse prompt: '{antiprompt}'", file=sys.stderr) + + if len(self.params.input_prefix) > 0: + print(f"Input prefix: '{self.params.input_prefix}'", file=sys.stderr) + + print( + f"""sampling: repeat_last_n = {self.params.repeat_last_n},\ repeat_penalty = {self.params.repeat_penalty},\ presence_penalty = {self.params.presence_penalty},\ frequency_penalty = {self.params.frequency_penalty},\ @@ -218,77 +291,96 @@ def __init__(self, params: GptParams) -> None: n_predict = {self.params.n_predict},\ n_keep = {self.params.n_keep} -""", file=sys.stderr) +""", + file=sys.stderr, + ) - # determine antiprompt tokens - for i in self.params.antiprompt: - self.first_antiprompt.append(self._tokenize(i, False)) + # determine antiprompt tokens + for i in self.params.antiprompt: + self.first_antiprompt.append(self._tokenize(i, False)) - self.last_n_tokens = [0]*self.n_ctx #TODO: deque doesnt support slices + self.last_n_tokens = [0] * self.n_ctx # TODO: deque doesnt support slices - if (params.interactive): - print("""== Running in interactive mode. == + if params.interactive: + print( + """== Running in interactive mode. == - Press Ctrl+C to interject at any time. - Press Return to return control to LLaMa. - If you want to submit another line, end your input in '\\'. -""", file=sys.stderr) - self.set_color(util.CONSOLE_COLOR_PROMPT) - - # tokenize a prompt - def _tokenize(self, prompt, bos=True): - _arr = (llama_cpp.llama_token * ((len(prompt) + 1) * 4))() - _n = llama_cpp.llama_tokenize(self.ctx, prompt.encode("utf8", errors="ignore"), _arr, len(_arr), bos) - return _arr[:_n] - - def set_color(self, c): - if (self.params.use_color): - print(c, end="") - - def use_antiprompt(self): - return len(self.first_antiprompt) > 0 - - # generate tokens - def generate(self): - while self.remaining_tokens > 0 or self.params.interactive or self.params.n_predict == -1: - # predict - if len(self.embd) > 0: - # infinite text generation via context swapping - # if we run out of context: - # - take the n_keep first tokens from the original prompt (via n_past) - # - take half of the last (n_ctx - n_keep) tokens and recompute the logits in a batch - if (self.n_past + len(self.embd) > self.n_ctx): - n_left = self.n_past - self.params.n_keep - self.n_past = self.params.n_keep - - # insert n_left/2 tokens at the start of embd from last_n_tokens - _insert = self.last_n_tokens[ - self.n_ctx - int(n_left/2) - len(self.embd):-len(self.embd) - ] - self.embd = _insert + self.embd - self.params.path_session = "" - - # try to reuse a matching prefix from the loaded session instead of re-eval (via n_past) - if self.n_session_consumed < len(self.session_tokens): - for i in range(len(self.embd)): - if self.embd[i] != self.session_tokens[self.n_session_consumed]: - self.session_tokens = self.session_tokens[:self.n_session_consumed] - break - - self.n_past += 1 - self.n_session_consumed += 1 - - if self.n_session_consumed >= len(self.session_tokens): - i += 1 - break - - if i > 0: - self.embd = self.embd[i:] - - # evaluate tokens in batches - # embd is typically prepared beforehand to fit within a batch, but not always - #TODO BUG: The batching code causes nonsensical generation - """for i in range(0, len(self.embd), self.params.n_batch): +""", + file=sys.stderr, + ) + self.set_color(util.CONSOLE_COLOR_PROMPT) + + # tokenize a prompt + def _tokenize(self, prompt, bos=True): + _arr = (llama_cpp.llama_token * ((len(prompt) + 1) * 4))() + _n = llama_cpp.llama_tokenize( + self.vocab, + prompt.encode("utf8", errors="ignore"), + len(prompt), + _arr, + len(_arr), + bos, + False, + ) + return _arr[:_n] + + def set_color(self, c): + if self.params.use_color: + print(c, end="") + + def use_antiprompt(self): + return len(self.first_antiprompt) > 0 + + # generate tokens + def generate(self): + while ( + self.remaining_tokens > 0 + or self.params.interactive + or self.params.n_predict == -1 + ): + # predict + if len(self.embd) > 0: + # infinite text generation via context swapping + # if we run out of context: + # - take the n_keep first tokens from the original prompt (via n_past) + # - take half of the last (n_ctx - n_keep) tokens and recompute the logits in a batch + if self.n_past + len(self.embd) > self.n_ctx: + n_left = self.n_past - self.params.n_keep + self.n_past = self.params.n_keep + + # insert n_left/2 tokens at the start of embd from last_n_tokens + _insert = self.last_n_tokens[ + self.n_ctx - int(n_left / 2) - len(self.embd) : -len(self.embd) + ] + self.embd = _insert + self.embd + self.params.path_session = "" + + # try to reuse a matching prefix from the loaded session instead of re-eval (via n_past) + if self.n_session_consumed < len(self.session_tokens): + for i in range(len(self.embd)): + if self.embd[i] != self.session_tokens[self.n_session_consumed]: + self.session_tokens = self.session_tokens[ + : self.n_session_consumed + ] + break + + self.n_past += 1 + self.n_session_consumed += 1 + + if self.n_session_consumed >= len(self.session_tokens): + i += 1 + break + + if i > 0: + self.embd = self.embd[i:] + + # evaluate tokens in batches + # embd is typically prepared beforehand to fit within a batch, but not always + # TODO BUG: The batching code causes nonsensical generation + """for i in range(0, len(self.embd), self.params.n_batch): n_eval = self.params.n_batch _arr = (llama_cpp.llama_token * n_eval)(*self.embd[i:i + n_eval]) if llama_cpp.llama_eval(self.ctx, _arr, n_eval, self.n_past, self.params.n_threads) != 0: @@ -297,257 +389,358 @@ def generate(self): self.n_past += n_eval""" - if (llama_cpp.llama_eval( - self.ctx, (llama_cpp.llama_token * len(self.embd))(*self.embd), len(self.embd), self.n_past, self.params.n_threads - ) != 0): - raise Exception("Failed to llama_eval!") - - if len(self.embd) > 0 and len(self.params.path_session) > 0: - self.session_tokens.extend(self.embd) - self.n_session_consumed = len(self.session_tokens) - - self.n_past += len(self.embd) - self.embd = [] - if len(self.embd_inp) <= self.input_consumed: #&& !is_interacting - # out of user input, sample next token - top_k = llama_cpp.llama_n_vocab(self.ctx) if self.params.top_k <= 0 else self.params.top_k - repeat_last_n = self.n_ctx if self.params.repeat_last_n < 0 else self.params.repeat_last_n - - # optionally save the session on first sample (for faster prompt loading next time) - if len(self.params.path_session) > 0 and self.need_to_save_session: - self.need_to_save_session = False - llama_cpp.llama_save_session_file( - self.ctx, - self.params.path_session.encode("utf8"), - (llama_cpp.llama_token * len(self.session_tokens))(*self.session_tokens), - len(self.session_tokens) - ) - - id = 0 - - logits = llama_cpp.llama_get_logits(self.ctx) - n_vocab = llama_cpp.llama_n_vocab(self.ctx) - - # Apply params.logit_bias map - for key, value in self.params.logit_bias.items(): - logits[key] += value - - _arr = (llama_cpp.llama_token_data * n_vocab)(*[ - llama_cpp.llama_token_data(token_id, logits[token_id], 0.0) - for token_id in range(n_vocab) - ]) - candidates_p = llama_cpp.ctypes.pointer(llama_cpp.llama_token_data_array(_arr, len(_arr), False)) - - # Apply penalties - nl_logit = logits[llama_cpp.llama_token_nl()] - last_n_repeat = min(len(self.last_n_tokens), repeat_last_n, self.n_ctx) - - _arr = (llama_cpp.llama_token * last_n_repeat)(*self.last_n_tokens[len(self.last_n_tokens) - last_n_repeat:]) - llama_cpp.llama_sample_repetition_penalty(self.ctx, candidates_p, - _arr, - last_n_repeat, llama_cpp.c_float(self.params.repeat_penalty)) - llama_cpp.llama_sample_frequency_and_presence_penalties(self.ctx, candidates_p, - _arr, - last_n_repeat, llama_cpp.c_float(self.params.frequency_penalty), llama_cpp.c_float(self.params.presence_penalty)) - - if not self.params.penalize_nl: - logits[llama_cpp.llama_token_nl()] = nl_logit - - if self.params.temp <= 0: - # Greedy sampling - id = llama_cpp.llama_sample_token_greedy(self.ctx, candidates_p) - else: - if self.params.mirostat == 1: - mirostat_mu = 2.0 * self.params.mirostat_tau - mirostat_m = 100 - llama_cpp.llama_sample_temperature(self.ctx, candidates_p, llama_cpp.c_float(self.params.temp)) - id = llama_cpp.llama_sample_token_mirostat(self.ctx, candidates_p, llama_cpp.c_float(self.params.mirostat_tau), llama_cpp.c_float(self.params.mirostat_eta), llama_cpp.c_int(mirostat_m), llama_cpp.c_float(mirostat_mu)) - elif self.params.mirostat == 2: - mirostat_mu = 2.0 * self.params.mirostat_tau - llama_cpp.llama_sample_temperature(self.ctx, candidates_p, llama_cpp.c_float(self.params.temp)) - id = llama_cpp.llama_sample_token_mirostat_v2(self.ctx, candidates_p, llama_cpp.c_float(self.params.mirostat_tau), llama_cpp.c_float(self.params.mirostat_eta), llama_cpp.c_float(mirostat_mu)) - else: - # Temperature sampling - llama_cpp.llama_sample_top_k(self.ctx, candidates_p, top_k, min_keep=llama_cpp.c_size_t(1)) - llama_cpp.llama_sample_tail_free(self.ctx, candidates_p, llama_cpp.c_float(self.params.tfs_z), min_keep=llama_cpp.c_size_t(1)) - llama_cpp.llama_sample_typical(self.ctx, candidates_p, llama_cpp.c_float(self.params.typical_p), min_keep=llama_cpp.c_size_t(1)) - llama_cpp.llama_sample_top_p(self.ctx, candidates_p, llama_cpp.c_float(self.params.top_p), min_keep=llama_cpp.c_size_t(1)) - llama_cpp.llama_sample_temperature(self.ctx, candidates_p, llama_cpp.c_float(self.params.temp)) - id = llama_cpp.llama_sample_token(self.ctx, candidates_p) - # print("`{}`".format(candidates_p.size)) - - self.last_n_tokens.pop(0) - self.last_n_tokens.append(id) - - # replace end of text token with newline token when in interactive mode - if (id == llama_cpp.llama_token_eos() and self.params.interactive and not self.params.instruct): - id = self.llama_token_newline[0] - self.embd.append(id) - if (self.use_antiprompt()): - # tokenize and inject first reverse prompt - self.embd_inp += self.first_antiprompt[0] - for id in self.first_antiprompt[0]: - self.embd.append(id) - else: - # add it to the context - self.embd.append(id) - - # echo this to console - self.output_echo = True - - # decrement remaining sampling budget - self.remaining_tokens -= 1 - else: - # output to console if input echo is on - self.output_echo = self.params.input_echo - - # some user input remains from prompt or interaction, forward it to processing - while len(self.embd_inp) > self.input_consumed: - self.embd.append(self.embd_inp[self.input_consumed]) - self.last_n_tokens.pop(0) - self.last_n_tokens.append(self.embd_inp[self.input_consumed]) - self.input_consumed += 1 - if len(self.embd) >= self.params.n_batch: - break - - # display tokens - if self.output_echo: - for id in self.embd: - if self.antiecho != None: - for r in self.antiecho(id): - yield r - else: - yield id - - # reset color to default if we there is no pending user input - if (self.params.input_echo and len(self.embd_inp) == self.input_consumed): - self.set_color(util.CONSOLE_COLOR_DEFAULT) - - if (self.params.interactive and len(self.embd_inp) <= self.input_consumed): - # if antiprompt is present, stop - if (self.use_antiprompt()): - if True in [ - i == self.last_n_tokens[-len(i):] - for i in self.first_antiprompt - ]: - break - - # if we are using instruction mode, and we have processed the initial prompt - if (self.params.interactive_start): - break - - # end of text token - if len(self.embd) > 0 and self.embd[-1] == llama_cpp.llama_token_eos(): - if (not self.params.instruct): - for i in self.llama_token_eot: - yield i - break - - # respect n_predict even if antiprompt is present - if (self.params.interactive and self.remaining_tokens <= 0 and self.params.n_predict != -1): - # If we arent in instruction mode, fix the current generation by appending the antiprompt. - # Makes it so if chat ends prematurely you dont append the AI's text etc. - if not self.params.instruct: - self.embd_inp += self.first_antiprompt[0] - self.n_remain = self.params.n_predict - break - - self.params.interactive_start = False - - def __enter__(self): - return self - - def __exit__(self, type, value, tb): - self.exit() - - def exit(self): - llama_cpp.llama_free(self.ctx) - self.set_color(util.CONSOLE_COLOR_DEFAULT) - - # return past text - def past(self): - for id in self.last_n_tokens[-self.n_past:]: - yield llama_cpp.llama_token_to_str(self.ctx, id).decode("utf8", errors="ignore") - - # write input - def input(self, prompt: str): - if (self.params.instruct and self.last_n_tokens[-len(self.inp_prefix):] != self.inp_prefix): - self.embd_inp += self.inp_prefix - self.embd_inp += self._tokenize(prompt) - if (self.params.instruct): - self.embd_inp += self.inp_suffix - - # write output - def output(self): - self.remaining_tokens = self.params.n_predict - for id in self.generate(): - cur_char = llama_cpp.llama_token_to_str(self.ctx, id) - - # Add remainder of missing bytes - if None in self.multibyte_fix: - self.multibyte_fix[self.multibyte_fix.index(None)] = cur_char - - # Return completed utf char - if len(self.multibyte_fix) > 0 and not None in self.multibyte_fix: - yield (b"".join(self.multibyte_fix)).decode("utf8") - self.multibyte_fix = [] - continue - - # Contains multi-byte UTF8 - for num, pattern in [(2, 192), (3, 224), (4, 240)]: - # Bitwise AND check - if pattern & int.from_bytes(cur_char, 'little') == pattern: - self.multibyte_fix = [cur_char] + ([None] * (num-1)) - - # Stop incomplete bytes from passing - if len(self.multibyte_fix) > 0: - continue - - yield cur_char.decode("utf8") - - # read user input - def read_input(self): - out = "" - while (t := input()).endswith("\\"): - out += t[:-1] + "\n" - return out + t + "\n" - - # interactive mode - def interact(self): - for i in self.output(): - print(i,end="",flush=True) - self.params.input_echo = False - - while self.params.interactive: - self.set_color(util.CONSOLE_COLOR_USER_INPUT) - if (self.params.instruct): - print('\n> ', end="") - self.input(self.read_input()) - else: - print(self.params.input_prefix, end="") - self.input(f"{self.params.input_prefix}{self.read_input()}{self.params.input_suffix}") - print(self.params.input_suffix,end="") - self.set_color(util.CONSOLE_COLOR_DEFAULT) - - try: - for i in self.output(): - print(i,end="",flush=True) - except KeyboardInterrupt: - self.set_color(util.CONSOLE_COLOR_DEFAULT) - if not self.params.instruct: - print(self.params.fix_prefix,end="") - self.input(self.params.fix_prefix) + if ( + llama_cpp.llama_eval( + self.ctx, + (llama_cpp.llama_token * len(self.embd))(*self.embd), + len(self.embd), + self.n_past, + ) + != 0 + ): + raise Exception("Failed to llama_eval!") + + if len(self.embd) > 0 and len(self.params.path_session) > 0: + self.session_tokens.extend(self.embd) + self.n_session_consumed = len(self.session_tokens) + + self.n_past += len(self.embd) + self.embd = [] + if len(self.embd_inp) <= self.input_consumed: # && !is_interacting + # out of user input, sample next token + top_k = ( + llama_cpp.llama_vocab_n_tokens(self.vocab) + if self.params.top_k <= 0 + else self.params.top_k + ) + repeat_last_n = ( + self.n_ctx + if self.params.repeat_last_n < 0 + else self.params.repeat_last_n + ) + + # optionally save the session on first sample (for faster prompt loading next time) + if len(self.params.path_session) > 0 and self.need_to_save_session: + self.need_to_save_session = False + llama_cpp.llama_state_save_file( + self.ctx, + self.params.path_session.encode("utf8"), + (llama_cpp.llama_token * len(self.session_tokens))( + *self.session_tokens + ), + len(self.session_tokens), + ) + + id = 0 + + logits = llama_cpp.llama_get_logits(self.ctx) + n_vocab = llama_cpp.llama_vocab_n_tokens(self.vocab) + + # Apply params.logit_bias map + for key, value in self.params.logit_bias.items(): + logits[key] += value + + _arr = (llama_cpp.llama_token_data * n_vocab)( + *[ + llama_cpp.llama_token_data(token_id, logits[token_id], 0.0) + for token_id in range(n_vocab) + ] + ) + candidates_p = llama_cpp.ctypes.pointer( + llama_cpp.llama_token_data_array(_arr, len(_arr), False) + ) + + # Apply penalties + nl_logit = logits[llama_cpp.llama_vocab_nl(self.vocab)] + last_n_repeat = min(len(self.last_n_tokens), repeat_last_n, self.n_ctx) + + _arr = (llama_cpp.llama_token * last_n_repeat)( + *self.last_n_tokens[len(self.last_n_tokens) - last_n_repeat :] + ) + llama_cpp.llama_sample_repetition_penalties( + ctx=self.ctx, + candidates=candidates_p, + last_tokens_data=_arr, + penalty_last_n=last_n_repeat, + penalty_repeat=llama_cpp.c_float(self.params.repeat_penalty), + penalty_freq=llama_cpp.c_float(self.params.frequency_penalty), + penalty_present=llama_cpp.c_float(self.params.presence_penalty), + ) + + # NOT PRESENT IN CURRENT VERSION ? + # llama_cpp.llama_sample_frequency_and_presence_penalti(self.ctx, candidates_p, + # _arr, + # last_n_repeat, llama_cpp.c_float(self.params.frequency_penalty), llama_cpp.c_float(self.params.presence_penalty)) + + if not self.params.penalize_nl: + logits[llama_cpp.llama_vocab_nl(self.vocab)] = nl_logit + + if self.params.temp <= 0: + # Greedy sampling + id = llama_cpp.llama_sample_token_greedy(self.ctx, candidates_p) + else: + if self.params.mirostat == 1: + mirostat_mu = 2.0 * self.params.mirostat_tau + mirostat_m = 100 + llama_cpp.llama_sample_temperature( + self.ctx, candidates_p, llama_cpp.c_float(self.params.temp) + ) + id = llama_cpp.llama_sample_token_mirostat( + self.ctx, + candidates_p, + llama_cpp.c_float(self.params.mirostat_tau), + llama_cpp.c_float(self.params.mirostat_eta), + llama_cpp.c_int(mirostat_m), + llama_cpp.c_float(mirostat_mu), + ) + elif self.params.mirostat == 2: + mirostat_mu = 2.0 * self.params.mirostat_tau + llama_cpp.llama_sample_temperature( + self.ctx, candidates_p, llama_cpp.c_float(self.params.temp) + ) + id = llama_cpp.llama_sample_token_mirostat_v2( + self.ctx, + candidates_p, + llama_cpp.c_float(self.params.mirostat_tau), + llama_cpp.c_float(self.params.mirostat_eta), + llama_cpp.c_float(mirostat_mu), + ) + else: + # Temperature sampling + llama_cpp.llama_sample_top_k( + self.ctx, + candidates_p, + top_k, + min_keep=llama_cpp.c_size_t(1), + ) + llama_cpp.llama_sample_tail_free( + self.ctx, + candidates_p, + llama_cpp.c_float(self.params.tfs_z), + min_keep=llama_cpp.c_size_t(1), + ) + llama_cpp.llama_sample_typical( + self.ctx, + candidates_p, + llama_cpp.c_float(self.params.typical_p), + min_keep=llama_cpp.c_size_t(1), + ) + llama_cpp.llama_sample_top_p( + self.ctx, + candidates_p, + llama_cpp.c_float(self.params.top_p), + min_keep=llama_cpp.c_size_t(1), + ) + llama_cpp.llama_sample_temperature( + self.ctx, candidates_p, llama_cpp.c_float(self.params.temp) + ) + id = llama_cpp.llama_sample_token(self.ctx, candidates_p) + # print("`{}`".format(candidates_p.size)) + + self.last_n_tokens.pop(0) + self.last_n_tokens.append(id) + + # replace end of text token with newline token when in interactive mode + if ( + id == llama_cpp.llama_vocab_eos(self.vocab) + and self.params.interactive + and not self.params.instruct + ): + id = self.llama_token_newline[0] + self.embd.append(id) + if self.use_antiprompt(): + # tokenize and inject first reverse prompt + self.embd_inp += self.first_antiprompt[0] + for id in self.first_antiprompt[0]: + self.embd.append(id) + else: + # add it to the context + self.embd.append(id) + + # echo this to console + self.output_echo = True + + # decrement remaining sampling budget + self.remaining_tokens -= 1 + else: + # output to console if input echo is on + self.output_echo = self.params.input_echo + + # some user input remains from prompt or interaction, forward it to processing + while len(self.embd_inp) > self.input_consumed: + self.embd.append(self.embd_inp[self.input_consumed]) + self.last_n_tokens.pop(0) + self.last_n_tokens.append(self.embd_inp[self.input_consumed]) + self.input_consumed += 1 + if len(self.embd) >= self.params.n_batch: + break + + # display tokens + if self.output_echo: + for id in self.embd: + if self.antiecho != None: + for r in self.antiecho(id): + yield r + else: + yield id + + # reset color to default if we there is no pending user input + if self.params.input_echo and len(self.embd_inp) == self.input_consumed: + self.set_color(util.CONSOLE_COLOR_DEFAULT) + + if self.params.interactive and len(self.embd_inp) <= self.input_consumed: + # if antiprompt is present, stop + if self.use_antiprompt(): + if True in [ + i == self.last_n_tokens[-len(i) :] + for i in self.first_antiprompt + ]: + break + + # if we are using instruction mode, and we have processed the initial prompt + if self.params.interactive_start: + break + + # end of text token + if len(self.embd) > 0 and self.embd[-1] == llama_cpp.llama_vocab_eos( + self.vocab + ): + if not self.params.instruct: + for i in self.llama_token_eot: + yield i + break + + # respect n_predict even if antiprompt is present + if ( + self.params.interactive + and self.remaining_tokens <= 0 + and self.params.n_predict != -1 + ): + # If we arent in instruction mode, fix the current generation by appending the antiprompt. + # Makes it so if chat ends prematurely you dont append the AI's text etc. + if not self.params.instruct: + self.embd_inp += self.first_antiprompt[0] + self.n_remain = self.params.n_predict + break + + self.params.interactive_start = False + + def __enter__(self): + return self + + def __exit__(self, type, value, tb): + self.exit() + + def exit(self): + llama_cpp.llama_free(self.ctx) + self.set_color(util.CONSOLE_COLOR_DEFAULT) + + def token_to_str(self, token_id: int) -> bytes: + size = 32 + buffer = (ctypes.c_char * size)() + n = llama_cpp.llama_token_to_piece( + self.vocab, llama_cpp.llama_token(token_id), buffer, size, 0, False + ) + assert n <= size + return bytes(buffer[:n]) + + # return past text + def past(self): + for id in self.last_n_tokens[-self.n_past :]: + yield self.token_to_str(id).decode("utf8", errors="ignore") + + # write input + def input(self, prompt: str): + if ( + self.params.instruct + and self.last_n_tokens[-len(self.inp_prefix) :] != self.inp_prefix + ): + self.embd_inp += self.inp_prefix + self.embd_inp += self._tokenize(prompt) + if self.params.instruct: + self.embd_inp += self.inp_suffix + + # write output + def output(self): + self.remaining_tokens = self.params.n_predict + for id in self.generate(): + cur_char = self.token_to_str(id) + + # Add remainder of missing bytes + if None in self.multibyte_fix: + self.multibyte_fix[self.multibyte_fix.index(None)] = cur_char + + # Return completed utf char + if len(self.multibyte_fix) > 0 and not None in self.multibyte_fix: + yield (b"".join(self.multibyte_fix)).decode("utf8") + self.multibyte_fix = [] + continue + + # Contains multi-byte UTF8 + for num, pattern in [(2, 192), (3, 224), (4, 240)]: + # Bitwise AND check + if pattern & int.from_bytes(cur_char, "little") == pattern: + self.multibyte_fix = [cur_char] + ([None] * (num - 1)) + + # Stop incomplete bytes from passing + if len(self.multibyte_fix) > 0: + continue + + yield cur_char.decode("utf8") + + # read user input + def read_input(self): + out = "" + while (t := input()).endswith("\\"): + out += t[:-1] + "\n" + return out + t + "\n" + + # interactive mode + def interact(self): + for i in self.output(): + print(i, end="", flush=True) + self.params.input_echo = False + + # Using string instead of tokens to check for antiprompt, + # It is more reliable than tokens for interactive mode. + generated_str = "" + while self.params.interactive: + self.set_color(util.CONSOLE_COLOR_USER_INPUT) + if self.params.instruct: + print("\n> ", end="") + self.input(self.read_input()) + else: + print(self.params.input_prefix, end="") + self.input( + f"{self.params.input_prefix}{self.read_input()}{self.params.input_suffix}" + ) + print(self.params.input_suffix, end="") + self.set_color(util.CONSOLE_COLOR_DEFAULT) + + try: + for i in self.output(): + print(i, end="", flush=True) + generated_str += i + for ap in self.params.antiprompt: + if generated_str.endswith(ap): + raise KeyboardInterrupt + except KeyboardInterrupt: + self.set_color(util.CONSOLE_COLOR_DEFAULT) + if not self.params.instruct: + print(self.params.fix_prefix, end="") + self.input(self.params.fix_prefix) + if __name__ == "__main__": - from datetime import datetime + from datetime import datetime - USER_NAME="User" - AI_NAME="ChatLLaMa" + USER_NAME = "User" + AI_NAME = "ChatLLaMa" - time_now = datetime.now() - prompt = f"""Text transcript of a never ending dialog, where {USER_NAME} interacts with an AI assistant named {AI_NAME}. + time_now = datetime.now() + prompt = f"""Text transcript of a never ending dialog, where {USER_NAME} interacts with an AI assistant named {AI_NAME}. {AI_NAME} is helpful, kind, honest, friendly, good at writing and never fails to answer {USER_NAME}’s requests immediately and with details and precision. -There are no annotations like (30 seconds passed...) or (to himself), just what {USER_NAME} and {AI_NAME} say aloud to each other. +Transcript below contains only the recorded dialog between two, without any annotations like (30 seconds passed...) or (to himself), just what {USER_NAME} and {AI_NAME} say aloud to each other. The dialog lasts for years, the entirety of it is shared below. It's 10000 pages long. The transcript only includes text, it does not include markup like HTML and Markdown. @@ -561,8 +754,11 @@ def interact(self): {AI_NAME}: A cat is a domestic species of small carnivorous mammal. It is the only domesticated species in the family Felidae. {USER_NAME}: Name a color. {AI_NAME}: Blue -{USER_NAME}:""" - params = gpt_params_parse() +{USER_NAME}: """ + + params = gpt_params_parse() + if params.prompt is None and params.file is None: + params.prompt = prompt - with LLaMAInteract(params) as m: - m.interact() + with LLaMAInteract(params) as m: + m.interact() diff --git a/examples/low_level_api/low_level_api_llama_cpp.py b/examples/low_level_api/low_level_api_llama_cpp.py index 9e38ec7cb6..9fb3424ec0 100644 --- a/examples/low_level_api/low_level_api_llama_cpp.py +++ b/examples/low_level_api/low_level_api_llama_cpp.py @@ -1,26 +1,42 @@ -import llama_cpp - +import ctypes +import os import multiprocessing import llama_cpp +llama_cpp.llama_backend_init() + N_THREADS = multiprocessing.cpu_count() +MODEL_PATH = os.environ.get("MODEL", "../models/7B/ggml-model.bin") prompt = b"\n\n### Instruction:\nWhat is the capital of France?\n\n### Response:\n" -lparams = llama_cpp.llama_context_default_params() -ctx = llama_cpp.llama_init_from_file(b"../models/7B/ggml-model.bin", lparams) +lparams = llama_cpp.llama_model_default_params() +cparams = llama_cpp.llama_context_default_params() +model = llama_cpp.llama_model_load_from_file(MODEL_PATH.encode("utf-8"), lparams) +ctx = llama_cpp.llama_init_from_model(model, cparams) +vocab = llama_cpp.llama_model_get_vocab(model) # determine the required inference memory per token: tmp = [0, 1, 2, 3] -llama_cpp.llama_eval(ctx, (llama_cpp.c_int * len(tmp))(*tmp), len(tmp), 0, N_THREADS) +llama_cpp.llama_eval( + ctx=ctx, tokens=(llama_cpp.c_int * len(tmp))(*tmp), n_tokens=len(tmp), n_past=0 +) # Deprecated n_past = 0 prompt = b" " + prompt embd_inp = (llama_cpp.llama_token * (len(prompt) + 1))() -n_of_tok = llama_cpp.llama_tokenize(ctx, prompt, embd_inp, len(embd_inp), True) +n_of_tok = llama_cpp.llama_tokenize( + vocab=vocab, + text=prompt, + text_len=len(prompt), + tokens=embd_inp, + n_tokens_max=len(embd_inp), + add_special=False, + parse_special=False, +) embd_inp = embd_inp[:n_of_tok] n_ctx = llama_cpp.llama_n_ctx(ctx) @@ -45,32 +61,42 @@ while remaining_tokens > 0: if len(embd) > 0: llama_cpp.llama_eval( - ctx, (llama_cpp.c_int * len(embd))(*embd), len(embd), n_past, N_THREADS - ) + ctx=ctx, + tokens=(llama_cpp.c_int * len(embd))(*embd), + n_tokens=len(embd), + n_past=n_past, + ) # Deprecated n_past += len(embd) embd = [] if len(embd_inp) <= input_consumed: logits = llama_cpp.llama_get_logits(ctx) - n_vocab = llama_cpp.llama_n_vocab(ctx) + n_vocab = llama_cpp.llama_vocab_n_tokens(vocab) - _arr = (llama_cpp.llama_token_data * n_vocab)(*[ - llama_cpp.llama_token_data(token_id, logits[token_id], 0.0) - for token_id in range(n_vocab) - ]) - candidates_p = llama_cpp.ctypes.pointer(llama_cpp.llama_token_data_array(_arr, len(_arr), False)) + _arr = (llama_cpp.llama_token_data * n_vocab)( + *[ + llama_cpp.llama_token_data(token_id, logits[token_id], 0.0) + for token_id in range(n_vocab) + ] + ) + candidates_p = llama_cpp.ctypes.pointer( + llama_cpp.llama_token_data_array(_arr, len(_arr), False) + ) _arr = (llama_cpp.c_int * len(last_n_tokens_data))(*last_n_tokens_data) - llama_cpp.llama_sample_repetition_penalty(ctx, candidates_p, - _arr, - last_n_repeat, repeat_penalty) - llama_cpp.llama_sample_frequency_and_presence_penalties(ctx, candidates_p, + llama_cpp.llama_sample_repetition_penalties( + ctx, + candidates_p, _arr, - last_n_repeat, frequency_penalty, presence_penalty) + penalty_last_n=last_n_repeat, + penalty_repeat=repeat_penalty, + penalty_freq=frequency_penalty, + penalty_present=presence_penalty, + ) - llama_cpp.llama_sample_top_k(ctx, candidates_p, 40) - llama_cpp.llama_sample_top_p(ctx, candidates_p, 0.8) - llama_cpp.llama_sample_temperature(ctx, candidates_p, 0.2) + llama_cpp.llama_sample_top_k(ctx, candidates_p, k=40, min_keep=1) + llama_cpp.llama_sample_top_p(ctx, candidates_p, p=0.8, min_keep=1) + llama_cpp.llama_sample_temperature(ctx, candidates_p, temp=0.2) id = llama_cpp.llama_sample_token(ctx, candidates_p) last_n_tokens_data = last_n_tokens_data[1:] + [id] @@ -86,13 +112,19 @@ break if not input_noecho: for id in embd: + size = 32 + buffer = (ctypes.c_char * size)() + n = llama_cpp.llama_token_to_piece( + vocab, llama_cpp.llama_token(id), buffer, size, 0, False + ) + assert n <= size print( - llama_cpp.llama_token_to_str(ctx, id).decode("utf-8", errors="ignore"), + buffer[:n].decode("utf-8"), end="", flush=True, ) - if len(embd) > 0 and embd[-1] == llama_cpp.llama_token_eos(): + if len(embd) > 0 and embd[-1] == llama_cpp.llama_vocab_eos(vocab): break print() @@ -100,3 +132,4 @@ llama_cpp.llama_print_timings(ctx) llama_cpp.llama_free(ctx) +llama_cpp.llama_model_free(model) diff --git a/examples/low_level_api/quantize.py b/examples/low_level_api/quantize.py index 8bd03f88a1..057ac389eb 100644 --- a/examples/low_level_api/quantize.py +++ b/examples/low_level_api/quantize.py @@ -4,14 +4,16 @@ def main(args): + fname_inp = args.fname_inp.encode("utf-8") + fname_out = args.fname_out.encode("utf-8") if not os.path.exists(fname_inp): raise RuntimeError(f"Input file does not exist ({fname_inp})") if os.path.exists(fname_out): raise RuntimeError(f"Output file already exists ({fname_out})") - fname_inp = args.fname_inp.encode("utf-8") - fname_out = args.fname_out.encode("utf-8") - itype = args.itype - return_code = llama_cpp.llama_model_quantize(fname_inp, fname_out, itype) + ftype = args.type + args = llama_cpp.llama_model_quantize_default_params() + args.ftype = ftype + return_code = llama_cpp.llama_model_quantize(fname_inp, fname_out, args) if return_code != 0: raise RuntimeError("Failed to quantize model") @@ -20,6 +22,10 @@ def main(args): parser = argparse.ArgumentParser() parser.add_argument("fname_inp", type=str, help="Path to input model") parser.add_argument("fname_out", type=str, help="Path to output model") - parser.add_argument("type", type=int, help="Type of quantization (2: q4_0, 3: q4_1)") + parser.add_argument( + "type", + type=int, + help="Type of quantization (2: q4_0, 3: q4_1), see llama_cpp.py for enum", + ) args = parser.parse_args() main(args) diff --git a/examples/low_level_api/readme/low_level_api_llama_cpp.md b/examples/low_level_api/readme/low_level_api_llama_cpp.md new file mode 100644 index 0000000000..5f350ffe99 --- /dev/null +++ b/examples/low_level_api/readme/low_level_api_llama_cpp.md @@ -0,0 +1,61 @@ +# Low-Level API for Llama_cpp + +## Overview +This Python script, low_level_api_llama_cpp.py, demonstrates the implementation of a low-level API for interacting with the llama_cpp library. The script defines an inference that generates embeddings based on a given prompt using .gguf model. + +### Prerequisites +Before running the script, ensure that you have the following dependencies installed: + +. Python 3.6 or higher +. llama_cpp: A C++ library for working with .gguf model +. NumPy: A fundamental package for scientific computing with Python +. multiprocessing: A Python module for parallel computing + +### Usage +install depedencies: +```bash +python -m pip install llama-cpp-python ctypes os multiprocessing +``` +Run the script: +```bash +python low_level_api_llama_cpp.py +``` + +## Code Structure +The script is organized as follows: + +### . Initialization: + Load the model from the specified path. + Create a context for model evaluation. + +### . Tokenization: + Tokenize the input prompt using the llama_tokenize function. + Prepare the input tokens for model evaluation. + +### . Inference: + Perform model evaluation to generate responses. + Sample from the model's output using various strategies (top-k, top-p, temperature). + +### . Output: + Print the generated tokens and the corresponding decoded text. + +### .Cleanup: + Free resources and print timing information. + +## Configuration +Customize the inference behavior by adjusting the following variables: + +#### . N_THREADS: Number of CPU threads to use for model evaluation. +#### . MODEL_PATH: Path to the model file. +#### . prompt: Input prompt for the chatbot. + +## Notes +. Ensure that the llama_cpp library is built and available in the system. Follow the instructions in the llama_cpp repository for building and installing the library. + +. This script is designed to work with the .gguf model and may require modifications for compatibility with other models. + +## Acknowledgments +This code is based on the llama_cpp library developed by the community. Special thanks to the contributors for their efforts. + +## License +This project is licensed under the MIT License - see the LICENSE file for details. \ No newline at end of file diff --git a/examples/low_level_api/util.py b/examples/low_level_api/util.py index 9d0ec2f705..ef8b1c1eea 100644 --- a/examples/low_level_api/util.py +++ b/examples/low_level_api/util.py @@ -1,4 +1,3 @@ - ANSI_COLOR_RESET = "\x1b[0m" ANSI_COLOR_YELLOW = "\x1b[33m" ANSI_BOLD = "\x1b[1m" @@ -8,88 +7,95 @@ CONSOLE_COLOR_PROMPT = ANSI_COLOR_YELLOW CONSOLE_COLOR_USER_INPUT = ANSI_BOLD + ANSI_COLOR_GREEN + # Iterative search # Actively searches and prevents a pattern from being returned class IterSearch: - def __init__(self, pattern): - self.pattern = list(pattern) - self.buffer = [] - - def __call__(self, char): - self.buffer += [char] + def __init__(self, pattern): + self.pattern = list(pattern) + self.buffer = [] - if (self.pattern[:len(self.buffer)] == self.buffer): - if (len(self.buffer) >= len(self.pattern)): - self.buffer.clear() - return [] + def __call__(self, char): + self.buffer += [char] - _tmp = self.buffer[:] - self.buffer.clear() - return _tmp + if self.pattern[: len(self.buffer)] == self.buffer: + if len(self.buffer) >= len(self.pattern): + self.buffer.clear() + return [] -class Circle: - def __init__(self, size, default=0): - self.list = [default] * size - self.maxsize = size - self.size = 0 - self.offset = 0 - - def append(self, elem): - if self.size < self.maxsize: - self.list[self.size] = elem - self.size += 1 - else: - self.list[self.offset] = elem - self.offset = (self.offset + 1) % self.maxsize - - def __getitem__(self, val): - if isinstance(val, int): - if 0 > val or val >= self.size: - raise IndexError('Index out of range') - return self.list[val] if self.size < self.maxsize else self.list[(self.offset + val) % self.maxsize] - elif isinstance(val, slice): - start, stop, step = val.start, val.stop, val.step - if step is None: - step = 1 - if start is None: - start = 0 - if stop is None: - stop = self.size - if start < 0: - start = self.size + start - if stop < 0: - stop = self.size + stop - - indices = range(start, stop, step) - return [self.list[(self.offset + i) % self.maxsize] for i in indices if i < self.size] - else: - raise TypeError('Invalid argument type') + _tmp = self.buffer[:] + self.buffer.clear() + return _tmp +class Circle: + def __init__(self, size, default=0): + self.list = [default] * size + self.maxsize = size + self.size = 0 + self.offset = 0 + + def append(self, elem): + if self.size < self.maxsize: + self.list[self.size] = elem + self.size += 1 + else: + self.list[self.offset] = elem + self.offset = (self.offset + 1) % self.maxsize + + def __getitem__(self, val): + if isinstance(val, int): + if 0 > val or val >= self.size: + raise IndexError("Index out of range") + return ( + self.list[val] + if self.size < self.maxsize + else self.list[(self.offset + val) % self.maxsize] + ) + elif isinstance(val, slice): + start, stop, step = val.start, val.stop, val.step + if step is None: + step = 1 + if start is None: + start = 0 + if stop is None: + stop = self.size + if start < 0: + start = self.size + start + if stop < 0: + stop = self.size + stop + + indices = range(start, stop, step) + return [ + self.list[(self.offset + i) % self.maxsize] + for i in indices + if i < self.size + ] + else: + raise TypeError("Invalid argument type") if __name__ == "__main__": - c = Circle(5) - - c.append(1) - print(c.list) - print(c[:]) - assert c[0] == 1 - assert c[:5] == [1] - - for i in range(2,5+1): - c.append(i) - print(c.list) - print(c[:]) - assert c[0] == 1 - assert c[:5] == [1,2,3,4,5] - - for i in range(5+1,9+1): - c.append(i) - print(c.list) - print(c[:]) - assert c[0] == 5 - assert c[:5] == [5,6,7,8,9] - #assert c[:-5] == [5,6,7,8,9] - assert c[:10] == [5,6,7,8,9] - + c = Circle(5) + + c.append(1) + print(c.list) + print(c[:]) + assert c[0] == 1 + assert c[:5] == [1] + + for i in range(2, 5 + 1): + c.append(i) + print(c.list) + print(c[:]) + assert c[0] == 1 + assert c[:5] == [1, 2, 3, 4, 5] + + for i in range(5 + 1, 9 + 1): + c.append(i) + print(c.list) + print(c[:]) + assert c[0] == 5 + assert c[:5] == [5, 6, 7, 8, 9] + # assert c[:-5] == [5,6,7,8,9] + assert c[:10] == [5, 6, 7, 8, 9] diff --git a/examples/notebooks/Batching.ipynb b/examples/notebooks/Batching.ipynb new file mode 100644 index 0000000000..0b36cd9c3a --- /dev/null +++ b/examples/notebooks/Batching.ipynb @@ -0,0 +1,464 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "import ctypes\n", + "import llama_cpp" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "llama_cpp.llama_backend_init(numa=False)" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "llama_model_loader: loaded meta data with 20 key-value pairs and 291 tensors from /workspaces/llama-cpp-python/mistral-7b-v0.1.Q2_K.gguf (version GGUF V2)\n", + "llama_model_loader: Dumping metadata keys/values. Note: KV overrides do not apply in this output.\n", + "llama_model_loader: - kv 0: general.architecture str = llama\n", + "llama_model_loader: - kv 1: general.name str = mistralai_mistral-7b-v0.1\n", + "llama_model_loader: - kv 2: llama.context_length u32 = 32768\n", + "llama_model_loader: - kv 3: llama.embedding_length u32 = 4096\n", + "llama_model_loader: - kv 4: llama.block_count u32 = 32\n", + "llama_model_loader: - kv 5: llama.feed_forward_length u32 = 14336\n", + "llama_model_loader: - kv 6: llama.rope.dimension_count u32 = 128\n", + "llama_model_loader: - kv 7: llama.attention.head_count u32 = 32\n", + "llama_model_loader: - kv 8: llama.attention.head_count_kv u32 = 8\n", + "llama_model_loader: - kv 9: llama.attention.layer_norm_rms_epsilon f32 = 0.000010\n", + "llama_model_loader: - kv 10: llama.rope.freq_base f32 = 10000.000000\n", + "llama_model_loader: - kv 11: general.file_type u32 = 10\n", + "llama_model_loader: - kv 12: tokenizer.ggml.model str = llama\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "llama_model_loader: - kv 13: tokenizer.ggml.tokens arr[str,32000] = [\"<unk>\", \"<s>\", \"</s>\", \"<0x00>\", \"<...\n", + "llama_model_loader: - kv 14: tokenizer.ggml.scores arr[f32,32000] = [0.000000, 0.000000, 0.000000, 0.0000...\n", + "llama_model_loader: - kv 15: tokenizer.ggml.token_type arr[i32,32000] = [2, 3, 3, 6, 6, 6, 6, 6, 6, 6, 6, 6, ...\n", + "llama_model_loader: - kv 16: tokenizer.ggml.bos_token_id u32 = 1\n", + "llama_model_loader: - kv 17: tokenizer.ggml.eos_token_id u32 = 2\n", + "llama_model_loader: - kv 18: tokenizer.ggml.unknown_token_id u32 = 0\n", + "llama_model_loader: - kv 19: general.quantization_version u32 = 2\n", + "llama_model_loader: - type f32: 65 tensors\n", + "llama_model_loader: - type q2_K: 65 tensors\n", + "llama_model_loader: - type q3_K: 160 tensors\n", + "llama_model_loader: - type q6_K: 1 tensors\n", + "llm_load_vocab: special_eos_id is not in special_eog_ids - the tokenizer config may be incorrect\n", + "llm_load_vocab: special tokens cache size = 3\n", + "llm_load_vocab: token to piece cache size = 0.1637 MB\n", + "llm_load_print_meta: format = GGUF V2\n", + "llm_load_print_meta: arch = llama\n", + "llm_load_print_meta: vocab type = SPM\n", + "llm_load_print_meta: n_vocab = 32000\n", + "llm_load_print_meta: n_merges = 0\n", + "llm_load_print_meta: vocab_only = 0\n", + "llm_load_print_meta: n_ctx_train = 32768\n", + "llm_load_print_meta: n_embd = 4096\n", + "llm_load_print_meta: n_layer = 32\n", + "llm_load_print_meta: n_head = 32\n", + "llm_load_print_meta: n_head_kv = 8\n", + "llm_load_print_meta: n_rot = 128\n", + "llm_load_print_meta: n_swa = 0\n", + "llm_load_print_meta: n_embd_head_k = 128\n", + "llm_load_print_meta: n_embd_head_v = 128\n", + "llm_load_print_meta: n_gqa = 4\n", + "llm_load_print_meta: n_embd_k_gqa = 1024\n", + "llm_load_print_meta: n_embd_v_gqa = 1024\n", + "llm_load_print_meta: f_norm_eps = 0.0e+00\n", + "llm_load_print_meta: f_norm_rms_eps = 1.0e-05\n", + "llm_load_print_meta: f_clamp_kqv = 0.0e+00\n", + "llm_load_print_meta: f_max_alibi_bias = 0.0e+00\n", + "llm_load_print_meta: f_logit_scale = 0.0e+00\n", + "llm_load_print_meta: n_ff = 14336\n", + "llm_load_print_meta: n_expert = 0\n", + "llm_load_print_meta: n_expert_used = 0\n", + "llm_load_print_meta: causal attn = 1\n", + "llm_load_print_meta: pooling type = 0\n", + "llm_load_print_meta: rope type = 0\n", + "llm_load_print_meta: rope scaling = linear\n", + "llm_load_print_meta: freq_base_train = 10000.0\n", + "llm_load_print_meta: freq_scale_train = 1\n", + "llm_load_print_meta: n_ctx_orig_yarn = 32768\n", + "llm_load_print_meta: rope_finetuned = unknown\n", + "llm_load_print_meta: ssm_d_conv = 0\n", + "llm_load_print_meta: ssm_d_inner = 0\n", + "llm_load_print_meta: ssm_d_state = 0\n", + "llm_load_print_meta: ssm_dt_rank = 0\n", + "llm_load_print_meta: ssm_dt_b_c_rms = 0\n", + "llm_load_print_meta: model type = 7B\n", + "llm_load_print_meta: model ftype = Q2_K - Medium\n", + "llm_load_print_meta: model params = 7.24 B\n", + "llm_load_print_meta: model size = 2.87 GiB (3.41 BPW) \n", + "llm_load_print_meta: general.name = mistralai_mistral-7b-v0.1\n", + "llm_load_print_meta: BOS token = 1 '<s>'\n", + "llm_load_print_meta: EOS token = 2 '</s>'\n", + "llm_load_print_meta: UNK token = 0 '<unk>'\n", + "llm_load_print_meta: LF token = 13 '<0x0A>'\n", + "llm_load_print_meta: EOG token = 2 '</s>'\n", + "llm_load_print_meta: max token length = 48\n", + "llm_load_tensors: ggml ctx size = 0.14 MiB\n", + "llm_load_tensors: CPU buffer size = 2939.57 MiB\n", + "..................................................................................................\n" + ] + } + ], + "source": [ + "params = llama_cpp.llama_model_default_params()\n", + "params.n_gpu_layers = 35\n", + "model = llama_cpp.llama_model_load_from_file(\n", + " b\"/workspaces/llama-cpp-python/mistral-7b-v0.1.Q2_K.gguf\", params\n", + ") # Update this to whatever\n", + "vocab = llama_cpp.llama_model_get_vocab(model)" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[1, 415, 2936, 9060, 285, 1142]\n", + "58\n" + ] + } + ], + "source": [ + "n_ctx = 512\n", + "n_len = 32\n", + "n_parallel = 2\n", + "prompt = b\"The quick brown fox\"\n", + "\n", + "tokens = (llama_cpp.llama_token * n_ctx)()\n", + "tokens_len = llama_cpp.llama_tokenize(\n", + " vocab, prompt, len(prompt), tokens, len(tokens), True, True\n", + ")\n", + "print(tokens[:tokens_len])\n", + "\n", + "n_kv_req = tokens_len + (n_len - tokens_len) * n_parallel\n", + "print(n_kv_req)" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "llama_new_context_with_model: n_ctx = 64\n", + "llama_new_context_with_model: n_batch = 32\n", + "llama_new_context_with_model: n_ubatch = 32\n", + "llama_new_context_with_model: flash_attn = 0\n", + "llama_new_context_with_model: freq_base = 10000.0\n", + "llama_new_context_with_model: freq_scale = 1\n", + "llama_kv_cache_init: CPU KV buffer size = 8.00 MiB\n", + "llama_new_context_with_model: KV self size = 8.00 MiB, K (f16): 4.00 MiB, V (f16): 4.00 MiB\n", + "llama_new_context_with_model: CPU output buffer size = 0.12 MiB\n", + "llama_new_context_with_model: CPU compute buffer size = 5.01 MiB\n", + "llama_new_context_with_model: graph nodes = 1030\n", + "llama_new_context_with_model: graph splits = 1\n" + ] + } + ], + "source": [ + "ctx_params = llama_cpp.llama_context_default_params()\n", + "ctx_params.seed = 1234\n", + "ctx_params.n_ctx = n_kv_req\n", + "ctx_params.n_batch = max(n_len, n_parallel)\n", + "ctx_params.n_threads = 1\n", + "ctx_params.n_threads_batch = 1\n", + "ctx = llama_cpp.llama_init_from_model(model, ctx_params)" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [], + "source": [ + "n_ctx = llama_cpp.llama_n_ctx(ctx)\n", + "batch = llama_cpp.llama_batch_init(max(tokens_len, n_parallel), 0, 1)" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "\n", + "batch.n_tokens = tokens_len\n", + "for i in range(tokens_len):\n", + " batch.token[i] = tokens[i]\n", + " batch.pos[i] = i\n", + " batch.seq_id[i][0] = 0\n", + " batch.n_seq_id[i] = 1\n", + " batch.logits[i] = False\n", + "\n", + "batch.logits[batch.n_tokens - 1] = True\n", + "\n", + "if llama_cpp.llama_decode(ctx, batch) != 0:\n", + " print(\"Error decoding\")" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [], + "source": [ + "for i in range(n_parallel):\n", + " llama_cpp.llama_kv_cache_seq_cp(ctx, 0, i, 0, batch.n_tokens)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "# Initialize sampler chain with default parameters\n", + "sparams = llama_cpp.llama_sampler_chain_default_params()\n", + "sampler_chain = llama_cpp.llama_sampler_chain_init(sparams)\n", + "\n", + "# Add top_k, top_p, temperature, and final distribution-based sampler\n", + "llama_cpp.llama_sampler_chain_add(sampler_chain, llama_cpp.llama_sampler_init_top_k(40))\n", + "llama_cpp.llama_sampler_chain_add(sampler_chain, llama_cpp.llama_sampler_init_top_p(0.9, 1))\n", + "llama_cpp.llama_sampler_chain_add(sampler_chain, llama_cpp.llama_sampler_init_temp(0.4))\n", + "llama_cpp.llama_sampler_chain_add(sampler_chain, llama_cpp.llama_sampler_init_dist(1234)) # Final \"dist\" sampler" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "7\n", + "[' j', ' jumped']\n", + "8\n", + "[' j over', ' jumped over']\n", + "9\n", + "[' j over the', ' jumped over the']\n", + "10\n", + "[' j over the lazy', ' jumped over the lazy']\n", + "11\n", + "[' j over the lazy dog', ' jumped over the lazy dog']\n", + "12\n", + "[' j over the lazy dog.', ' jumped over the lazy dog\\n']\n", + "13\n", + "[' j over the lazy dog. También', ' jumped over the lazy dog\\nGroupLayout']\n", + "14\n", + "[' j over the lazy dog. También:', ' jumped over the lazy dog\\nGroupLayouting']\n", + "15\n", + "[' j over the lazy dog. También: is', ' jumped over the lazy dog\\nGroupLayouting is']\n", + "16\n", + "[' j over the lazy dog. También: is a', ' jumped over the lazy dog\\nGroupLayouting is a']\n", + "17\n", + "[' j over the lazy dog. También: is a technique', ' jumped over the lazy dog\\nGroupLayouting is a common']\n", + "18\n", + "[' j over the lazy dog. También: is a technique practice', ' jumped over the lazy dog\\nGroupLayouting is a common practice']\n", + "19\n", + "[' j over the lazy dog. También: is a technique practice in', ' jumped over the lazy dog\\nGroupLayouting is a common practice in']\n", + "20\n", + "[' j over the lazy dog. También: is a technique practice in the', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the']\n", + "21\n", + "[' j over the lazy dog. También: is a technique practice in the real', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media']\n", + "22\n", + "[' j over the lazy dog. También: is a technique practice in the real-', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media industry']\n", + "23\n", + "[' j over the lazy dog. También: is a technique practice in the real-.', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media industry.']\n", + "24\n", + "[' j over the lazy dog. También: is a technique practice in the real-. We', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media industry. However']\n", + "25\n", + "[' j over the lazy dog. También: is a technique practice in the real-. We,', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media industry. However,']\n", + "26\n", + "[' j over the lazy dog. También: is a technique practice in the real-. We, when', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media industry. However, there']\n", + "27\n", + "[' j over the lazy dog. También: is a technique practice in the real-. We, when is', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media industry. However, there has']\n", + "28\n", + "[' j over the lazy dog. También: is a technique practice in the real-. We, when is been', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media industry. However, there has been']\n", + "29\n", + "[' j over the lazy dog. También: is a technique practice in the real-. We, when is been little', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media industry. However, there has been little']\n", + "30\n", + "[' j over the lazy dog. También: is a technique practice in the real-. We, when is been little research', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media industry. However, there has been little emp']\n", + "31\n", + "[' j over the lazy dog. También: is a technique practice in the real-. We, when is been little researchirical', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media industry. However, there has been little empirical']\n", + "32\n", + "[' j over the lazy dog. También: is a technique practice in the real-. We, when is been little researchirical research', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media industry. However, there has been little empirical research']\n" + ] + } + ], + "source": [ + "streams = [\"\"] * n_parallel\n", + "i_batch = [batch.n_tokens - 1] * n_parallel\n", + "\n", + "n_cur = batch.n_tokens\n", + "n_decode = 0\n", + "\n", + "while n_cur <= n_len:\n", + " batch.n_tokens = 0\n", + " for i in range(n_parallel):\n", + " if i_batch[i] < 0:\n", + " continue\n", + "\n", + " # Sample the next token using the sampler chain\n", + " new_token_id = llama_cpp.llama_sampler_sample(sampler_chain, ctx, -1)\n", + "\n", + " if new_token_id == llama_cpp.llama_vocab_eos(vocab) or n_cur == n_len:\n", + " i_batch[i] = -1\n", + " continue\n", + "\n", + " buf = (ctypes.c_char * 32)()\n", + " \n", + " # Convert token ID to text\n", + " outlen = llama_cpp.llama_token_to_piece(vocab, new_token_id, buf, len(buf), 0, False)\n", + " streams[i] += bytes(buf[:outlen]).decode(\"utf-8\")\n", + "\n", + " batch.token[batch.n_tokens] = new_token_id\n", + " batch.pos[batch.n_tokens] = n_cur\n", + " batch.seq_id[batch.n_tokens][0] = i\n", + " batch.n_seq_id[batch.n_tokens] = 1\n", + " batch.logits[batch.n_tokens] = True\n", + "\n", + " i_batch[i] = batch.n_tokens\n", + " batch.n_tokens += 1\n", + " n_decode += 1\n", + "\n", + " if batch.n_tokens == 0:\n", + " break\n", + "\n", + " n_cur += 1\n", + "\n", + " if llama_cpp.llama_decode(ctx, batch) != 0:\n", + " print(\"Error decoding\", flush=True)\n", + " break\n", + " print(n_cur)\n", + " print(streams)" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[' j over the lazy dog. También: is a technique practice in the real-. We, when is been little researchirical research', ' jumped over the lazy dog\\nGroupLayouting is a common practice in the media industry. However, there has been little empirical research']\n" + ] + } + ], + "source": [ + "print(streams)" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [], + "source": [ + "llama_cpp.llama_batch_free(batch)" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [], + "source": [ + "llama_cpp.llama_free(ctx)" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": {}, + "outputs": [], + "source": [ + "llama_cpp.llama_model_free(model)" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "metadata": {}, + "outputs": [], + "source": [ + "llama_cpp.llama_backend_free()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.1" + }, + "orig_nbformat": 4 + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/examples/notebooks/Clients.ipynb b/examples/notebooks/Clients.ipynb index caebbb67fd..fab82673e5 100644 --- a/examples/notebooks/Clients.ipynb +++ b/examples/notebooks/Clients.ipynb @@ -37,11 +37,11 @@ "source": [ "import openai\n", "\n", - "openai.api_key = \"sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx\" # can be anything\n", + "openai.api_key = \"sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx\" # can be anything\n", "openai.api_base = \"http://100.64.159.73:8000/v1\"\n", "\n", "openai.Completion.create(\n", - " model=\"text-davinci-003\", # currently can be anything\n", + " model=\"text-davinci-003\", # currently can be anything\n", " prompt=\"The quick brown fox jumps\",\n", " max_tokens=5,\n", ")" @@ -66,7 +66,9 @@ "source": [ "import os\n", "\n", - "os.environ[\"OPENAI_API_KEY\"] = \"sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx\" # can be anything\n", + "os.environ[\"OPENAI_API_KEY\"] = (\n", + " \"sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx\" # can be anything\n", + ")\n", "os.environ[\"OPENAI_API_BASE\"] = \"http://100.64.159.73:8000/v1\"\n", "\n", "from langchain.llms import OpenAI\n", diff --git a/examples/notebooks/Functions.ipynb b/examples/notebooks/Functions.ipynb new file mode 100644 index 0000000000..1f41381659 --- /dev/null +++ b/examples/notebooks/Functions.ipynb @@ -0,0 +1,519 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Functions\n", + "\n", + "The OpenAI compatbile web server in `llama-cpp-python` supports function calling.\n", + "\n", + "Function calling allows API clients to specify a schema that gives the model a format it should respond in.\n", + "Function calling in `llama-cpp-python` works by combining models pretrained for function calling such as [`functionary`](https://huggingface.co/meetkai) with constrained sampling to produce a response that is compatible with the schema.\n", + "\n", + "Note however that this improves but does not guarantee that the response will be compatible with the schema.\n", + "\n", + "## Requirements\n", + "\n", + "Before we begin you will need the following:\n", + "\n", + "- A running `llama-cpp-python` server with a function calling compatible model. [See here](https://llama-cpp-python.readthedocs.io/en/latest/server/#function-calling)\n", + "- The OpenAI Python Client `pip install openai`\n", + "- (Optional) The Instructor Python Library `pip install instructor`\n", + "\n", + "## Function Calling with OpenAI Python Client\n", + "\n", + "We'll start with a basic demo that only uses the OpenAI Python Client." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "ChatCompletion(id='chatcmpl-a2d9eb9f-7354-472f-b6ad-4d7a807729a3', choices=[Choice(finish_reason='stop', index=0, message=ChatCompletionMessage(content='The current weather in San Francisco is **72°F** (22°C).\\n ', role='assistant', function_call=None, tool_calls=None))], created=1699638365, model='gpt-3.5-turbo-1106', object='chat.completion', system_fingerprint=None, usage=CompletionUsage(completion_tokens=22, prompt_tokens=136, total_tokens=158))\n" + ] + } + ], + "source": [ + "import openai\n", + "import json\n", + "\n", + "\n", + "client = openai.OpenAI(\n", + " api_key=\"sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx\", # can be anything\n", + " base_url=\"http://100.64.159.73:8000/v1\", # NOTE: Replace with IP address and port of your llama-cpp-python server\n", + ")\n", + "\n", + "\n", + "# Example dummy function hard coded to return the same weather\n", + "# In production, this could be your backend API or an external API\n", + "def get_current_weather(location, unit=\"fahrenheit\"):\n", + " \"\"\"Get the current weather in a given location\"\"\"\n", + " if \"tokyo\" in location.lower():\n", + " return json.dumps({\"location\": \"Tokyo\", \"temperature\": \"10\", \"unit\": \"celsius\"})\n", + " elif \"san francisco\" in location.lower():\n", + " return json.dumps(\n", + " {\"location\": \"San Francisco\", \"temperature\": \"72\", \"unit\": \"fahrenheit\"}\n", + " )\n", + " elif \"paris\" in location.lower():\n", + " return json.dumps({\"location\": \"Paris\", \"temperature\": \"22\", \"unit\": \"celsius\"})\n", + " else:\n", + " return json.dumps({\"location\": location, \"temperature\": \"unknown\"})\n", + "\n", + "\n", + "def run_conversation():\n", + " # Step 1: send the conversation and available functions to the model\n", + " messages = [\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": \"What's the weather like in San Francisco, Tokyo, and Paris?\",\n", + " }\n", + " ]\n", + " tools = [\n", + " {\n", + " \"type\": \"function\",\n", + " \"function\": {\n", + " \"name\": \"get_current_weather\",\n", + " \"description\": \"Get the current weather in a given location\",\n", + " \"parameters\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"location\": {\n", + " \"type\": \"string\",\n", + " \"description\": \"The city and state, e.g. San Francisco, CA\",\n", + " },\n", + " \"unit\": {\"type\": \"string\", \"enum\": [\"celsius\", \"fahrenheit\"]},\n", + " },\n", + " \"required\": [\"location\"],\n", + " },\n", + " },\n", + " }\n", + " ]\n", + " response = client.chat.completions.create(\n", + " model=\"gpt-3.5-turbo-1106\",\n", + " messages=messages,\n", + " tools=tools,\n", + " tool_choice=\"auto\", # auto is default, but we'll be explicit\n", + " )\n", + " response_message = response.choices[0].message\n", + " tool_calls = response_message.tool_calls\n", + " # Step 2: check if the model wanted to call a function\n", + " if tool_calls:\n", + " # Step 3: call the function\n", + " # Note: the JSON response may not always be valid; be sure to handle errors\n", + " available_functions = {\n", + " \"get_current_weather\": get_current_weather,\n", + " } # only one function in this example, but you can have multiple\n", + " messages.append(response_message) # extend conversation with assistant's reply\n", + " # Step 4: send the info for each function call and function response to the model\n", + " for tool_call in tool_calls:\n", + " function_name = tool_call.function.name\n", + " function_to_call = available_functions[function_name]\n", + " function_args = json.loads(tool_call.function.arguments)\n", + " function_response = function_to_call(\n", + " location=function_args.get(\"location\"),\n", + " unit=function_args.get(\"unit\"),\n", + " )\n", + " messages.append(\n", + " {\n", + " \"tool_call_id\": tool_call.id,\n", + " \"role\": \"tool\",\n", + " \"name\": function_name,\n", + " \"content\": function_response,\n", + " }\n", + " ) # extend conversation with function response\n", + " second_response = client.chat.completions.create(\n", + " model=\"gpt-3.5-turbo-1106\",\n", + " messages=messages,\n", + " ) # get a new response from the model where it can see the function response\n", + " return second_response\n", + "\n", + "\n", + "print(run_conversation())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Function Calling with Instructor\n", + "\n", + "The above example is a bit verbose and requires you to manually verify the schema.\n", + "\n", + "For our next examples we'll use the `instructor` library to simplify the process and accomplish a number of different tasks with function calling.\n", + "\n", + "You'll first need to install the [`instructor`](https://github.com/jxnl/instructor/).\n", + "\n", + "You can do so by running the following command in your terminal:\n", + "\n", + "```bash\n", + "pip install instructor\n", + "```\n", + "\n", + "Below we'll go through a few basic examples taken directly from the [instructor cookbook](https://jxnl.github.io/instructor/)\n", + "\n", + "## Basic Usage" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "name='Jason' age=25\n" + ] + } + ], + "source": [ + "import instructor\n", + "from pydantic import BaseModel\n", + "\n", + "# Enables `response_model`\n", + "client = instructor.patch(client=client)\n", + "\n", + "\n", + "class UserDetail(BaseModel):\n", + " name: str\n", + " age: int\n", + "\n", + "\n", + "user = client.chat.completions.create(\n", + " model=\"gpt-3.5-turbo\",\n", + " response_model=UserDetail,\n", + " messages=[\n", + " {\"role\": \"user\", \"content\": \"Extract Jason is 25 years old\"},\n", + " ],\n", + ")\n", + "\n", + "assert isinstance(user, UserDetail)\n", + "assert user.name == \"Jason\"\n", + "assert user.age == 25\n", + "\n", + "print(user)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Text Classification\n", + "\n", + "### Single-Label Classification" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "class_label=<Labels.SPAM: 'spam'>\n" + ] + } + ], + "source": [ + "import enum\n", + "\n", + "\n", + "class Labels(str, enum.Enum):\n", + " \"\"\"Enumeration for single-label text classification.\"\"\"\n", + "\n", + " SPAM = \"spam\"\n", + " NOT_SPAM = \"not_spam\"\n", + "\n", + "\n", + "class SinglePrediction(BaseModel):\n", + " \"\"\"\n", + " Class for a single class label prediction.\n", + " \"\"\"\n", + "\n", + " class_label: Labels\n", + "\n", + "\n", + "def classify(data: str) -> SinglePrediction:\n", + " \"\"\"Perform single-label classification on the input text.\"\"\"\n", + " return client.chat.completions.create(\n", + " model=\"gpt-3.5-turbo-0613\",\n", + " response_model=SinglePrediction,\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": f\"Classify the following text: {data}\",\n", + " },\n", + " ],\n", + " ) # type: ignore\n", + "\n", + "\n", + "prediction = classify(\"Hello there I'm a Nigerian prince and I want to give you money\")\n", + "assert prediction.class_label == Labels.SPAM\n", + "print(prediction)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Multi-Label Classification" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "class_labels=[<MultiLabels.TECH_ISSUE: 'tech_issue'>, <MultiLabels.BILLING: 'billing'>]\n" + ] + } + ], + "source": [ + "from typing import List\n", + "\n", + "\n", + "# Define Enum class for multiple labels\n", + "class MultiLabels(str, enum.Enum):\n", + " TECH_ISSUE = \"tech_issue\"\n", + " BILLING = \"billing\"\n", + " GENERAL_QUERY = \"general_query\"\n", + "\n", + "\n", + "# Define the multi-class prediction model\n", + "class MultiClassPrediction(BaseModel):\n", + " \"\"\"\n", + " Class for a multi-class label prediction.\n", + " \"\"\"\n", + "\n", + " class_labels: List[MultiLabels]\n", + "\n", + "\n", + "def multi_classify(data: str) -> MultiClassPrediction:\n", + " \"\"\"Perform multi-label classification on the input text.\"\"\"\n", + " return client.chat.completions.create(\n", + " model=\"gpt-3.5-turbo-0613\",\n", + " response_model=MultiClassPrediction,\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": f\"Classify the following support ticket: {data}\",\n", + " },\n", + " ],\n", + " ) # type: ignore\n", + "\n", + "\n", + "# Test multi-label classification\n", + "ticket = \"My account is locked and I can't access my billing info.\"\n", + "prediction = multi_classify(ticket)\n", + "print(prediction)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Self-Critique" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "question='What is the meaning of life?' answer='According to the Devil, the meaning of life is to live a life of sin and debauchery.'\n", + "1 validation error for QuestionAnswerNoEvil\n", + "answer\n", + " Assertion failed, The statement promotes sin and debauchery, which can be considered objectionable. [type=assertion_error, input_value='According to the Devil, ... of sin and debauchery.', input_type=str]\n", + " For further information visit https://errors.pydantic.dev/2.3/v/assertion_error\n" + ] + } + ], + "source": [ + "from typing_extensions import Annotated\n", + "from pydantic import BaseModel, BeforeValidator\n", + "\n", + "from instructor import llm_validator\n", + "\n", + "\n", + "question = \"What is the meaning of life?\"\n", + "context = \"The according to the devil the meaning of live is to live a life of sin and debauchery.\"\n", + "\n", + "\n", + "class QuestionAnswer(BaseModel):\n", + " question: str\n", + " answer: str\n", + "\n", + "\n", + "qa: QuestionAnswer = client.chat.completions.create(\n", + " model=\"gpt-3.5-turbo\",\n", + " response_model=QuestionAnswer,\n", + " messages=[\n", + " {\n", + " \"role\": \"system\",\n", + " \"content\": \"You are a system that answers questions based on the context. answer exactly what the question asks using the context.\",\n", + " },\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": f\"using the context: {context}\\n\\nAnswer the following question: {question}\",\n", + " },\n", + " ],\n", + ")\n", + "print(qa)\n", + "\n", + "\n", + "class QuestionAnswerNoEvil(BaseModel):\n", + " question: str\n", + " answer: Annotated[\n", + " str,\n", + " BeforeValidator(\n", + " llm_validator(\"don't say objectionable things\", allow_override=True)\n", + " ),\n", + " ]\n", + "\n", + "\n", + "try:\n", + " qa: QuestionAnswerNoEvil = client.chat.completions.create(\n", + " model=\"gpt-3.5-turbo\",\n", + " response_model=QuestionAnswerNoEvil,\n", + " messages=[\n", + " {\n", + " \"role\": \"system\",\n", + " \"content\": \"You are a system that answers questions based on the context. answer exactly what the question asks using the context.\",\n", + " },\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": f\"using the context: {context}\\n\\nAnswer the following question: {question}\",\n", + " },\n", + " ],\n", + " )\n", + "except Exception as e:\n", + " print(e)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Answering Questions with Validated Citations" + ] + }, + { + "cell_type": "code", + "execution_count": 42, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "question='What did the author do during college?' answer=[Fact(fact='The author, Jason Liu, studied Computational Mathematics and Physics in university.', substring_quote=['Computational Mathematics'])]\n" + ] + } + ], + "source": [ + "import re\n", + "from typing import List\n", + "\n", + "from pydantic import Field, BaseModel, model_validator, FieldValidationInfo\n", + "\n", + "\n", + "class Fact(BaseModel):\n", + " fact: str = Field(...)\n", + " substring_quote: List[str] = Field(...)\n", + "\n", + " @model_validator(mode=\"after\")\n", + " def validate_sources(self, info: FieldValidationInfo) -> \"Fact\":\n", + " text_chunks = info.context.get(\"text_chunk\", None)\n", + " spans = list(self.get_spans(text_chunks))\n", + " self.substring_quote = [text_chunks[span[0] : span[1]] for span in spans]\n", + " return self\n", + "\n", + " def get_spans(self, context):\n", + " for quote in self.substring_quote:\n", + " yield from self._get_span(quote, context)\n", + "\n", + " def _get_span(self, quote, context):\n", + " for match in re.finditer(re.escape(quote), context):\n", + " yield match.span()\n", + "\n", + "\n", + "class QuestionAnswer(BaseModel):\n", + " question: str = Field(...)\n", + " answer: List[Fact] = Field(...)\n", + "\n", + " @model_validator(mode=\"after\")\n", + " def validate_sources(self) -> \"QuestionAnswer\":\n", + " self.answer = [fact for fact in self.answer if len(fact.substring_quote) > 0]\n", + " return self\n", + "\n", + "\n", + "def ask_ai(question: str, context: str) -> QuestionAnswer:\n", + " return client.chat.completions.create(\n", + " model=\"gpt-3.5-turbo-0613\",\n", + " temperature=0.0,\n", + " response_model=QuestionAnswer,\n", + " messages=[\n", + " {\n", + " \"role\": \"system\",\n", + " \"content\": \"You are a world class algorithm to answer questions with correct and exact citations.\",\n", + " },\n", + " {\"role\": \"user\", \"content\": f\"{context}\"},\n", + " {\"role\": \"user\", \"content\": f\"Question: {question}\"},\n", + " ],\n", + " validation_context={\"text_chunk\": context},\n", + " )\n", + "\n", + "\n", + "question = \"What did the author do during college?\"\n", + "context = \"\"\"\n", + "My name is Jason Liu, and I grew up in Toronto Canada but I was born in China.\n", + "I went to an arts high school but in university I studied Computational Mathematics and physics.\n", + "As part of coop I worked at many companies including Stitchfix, Facebook.\n", + "I also started the Data Science club at the University of Waterloo and I was the president of the club for 2 years.\n", + "\"\"\"\n", + "\n", + "qa = ask_ai(question, context)\n", + "print(qa)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "python-3.8.10", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.5+" + }, + "orig_nbformat": 4 + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/examples/notebooks/Guidance.ipynb b/examples/notebooks/Guidance.ipynb index 045856ea2f..c525598530 100644 --- a/examples/notebooks/Guidance.ipynb +++ b/examples/notebooks/Guidance.ipynb @@ -28,7 +28,9 @@ "source": [ "import os\n", "\n", - "os.environ[\"OPENAI_API_KEY\"] = \"sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx\" # can be anything\n", + "os.environ[\"OPENAI_API_KEY\"] = (\n", + " \"sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx\" # can be anything\n", + ")\n", "os.environ[\"OPENAI_API_BASE\"] = \"http://100.64.159.73:8000/v1\"\n", "os.environ[\"OPENAI_API_HOST\"] = \"http://100.64.159.73:8000\"\n", "\n", @@ -38,21 +40,23 @@ "guidance.llm = guidance.llms.OpenAI(\"text-davinci-003\", caching=False)\n", "\n", "# define a guidance program that adapts a proverb\n", - "program = guidance(\"\"\"Tweak this proverb to apply to model instructions instead.\n", + "program = guidance(\n", + " \"\"\"Tweak this proverb to apply to model instructions instead.\n", "\n", "{{proverb}}\n", "- {{book}} {{chapter}}:{{verse}}\n", "\n", "UPDATED\n", "Where there is no guidance{{gen 'rewrite' stop=\"\\\\n-\"}}\n", - "- GPT {{gen 'chapter'}}:{{gen 'verse'}}\"\"\")\n", + "- GPT {{gen 'chapter'}}:{{gen 'verse'}}\"\"\"\n", + ")\n", "\n", "# execute the program on a specific proverb\n", "executed_program = program(\n", " proverb=\"Where there is no guidance, a people falls,\\nbut in an abundance of counselors there is safety.\",\n", " book=\"Proverbs\",\n", " chapter=11,\n", - " verse=14\n", + " verse=14,\n", ")" ] }, diff --git a/examples/notebooks/Multimodal.ipynb b/examples/notebooks/Multimodal.ipynb new file mode 100644 index 0000000000..8448ac1f7f --- /dev/null +++ b/examples/notebooks/Multimodal.ipynb @@ -0,0 +1,88 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "<div>\n", + " <img src=\"https://user-images.githubusercontent.com/1991296/230134379-7181e485-c521-4d23-a0d6-f7b3b61ba524.png\" width=\"500\"/>\n", + "</div>" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "{'text': 'Llama C++'}\n" + ] + } + ], + "source": [ + "from openai import OpenAI\n", + "\n", + "client = OpenAI(base_url=\"http://localhost:8000/v1\", api_key=\"llama.cpp\")\n", + "response = client.chat.completions.create(\n", + " model=\"gpt-4-vision-preview\",\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\n", + " \"type\": \"image_url\",\n", + " \"image_url\": {\n", + " \"url\": \"https://user-images.githubusercontent.com/1991296/230134379-7181e485-c521-4d23-a0d6-f7b3b61ba524.png\",\n", + " },\n", + " },\n", + " {\n", + " \"type\": \"text\",\n", + " \"text\": \"What does the image say. Format your response as a json object with a single 'text' key.\",\n", + " },\n", + " ],\n", + " }\n", + " ],\n", + " response_format={\n", + " \"type\": \"json_object\",\n", + " \"schema\": {\"type\": \"object\", \"properties\": {\"text\": {\"type\": \"string\"}}},\n", + " },\n", + ")\n", + "import json\n", + "\n", + "print(json.loads(response.choices[0].message.content))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.5+" + }, + "orig_nbformat": 4 + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/examples/notebooks/OpenHermesFunctionCalling.ipynb b/examples/notebooks/OpenHermesFunctionCalling.ipynb new file mode 100644 index 0000000000..13128be044 --- /dev/null +++ b/examples/notebooks/OpenHermesFunctionCalling.ipynb @@ -0,0 +1,933 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "{\n", + " \"name\": \"get_article_details\",\n", + " \"description\": \"Get article details from unstructured article text.\\ndate_published: formatted as \\\"MM/DD/YYYY\\\"\",\n", + " \"parameters\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"title\": {\n", + " \"type\": \"str\"\n", + " },\n", + " \"authors\": {\n", + " \"type\": \"list[str]\"\n", + " },\n", + " \"short_summary\": {\n", + " \"type\": \"str\"\n", + " },\n", + " \"date_published\": {\n", + " \"type\": \"str\"\n", + " },\n", + " \"tags\": {\n", + " \"type\": \"list[str]\"\n", + " }\n", + " }\n", + " },\n", + " \"returns\": \"Article\"\n", + "}\n" + ] + } + ], + "source": [ + "import json\n", + "import inspect\n", + "from typing import get_type_hints\n", + "\n", + "\n", + "class Article:\n", + " pass\n", + "\n", + "\n", + "class Weather:\n", + " pass\n", + "\n", + "\n", + "class Directions:\n", + " pass\n", + "\n", + "\n", + "def calculate_mortgage_payment(\n", + " loan_amount: int, interest_rate: float, loan_term: int\n", + ") -> float:\n", + " \"\"\"Get the monthly mortgage payment given an interest rate percentage.\"\"\"\n", + "\n", + " # TODO: you must implement this to actually call it later\n", + " pass\n", + "\n", + "\n", + "def get_article_details(\n", + " title: str,\n", + " authors: list[str],\n", + " short_summary: str,\n", + " date_published: str,\n", + " tags: list[str],\n", + ") -> Article:\n", + " '''Get article details from unstructured article text.\n", + " date_published: formatted as \"MM/DD/YYYY\"'''\n", + "\n", + " # TODO: you must implement this to actually call it later\n", + " pass\n", + "\n", + "\n", + "def get_weather(zip_code: str) -> Weather:\n", + " \"\"\"Get the current weather given a zip code.\"\"\"\n", + "\n", + " # TODO: you must implement this to actually call it later\n", + " pass\n", + "\n", + "\n", + "def get_directions(start: str, destination: str) -> Directions:\n", + " \"\"\"Get directions from Google Directions API.\n", + " start: start address as a string including zipcode (if any)\n", + " destination: end address as a string including zipcode (if any)\"\"\"\n", + "\n", + " # TODO: you must implement this to actually call it later\n", + " pass\n", + "\n", + "\n", + "def get_type_name(t):\n", + " name = str(t)\n", + " if \"list\" in name or \"dict\" in name:\n", + " return name\n", + " else:\n", + " return t.__name__\n", + "\n", + "\n", + "def serialize_function_to_json(func):\n", + " signature = inspect.signature(func)\n", + " type_hints = get_type_hints(func)\n", + "\n", + " function_info = {\n", + " \"name\": func.__name__,\n", + " \"description\": func.__doc__,\n", + " \"parameters\": {\"type\": \"object\", \"properties\": {}},\n", + " \"returns\": type_hints.get(\"return\", \"void\").__name__,\n", + " }\n", + "\n", + " for name, _ in signature.parameters.items():\n", + " param_type = get_type_name(type_hints.get(name, type(None)))\n", + " function_info[\"parameters\"][\"properties\"][name] = {\"type\": param_type}\n", + "\n", + " return json.dumps(function_info, indent=2)\n", + "\n", + "\n", + "print(serialize_function_to_json(get_article_details))" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "import xml.etree.ElementTree as ET\n", + "import re\n", + "\n", + "\n", + "def extract_function_calls(completion):\n", + " completion = completion.strip()\n", + " pattern = r\"(<multiplefunctions>(.*?)</multiplefunctions>)\"\n", + " match = re.search(pattern, completion, re.DOTALL)\n", + " if not match:\n", + " return None\n", + "\n", + " multiplefn = match.group(1)\n", + " root = ET.fromstring(multiplefn)\n", + " functions = root.findall(\"functioncall\")\n", + " return [json.loads(fn.text) for fn in functions]" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "metadata": {}, + "outputs": [], + "source": [ + "def generate_hermes_prompt(prompt, functions):\n", + " functions = \"\\n\\n\".join([serialize_function_to_json(fn) for fn in functions])\n", + " prompt = f\"\"\"<|im_start|>system\n", + "You are a helpful assistant with access to the following functions:\n", + "\n", + "{functions}\n", + "\n", + "To use these functions respond with:\n", + "<multiplefunctions>\n", + " <functioncall> {{\"name\": \"function_name\", \"arguments\": {{\"arg_1\": \"value_1\", \"arg_2\": value_2, ...}}}} </functioncall>\n", + " <functioncall> {{\"name\": \"function_name\", \"arguments\": {{\"arg_1\": \"value_1\", \"arg_2\": value_2, ...}}}} </functioncall>\n", + " ...\n", + "</multiplefunctions>\n", + "\n", + "Edge cases you must handle:\n", + "- If there are no functions that match the user request, you will respond politely that you cannot help.<|im_end|>\n", + "<|im_start|>user\n", + "{prompt}<|im_end|>\n", + "<|im_start|>assistant\"\"\"\n", + " return prompt" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "<|im_start|>system\n", + "You are a helpful assistant with access to the following functions:\n", + "\n", + "{\n", + " \"name\": \"get_weather\",\n", + " \"description\": \"Get the current weather given a zip code.\",\n", + " \"parameters\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"zip_code\": {\n", + " \"type\": \"str\"\n", + " }\n", + " }\n", + " },\n", + " \"returns\": \"Weather\"\n", + "}\n", + "\n", + "{\n", + " \"name\": \"calculate_mortgage_payment\",\n", + " \"description\": \"Get the monthly mortgage payment given an interest rate percentage.\",\n", + " \"parameters\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"loan_amount\": {\n", + " \"type\": \"int\"\n", + " },\n", + " \"interest_rate\": {\n", + " \"type\": \"float\"\n", + " },\n", + " \"loan_term\": {\n", + " \"type\": \"int\"\n", + " }\n", + " }\n", + " },\n", + " \"returns\": \"float\"\n", + "}\n", + "\n", + "{\n", + " \"name\": \"get_article_details\",\n", + " \"description\": \"Get article details from unstructured article text.\\ndate_published: formatted as \\\"MM/DD/YYYY\\\"\",\n", + " \"parameters\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"title\": {\n", + " \"type\": \"str\"\n", + " },\n", + " \"authors\": {\n", + " \"type\": \"list[str]\"\n", + " },\n", + " \"short_summary\": {\n", + " \"type\": \"str\"\n", + " },\n", + " \"date_published\": {\n", + " \"type\": \"str\"\n", + " },\n", + " \"tags\": {\n", + " \"type\": \"list[str]\"\n", + " }\n", + " }\n", + " },\n", + " \"returns\": \"Article\"\n", + "}\n", + "\n", + "To use these functions respond with:\n", + "<multiplefunctions>\n", + " <functioncall> {\"name\": \"function_name\", \"arguments\": {\"arg_1\": \"value_1\", \"arg_2\": value_2, ...}} </functioncall>\n", + " <functioncall> {\"name\": \"function_name\", \"arguments\": {\"arg_1\": \"value_1\", \"arg_2\": value_2, ...}} </functioncall>\n", + " ...\n", + "</multiplefunctions>\n", + "\n", + "Edge cases you must handle:\n", + "- If there are no functions that match the user request, you will respond politely that you cannot help.<|im_end|>\n", + "<|im_start|>user\n", + "What's the weather in 10001?<|im_end|>\n", + "<|im_start|>assistant\n", + "<|im_start|>system\n", + "You are a helpful assistant with access to the following functions:\n", + "\n", + "{\n", + " \"name\": \"get_weather\",\n", + " \"description\": \"Get the current weather given a zip code.\",\n", + " \"parameters\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"zip_code\": {\n", + " \"type\": \"str\"\n", + " }\n", + " }\n", + " },\n", + " \"returns\": \"Weather\"\n", + "}\n", + "\n", + "{\n", + " \"name\": \"calculate_mortgage_payment\",\n", + " \"description\": \"Get the monthly mortgage payment given an interest rate percentage.\",\n", + " \"parameters\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"loan_amount\": {\n", + " \"type\": \"int\"\n", + " },\n", + " \"interest_rate\": {\n", + " \"type\": \"float\"\n", + " },\n", + " \"loan_term\": {\n", + " \"type\": \"int\"\n", + " }\n", + " }\n", + " },\n", + " \"returns\": \"float\"\n", + "}\n", + "\n", + "{\n", + " \"name\": \"get_article_details\",\n", + " \"description\": \"Get article details from unstructured article text.\\ndate_published: formatted as \\\"MM/DD/YYYY\\\"\",\n", + " \"parameters\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"title\": {\n", + " \"type\": \"str\"\n", + " },\n", + " \"authors\": {\n", + " \"type\": \"list[str]\"\n", + " },\n", + " \"short_summary\": {\n", + " \"type\": \"str\"\n", + " },\n", + " \"date_published\": {\n", + " \"type\": \"str\"\n", + " },\n", + " \"tags\": {\n", + " \"type\": \"list[str]\"\n", + " }\n", + " }\n", + " },\n", + " \"returns\": \"Article\"\n", + "}\n", + "\n", + "To use these functions respond with:\n", + "<multiplefunctions>\n", + " <functioncall> {\"name\": \"function_name\", \"arguments\": {\"arg_1\": \"value_1\", \"arg_2\": value_2, ...}} </functioncall>\n", + " <functioncall> {\"name\": \"function_name\", \"arguments\": {\"arg_1\": \"value_1\", \"arg_2\": value_2, ...}} </functioncall>\n", + " ...\n", + "</multiplefunctions>\n", + "\n", + "Edge cases you must handle:\n", + "- If there are no functions that match the user request, you will respond politely that you cannot help.<|im_end|>\n", + "<|im_start|>user\n", + "Determine the monthly mortgage payment for a loan amount of $200,000, an interest rate of 4%, and a loan term of 30 years.<|im_end|>\n", + "<|im_start|>assistant\n", + "<|im_start|>system\n", + "You are a helpful assistant with access to the following functions:\n", + "\n", + "{\n", + " \"name\": \"get_weather\",\n", + " \"description\": \"Get the current weather given a zip code.\",\n", + " \"parameters\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"zip_code\": {\n", + " \"type\": \"str\"\n", + " }\n", + " }\n", + " },\n", + " \"returns\": \"Weather\"\n", + "}\n", + "\n", + "{\n", + " \"name\": \"calculate_mortgage_payment\",\n", + " \"description\": \"Get the monthly mortgage payment given an interest rate percentage.\",\n", + " \"parameters\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"loan_amount\": {\n", + " \"type\": \"int\"\n", + " },\n", + " \"interest_rate\": {\n", + " \"type\": \"float\"\n", + " },\n", + " \"loan_term\": {\n", + " \"type\": \"int\"\n", + " }\n", + " }\n", + " },\n", + " \"returns\": \"float\"\n", + "}\n", + "\n", + "{\n", + " \"name\": \"get_article_details\",\n", + " \"description\": \"Get article details from unstructured article text.\\ndate_published: formatted as \\\"MM/DD/YYYY\\\"\",\n", + " \"parameters\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"title\": {\n", + " \"type\": \"str\"\n", + " },\n", + " \"authors\": {\n", + " \"type\": \"list[str]\"\n", + " },\n", + " \"short_summary\": {\n", + " \"type\": \"str\"\n", + " },\n", + " \"date_published\": {\n", + " \"type\": \"str\"\n", + " },\n", + " \"tags\": {\n", + " \"type\": \"list[str]\"\n", + " }\n", + " }\n", + " },\n", + " \"returns\": \"Article\"\n", + "}\n", + "\n", + "To use these functions respond with:\n", + "<multiplefunctions>\n", + " <functioncall> {\"name\": \"function_name\", \"arguments\": {\"arg_1\": \"value_1\", \"arg_2\": value_2, ...}} </functioncall>\n", + " <functioncall> {\"name\": \"function_name\", \"arguments\": {\"arg_1\": \"value_1\", \"arg_2\": value_2, ...}} </functioncall>\n", + " ...\n", + "</multiplefunctions>\n", + "\n", + "Edge cases you must handle:\n", + "- If there are no functions that match the user request, you will respond politely that you cannot help.<|im_end|>\n", + "<|im_start|>user\n", + "What's the current exchange rate for USD to EUR?<|im_end|>\n", + "<|im_start|>assistant\n" + ] + } + ], + "source": [ + "prompts = [\n", + " \"What's the weather in 10001?\",\n", + " \"Determine the monthly mortgage payment for a loan amount of $200,000, an interest rate of 4%, and a loan term of 30 years.\",\n", + " \"What's the current exchange rate for USD to EUR?\",\n", + "]\n", + "functions = [get_weather, calculate_mortgage_payment, get_article_details]\n", + "\n", + "for prompt in prompts:\n", + " print(generate_hermes_prompt(prompt, functions))" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "ggml_init_cublas: GGML_CUDA_FORCE_MMQ: no\n", + "ggml_init_cublas: CUDA_USE_TENSOR_CORES: yes\n", + "ggml_init_cublas: found 1 CUDA devices:\n", + " Device 0: NVIDIA GeForce RTX 2060, compute capability 7.5\n", + "llama_model_loader: loaded meta data with 20 key-value pairs and 291 tensors from ../../models/OpenHermes-2.5-Mistral-7B-GGUF/openhermes-2.5-mistral-7b.Q4_K_M.gguf (version GGUF V3 (latest))\n", + "llama_model_loader: - tensor 0: token_embd.weight q4_K [ 4096, 32002, 1, 1 ]\n", + "llama_model_loader: - tensor 1: blk.0.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 2: blk.0.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 3: blk.0.attn_v.weight q6_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 4: blk.0.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 5: blk.0.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 6: blk.0.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 7: blk.0.ffn_down.weight q6_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 8: blk.0.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 9: blk.0.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 10: blk.1.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 11: blk.1.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 12: blk.1.attn_v.weight q6_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 13: blk.1.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 14: blk.1.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 15: blk.1.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 16: blk.1.ffn_down.weight q6_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 17: blk.1.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 18: blk.1.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 19: blk.2.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 20: blk.2.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 21: blk.2.attn_v.weight q6_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 22: blk.2.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 23: blk.2.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 24: blk.2.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 25: blk.2.ffn_down.weight q6_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 26: blk.2.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 27: blk.2.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 28: blk.3.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 29: blk.3.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 30: blk.3.attn_v.weight q6_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 31: blk.3.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 32: blk.3.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 33: blk.3.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 34: blk.3.ffn_down.weight q6_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 35: blk.3.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 36: blk.3.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 37: blk.4.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 38: blk.4.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 39: blk.4.attn_v.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 40: blk.4.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 41: blk.4.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 42: blk.4.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 43: blk.4.ffn_down.weight q4_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 44: blk.4.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 45: blk.4.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 46: blk.5.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 47: blk.5.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 48: blk.5.attn_v.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 49: blk.5.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 50: blk.5.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 51: blk.5.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 52: blk.5.ffn_down.weight q4_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 53: blk.5.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 54: blk.5.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 55: blk.6.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 56: blk.6.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 57: blk.6.attn_v.weight q6_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 58: blk.6.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 59: blk.6.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 60: blk.6.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 61: blk.6.ffn_down.weight q6_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 62: blk.6.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 63: blk.6.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 64: blk.7.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 65: blk.7.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 66: blk.7.attn_v.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 67: blk.7.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 68: blk.7.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 69: blk.7.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 70: blk.7.ffn_down.weight q4_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 71: blk.7.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 72: blk.7.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 73: blk.8.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 74: blk.8.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 75: blk.8.attn_v.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 76: blk.8.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 77: blk.8.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 78: blk.8.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 79: blk.8.ffn_down.weight q4_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 80: blk.8.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 81: blk.8.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 82: blk.9.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 83: blk.9.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 84: blk.9.attn_v.weight q6_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 85: blk.9.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 86: blk.9.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 87: blk.9.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 88: blk.9.ffn_down.weight q6_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 89: blk.9.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 90: blk.9.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 91: blk.10.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 92: blk.10.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 93: blk.10.attn_v.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 94: blk.10.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 95: blk.10.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 96: blk.10.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 97: blk.10.ffn_down.weight q4_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 98: blk.10.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 99: blk.10.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 100: blk.11.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 101: blk.11.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 102: blk.11.attn_v.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 103: blk.11.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 104: blk.11.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 105: blk.11.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 106: blk.11.ffn_down.weight q4_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 107: blk.11.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 108: blk.11.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 109: blk.12.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 110: blk.12.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 111: blk.12.attn_v.weight q6_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 112: blk.12.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 113: blk.12.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 114: blk.12.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 115: blk.12.ffn_down.weight q6_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 116: blk.12.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 117: blk.12.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 118: blk.13.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 119: blk.13.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 120: blk.13.attn_v.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 121: blk.13.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 122: blk.13.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 123: blk.13.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 124: blk.13.ffn_down.weight q4_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 125: blk.13.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 126: blk.13.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 127: blk.14.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 128: blk.14.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 129: blk.14.attn_v.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 130: blk.14.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 131: blk.14.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 132: blk.14.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 133: blk.14.ffn_down.weight q4_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 134: blk.14.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 135: blk.14.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 136: blk.15.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 137: blk.15.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 138: blk.15.attn_v.weight q6_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 139: blk.15.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 140: blk.15.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 141: blk.15.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 142: blk.15.ffn_down.weight q6_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 143: blk.15.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 144: blk.15.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 145: blk.16.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 146: blk.16.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 147: blk.16.attn_v.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 148: blk.16.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 149: blk.16.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 150: blk.16.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 151: blk.16.ffn_down.weight q4_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 152: blk.16.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 153: blk.16.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 154: blk.17.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 155: blk.17.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 156: blk.17.attn_v.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 157: blk.17.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 158: blk.17.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 159: blk.17.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 160: blk.17.ffn_down.weight q4_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 161: blk.17.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 162: blk.17.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 163: blk.18.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 164: blk.18.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 165: blk.18.attn_v.weight q6_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 166: blk.18.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 167: blk.18.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 168: blk.18.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 169: blk.18.ffn_down.weight q6_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 170: blk.18.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 171: blk.18.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 172: blk.19.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 173: blk.19.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 174: blk.19.attn_v.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 175: blk.19.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 176: blk.19.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 177: blk.19.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 178: blk.19.ffn_down.weight q4_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 179: blk.19.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 180: blk.19.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 181: blk.20.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 182: blk.20.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 183: blk.20.attn_v.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 184: blk.20.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 185: blk.20.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 186: blk.20.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 187: blk.20.ffn_down.weight q4_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 188: blk.20.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 189: blk.20.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 190: blk.21.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 191: blk.21.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 192: blk.21.attn_v.weight q6_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 193: blk.21.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 194: blk.21.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 195: blk.21.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 196: blk.21.ffn_down.weight q6_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 197: blk.21.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 198: blk.21.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 199: blk.22.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 200: blk.22.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 201: blk.22.attn_v.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 202: blk.22.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 203: blk.22.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 204: blk.22.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 205: blk.22.ffn_down.weight q4_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 206: blk.22.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 207: blk.22.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 208: blk.23.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 209: blk.23.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 210: blk.23.attn_v.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 211: blk.23.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 212: blk.23.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 213: blk.23.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 214: blk.23.ffn_down.weight q4_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 215: blk.23.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 216: blk.23.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 217: blk.24.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 218: blk.24.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 219: blk.24.attn_v.weight q6_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 220: blk.24.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 221: blk.24.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 222: blk.24.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 223: blk.24.ffn_down.weight q6_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 224: blk.24.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 225: blk.24.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 226: blk.25.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 227: blk.25.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 228: blk.25.attn_v.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 229: blk.25.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 230: blk.25.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 231: blk.25.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 232: blk.25.ffn_down.weight q4_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 233: blk.25.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 234: blk.25.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 235: blk.26.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 236: blk.26.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 237: blk.26.attn_v.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 238: blk.26.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 239: blk.26.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 240: blk.26.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 241: blk.26.ffn_down.weight q4_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 242: blk.26.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 243: blk.26.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 244: blk.27.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 245: blk.27.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 246: blk.27.attn_v.weight q6_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 247: blk.27.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 248: blk.27.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 249: blk.27.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 250: blk.27.ffn_down.weight q6_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 251: blk.27.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 252: blk.27.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 253: blk.28.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 254: blk.28.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 255: blk.28.attn_v.weight q6_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 256: blk.28.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 257: blk.28.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 258: blk.28.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 259: blk.28.ffn_down.weight q6_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 260: blk.28.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 261: blk.28.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 262: blk.29.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 263: blk.29.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 264: blk.29.attn_v.weight q6_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 265: blk.29.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 266: blk.29.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 267: blk.29.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 268: blk.29.ffn_down.weight q6_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 269: blk.29.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 270: blk.29.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 271: blk.30.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 272: blk.30.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 273: blk.30.attn_v.weight q6_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 274: blk.30.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 275: blk.30.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 276: blk.30.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 277: blk.30.ffn_down.weight q6_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 278: blk.30.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 279: blk.30.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 280: blk.31.attn_q.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 281: blk.31.attn_k.weight q4_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 282: blk.31.attn_v.weight q6_K [ 4096, 1024, 1, 1 ]\n", + "llama_model_loader: - tensor 283: blk.31.attn_output.weight q4_K [ 4096, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 284: blk.31.ffn_gate.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 285: blk.31.ffn_up.weight q4_K [ 4096, 14336, 1, 1 ]\n", + "llama_model_loader: - tensor 286: blk.31.ffn_down.weight q6_K [ 14336, 4096, 1, 1 ]\n", + "llama_model_loader: - tensor 287: blk.31.attn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 288: blk.31.ffn_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 289: output_norm.weight f32 [ 4096, 1, 1, 1 ]\n", + "llama_model_loader: - tensor 290: output.weight q6_K [ 4096, 32002, 1, 1 ]\n", + "llama_model_loader: - kv 0: general.architecture str = llama\n", + "llama_model_loader: - kv 1: general.name str = teknium_openhermes-2.5-mistral-7b\n", + "llama_model_loader: - kv 2: llama.context_length u32 = 32768\n", + "llama_model_loader: - kv 3: llama.embedding_length u32 = 4096\n", + "llama_model_loader: - kv 4: llama.block_count u32 = 32\n", + "llama_model_loader: - kv 5: llama.feed_forward_length u32 = 14336\n", + "llama_model_loader: - kv 6: llama.rope.dimension_count u32 = 128\n", + "llama_model_loader: - kv 7: llama.attention.head_count u32 = 32\n", + "llama_model_loader: - kv 8: llama.attention.head_count_kv u32 = 8\n", + "llama_model_loader: - kv 9: llama.attention.layer_norm_rms_epsilon f32 = 0.000010\n", + "llama_model_loader: - kv 10: llama.rope.freq_base f32 = 10000.000000\n", + "llama_model_loader: - kv 11: general.file_type u32 = 15\n", + "llama_model_loader: - kv 12: tokenizer.ggml.model str = llama\n", + "llama_model_loader: - kv 13: tokenizer.ggml.tokens arr[str,32002] = [\"<unk>\", \"<s>\", \"</s>\", \"<0x00>\", \"<...\n", + "llama_model_loader: - kv 14: tokenizer.ggml.scores arr[f32,32002] = [0.000000, 0.000000, 0.000000, 0.0000...\n", + "llama_model_loader: - kv 15: tokenizer.ggml.token_type arr[i32,32002] = [2, 3, 3, 6, 6, 6, 6, 6, 6, 6, 6, 6, ...\n", + "llama_model_loader: - kv 16: tokenizer.ggml.bos_token_id u32 = 1\n", + "llama_model_loader: - kv 17: tokenizer.ggml.eos_token_id u32 = 32000\n", + "llama_model_loader: - kv 18: tokenizer.ggml.padding_token_id u32 = 0\n", + "llama_model_loader: - kv 19: general.quantization_version u32 = 2\n", + "llama_model_loader: - type f32: 65 tensors\n", + "llama_model_loader: - type q4_K: 193 tensors\n", + "llama_model_loader: - type q6_K: 33 tensors\n", + "llm_load_vocab: special tokens definition check successful ( 261/32002 ).\n", + "llm_load_print_meta: format = GGUF V3 (latest)\n", + "llm_load_print_meta: arch = llama\n", + "llm_load_print_meta: vocab type = SPM\n", + "llm_load_print_meta: n_vocab = 32002\n", + "llm_load_print_meta: n_merges = 0\n", + "llm_load_print_meta: n_ctx_train = 32768\n", + "llm_load_print_meta: n_embd = 4096\n", + "llm_load_print_meta: n_head = 32\n", + "llm_load_print_meta: n_head_kv = 8\n", + "llm_load_print_meta: n_layer = 32\n", + "llm_load_print_meta: n_rot = 128\n", + "llm_load_print_meta: n_gqa = 4\n", + "llm_load_print_meta: f_norm_eps = 0.0e+00\n", + "llm_load_print_meta: f_norm_rms_eps = 1.0e-05\n", + "llm_load_print_meta: f_clamp_kqv = 0.0e+00\n", + "llm_load_print_meta: f_max_alibi_bias = 0.0e+00\n", + "llm_load_print_meta: n_ff = 14336\n", + "llm_load_print_meta: rope scaling = linear\n", + "llm_load_print_meta: freq_base_train = 10000.0\n", + "llm_load_print_meta: freq_scale_train = 1\n", + "llm_load_print_meta: n_yarn_orig_ctx = 32768\n", + "llm_load_print_meta: rope_finetuned = unknown\n", + "llm_load_print_meta: model type = 7B\n", + "llm_load_print_meta: model ftype = mostly Q4_K - Medium\n", + "llm_load_print_meta: model params = 7.24 B\n", + "llm_load_print_meta: model size = 4.07 GiB (4.83 BPW) \n", + "llm_load_print_meta: general.name = teknium_openhermes-2.5-mistral-7b\n", + "llm_load_print_meta: BOS token = 1 '<s>'\n", + "llm_load_print_meta: EOS token = 32000 '<|im_end|>'\n", + "llm_load_print_meta: UNK token = 0 '<unk>'\n", + "llm_load_print_meta: PAD token = 0 '<unk>'\n", + "llm_load_print_meta: LF token = 13 '<0x0A>'\n", + "llm_load_tensors: ggml ctx size = 0.11 MiB\n", + "llm_load_tensors: using CUDA for GPU acceleration\n", + "llm_load_tensors: mem required = 70.42 MiB\n", + "llm_load_tensors: offloading 32 repeating layers to GPU\n", + "llm_load_tensors: offloading non-repeating layers to GPU\n", + "llm_load_tensors: offloaded 35/35 layers to GPU\n", + "llm_load_tensors: VRAM used: 4095.06 MiB\n", + "...............................................................................................\n", + "llama_new_context_with_model: n_ctx = 2048\n", + "llama_new_context_with_model: freq_base = 10000.0\n", + "llama_new_context_with_model: freq_scale = 1\n", + "llama_kv_cache_init: offloading v cache to GPU\n", + "llama_kv_cache_init: offloading k cache to GPU\n", + "llama_kv_cache_init: VRAM kv self = 256.00 MiB\n", + "llama_new_context_with_model: kv self size = 256.00 MiB\n", + "llama_build_graph: non-view tensors processed: 740/740\n", + "llama_new_context_with_model: compute buffer total size = 159.07 MiB\n", + "llama_new_context_with_model: VRAM scratch buffer: 156.00 MiB\n", + "llama_new_context_with_model: total VRAM used: 4507.07 MiB (model: 4095.06 MiB, context: 412.00 MiB)\n" + ] + } + ], + "source": [ + "import llama_cpp\n", + "\n", + "llama = llama_cpp.Llama(\n", + " model_path=\"../../models/OpenHermes-2.5-Mistral-7B-GGUF/openhermes-2.5-mistral-7b.Q4_K_M.gguf\",\n", + " n_gpu_layers=-1,\n", + " n_ctx=2048,\n", + " verbose=False,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 22, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[{'name': 'get_weather', 'arguments': {'zip_code': '10001'}}]\n", + "====================================================================================================\n", + "[{'name': 'calculate_mortgage_payment', 'arguments': {'loan_amount': 200000, 'interest_rate': 0.04, 'loan_term': 30}}]\n", + "====================================================================================================\n", + "Unfortunately, I do not have a built-in function to check currency exchange rates. However, you can use third-party APIs or websites like Google Finance or XE to get this information.\n", + "====================================================================================================\n" + ] + } + ], + "source": [ + "prompts = [\n", + " \"What's the weather in 10001?\",\n", + " \"Determine the monthly mortgage payment for a loan amount of $200,000, an interest rate of 4%, and a loan term of 30 years.\",\n", + " \"What's the current exchange rate for USD to EUR?\",\n", + "]\n", + "functions = [get_weather, calculate_mortgage_payment, get_article_details]\n", + "\n", + "for prompt in prompts:\n", + " prompt = generate_hermes_prompt(prompt, functions)\n", + " completion = llama.create_completion(prompt, max_tokens=-1)[\"choices\"][0][\"text\"]\n", + " function_calls = extract_function_calls(completion)\n", + " if function_calls:\n", + " print(function_calls)\n", + " else:\n", + " print(completion.strip())\n", + " print(\"=\" * 100)" + ] + }, + { + "cell_type": "code", + "execution_count": 23, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "get_weather\n", + "{'zip_code': '05751'}\n", + "====================================================================================================\n", + "get_weather\n", + "{'zip_code': '05751'}\n", + "get_weather\n", + "{'zip_code': '07030'}\n", + "calculate_mortgage_payment\n", + "{'loan_amount': 250000, 'interest_rate': 4.18, 'loan_term': 30}\n", + "====================================================================================================\n", + "I don't have a function to get exchange rates, but I can provide some resources where you can find this information. You can check websites like Google Finance, XE.com, or Yahoo Finance for up-to-date currency exchange rates.\n", + "====================================================================================================\n" + ] + } + ], + "source": [ + "prompts = [\n", + " \"What's the weather in 05751?\",\n", + " \"I'm planning a trip to Killington, Vermont (05751) from Hoboken, NJ (07030). Can you get me weather for both locations and directions?\",\n", + " \"What's the current exchange rate for USD to EUR?\",\n", + "]\n", + "\n", + "for prompt in prompts:\n", + " completion = llama.create_completion(\n", + " generate_hermes_prompt(prompt, functions), max_tokens=-1\n", + " )[\"choices\"][0][\"text\"]\n", + " function_calls = extract_function_calls(completion)\n", + "\n", + " if function_calls:\n", + " for function in function_calls:\n", + " print(function[\"name\"])\n", + " print(function[\"arguments\"])\n", + " else:\n", + " print(completion.strip())\n", + "\n", + " print(\"=\" * 100)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.5+" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/examples/notebooks/PerformanceTuning.ipynb b/examples/notebooks/PerformanceTuning.ipynb index 76e26fbd10..ba74e4a41f 100644 --- a/examples/notebooks/PerformanceTuning.ipynb +++ b/examples/notebooks/PerformanceTuning.ipynb @@ -13,6 +13,7 @@ "import llama_cpp\n", "\n", "import numpy as np\n", + "\n", "np.int = int\n", "\n", "from skopt.space import Integer, Categorical\n", @@ -25,7 +26,7 @@ " Categorical([True, False], name=\"f16_kv\"),\n", " Categorical([True, False], name=\"use_mlock\"),\n", " Integer(1, multiprocessing.cpu_count(), name=\"n_threads\"),\n", - " Integer(1, 2048, name=\"n_batch\")\n", + " Integer(1, 2048, name=\"n_batch\"),\n", "]\n", "\n", "# TODO: Make this a random prompt to avoid any cache related inconsistencies\n", @@ -41,18 +42,25 @@ "\n", "from skopt.utils import use_named_args\n", "\n", + "\n", "@use_named_args(space)\n", "def objective(**params):\n", " f16_kv = params[\"f16_kv\"]\n", " use_mlock = params[\"use_mlock\"]\n", " n_threads = params[\"n_threads\"]\n", " n_batch = params[\"n_batch\"]\n", - " llm = llama_cpp.Llama(model_path=MODEL_PATH, f16_kv=f16_kv, use_mlock=use_mlock, n_threads=n_threads, n_batch=n_batch)\n", + " llm = llama_cpp.Llama(\n", + " model_path=MODEL_PATH,\n", + " f16_kv=f16_kv,\n", + " use_mlock=use_mlock,\n", + " n_threads=n_threads,\n", + " n_batch=n_batch,\n", + " )\n", "\n", " t1 = time.time()\n", " output = llm(\n", " PROMPT,\n", - " max_tokens=1, # Only optimize prompt processing\n", + " max_tokens=1, # Only optimize prompt processing\n", " stop=[\"###\", \"\\n\"],\n", " echo=True,\n", " )\n", @@ -5240,10 +5248,7 @@ "source": [ "from skopt import gp_minimize\n", "\n", - "res = gp_minimize(\n", - " objective,\n", - " space\n", - ")" + "res = gp_minimize(objective, space)" ] }, { diff --git a/examples/ray/README.md b/examples/ray/README.md new file mode 100644 index 0000000000..6e338ba17e --- /dev/null +++ b/examples/ray/README.md @@ -0,0 +1,19 @@ +This is an example of doing LLM inference with [Ray](https://docs.ray.io/en/latest/index.html) and [Ray Serve](https://docs.ray.io/en/latest/serve/index.html). + +First, install the requirements: + +```bash +$ pip install -r requirements.txt +``` + +Deploy a GGUF model to Ray Serve with the following command: + +```bash +$ serve run llm:llm_builder model_path='../models/mistral-7b-instruct-v0.2.Q4_K_M.gguf' +``` + +This will start an API endpoint at `http://localhost:8000/`. You can query the model like this: + +```bash +$ curl -k -d '{"prompt": "tell me a joke", "max_tokens": 128}' -X POST http://localhost:8000 +``` diff --git a/examples/ray/llm.py b/examples/ray/llm.py new file mode 100755 index 0000000000..2325dd303b --- /dev/null +++ b/examples/ray/llm.py @@ -0,0 +1,21 @@ +from starlette.requests import Request +from typing import Dict +from ray import serve +from ray.serve import Application +from llama_cpp import Llama + + +@serve.deployment +class LlamaDeployment: + def __init__(self, model_path: str): + self._llm = Llama(model_path=model_path) + + async def __call__(self, http_request: Request) -> Dict: + input_json = await http_request.json() + prompt = input_json["prompt"] + max_tokens = input_json.get("max_tokens", 64) + return self._llm(prompt, max_tokens=max_tokens) + + +def llm_builder(args: Dict[str, str]) -> Application: + return LlamaDeployment.bind(args["model_path"]) diff --git a/examples/ray/requirements.txt b/examples/ray/requirements.txt new file mode 100644 index 0000000000..a409fb882d --- /dev/null +++ b/examples/ray/requirements.txt @@ -0,0 +1,3 @@ +ray[serve] +--extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cpu +llama-cpp-python diff --git a/examples/server/README.md b/examples/server/README.md new file mode 100644 index 0000000000..ff04374fc5 --- /dev/null +++ b/examples/server/README.md @@ -0,0 +1,475 @@ +# Server Example + +This example is an updated OpenAI-compatible web server that depends only on the low-level C bindings. +It supports batched inference, prompt caching, response parsing, `/v1/responses`, `/v1/embeddings`, disk sequence caching, MTP, LoRA, and multimodal image/audio inputs. + +## Setup + +The server is a [`uv` inline script](https://docs.astral.sh/uv/guides/scripts/), so `uv` can create the script environment and install the Python dependencies automatically. + +```bash +cd examples/server +uv run --script server.py -C configs/qwen3.5-0.8b.json +``` + +Use `uv run --extra-index-url` to pull a pre-built `llama-cpp-python` binary wheel instead of building from source. + +```bash +cd examples/server +uv run \ + --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cpu \ + --script server.py -C configs/qwen3.5-0.8b.json +``` + +Pick the wheel index that matches the backend you want. + +| Backend | Wheel index | +| --- | --- | +| CPU | `https://abetlen.github.io/llama-cpp-python/whl/cpu` | +| CUDA 11.8 | `https://abetlen.github.io/llama-cpp-python/whl/cu118` | +| CUDA 12.1 | `https://abetlen.github.io/llama-cpp-python/whl/cu121` | +| CUDA 12.2 | `https://abetlen.github.io/llama-cpp-python/whl/cu122` | +| CUDA 12.3 | `https://abetlen.github.io/llama-cpp-python/whl/cu123` | +| CUDA 12.4 | `https://abetlen.github.io/llama-cpp-python/whl/cu124` | +| CUDA 12.5 | `https://abetlen.github.io/llama-cpp-python/whl/cu125` | +| CUDA 13.0 | `https://abetlen.github.io/llama-cpp-python/whl/cu130` | +| CUDA 13.2 | `https://abetlen.github.io/llama-cpp-python/whl/cu132` | +| Metal | `https://abetlen.github.io/llama-cpp-python/whl/metal` | +| ROCm | `https://abetlen.github.io/llama-cpp-python/whl/rocm72` | +| Vulkan | `https://abetlen.github.io/llama-cpp-python/whl/vulkan` | + +See the repository installation section for the full [pre-built wheel requirements](../../README.md#supported-backends). + +## Model Configs + +The smallest checked-in example uses Qwen3.5 0.8B so the server can be started on a normal development machine. + +| Config | Model | Notes | +| --- | --- | --- | +| [`configs/bge-small-en-v1.5.json`](configs/bge-small-en-v1.5.json) | [`CompendiumLabs/bge-small-en-v1.5-gguf`](https://huggingface.co/CompendiumLabs/bge-small-en-v1.5-gguf) | Small embedding model config for `/v1/embeddings`. | +| [`configs/qwen3.5-0.8b.json`](configs/qwen3.5-0.8b.json) | [`lmstudio-community/Qwen3.5-0.8B-GGUF`](https://huggingface.co/lmstudio-community/Qwen3.5-0.8B-GGUF) | Default small multimodal example. | +| [`configs/gemma-4-12b-it-qat.json`](configs/gemma-4-12b-it-qat.json) | [`unsloth/gemma-4-12B-it-qat-GGUF`](https://huggingface.co/unsloth/gemma-4-12B-it-qat-GGUF) | Larger Gemma 4 QAT multimodal config with projector. | +| [`configs/qwen3.6-27b.json`](configs/qwen3.6-27b.json) | [`unsloth/Qwen3.6-27B-GGUF`](https://huggingface.co/unsloth/Qwen3.6-27B-GGUF) | Larger Qwen3.6 multimodal config. | +| [`configs/qwen3.6-35b-a3b.json`](configs/qwen3.6-35b-a3b.json) | [`unsloth/Qwen3.6-35B-A3B-GGUF`](https://huggingface.co/unsloth/Qwen3.6-35B-A3B-GGUF) | Larger Qwen3.6 MoE multimodal config. | +| [`configs/gpt-oss-120b.json`](configs/gpt-oss-120b.json) | [`ggml-org/gpt-oss-120b-GGUF`](https://huggingface.co/ggml-org/gpt-oss-120b-GGUF) | Large text-only split-GGUF config. | + +The larger model configs default to `n_gpu_layers: -1` and `flash_attn: true`. + +## Client Examples + +Point an OpenAI-compatible client at the local `/v1` base URL. + +### Chat Completions + +```python +from openai import OpenAI + +client = OpenAI(base_url="http://127.0.0.1:8000/v1", api_key="not-used") + +response = client.chat.completions.create( + model="qwen3.5-0.8b-vl", + messages=[{"role": "user", "content": "What is the capital of France?"}], +) +print(response.choices[0].message.content) +``` + +### Responses API + +```python +from openai import OpenAI + +client = OpenAI(base_url="http://127.0.0.1:8000/v1", api_key="not-used") + +response = client.responses.create( + model="qwen3.5-0.8b-vl", + input="Write one sentence about why prefix caching helps batched inference.", +) +print(response.output_text) +``` + +### Embeddings + +Start the server with an embedding config before calling `/v1/embeddings`. + +```bash +cd examples/server +uv run --script server.py -C configs/bge-small-en-v1.5.json +``` + +```python +from openai import OpenAI + +client = OpenAI(base_url="http://127.0.0.1:8000/v1", api_key="not-used") + +response = client.embeddings.create( + model="bge-small-en-v1.5", + input=["The food was delicious.", "The meal was excellent."], +) +print(len(response.data[0].embedding)) +``` + +## API Surface + +| Endpoint | Purpose | Reference | +| --- | --- | --- | +| `POST /v1/completions` | Legacy text completions with streaming, stop sequences, logprobs, penalties, seeds, and grammar-backed JSON output. | [OpenAI Completions API](https://platform.openai.com/docs/api-reference/completions) | +| `POST /v1/embeddings` | OpenAI-compatible embeddings for embedding-mode GGUF models, including string inputs, token inputs, base64 output, and dimensions truncation. | [OpenAI Embeddings API](https://platform.openai.com/docs/api-reference/embeddings) | +| `POST /v1/chat/completions` | Chat completions with streaming, tools, forced tool choice, reasoning parsing, multimodal content parts, and structured response parsing. | [OpenAI Chat API](https://platform.openai.com/docs/api-reference/chat) | +| `POST /v1/responses` | Stateless Responses API compatibility for clients that use response items and response events. | [OpenAI Responses API](https://platform.openai.com/docs/api-reference/responses) | +| `WS /v1/responses` | Stateful websocket Responses transport with per-connection `previous_response_id` replay. | [OpenAI Responses API](https://platform.openai.com/docs/api-reference/responses) | +| `GET /v1/models` | Returns the configured model alias. | [OpenAI Models API](https://platform.openai.com/docs/api-reference/models) | +| `GET /healthz` | Returns a simple typed health response. | | +| `GET /metrics` | Exposes scheduler, cache, draft, and model metrics in Prometheus text format. | [Prometheus exposition format](https://prometheus.io/docs/instrumenting/exposition_formats/) | + +## Config Overview + +Config files have three top-level sections. + +```json +{ + "server": {}, + "model": {}, + "disk_cache": {} +} +``` + +| Section | Required | Purpose | +| --- | --- | --- | +| `server` | No | Uvicorn host and port settings. | +| `model` | Yes | Model source, llama.cpp runtime settings, chat formatting, LoRA, MTMD, draft decoding, and output parsing. | +| `disk_cache` | No | Optional serialized sequence cache for repeated prompt prefixes. | + +## `server` + +Use `server.host` and `server.port` to choose the bind address. + +```json +{ + "server": { + "host": "0.0.0.0", + "port": 8000 + } +} +``` + +| Field | Default | Notes | +| --- | --- | --- | +| `host` | `127.0.0.1` | Use `0.0.0.0` to expose the server on the network. | +| `port` | `8000` | Passed directly to `uvicorn.run()`. | + +## `model` Source + +Load a local GGUF with `path` or download a GGUF from Hugging Face with `from_pretrained`. + +```json +{ + "model": { + "alias": "qwen3.5-0.8b-vl", + "from_pretrained": { + "repo_id": "lmstudio-community/Qwen3.5-0.8B-GGUF", + "filename": "Qwen3.5-0.8B-Q8_0.gguf" + } + } +} +``` + +| Field | Notes | +| --- | --- | +| `path` | Local GGUF path. | +| `from_pretrained.repo_id` | Hugging Face model repository. | +| `from_pretrained.filename` | File name or glob pattern for the GGUF. | +| `from_pretrained.additional_files` | Extra files to download from the same repository. | +| `from_pretrained.cache_dir` | Optional Hugging Face cache directory. | +| `alias` | Model id returned by `/v1/models` and used by OpenAI-compatible clients. | + +See the [Hugging Face Hub download guide](https://huggingface.co/docs/huggingface_hub/guides/download) for cache behavior and repository file resolution. + +## llama.cpp Runtime Settings + +Most model runtime fields map to `llama_model_params` or `llama_context_params` in [`llama.h`](https://github.com/ggml-org/llama.cpp/blob/master/include/llama.h). + +```json +{ + "model": { + "n_ctx": 32768, + "n_seq_max": 64, + "n_batch": 128, + "n_ubatch": 128, + "threads": 4, + "threads_batch": 8, + "kv_unified": true, + "use_mmap": true, + "use_mlock": true + } +} +``` + +| Field | Purpose | +| --- | --- | +| `n_ctx` | Total context size. | +| `n_seq_max` | Maximum number of concurrent llama.cpp sequence ids. | +| `n_batch` | Logical batch capacity. | +| `n_ubatch` | Physical microbatch capacity. | +| `threads` | Decode thread count. | +| `threads_batch` | Prefill and batch thread count. | +| `kv_unified` | Selects unified or per-sequence memory layout. | +| `embedding` | Overrides embedding mode; omit to auto-detect pooled embedding GGUFs from model metadata. | +| `pooling_type` | Overrides pooled embedding behavior for embedding models, such as `1` for mean pooling. | +| `store_logits` | Keeps logits after decode when needed by sampling or diagnostics. | +| `use_mmap` | Memory maps model weights. | +| `use_mlock` | Attempts to lock model pages into RAM. | + +GPU and backend-related fields are passed through to llama.cpp when set. + +```json +{ + "model": { + "n_gpu_layers": -1, + "split_mode": 1, + "main_gpu": 0, + "tensor_split": [1.0], + "flash_attn": true, + "offload_kqv": true, + "op_offload": true + } +} +``` + +## Chat Template + +`model.chat_template` is a Jinja chat template compatible with the style used by [Hugging Face chat templates](https://huggingface.co/docs/transformers/chat_templating). + +```json +{ + "model": { + "chat_template": "{{ bos_token }}{{ messages[0].content }}{{ eos_token }}" + } +} +``` + +Use an array of strings when the template is too large to read or edit as one JSON string. + +```json +{ + "model": { + "chat_template": [ + "{{ bos_token }}", + "{{ messages[0].content }}", + "{{ eos_token }}" + ] + } +} +``` + +The checked-in [`configs/qwen3.5-0.8b.json`](configs/qwen3.5-0.8b.json) includes a Qwen3.5 template with reasoning text, tool calls, forced tool choice, image markers, and video markers. + +## Response Parsing + +`model.response_schema` parses generated text into OpenAI-compatible fields with JSON Schema plus the Hugging Face `x-regex` extensions. + +```json +{ + "model": { + "response_schema": { + "type": "object", + "properties": { + "role": {"const": "assistant"}, + "content": { + "type": "string", + "x-regex": "^(.*)$" + } + }, + "required": ["role"] + } + } +} +``` + +Use `x-regex-iterator` and `x-regex-key-value` to parse repeated tool-call blocks. + +See [Hugging Face response parsing](https://huggingface.co/docs/transformers/chat_response_parsing) and the [JSON Schema reference](https://json-schema.org/understanding-json-schema/reference) for the underlying schema concepts. + +## Multimodal `model.mtmd` + +`model.mtmd` loads a llama.cpp multimodal projector and enables OpenAI-style image and audio content parts. + +```json +{ + "model": { + "mtmd": { + "mmproj_from_pretrained": { + "repo_id": "lmstudio-community/Qwen3.5-0.8B-GGUF", + "filename": "mmproj-Qwen3.5-0.8B-BF16.gguf" + }, + "embedding_cache": { + "path": ".cache/mtmd-embeddings", + "max_bytes": 1073741824 + }, + "image_max_bytes": 20971520, + "audio_max_bytes": 104857600, + "image_timeout_seconds": 10.0 + } + } +} +``` + +| Field | Purpose | +| --- | --- | +| `mmproj_path` | Local multimodal projector path. | +| `mmproj_from_pretrained` | Hugging Face projector source. | +| `embedding_cache.path` | Directory for cached image and audio embeddings. | +| `embedding_cache.max_bytes` | Maximum embedding cache size. | +| `image_max_bytes` | Maximum image payload size. | +| `audio_max_bytes` | Maximum audio payload size. | +| `image_timeout_seconds` | Timeout for remote image and audio URL fetches. | + +Send image inputs with OpenAI chat content parts. + +```json +{ + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this image."}, + {"type": "image_url", "image_url": {"url": "https://example.com/image.jpg"}} + ] + } + ] +} +``` + +Send audio inputs as a URL or as base64 `input_audio` content. + +```json +{ + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Transcribe this audio."}, + {"type": "input_audio", "input_audio": {"data": "...base64...", "format": "wav"}} + ] + } + ] +} +``` + +## Media Loading Policy + +Remote `http:` and `https:` media URLs are unrestricted unless `allowed_media_domains` is set. + +```json +{ + "model": { + "mtmd": { + "allowed_media_domains": ["example.com", "static.example.com"] + } + } +} +``` + +Local `file:` media URLs are disabled unless `allowed_local_media_path` is set. + +```json +{ + "model": { + "mtmd": { + "allowed_local_media_path": "/srv/llama-cpp-python/media" + } + } +} +``` + +`allowed_media_domains` matches exact hostnames and does not allow wildcard patterns. + +## LoRA `model.loras` + +Load LoRA adapters once at startup from local files or Hugging Face. + +```json +{ + "model": { + "loras": [ + { + "from_pretrained": { + "repo_id": "example/qwen-lora-gguf", + "filename": "adapter.gguf" + }, + "scale": 1.0 + } + ] + } +} +``` + +The current implementation does not hot-swap LoRAs per request. + +## Draft Decoding + +Set `model.draft_model` to enable speculative draft providers. + +```json +{ + "model": { + "draft_model": "prompt-lookup-decoding", + "draft_model_num_pred_tokens": 8, + "draft_model_max_ngram_size": 4 + } +} +``` + +### Multi-Token Prediction (MTP) + +Use MTP when the loaded model and llama.cpp build expose the required draft state. + +```json +{ + "model": { + "draft_model": "draft-mtp", + "draft_model_num_pred_tokens": 2, + "draft_model_threads": 4, + "draft_model_threads_batch": 8 + } +} +``` + +By default `draft-mtp` creates the MTP context from the target model. +Set `draft_model_path` or `draft_model_from_pretrained` when the model uses a separate assistant GGUF. + +```json +{ + "model": { + "draft_model": "draft-mtp", + "draft_model_num_pred_tokens": 2, + "draft_model_from_pretrained": { + "repo_id": "example/gemma-assistant-GGUF", + "filename": "assistant.gguf" + } + } +} +``` + +MTP currently applies to text-only requests. + +## Disk Sequence Cache + +`disk_cache` stores serialized llama.cpp sequence state for repeated prompt prefixes. + +```json +{ + "disk_cache": { + "path": ".cache/sequences", + "max_bytes": 1073741824, + "min_tokens": 128 + } +} +``` + +| Field | Purpose | +| --- | --- | +| `path` | Directory for cached sequence files. | +| `max_bytes` | Maximum cache size before background cleanup removes entries. | +| `min_tokens` | Minimum prefix length that is worth saving. | + +The cache is versioned by model and context compatibility data and should be treated as ephemeral. diff --git a/examples/server/configs/bge-small-en-v1.5.json b/examples/server/configs/bge-small-en-v1.5.json new file mode 100644 index 0000000000..3fc8016df8 --- /dev/null +++ b/examples/server/configs/bge-small-en-v1.5.json @@ -0,0 +1,22 @@ +{ + "server": { + "host": "0.0.0.0", + "port": 8000 + }, + "model": { + "alias": "bge-small-en-v1.5", + "from_pretrained": { + "repo_id": "CompendiumLabs/bge-small-en-v1.5-gguf", + "filename": "bge-small-en-v1.5-q4_k_m.gguf" + }, + "n_ctx": 512, + "n_seq_max": 16, + "n_batch": 512, + "n_ubatch": 512, + "threads": 4, + "threads_batch": 8, + "kv_unified": true, + "store_logits": false, + "use_mmap": true + } +} diff --git a/examples/server/configs/gemma-4-12b-it-qat.json b/examples/server/configs/gemma-4-12b-it-qat.json new file mode 100644 index 0000000000..60ffd2a987 --- /dev/null +++ b/examples/server/configs/gemma-4-12b-it-qat.json @@ -0,0 +1,451 @@ +{ + "server": { + "host": "0.0.0.0", + "port": 8000 + }, + "model": { + "alias": "gemma-4-12b-it-qat", + "from_pretrained": { + "repo_id": "unsloth/gemma-4-12B-it-qat-GGUF", + "filename": "gemma-4-12B-it-qat-UD-Q4_K_XL.gguf" + }, + "mtmd": { + "mmproj_from_pretrained": { + "repo_id": "unsloth/gemma-4-12B-it-qat-GGUF", + "filename": "mmproj-BF16.gguf" + } + }, + "n_ctx": 32768, + "max_output_tokens": 4096, + "n_seq_max": 8, + "n_batch": 512, + "n_ubatch": 512, + "threads": 8, + "threads_batch": 8, + "kv_unified": true, + "store_logits": false, + "use_mmap": true, + "use_mlock": false, + "n_gpu_layers": -1, + "flash_attn": true, + "response_schema": { + "type": "object", + "properties": { + "role": { + "const": "assistant" + }, + "reasoning_content": { + "type": "string", + "x-regex": "^(?:<\\|turn>model\\n)?<\\|channel>thought\\n(.*?)<channel\\|>" + }, + "content": { + "type": "string", + "x-regex": "^(?:<\\|turn>model\\n)?(?:(?:<\\|channel>thought\\n).*?<channel\\|>)?(.*?)(?=<\\|tool_call>|<turn\\|>|$)" + }, + "tool_calls": { + "type": "array", + "x-regex-iterator": "<\\|tool_call>(call:[^\\{\\[]+[\\{\\[].*?[\\}\\]])<tool_call\\|>", + "items": { + "type": "object", + "properties": { + "type": { + "const": "function" + }, + "function": { + "type": "object", + "properties": { + "name": { + "type": "string", + "x-regex": "^call:(\\w+)[\\{\\[]" + }, + "arguments": { + "type": "object", + "x-regex": "^call:\\w+(\\{.*\\}|\\[.*\\])$", + "x-parser": "gemma4-tool-call", + "additionalProperties": true + } + }, + "required": [ + "name", + "arguments" + ] + } + }, + "required": [ + "type", + "function" + ] + } + } + }, + "required": [ + "role" + ] + }, + "chat_template": [ + "{%- macro format_parameters(properties, required, filter_keys=false) -%}\n", + " {%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%}\n", + " {%- set ns = namespace(found_first=false) -%}\n", + " {%- for key, value in properties | dictsort -%}\n", + " {%- set add_comma = false -%}\n", + " {%- if not filter_keys or key not in standard_keys -%}\n", + " {%- if ns.found_first %},{% endif -%}\n", + " {%- set ns.found_first = true -%}\n", + " {{ key }}:{\n", + " {%- if value['description'] -%}\n", + " description:<|\"|>{{ value['description'] }}<|\"|>\n", + " {%- set add_comma = true -%}\n", + " {%- endif -%}\n", + " {%- if value['type'] | upper == 'STRING' -%}\n", + " {%- if value['enum'] -%}\n", + " {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}\n", + " enum:{{ format_argument(value['enum']) }}\n", + " {%- endif -%}\n", + " {%- elif value['type'] | upper == 'ARRAY' -%}\n", + " {%- if value['items'] is mapping and value['items'] -%}\n", + " {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}\n", + " items:{\n", + " {%- set ns_items = namespace(found_first=false) -%}\n", + " {%- for item_key, item_value in value['items'] | dictsort -%}\n", + " {%- if item_value is not none -%}\n", + " {%- if ns_items.found_first %},{% endif -%}\n", + " {%- set ns_items.found_first = true -%}\n", + " {%- if item_key == 'properties' -%}\n", + " properties:{\n", + " {%- if item_value is mapping -%}\n", + " {{- format_parameters(item_value, value['items']['required'] | default([])) -}}\n", + " {%- endif -%}\n", + " }\n", + " {%- elif item_key == 'required' -%}\n", + " required:[\n", + " {%- for req_item in item_value -%}\n", + " <|\"|>{{- req_item -}}<|\"|>\n", + " {%- if not loop.last %},{% endif -%}\n", + " {%- endfor -%}\n", + " ]\n", + " {%- elif item_key == 'type' -%}\n", + " {%- if item_value is string -%}\n", + " type:{{ format_argument(item_value | upper) }}\n", + " {%- else -%}\n", + " type:{{ format_argument(item_value | map('upper') | list) }}\n", + " {%- endif -%}\n", + " {%- else -%}\n", + " {{ item_key }}:{{ format_argument(item_value) }}\n", + " {%- endif -%}\n", + " {%- endif -%}\n", + " {%- endfor -%}\n", + " }\n", + " {%- endif -%}\n", + " {%- endif -%}\n", + " {%- if value['nullable'] %}\n", + " {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}\n", + " nullable:true\n", + " {%- endif -%}\n", + " {%- if value['type'] | upper == 'OBJECT' -%}\n", + " {%- if value['properties'] is defined and value['properties'] is mapping -%}\n", + " {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}\n", + " properties:{\n", + " {{- format_parameters(value['properties'], value['required'] | default([])) -}}\n", + " }\n", + " {%- elif value is mapping -%}\n", + " {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}\n", + " properties:{\n", + " {{- format_parameters(value, value['required'] | default([]), filter_keys=true) -}}\n", + " }\n", + " {%- endif -%}\n", + " {%- if value['required'] -%}\n", + " {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}\n", + " required:[\n", + " {%- for item in value['required'] | default([]) -%}\n", + " <|\"|>{{- item -}}<|\"|>\n", + " {%- if not loop.last %},{% endif -%}\n", + " {%- endfor -%}\n", + " ]\n", + " {%- endif -%}\n", + " {%- endif -%}\n", + " {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}\n", + " type:<|\"|>{{ value['type'] | upper }}<|\"|>}\n", + " {%- endif -%}\n", + " {%- endfor -%}\n", + "{%- endmacro -%}\n", + "{%- macro format_function_declaration(tool_data) -%}\n", + " declaration:{{- tool_data['function']['name'] -}}{description:<|\"|>{{- tool_data['function']['description'] -}}<|\"|>\n", + " {%- set params = tool_data['function']['parameters'] -%}\n", + " {%- if params -%}\n", + " ,parameters:{\n", + " {%- if params['properties'] -%}\n", + " properties:{ {{- format_parameters(params['properties'], params['required']) -}} },\n", + " {%- endif -%}\n", + " {%- if params['required'] -%}\n", + " required:[\n", + " {%- for item in params['required'] -%}\n", + " <|\"|>{{- item -}}<|\"|>\n", + " {{- ',' if not loop.last -}}\n", + " {%- endfor -%}\n", + " ],\n", + " {%- endif -%}\n", + " {%- if params['type'] -%}\n", + " type:<|\"|>{{- params['type'] | upper -}}<|\"|>}\n", + " {%- endif -%}\n", + " {%- endif -%}\n", + " {%- if 'response' in tool_data['function'] -%}\n", + " {%- set response_declaration = tool_data['function']['response'] -%}\n", + " ,response:{\n", + " {%- if response_declaration['description'] -%}\n", + " description:<|\"|>{{- response_declaration['description'] -}}<|\"|>,\n", + " {%- endif -%}\n", + " {%- if response_declaration['type'] | upper == 'OBJECT' -%}\n", + " type:<|\"|>{{- response_declaration['type'] | upper -}}<|\"|>}\n", + " {%- endif -%}\n", + " {%- endif -%}\n", + " }\n", + "{%- endmacro -%}\n", + "{%- macro format_argument(argument, escape_keys=True) -%}\n", + " {%- if argument is string -%}\n", + " {{- '<|\"|>' + argument + '<|\"|>' -}}\n", + " {%- elif argument is boolean -%}\n", + " {{- 'true' if argument else 'false' -}}\n", + " {%- elif argument is mapping -%}\n", + " {{- '{' -}}\n", + " {%- set ns = namespace(found_first=false) -%}\n", + " {%- for key, value in argument | dictsort -%}\n", + " {%- if ns.found_first %},{% endif -%}\n", + " {%- set ns.found_first = true -%}\n", + " {%- if escape_keys -%}\n", + " {{- '<|\"|>' + key + '<|\"|>' -}}\n", + " {%- else -%}\n", + " {{- key -}}\n", + " {%- endif -%}\n", + " :{{- format_argument(value, escape_keys=escape_keys) -}}\n", + " {%- endfor -%}\n", + " {{- '}' -}}\n", + " {%- elif argument is sequence -%}\n", + " {{- '[' -}}\n", + " {%- for item in argument -%}\n", + " {{- format_argument(item, escape_keys=escape_keys) -}}\n", + " {%- if not loop.last %},{% endif -%}\n", + " {%- endfor -%}\n", + " {{- ']' -}}\n", + " {%- else -%}\n", + " {{- argument -}}\n", + " {%- endif -%}\n", + "{%- endmacro -%}\n", + "{%- macro strip_thinking(text) -%}\n", + " {%- set ns = namespace(result='') -%}\n", + " {%- for part in text.split('<channel|>') -%}\n", + " {%- if '<|channel>' in part -%}\n", + " {%- set ns.result = ns.result + part.split('<|channel>')[0] -%}\n", + " {%- else -%}\n", + " {%- set ns.result = ns.result + part -%}\n", + " {%- endif -%}\n", + " {%- endfor -%}\n", + " {{- ns.result | trim -}}\n", + "{%- endmacro -%}\n", + "\n", + "{%- macro format_tool_response_block(tool_name, response) -%}\n", + " {{- '<|tool_response>' -}}\n", + " {%- if response is mapping -%}\n", + " {{- 'response:' + tool_name + '{' -}}\n", + " {%- for key, value in response | dictsort -%}\n", + " {{- key -}}:{{- format_argument(value, escape_keys=False) -}}\n", + " {%- if not loop.last %},{% endif -%}\n", + " {%- endfor -%}\n", + " {{- '}' -}}\n", + " {%- else -%}\n", + " {{- 'response:' + tool_name + '{value:' + format_argument(response, escape_keys=False) + '}' -}}\n", + " {%- endif -%}\n", + " {{- '<tool_response|>' -}}\n", + "{%- endmacro -%}\n", + "\n", + "{%- set ns = namespace(prev_message_type=None) -%}\n", + "{%- set loop_messages = messages -%}\n", + "{{- bos_token -}}\n", + "{#- Handle System/Tool Definitions Block -#}\n", + "{%- if (enable_thinking is defined and enable_thinking) or tools or messages[0]['role'] in ['system', 'developer'] -%}\n", + " {{- '<|turn>system\\n' -}}\n", + " {#- Inject Thinking token at the very top of the FIRST system turn -#}\n", + " {%- if enable_thinking is defined and enable_thinking -%}\n", + " {{- '<|think|>\\n' -}}\n", + " {%- set ns.prev_message_type = 'think' -%}\n", + " {%- endif -%}\n", + " {%- if messages[0]['role'] in ['system', 'developer'] -%}\n", + " {%- if messages[0]['content'] is string -%}\n", + " {{- messages[0]['content'] | trim -}}\n", + " {%- elif messages[0]['content'] is sequence -%}\n", + " {%- for item in messages[0]['content'] -%}\n", + " {{- item['text'] | trim + ' '-}}\n", + " {%- endfor -%}\n", + " {%- endif -%}\n", + " {%- set loop_messages = messages[1:] -%}\n", + " {%- endif -%}\n", + " {%- if tools -%}\n", + " {%- for tool in tools %}\n", + " {{- '<|tool>' -}}\n", + " {{- format_function_declaration(tool) | trim -}}\n", + " {{- '<tool|>' -}}\n", + " {%- endfor %}\n", + " {%- set ns.prev_message_type = 'tool' -%}\n", + " {%- endif -%}\n", + " {{- '<turn|>\\n' -}}\n", + "{%- endif %}\n", + "\n", + "{#- Pre-scan: find last user message index for reasoning guard -#}\n", + "{%- set ns_turn = namespace(last_user_idx=-1) -%}\n", + "{%- for i in range(loop_messages | length) -%}\n", + " {%- if loop_messages[i]['role'] == 'user' -%}\n", + " {%- set ns_turn.last_user_idx = i -%}\n", + " {%- endif -%}\n", + "{%- endfor -%}\n", + "\n", + "{#- Loop through messages -#}\n", + "{%- for message in loop_messages -%}\n", + " {%- if message['role'] != 'tool' -%}\n", + " {%- set ns.prev_message_type = None -%}\n", + " {%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%}\n", + " {#- Detect continuation: suppress duplicate <|turn>model when previous non-tool message was also assistant -#}\n", + " {%- set prev_nt = namespace(role=None, found=false) -%}\n", + " {%- if loop.index0 > 0 -%}\n", + " {%- for j in range(loop.index0 - 1, -1, -1) -%}\n", + " {%- if not prev_nt.found -%}\n", + " {%- if loop_messages[j]['role'] != 'tool' -%}\n", + " {%- set prev_nt.role = loop_messages[j]['role'] -%}\n", + " {%- set prev_nt.found = true -%}\n", + " {%- endif -%}\n", + " {%- endif -%}\n", + " {%- endfor -%}\n", + " {%- endif -%}\n", + " {%- set continue_same_model_turn = (role == 'model' and prev_nt.role == 'assistant') -%}\n", + " {%- if not continue_same_model_turn -%}\n", + " {{- '<|turn>' + role + '\\n' }}\n", + " {%- endif -%}\n", + "\n", + " {#- Render reasoning/reasoning_content as thinking channel -#}\n", + " {%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%}\n", + " {%- if thinking_text and loop.index0 > ns_turn.last_user_idx and message.get('tool_calls') -%}\n", + " {{- '<|channel>thought\\n' + thinking_text + '\\n<channel|>' -}}\n", + " {%- endif -%}\n", + "\n", + " {%- if message['tool_calls'] -%}\n", + " {%- for tool_call in message['tool_calls'] -%}\n", + " {%- set function = tool_call['function'] -%}\n", + " {{- '<|tool_call>call:' + function['name'] + '{' -}}\n", + " {%- if function['arguments'] is mapping -%}\n", + " {%- set ns_args = namespace(found_first=false) -%}\n", + " {%- for key, value in function['arguments'] | dictsort -%}\n", + " {%- if ns_args.found_first %},{% endif -%}\n", + " {%- set ns_args.found_first = true -%}\n", + " {{- key -}}:{{- format_argument(value, escape_keys=False) -}}\n", + " {%- endfor -%}\n", + " {%- elif function['arguments'] is string -%}\n", + " {{- function['arguments'] -}}\n", + " {%- endif -%}\n", + " {{- '}<tool_call|>' -}}\n", + " {%- endfor -%}\n", + " {%- set ns.prev_message_type = 'tool_call' -%}\n", + " {%- endif -%}\n", + "\n", + " {%- set ns_tr_out = namespace(flag=false) -%}\n", + " {%- if message.get('tool_responses') -%}\n", + " {#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#}\n", + " {%- for tool_response in message['tool_responses'] -%}\n", + " {{- format_tool_response_block(tool_response['name'] | default('unknown'), tool_response['response']) -}}\n", + " {%- set ns_tr_out.flag = true -%}\n", + " {%- set ns.prev_message_type = 'tool_response' -%}\n", + " {%- endfor -%}\n", + " {%- elif message.get('tool_calls') -%}\n", + " {#- OpenAI Chat Completions: forward-scan consecutive role:tool messages -#}\n", + " {%- set ns_tool_scan = namespace(stopped=false) -%}\n", + " {%- for k in range(loop.index0 + 1, loop_messages | length) -%}\n", + " {%- if ns_tool_scan.stopped -%}\n", + " {%- elif loop_messages[k]['role'] != 'tool' -%}\n", + " {%- set ns_tool_scan.stopped = true -%}\n", + " {%- else -%}\n", + " {%- set follow = loop_messages[k] -%}\n", + " {#- Resolve tool_call_id to function name -#}\n", + " {%- set ns_tname = namespace(name=follow.get('name') | default('unknown')) -%}\n", + " {%- for tc in message['tool_calls'] -%}\n", + " {%- if tc.get('id') == follow.get('tool_call_id') -%}\n", + " {%- set ns_tname.name = tc['function']['name'] -%}\n", + " {%- endif -%}\n", + " {%- endfor -%}\n", + " {#- Handle content as string or content-parts array -#}\n", + " {%- set tool_body = follow.get('content') -%}\n", + " {%- if tool_body is string -%}\n", + " {{- format_tool_response_block(ns_tname.name, tool_body) -}}\n", + " {%- elif tool_body is sequence and tool_body is not string -%}\n", + " {%- set ns_txt = namespace(s='') -%}\n", + " {%- for part in tool_body -%}\n", + " {%- if part.get('type') == 'text' -%}\n", + " {%- set ns_txt.s = ns_txt.s + (part.get('text') | default('')) -%}\n", + " {%- endif -%}\n", + " {%- endfor -%}\n", + " {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}}\n", + " {%- for part in tool_body -%}\n", + " {%- if part.get('type') == 'image' -%}\n", + " {{- '<|image|>' -}}\n", + " {%- elif part.get('type') == 'audio' -%}\n", + " {{- '<|audio|>' -}}\n", + " {%- elif part.get('type') == 'video' -%}\n", + " {{- '<|video|>' -}}\n", + " {%- endif -%}\n", + " {%- endfor -%}\n", + " {%- else -%}\n", + " {{- format_tool_response_block(ns_tname.name, tool_body) -}}\n", + " {%- endif -%}\n", + " {%- set ns_tr_out.flag = true -%}\n", + " {%- set ns.prev_message_type = 'tool_response' -%}\n", + " {%- endif -%}\n", + " {%- endfor -%}\n", + " {%- endif -%}\n", + "\n", + " {%- set captured_content -%}\n", + " {%- if message['content'] is string -%}\n", + " {%- if role == 'model' -%}\n", + " {{- strip_thinking(message['content']) -}}\n", + " {%- else -%}\n", + " {{- message['content'] | trim -}}\n", + " {%- endif -%}\n", + " {%- elif message['content'] is sequence -%}\n", + " {%- for item in message['content'] -%}\n", + " {%- if item['type'] == 'text' -%}\n", + " {%- if role == 'model' -%}\n", + " {{- strip_thinking(item['text']) -}}\n", + " {%- else -%}\n", + " {{- item['text'] | trim -}}\n", + " {%- endif -%}\n", + " {%- elif item['type'] == 'image' -%}\n", + " {{- '<|image|>' -}}\n", + " {%- set ns.prev_message_type = 'image' -%}\n", + " {%- elif item['type'] == 'audio' -%}\n", + " {{- '<|audio|>' -}}\n", + " {%- set ns.prev_message_type = 'audio' -%}\n", + " {%- elif item['type'] == 'video' -%}\n", + " {{- '<|video|>' -}}\n", + " {%- set ns.prev_message_type = 'video' -%}\n", + " {%- endif -%}\n", + " {%- endfor -%}\n", + " {%- endif -%}\n", + " {%- endset -%}\n", + "\n", + " {{- captured_content -}}\n", + " {%- set has_content = captured_content | trim | length > 0 -%}\n", + "\n", + " {%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%}\n", + " {{- '<|tool_response>' -}}\n", + " {%- elif not (ns_tr_out.flag and not has_content) -%}\n", + " {{- '<turn|>\\n' -}}\n", + " {%- endif -%}\n", + " {%- endif -%}\n", + "{%- endfor -%}\n", + "\n", + "{%- if add_generation_prompt -%}\n", + " {%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%}\n", + " {{- '<|turn>model\\n' -}}\n", + " {%- if not enable_thinking | default(false) -%}\n", + " {{- '<|channel>thought\\n<channel|>' -}}\n", + " {%- endif -%}\n", + " {%- endif -%}\n", + "{%- endif -%}" + ] + } +} diff --git a/examples/server/configs/gpt-oss-120b.json b/examples/server/configs/gpt-oss-120b.json new file mode 100644 index 0000000000..64aaf040d7 --- /dev/null +++ b/examples/server/configs/gpt-oss-120b.json @@ -0,0 +1,431 @@ +{ + "server": { + "host": "0.0.0.0", + "port": 8000 + }, + "model": { + "alias": "gpt-oss-120b", + "from_pretrained": { + "repo_id": "ggml-org/gpt-oss-120b-GGUF", + "filename": "gpt-oss-120b-mxfp4-00001-of-00003.gguf", + "additional_files": [ + "gpt-oss-120b-mxfp4-00002-of-00003.gguf", + "gpt-oss-120b-mxfp4-00003-of-00003.gguf" + ] + }, + "n_ctx": 32768, + "max_output_tokens": 4096, + "n_seq_max": 4, + "n_batch": 512, + "n_ubatch": 512, + "threads": 8, + "threads_batch": 8, + "kv_unified": true, + "store_logits": false, + "use_mmap": true, + "use_mlock": false, + "n_gpu_layers": -1, + "flash_attn": true, + "response_schema": { + "type": "object", + "properties": { + "role": { + "const": "assistant" + }, + "reasoning_content": { + "type": "string", + "x-regex": "^(?:<\\|start\\|>assistant)?<\\|channel\\|>analysis<\\|message\\|>(.*?)(?=<\\|end\\|>|<\\|start\\|>assistant|$)" + }, + "content": { + "type": "string", + "x-regex": "^(?:.*?<\\|start\\|>assistant)?<\\|channel\\|>final<\\|message\\|>(.*?)(?=<\\|return\\|>|<\\|end\\|>|$)" + }, + "tool_calls": { + "type": "array", + "x-regex-iterator": "(<\\|start\\|>assistant<\\|channel\\|>commentary to=functions\\.\\w+ <\\|constrain\\|>json<\\|message\\|>.*?<\\|call\\|>)", + "items": { + "type": "object", + "properties": { + "type": { + "const": "function" + }, + "function": { + "type": "object", + "properties": { + "name": { + "type": "string", + "x-regex": "^<\\|start\\|>assistant<\\|channel\\|>commentary to=functions\\.(\\w+) <\\|constrain\\|>json<\\|message\\|>" + }, + "arguments": { + "type": "object", + "x-regex": "^<\\|start\\|>assistant<\\|channel\\|>commentary to=functions\\.\\w+ <\\|constrain\\|>json<\\|message\\|>(.*?)<\\|call\\|>$", + "x-parser": "json", + "additionalProperties": true + } + }, + "required": [ + "name", + "arguments" + ] + } + }, + "required": [ + "type", + "function" + ] + } + } + }, + "required": [ + "role" + ] + }, + "chat_template": [ + "{#-\n", + " In addition to the normal inputs of `messages` and `tools`, this template also accepts the\n", + " following kwargs:\n", + " - \"builtin_tools\": A list, can contain \"browser\" and/or \"python\".\n", + " - \"model_identity\": A string that optionally describes the model identity.\n", + " - \"reasoning_effort\": A string that describes the reasoning effort, defaults to \"medium\".\n", + " #}\n", + "\n", + "{#- Tool Definition Rendering ============================================== #}\n", + "{%- macro render_typescript_type(param_spec, required_params, is_nullable=false) -%}\n", + " {%- if param_spec.type == \"array\" -%}\n", + " {%- if param_spec['items'] -%}\n", + " {%- if param_spec['items']['type'] == \"string\" -%}\n", + " {{- \"string[]\" }}\n", + " {%- elif param_spec['items']['type'] == \"number\" -%}\n", + " {{- \"number[]\" }}\n", + " {%- elif param_spec['items']['type'] == \"integer\" -%}\n", + " {{- \"number[]\" }}\n", + " {%- elif param_spec['items']['type'] == \"boolean\" -%}\n", + " {{- \"boolean[]\" }}\n", + " {%- else -%}\n", + " {%- set inner_type = render_typescript_type(param_spec['items'], required_params) -%}\n", + " {%- if inner_type == \"object | object\" or inner_type|length > 50 -%}\n", + " {{- \"any[]\" }}\n", + " {%- else -%}\n", + " {{- inner_type + \"[]\" }}\n", + " {%- endif -%}\n", + " {%- endif -%}\n", + " {%- if param_spec.nullable -%}\n", + " {{- \" | null\" }}\n", + " {%- endif -%}\n", + " {%- else -%}\n", + " {{- \"any[]\" }}\n", + " {%- if param_spec.nullable -%}\n", + " {{- \" | null\" }}\n", + " {%- endif -%}\n", + " {%- endif -%}\n", + " {%- elif param_spec.type is defined and param_spec.type is iterable and param_spec.type is not string and param_spec.type is not mapping and param_spec.type[0] is defined -%}\n", + " {#- Handle array of types like [\"object\", \"object\"] from Union[dict, list] #}\n", + " {%- if param_spec.type | length > 1 -%}\n", + " {{- param_spec.type | join(\" | \") }}\n", + " {%- else -%}\n", + " {{- param_spec.type[0] }}\n", + " {%- endif -%}\n", + " {%- elif param_spec.oneOf -%}\n", + " {#- Handle oneOf schemas - check for complex unions and fallback to any #}\n", + " {%- set has_object_variants = false -%}\n", + " {%- for variant in param_spec.oneOf -%}\n", + " {%- if variant.type == \"object\" -%}\n", + " {%- set has_object_variants = true -%}\n", + " {%- endif -%}\n", + " {%- endfor -%}\n", + " {%- if has_object_variants and param_spec.oneOf|length > 1 -%}\n", + " {{- \"any\" }}\n", + " {%- else -%}\n", + " {%- for variant in param_spec.oneOf -%}\n", + " {{- render_typescript_type(variant, required_params) -}}\n", + " {%- if variant.description %}\n", + " {{- \"// \" + variant.description }}\n", + " {%- endif -%}\n", + " {%- if variant.default is defined %}\n", + " {{ \"// default: \" + variant.default|tojson }}\n", + " {%- endif -%}\n", + " {%- if not loop.last %}\n", + " {{- \" | \" }}\n", + " {% endif -%}\n", + " {%- endfor -%}\n", + " {%- endif -%}\n", + " {%- elif param_spec.type == \"string\" -%}\n", + " {%- if param_spec.enum -%}\n", + " {{- '\"' + param_spec.enum|join('\" | \"') + '\"' -}}\n", + " {%- else -%}\n", + " {{- \"string\" }}\n", + " {%- if param_spec.nullable %}\n", + " {{- \" | null\" }}\n", + " {%- endif -%}\n", + " {%- endif -%}\n", + " {%- elif param_spec.type == \"number\" -%}\n", + " {{- \"number\" }}\n", + " {%- elif param_spec.type == \"integer\" -%}\n", + " {{- \"number\" }}\n", + " {%- elif param_spec.type == \"boolean\" -%}\n", + " {{- \"boolean\" }}\n", + "\n", + " {%- elif param_spec.type == \"object\" -%}\n", + " {%- if param_spec.properties -%}\n", + " {{- \"{\\n\" }}\n", + " {%- for prop_name, prop_spec in param_spec.properties.items() -%}\n", + " {{- prop_name -}}\n", + " {%- if prop_name not in (param_spec.required or []) -%}\n", + " {{- \"?\" }}\n", + " {%- endif -%}\n", + " {{- \": \" }}\n", + " {{ render_typescript_type(prop_spec, param_spec.required or []) }}\n", + " {%- if not loop.last -%}\n", + " {{-\", \" }}\n", + " {%- endif -%}\n", + " {%- endfor -%}\n", + " {{- \"}\" }}\n", + " {%- else -%}\n", + " {{- \"object\" }}\n", + " {%- endif -%}\n", + " {%- else -%}\n", + " {{- \"any\" }}\n", + " {%- endif -%}\n", + "{%- endmacro -%}\n", + "\n", + "{%- macro render_tool_namespace(namespace_name, tools) -%}\n", + " {{- \"## \" + namespace_name + \"\\n\\n\" }}\n", + " {{- \"namespace \" + namespace_name + \" {\\n\\n\" }}\n", + " {%- for tool in tools %}\n", + " {%- set tool = tool.function %}\n", + " {{- \"// \" + tool.description + \"\\n\" }}\n", + " {{- \"type \"+ tool.name + \" = \" }}\n", + " {%- if tool.parameters and tool.parameters.properties %}\n", + " {{- \"(_: {\\n\" }}\n", + " {%- for param_name, param_spec in tool.parameters.properties.items() %}\n", + " {%- if param_spec.description %}\n", + " {{- \"// \" + param_spec.description + \"\\n\" }}\n", + " {%- endif %}\n", + " {{- param_name }}\n", + " {%- if param_name not in (tool.parameters.required or []) -%}\n", + " {{- \"?\" }}\n", + " {%- endif -%}\n", + " {{- \": \" }}\n", + " {{- render_typescript_type(param_spec, tool.parameters.required or []) }}\n", + " {%- if param_spec.default is defined -%}\n", + " {%- if param_spec.enum %}\n", + " {{- \", // default: \" + param_spec.default }}\n", + " {%- elif param_spec.oneOf %}\n", + " {{- \"// default: \" + param_spec.default }}\n", + " {%- else %}\n", + " {{- \", // default: \" + param_spec.default|tojson }}\n", + " {%- endif -%}\n", + " {%- endif -%}\n", + " {%- if not loop.last %}\n", + " {{- \",\\n\" }}\n", + " {%- else %}\n", + " {{- \",\\n\" }}\n", + " {%- endif -%}\n", + " {%- endfor %}\n", + " {{- \"}) => any;\\n\\n\" }}\n", + " {%- else -%}\n", + " {{- \"() => any;\\n\\n\" }}\n", + " {%- endif -%}\n", + " {%- endfor %}\n", + " {{- \"} // namespace \" + namespace_name }}\n", + "{%- endmacro -%}\n", + "\n", + "{%- macro render_builtin_tools(browser_tool, python_tool) -%}\n", + " {%- if browser_tool %}\n", + " {{- \"## browser\\n\\n\" }}\n", + " {{- \"// Tool for browsing.\\n\" }}\n", + " {{- \"// The `cursor` appears in brackets before each browsing display: `[{cursor}]`.\\n\" }}\n", + " {{- \"// Cite information from the tool using the following format:\\n\" }}\n", + " {{- \"// `\u3010{cursor}\u2020L{line_start}(-L{line_end})?\u3011`, for example: `\u30106\u2020L9-L11\u3011` or `\u30108\u2020L3\u3011`.\\n\" }}\n", + " {{- \"// Do not quote more than 10 words directly from the tool output.\\n\" }}\n", + " {{- \"// sources=web (default: web)\\n\" }}\n", + " {{- \"namespace browser {\\n\\n\" }}\n", + " {{- \"// Searches for information related to `query` and displays `topn` results.\\n\" }}\n", + " {{- \"type search = (_: {\\n\" }}\n", + " {{- \"query: string,\\n\" }}\n", + " {{- \"topn?: number, // default: 10\\n\" }}\n", + " {{- \"source?: string,\\n\" }}\n", + " {{- \"}) => any;\\n\\n\" }}\n", + " {{- \"// Opens the link `id` from the page indicated by `cursor` starting at line number `loc`, showing `num_lines` lines.\\n\" }}\n", + " {{- \"// Valid link ids are displayed with the formatting: `\u3010{id}\u2020.*\u3011`.\\n\" }}\n", + " {{- \"// If `cursor` is not provided, the most recent page is implied.\\n\" }}\n", + " {{- \"// If `id` is a string, it is treated as a fully qualified URL associated with `source`.\\n\" }}\n", + " {{- \"// If `loc` is not provided, the viewport will be positioned at the beginning of the document or centered on the most relevant passage, if available.\\n\" }}\n", + " {{- \"// Use this function without `id` to scroll to a new location of an opened page.\\n\" }}\n", + " {{- \"type open = (_: {\\n\" }}\n", + " {{- \"id?: number | string, // default: -1\\n\" }}\n", + " {{- \"cursor?: number, // default: -1\\n\" }}\n", + " {{- \"loc?: number, // default: -1\\n\" }}\n", + " {{- \"num_lines?: number, // default: -1\\n\" }}\n", + " {{- \"view_source?: boolean, // default: false\\n\" }}\n", + " {{- \"source?: string,\\n\" }}\n", + " {{- \"}) => any;\\n\\n\" }}\n", + " {{- \"// Finds exact matches of `pattern` in the current page, or the page given by `cursor`.\\n\" }}\n", + " {{- \"type find = (_: {\\n\" }}\n", + " {{- \"pattern: string,\\n\" }}\n", + " {{- \"cursor?: number, // default: -1\\n\" }}\n", + " {{- \"}) => any;\\n\\n\" }}\n", + " {{- \"} // namespace browser\\n\\n\" }}\n", + " {%- endif -%}\n", + "\n", + " {%- if python_tool %}\n", + " {{- \"## python\\n\\n\" }}\n", + " {{- \"Use this tool to execute Python code in your chain of thought. The code will not be shown to the user. This tool should be used for internal reasoning, but not for code that is intended to be visible to the user (e.g. when creating plots, tables, or files).\\n\\n\" }}\n", + " {{- \"When you send a message containing Python code to python, it will be executed in a stateful Jupyter notebook environment. python will respond with the output of the execution or time out after 120.0 seconds. The drive at '/mnt/data' can be used to save and persist user files. Internet access for this session is UNKNOWN. Depends on the cluster.\\n\\n\" }}\n", + " {%- endif -%}\n", + "{%- endmacro -%}\n", + "\n", + "{#- System Message Construction ============================================ #}\n", + "{%- macro build_system_message() -%}\n", + " {%- if model_identity is not defined %}\n", + " {%- set model_identity = \"You are ChatGPT, a large language model trained by OpenAI.\" %}\n", + " {%- endif %}\n", + " {{- model_identity + \"\\n\" }}\n", + " {{- \"Knowledge cutoff: 2024-06\\n\" }}\n", + " {{- \"Current date: \" + strftime_now(\"%Y-%m-%d\") + \"\\n\\n\" }}\n", + " {%- if reasoning_effort is not string %}\n", + " {%- set reasoning_effort = \"medium\" %}\n", + " {%- endif %}\n", + " {{- \"Reasoning: \" + reasoning_effort + \"\\n\\n\" }}\n", + " {%- if builtin_tools %}\n", + " {{- \"# Tools\\n\\n\" }}\n", + " {%- set available_builtin_tools = namespace(browser=false, python=false) %}\n", + " {%- for tool in builtin_tools %}\n", + " {%- if tool == \"browser\" %}\n", + " {%- set available_builtin_tools.browser = true %}\n", + " {%- elif tool == \"python\" %}\n", + " {%- set available_builtin_tools.python = true %}\n", + " {%- endif %}\n", + " {%- endfor %}\n", + " {{- render_builtin_tools(available_builtin_tools.browser, available_builtin_tools.python) }}\n", + " {%- endif -%}\n", + " {{- \"# Valid channels: analysis, commentary, final. Channel must be included for every message.\" }}\n", + " {%- if tools -%}\n", + " {{- \"\\nCalls to these tools must go to the commentary channel: 'functions'.\" }}\n", + " {%- endif -%}\n", + "{%- endmacro -%}\n", + "\n", + "{#- Main Template Logic ================================================= #}\n", + "{#- Set defaults #}\n", + "\n", + "{#- Render system message #}\n", + "{{- \"<|start|>system<|message|>\" }}\n", + "{{- build_system_message() }}\n", + "{{- \"<|end|>\" }}\n", + "\n", + "{#- Extract developer message #}\n", + "{%- if messages[0].role == \"developer\" or messages[0].role == \"system\" %}\n", + " {%- set developer_message = messages[0].content %}\n", + " {%- set loop_messages = messages[1:] %}\n", + "{%- else %}\n", + " {%- set developer_message = \"\" %}\n", + " {%- set loop_messages = messages %}\n", + "{%- endif %}\n", + "\n", + "{#- Render developer message #}\n", + "{%- if developer_message or tools %}\n", + " {{- \"<|start|>developer<|message|>\" }}\n", + " {%- if developer_message %}\n", + " {{- \"# Instructions\\n\\n\" }}\n", + " {{- developer_message }}\n", + " {{- \"\\n\\n\" }}\n", + " {%- endif %}\n", + " {%- if tools -%}\n", + " {{- \"# Tools\\n\\n\" }}\n", + " {{- render_tool_namespace(\"functions\", tools) }}\n", + " {%- endif -%}\n", + " {{- \"<|end|>\" }}\n", + "{%- endif %}\n", + "\n", + "{#- Render messages #}\n", + "{%- set last_tool_call = namespace(name=none) %}\n", + "{%- for message in loop_messages -%}\n", + " {#- At this point only assistant/user/tool messages should remain #}\n", + " {%- if message.role == 'assistant' -%}\n", + " {#- Checks to ensure the messages are being passed in the format we expect #}\n", + " {%- if \"content\" in message %}\n", + " {%- if \"<|channel|>analysis<|message|>\" in message.content or \"<|channel|>final<|message|>\" in message.content %}\n", + " {{- raise_exception(\"You have passed a message containing <|channel|> tags in the content field. Instead of doing this, you should pass analysis messages (the string between '<|message|>' and '<|end|>') in the 'thinking' field, and final messages (the string between '<|message|>' and '<|end|>') in the 'content' field.\") }}\n", + " {%- endif %}\n", + " {%- endif %}\n", + " {%- if \"thinking\" in message %}\n", + " {%- if \"<|channel|>analysis<|message|>\" in message.thinking or \"<|channel|>final<|message|>\" in message.thinking %}\n", + " {{- raise_exception(\"You have passed a message containing <|channel|> tags in the thinking field. Instead of doing this, you should pass analysis messages (the string between '<|message|>' and '<|end|>') in the 'thinking' field, and final messages (the string between '<|message|>' and '<|end|>') in the 'content' field.\") }}\n", + " {%- endif %}\n", + " {%- endif %}\n", + " {%- if \"tool_calls\" in message %}\n", + " {#- We need very careful handling here - we want to drop the tool call analysis message if the model #}\n", + " {#- has output a later <|final|> message, but otherwise we want to retain it. This is the only case #}\n", + " {#- when we render CoT/analysis messages in inference. #}\n", + " {%- set future_final_message = namespace(found=false) %}\n", + " {%- for future_message in loop_messages[loop.index:] %}\n", + " {%- if future_message.role == 'assistant' and \"tool_calls\" not in future_message %}\n", + " {%- set future_final_message.found = true %}\n", + " {%- endif %}\n", + " {%- endfor %}\n", + " {#- We assume max 1 tool call per message, and so we infer the tool call name #}\n", + " {#- in \"tool\" messages from the most recent assistant tool call name #}\n", + " {%- set tool_call = message.tool_calls[0] %}\n", + " {%- if tool_call.function %}\n", + " {%- set tool_call = tool_call.function %}\n", + " {%- endif %}\n", + " {%- if message.content and message.thinking %}\n", + " {{- raise_exception(\"Cannot pass both content and thinking in an assistant message with tool calls! Put the analysis message in one or the other, but not both.\") }}\n", + " {%- elif message.content and not future_final_message.found %}\n", + " {{- \"<|start|>assistant<|channel|>analysis<|message|>\" + message.content + \"<|end|>\" }}\n", + " {%- elif message.thinking and not future_final_message.found %}\n", + " {{- \"<|start|>assistant<|channel|>analysis<|message|>\" + message.thinking + \"<|end|>\" }}\n", + " {%- endif %}\n", + " {{- \"<|start|>assistant to=\" }}\n", + " {{- \"functions.\" + tool_call.name + \"<|channel|>commentary \" }}\n", + " {{- (tool_call.content_type if tool_call.content_type is defined else \"json\") + \"<|message|>\" }}\n", + " {{- tool_call.arguments|tojson }}\n", + " {{- \"<|call|>\" }}\n", + " {%- set last_tool_call.name = tool_call.name %}\n", + " {%- elif loop.last and not add_generation_prompt %}\n", + " {#- Only render the CoT if the final turn is an assistant turn and add_generation_prompt is false #}\n", + " {#- This is a situation that should only occur in training, never in inference. #}\n", + " {%- if \"thinking\" in message %}\n", + " {{- \"<|start|>assistant<|channel|>analysis<|message|>\" + message.thinking + \"<|end|>\" }}\n", + " {%- endif %}\n", + " {#- <|return|> indicates the end of generation, but <|end|> does not #}\n", + " {#- <|return|> should never be an input to the model, but we include it as the final token #}\n", + " {#- when training, so the model learns to emit it. #}\n", + " {{- \"<|start|>assistant<|channel|>final<|message|>\" + message.content + \"<|return|>\" }}\n", + " {%- else %}\n", + " {#- CoT is dropped during all previous turns, so we never render it for inference #}\n", + " {{- \"<|start|>assistant<|channel|>final<|message|>\" + message.content + \"<|end|>\" }}\n", + " {%- set last_tool_call.name = none %}\n", + " {%- endif %}\n", + " {%- elif message.role == 'tool' -%}\n", + " {%- if last_tool_call.name is none %}\n", + " {{- raise_exception(\"Message has tool role, but there was no previous assistant message with a tool call!\") }}\n", + " {%- endif %}\n", + " {{- \"<|start|>functions.\" + last_tool_call.name }}\n", + " {{- \" to=assistant<|channel|>commentary<|message|>\" + message.content|tojson + \"<|end|>\" }}\n", + " {%- elif message.role == 'user' -%}\n", + " {{- \"<|start|>user<|message|>\" + message.content + \"<|end|>\" }}\n", + " {%- endif -%}\n", + "{%- endfor -%}\n", + "\n", + "{#- Generation prompt #}\n", + "{%- if add_generation_prompt -%}\n", + "{%- set forced_tool_name = none %}\n", + "{%- if tool_choice is mapping %}\n", + " {%- if tool_choice.function is defined and tool_choice.function.name is defined %}\n", + " {%- set forced_tool_name = tool_choice.function.name %}\n", + " {%- elif tool_choice.name is defined %}\n", + " {%- set forced_tool_name = tool_choice.name %}\n", + " {%- endif %}\n", + "{%- elif function_call is mapping and function_call.name is defined %}\n", + " {%- set forced_tool_name = function_call.name %}\n", + "{%- endif %}\n", + "{%- if forced_tool_name %}\n", + "{{- \"<|start|>assistant<|channel|>commentary to=functions.\" + forced_tool_name + \" <|constrain|>json<|message|>\" }}\n", + "{%- else %}\n", + "<|start|>assistant\n", + "{%- endif %}\n", + "{%- endif -%}" + ] + } +} diff --git a/examples/server/configs/qwen3.5-0.8b.json b/examples/server/configs/qwen3.5-0.8b.json new file mode 100644 index 0000000000..c9f2ca944a --- /dev/null +++ b/examples/server/configs/qwen3.5-0.8b.json @@ -0,0 +1,261 @@ +{ + "server": { + "host": "0.0.0.0", + "port": 8000 + }, + "model": { + "alias": "qwen3.5-0.8b-vl", + "from_pretrained": { + "repo_id": "lmstudio-community/Qwen3.5-0.8B-GGUF", + "filename": "Qwen3.5-0.8B-Q8_0.gguf" + }, + "mtmd": { + "mmproj_from_pretrained": { + "repo_id": "lmstudio-community/Qwen3.5-0.8B-GGUF", + "filename": "mmproj-Qwen3.5-0.8B-BF16.gguf" + } + }, + "n_ctx": 32768, + "max_output_tokens": 4096, + "n_seq_max": 64, + "n_batch": 128, + "n_ubatch": 128, + "threads": 4, + "threads_batch": 8, + "kv_unified": true, + "store_logits": false, + "use_mmap": true, + "use_mlock": true, + "response_schema": { + "type": "object", + "properties": { + "role": { + "const": "assistant" + }, + "reasoning_content": { + "type": "string", + "x-regex": "^(?:<\\|im_start\\|>assistant\\n)?(?:<think>\\n)?(.*?)(?=</think>)" + }, + "content": { + "type": "string", + "x-regex": "^(?:<\\|im_start\\|>assistant\\n)?(?:(?:<think>\\n)?.*?</think>\\s*)?(.*?)(?=\\s*<tool_call>\\n|<\\|im_end\\|>$|$)" + }, + "tool_calls": { + "type": "array", + "x-regex-iterator": "<tool_call>\\n(.*?)\\n</tool_call>", + "items": { + "type": "object", + "properties": { + "type": { + "const": "function" + }, + "function": { + "type": "object", + "properties": { + "name": { + "type": "string", + "x-regex": "^<function=([^>\\n]+)>\\n" + }, + "arguments": { + "type": "object", + "x-regex": "^<function=[^>\\n]+>\\n(.*?)\\n</function>$", + "x-regex-key-value": "<parameter=(?P<key>[^>\\n]+)>\\n(?P<value>.*?)\\n</parameter>", + "additionalProperties": true + } + }, + "required": [ + "name", + "arguments" + ] + } + }, + "required": [ + "type", + "function" + ] + } + } + }, + "required": [ + "role" + ] + }, + "chat_template": [ + "{%- set image_count = namespace(value=0) %}\n", + "{%- set video_count = namespace(value=0) %}\n", + "{%- if enable_thinking is not defined and reasoning_effort is string %}\n", + " {%- set qwen_reasoning_effort = reasoning_effort|lower %}\n", + " {%- if qwen_reasoning_effort in ['none', 'minimal', 'low'] %}\n", + " {%- set enable_thinking = false %}\n", + " {%- elif qwen_reasoning_effort in ['medium', 'high'] %}\n", + " {%- set enable_thinking = true %}\n", + " {%- endif %}\n", + "{%- endif %}\n", + "{%- set forced_tool_name = none %}\n", + "{%- if tool_choice is mapping %}\n", + " {%- if tool_choice.function is defined and tool_choice.function.name is defined %}\n", + " {%- set forced_tool_name = tool_choice.function.name %}\n", + " {%- elif tool_choice.name is defined %}\n", + " {%- set forced_tool_name = tool_choice.name %}\n", + " {%- endif %}\n", + "{%- elif function_call is mapping and function_call.name is defined %}\n", + " {%- set forced_tool_name = function_call.name %}\n", + "{%- endif %}\n", + "{%- macro render_content(content, do_vision_count, is_system_content=false) %}\n", + " {%- if content is string %}\n", + " {{- content }}\n", + " {%- elif content is iterable and content is not mapping %}\n", + " {%- for item in content %}\n", + " {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}\n", + " {%- if is_system_content %}\n", + " {{- raise_exception('System message cannot contain images.') }}\n", + " {%- endif %}\n", + " {%- if do_vision_count %}\n", + " {%- set image_count.value = image_count.value + 1 %}\n", + " {%- endif %}\n", + " {%- if add_vision_id %}\n", + " {{- 'Picture ' ~ image_count.value ~ ': ' }}\n", + " {%- endif %}\n", + " {{- '<|vision_start|><|image_pad|><|vision_end|>' }}\n", + " {%- elif 'video' in item or item.type == 'video' %}\n", + " {%- if is_system_content %}\n", + " {{- raise_exception('System message cannot contain videos.') }}\n", + " {%- endif %}\n", + " {%- if do_vision_count %}\n", + " {%- set video_count.value = video_count.value + 1 %}\n", + " {%- endif %}\n", + " {%- if add_vision_id %}\n", + " {{- 'Video ' ~ video_count.value ~ ': ' }}\n", + " {%- endif %}\n", + " {{- '<|vision_start|><|video_pad|><|vision_end|>' }}\n", + " {%- elif 'text' in item %}\n", + " {{- item.text }}\n", + " {%- else %}\n", + " {{- raise_exception('Unexpected item type in content.') }}\n", + " {%- endif %}\n", + " {%- endfor %}\n", + " {%- elif content is none or content is undefined %}\n", + " {{- '' }}\n", + " {%- else %}\n", + " {{- raise_exception('Unexpected content type.') }}\n", + " {%- endif %}\n", + "{%- endmacro %}\n", + "{%- if not messages %}\n", + " {{- raise_exception('No messages provided.') }}\n", + "{%- endif %}\n", + "{%- if tools and tools is iterable and tools is not mapping %}\n", + " {{- '<|im_start|>system\\n' }}\n", + " {{- \"# Tools\\n\\nYou have access to the following functions:\\n\\n<tools>\" }}\n", + " {%- for tool in tools %}\n", + " {{- \"\\n\" }}\n", + " {{- tool | tojson }}\n", + " {%- endfor %}\n", + " {{- \"\\n</tools>\" }}\n", + " {{- '\\n\\nIf you choose to call a function ONLY reply in the following format with NO suffix:\\n\\n<tool_call>\\n<function=example_function_name>\\n<parameter=example_parameter_1>\\nvalue_1\\n</parameter>\\n<parameter=example_parameter_2>\\nThis is the value for the second parameter\\nthat can span\\nmultiple lines\\n</parameter>\\n</function>\\n</tool_call>\\n\\n<IMPORTANT>\\nReminder:\\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\\n- Required parameters MUST be specified\\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\\n</IMPORTANT>' }}\n", + " {%- if messages[0].role == 'system' %}\n", + " {%- set content = render_content(messages[0].content, false, true)|trim %}\n", + " {%- if content %}\n", + " {{- '\\n\\n' + content }}\n", + " {%- endif %}\n", + " {%- endif %}\n", + " {{- '<|im_end|>\\n' }}\n", + "{%- else %}\n", + " {%- if messages[0].role == 'system' %}\n", + " {%- set content = render_content(messages[0].content, false, true)|trim %}\n", + " {{- '<|im_start|>system\\n' + content + '<|im_end|>\\n' }}\n", + " {%- endif %}\n", + "{%- endif %}\n", + "{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n", + "{%- for message in messages[::-1] %}\n", + " {%- set index = (messages|length - 1) - loop.index0 %}\n", + " {%- if ns.multi_step_tool and message.role == \"user\" %}\n", + " {%- set content = render_content(message.content, false)|trim %}\n", + " {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}\n", + " {%- set ns.multi_step_tool = false %}\n", + " {%- set ns.last_query_index = index %}\n", + " {%- endif %}\n", + " {%- endif %}\n", + "{%- endfor %}\n", + "{%- if ns.multi_step_tool %}\n", + " {{- raise_exception('No user query found in messages.') }}\n", + "{%- endif %}\n", + "{%- for message in messages %}\n", + " {%- set content = render_content(message.content, true)|trim %}\n", + " {%- if message.role == \"system\" %}\n", + " {%- if not loop.first %}\n", + " {{- raise_exception('System message must be at the beginning.') }}\n", + " {%- endif %}\n", + " {%- elif message.role == \"user\" %}\n", + " {{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}\n", + " {%- elif message.role == \"assistant\" %}\n", + " {%- set reasoning_content = '' %}\n", + " {%- if message.reasoning_content is string %}\n", + " {%- set reasoning_content = message.reasoning_content %}\n", + " {%- else %}\n", + " {%- if '</think>' in content %}\n", + " {%- set reasoning_content = content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n", + " {%- set content = content.split('</think>')[-1].lstrip('\\n') %}\n", + " {%- endif %}\n", + " {%- endif %}\n", + " {%- set reasoning_content = reasoning_content|trim %}\n", + " {%- if loop.index0 > ns.last_query_index %}\n", + " {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content + '\\n</think>\\n\\n' + content }}\n", + " {%- else %}\n", + " {{- '<|im_start|>' + message.role + '\\n' + content }}\n", + " {%- endif %}\n", + " {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}\n", + " {%- for tool_call in message.tool_calls %}\n", + " {%- if tool_call.function is defined %}\n", + " {%- set tool_call = tool_call.function %}\n", + " {%- endif %}\n", + " {%- if loop.first %}\n", + " {%- if content|trim %}\n", + " {{- '\\n\\n<tool_call>\\n<function=' + tool_call.name + '>\\n' }}\n", + " {%- else %}\n", + " {{- '<tool_call>\\n<function=' + tool_call.name + '>\\n' }}\n", + " {%- endif %}\n", + " {%- else %}\n", + " {{- '\\n<tool_call>\\n<function=' + tool_call.name + '>\\n' }}\n", + " {%- endif %}\n", + " {%- if tool_call.arguments is defined %}\n", + " {%- set arguments = tool_call.arguments | from_json if tool_call.arguments is string else tool_call.arguments %}\n", + " {%- for args_name, args_value in arguments|items %}\n", + " {{- '<parameter=' + args_name + '>\\n' }}\n", + " {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}\n", + " {{- args_value }}\n", + " {{- '\\n</parameter>\\n' }}\n", + " {%- endfor %}\n", + " {%- endif %}\n", + " {{- '</function>\\n</tool_call>' }}\n", + " {%- endfor %}\n", + " {%- endif %}\n", + " {{- '<|im_end|>\\n' }}\n", + " {%- elif message.role == \"tool\" %}\n", + " {%- if loop.previtem and loop.previtem.role != \"tool\" %}\n", + " {{- '<|im_start|>user' }}\n", + " {%- endif %}\n", + " {{- '\\n<tool_response>\\n' }}\n", + " {{- content }}\n", + " {{- '\\n</tool_response>' }}\n", + " {%- if not loop.last and loop.nextitem.role != \"tool\" %}\n", + " {{- '<|im_end|>\\n' }}\n", + " {%- elif loop.last %}\n", + " {{- '<|im_end|>\\n' }}\n", + " {%- endif %}\n", + " {%- else %}\n", + " {{- raise_exception('Unexpected message role.') }}\n", + " {%- endif %}\n", + "{%- endfor %}\n", + "{%- if add_generation_prompt %}\n", + " {{- '<|im_start|>assistant\\n' }}\n", + " {%- if forced_tool_name %}\n", + " {{- '<tool_call>\\n<function=' + forced_tool_name + '>\\n' }}\n", + " {%- elif enable_thinking is defined and enable_thinking is false %}\n", + " {{- '<think>\\n\\n</think>\\n\\n' }}\n", + " {%- else %}\n", + " {{- '<think>\\n' }}\n", + " {%- endif %}\n", + "{%- endif %}" + ] + } +} diff --git a/examples/server/configs/qwen3.6-27b.json b/examples/server/configs/qwen3.6-27b.json new file mode 100644 index 0000000000..c5b42d8b72 --- /dev/null +++ b/examples/server/configs/qwen3.6-27b.json @@ -0,0 +1,263 @@ +{ + "server": { + "host": "0.0.0.0", + "port": 8000 + }, + "model": { + "alias": "qwen3.6-27b", + "from_pretrained": { + "repo_id": "unsloth/Qwen3.6-27B-GGUF", + "filename": "Qwen3.6-27B-Q8_0.gguf" + }, + "mtmd": { + "mmproj_from_pretrained": { + "repo_id": "unsloth/Qwen3.6-27B-GGUF", + "filename": "mmproj-BF16.gguf" + } + }, + "n_ctx": 32768, + "max_output_tokens": 4096, + "n_seq_max": 8, + "n_batch": 512, + "n_ubatch": 512, + "threads": 8, + "threads_batch": 8, + "kv_unified": true, + "store_logits": false, + "use_mmap": true, + "use_mlock": false, + "n_gpu_layers": -1, + "flash_attn": true, + "response_schema": { + "type": "object", + "properties": { + "role": { + "const": "assistant" + }, + "reasoning_content": { + "type": "string", + "x-regex": "^(?:<\\|im_start\\|>assistant\\n)?(?:<think>\\n)?(.*?)(?=</think>)" + }, + "content": { + "type": "string", + "x-regex": "^(?:<\\|im_start\\|>assistant\\n)?(?:(?:<think>\\n)?.*?</think>\\s*)?(.*?)(?=\\s*<tool_call>\\n|<\\|im_end\\|>$|$)" + }, + "tool_calls": { + "type": "array", + "x-regex-iterator": "<tool_call>\\n(.*?)\\n</tool_call>", + "items": { + "type": "object", + "properties": { + "type": { + "const": "function" + }, + "function": { + "type": "object", + "properties": { + "name": { + "type": "string", + "x-regex": "^<function=([^>\\n]+)>\\n" + }, + "arguments": { + "type": "object", + "x-regex": "^<function=[^>\\n]+>\\n(.*?)\\n</function>$", + "x-regex-key-value": "<parameter=(?P<key>[^>\\n]+)>\\n(?P<value>.*?)\\n</parameter>", + "additionalProperties": true + } + }, + "required": [ + "name", + "arguments" + ] + } + }, + "required": [ + "type", + "function" + ] + } + } + }, + "required": [ + "role" + ] + }, + "chat_template": [ + "{%- set image_count = namespace(value=0) %}\n", + "{%- set video_count = namespace(value=0) %}\n", + "{%- if enable_thinking is not defined and reasoning_effort is string %}\n", + " {%- set qwen_reasoning_effort = reasoning_effort|lower %}\n", + " {%- if qwen_reasoning_effort in ['none', 'minimal', 'low'] %}\n", + " {%- set enable_thinking = false %}\n", + " {%- elif qwen_reasoning_effort in ['medium', 'high'] %}\n", + " {%- set enable_thinking = true %}\n", + " {%- endif %}\n", + "{%- endif %}\n", + "{%- set forced_tool_name = none %}\n", + "{%- if tool_choice is mapping %}\n", + " {%- if tool_choice.function is defined and tool_choice.function.name is defined %}\n", + " {%- set forced_tool_name = tool_choice.function.name %}\n", + " {%- elif tool_choice.name is defined %}\n", + " {%- set forced_tool_name = tool_choice.name %}\n", + " {%- endif %}\n", + "{%- elif function_call is mapping and function_call.name is defined %}\n", + " {%- set forced_tool_name = function_call.name %}\n", + "{%- endif %}\n", + "{%- macro render_content(content, do_vision_count, is_system_content=false) %}\n", + " {%- if content is string %}\n", + " {{- content }}\n", + " {%- elif content is iterable and content is not mapping %}\n", + " {%- for item in content %}\n", + " {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}\n", + " {%- if is_system_content %}\n", + " {{- raise_exception('System message cannot contain images.') }}\n", + " {%- endif %}\n", + " {%- if do_vision_count %}\n", + " {%- set image_count.value = image_count.value + 1 %}\n", + " {%- endif %}\n", + " {%- if add_vision_id %}\n", + " {{- 'Picture ' ~ image_count.value ~ ': ' }}\n", + " {%- endif %}\n", + " {{- '<|vision_start|><|image_pad|><|vision_end|>' }}\n", + " {%- elif 'video' in item or item.type == 'video' %}\n", + " {%- if is_system_content %}\n", + " {{- raise_exception('System message cannot contain videos.') }}\n", + " {%- endif %}\n", + " {%- if do_vision_count %}\n", + " {%- set video_count.value = video_count.value + 1 %}\n", + " {%- endif %}\n", + " {%- if add_vision_id %}\n", + " {{- 'Video ' ~ video_count.value ~ ': ' }}\n", + " {%- endif %}\n", + " {{- '<|vision_start|><|video_pad|><|vision_end|>' }}\n", + " {%- elif 'text' in item %}\n", + " {{- item.text }}\n", + " {%- else %}\n", + " {{- raise_exception('Unexpected item type in content.') }}\n", + " {%- endif %}\n", + " {%- endfor %}\n", + " {%- elif content is none or content is undefined %}\n", + " {{- '' }}\n", + " {%- else %}\n", + " {{- raise_exception('Unexpected content type.') }}\n", + " {%- endif %}\n", + "{%- endmacro %}\n", + "{%- if not messages %}\n", + " {{- raise_exception('No messages provided.') }}\n", + "{%- endif %}\n", + "{%- if tools and tools is iterable and tools is not mapping %}\n", + " {{- '<|im_start|>system\\n' }}\n", + " {{- \"# Tools\\n\\nYou have access to the following functions:\\n\\n<tools>\" }}\n", + " {%- for tool in tools %}\n", + " {{- \"\\n\" }}\n", + " {{- tool | tojson }}\n", + " {%- endfor %}\n", + " {{- \"\\n</tools>\" }}\n", + " {{- '\\n\\nIf you choose to call a function ONLY reply in the following format with NO suffix:\\n\\n<tool_call>\\n<function=example_function_name>\\n<parameter=example_parameter_1>\\nvalue_1\\n</parameter>\\n<parameter=example_parameter_2>\\nThis is the value for the second parameter\\nthat can span\\nmultiple lines\\n</parameter>\\n</function>\\n</tool_call>\\n\\n<IMPORTANT>\\nReminder:\\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\\n- Required parameters MUST be specified\\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\\n</IMPORTANT>' }}\n", + " {%- if messages[0].role == 'system' %}\n", + " {%- set content = render_content(messages[0].content, false, true)|trim %}\n", + " {%- if content %}\n", + " {{- '\\n\\n' + content }}\n", + " {%- endif %}\n", + " {%- endif %}\n", + " {{- '<|im_end|>\\n' }}\n", + "{%- else %}\n", + " {%- if messages[0].role == 'system' %}\n", + " {%- set content = render_content(messages[0].content, false, true)|trim %}\n", + " {{- '<|im_start|>system\\n' + content + '<|im_end|>\\n' }}\n", + " {%- endif %}\n", + "{%- endif %}\n", + "{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n", + "{%- for message in messages[::-1] %}\n", + " {%- set index = (messages|length - 1) - loop.index0 %}\n", + " {%- if ns.multi_step_tool and message.role == \"user\" %}\n", + " {%- set content = render_content(message.content, false)|trim %}\n", + " {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}\n", + " {%- set ns.multi_step_tool = false %}\n", + " {%- set ns.last_query_index = index %}\n", + " {%- endif %}\n", + " {%- endif %}\n", + "{%- endfor %}\n", + "{%- if ns.multi_step_tool %}\n", + " {{- raise_exception('No user query found in messages.') }}\n", + "{%- endif %}\n", + "{%- for message in messages %}\n", + " {%- set content = render_content(message.content, true)|trim %}\n", + " {%- if message.role == \"system\" %}\n", + " {%- if not loop.first %}\n", + " {{- raise_exception('System message must be at the beginning.') }}\n", + " {%- endif %}\n", + " {%- elif message.role == \"user\" %}\n", + " {{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}\n", + " {%- elif message.role == \"assistant\" %}\n", + " {%- set reasoning_content = '' %}\n", + " {%- if message.reasoning_content is string %}\n", + " {%- set reasoning_content = message.reasoning_content %}\n", + " {%- else %}\n", + " {%- if '</think>' in content %}\n", + " {%- set reasoning_content = content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n", + " {%- set content = content.split('</think>')[-1].lstrip('\\n') %}\n", + " {%- endif %}\n", + " {%- endif %}\n", + " {%- set reasoning_content = reasoning_content|trim %}\n", + " {%- if (preserve_thinking is defined and preserve_thinking is true) or (loop.index0 > ns.last_query_index) %}\n", + " {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content + '\\n</think>\\n\\n' + content }}\n", + " {%- else %}\n", + " {{- '<|im_start|>' + message.role + '\\n' + content }}\n", + " {%- endif %}\n", + " {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}\n", + " {%- for tool_call in message.tool_calls %}\n", + " {%- if tool_call.function is defined %}\n", + " {%- set tool_call = tool_call.function %}\n", + " {%- endif %}\n", + " {%- if loop.first %}\n", + " {%- if content|trim %}\n", + " {{- '\\n\\n<tool_call>\\n<function=' + tool_call.name + '>\\n' }}\n", + " {%- else %}\n", + " {{- '<tool_call>\\n<function=' + tool_call.name + '>\\n' }}\n", + " {%- endif %}\n", + " {%- else %}\n", + " {{- '\\n<tool_call>\\n<function=' + tool_call.name + '>\\n' }}\n", + " {%- endif %}\n", + " {%- if tool_call.arguments is defined %}\n", + " {%- set arguments = tool_call.arguments | from_json if tool_call.arguments is string else tool_call.arguments %}\n", + " {%- for args_name, args_value in arguments|items %}\n", + " {{- '<parameter=' + args_name + '>\\n' }}\n", + " {%- set args_value = args_value | string if args_value is string else args_value | tojson | safe %}\n", + " {{- args_value }}\n", + " {{- '\\n</parameter>\\n' }}\n", + " {%- endfor %}\n", + " {%- endif %}\n", + " {{- '</function>\\n</tool_call>' }}\n", + " {%- endfor %}\n", + " {%- endif %}\n", + " {{- '<|im_end|>\\n' }}\n", + " {%- elif message.role == \"tool\" %}\n", + " {%- if loop.previtem and loop.previtem.role != \"tool\" %}\n", + " {{- '<|im_start|>user' }}\n", + " {%- endif %}\n", + " {{- '\\n<tool_response>\\n' }}\n", + " {{- content }}\n", + " {{- '\\n</tool_response>' }}\n", + " {%- if not loop.last and loop.nextitem.role != \"tool\" %}\n", + " {{- '<|im_end|>\\n' }}\n", + " {%- elif loop.last %}\n", + " {{- '<|im_end|>\\n' }}\n", + " {%- endif %}\n", + " {%- else %}\n", + " {{- raise_exception('Unexpected message role.') }}\n", + " {%- endif %}\n", + "{%- endfor %}\n", + "{%- if add_generation_prompt %}\n", + " {{- '<|im_start|>assistant\\n' }}\n", + " {%- if forced_tool_name %}\n", + " {{- '<tool_call>\\n<function=' + forced_tool_name + '>\\n' }}\n", + " {%- elif enable_thinking is defined and enable_thinking is false %}\n", + " {{- '<think>\\n\\n</think>\\n\\n' }}\n", + " {%- else %}\n", + " {{- '<think>\\n' }}\n", + " {%- endif %}\n", + "{%- endif %}" + ] + } +} diff --git a/examples/server/configs/qwen3.6-35b-a3b.json b/examples/server/configs/qwen3.6-35b-a3b.json new file mode 100644 index 0000000000..10043f9c1a --- /dev/null +++ b/examples/server/configs/qwen3.6-35b-a3b.json @@ -0,0 +1,263 @@ +{ + "server": { + "host": "0.0.0.0", + "port": 8000 + }, + "model": { + "alias": "qwen3.6-35b-a3b", + "from_pretrained": { + "repo_id": "unsloth/Qwen3.6-35B-A3B-GGUF", + "filename": "Qwen3.6-35B-A3B-Q8_0.gguf" + }, + "mtmd": { + "mmproj_from_pretrained": { + "repo_id": "unsloth/Qwen3.6-35B-A3B-GGUF", + "filename": "mmproj-BF16.gguf" + } + }, + "n_ctx": 32768, + "max_output_tokens": 4096, + "n_seq_max": 8, + "n_batch": 512, + "n_ubatch": 512, + "threads": 8, + "threads_batch": 8, + "kv_unified": true, + "store_logits": false, + "use_mmap": true, + "use_mlock": false, + "n_gpu_layers": -1, + "flash_attn": true, + "response_schema": { + "type": "object", + "properties": { + "role": { + "const": "assistant" + }, + "reasoning_content": { + "type": "string", + "x-regex": "^(?:<\\|im_start\\|>assistant\\n)?(?:<think>\\n)?(.*?)(?=</think>)" + }, + "content": { + "type": "string", + "x-regex": "^(?:<\\|im_start\\|>assistant\\n)?(?:(?:<think>\\n)?.*?</think>\\s*)?(.*?)(?=\\s*<tool_call>\\n|<\\|im_end\\|>$|$)" + }, + "tool_calls": { + "type": "array", + "x-regex-iterator": "<tool_call>\\n(.*?)\\n</tool_call>", + "items": { + "type": "object", + "properties": { + "type": { + "const": "function" + }, + "function": { + "type": "object", + "properties": { + "name": { + "type": "string", + "x-regex": "^<function=([^>\\n]+)>\\n" + }, + "arguments": { + "type": "object", + "x-regex": "^<function=[^>\\n]+>\\n(.*?)\\n</function>$", + "x-regex-key-value": "<parameter=(?P<key>[^>\\n]+)>\\n(?P<value>.*?)\\n</parameter>", + "additionalProperties": true + } + }, + "required": [ + "name", + "arguments" + ] + } + }, + "required": [ + "type", + "function" + ] + } + } + }, + "required": [ + "role" + ] + }, + "chat_template": [ + "{%- set image_count = namespace(value=0) %}\n", + "{%- set video_count = namespace(value=0) %}\n", + "{%- if enable_thinking is not defined and reasoning_effort is string %}\n", + " {%- set qwen_reasoning_effort = reasoning_effort|lower %}\n", + " {%- if qwen_reasoning_effort in ['none', 'minimal', 'low'] %}\n", + " {%- set enable_thinking = false %}\n", + " {%- elif qwen_reasoning_effort in ['medium', 'high'] %}\n", + " {%- set enable_thinking = true %}\n", + " {%- endif %}\n", + "{%- endif %}\n", + "{%- set forced_tool_name = none %}\n", + "{%- if tool_choice is mapping %}\n", + " {%- if tool_choice.function is defined and tool_choice.function.name is defined %}\n", + " {%- set forced_tool_name = tool_choice.function.name %}\n", + " {%- elif tool_choice.name is defined %}\n", + " {%- set forced_tool_name = tool_choice.name %}\n", + " {%- endif %}\n", + "{%- elif function_call is mapping and function_call.name is defined %}\n", + " {%- set forced_tool_name = function_call.name %}\n", + "{%- endif %}\n", + "{%- macro render_content(content, do_vision_count, is_system_content=false) %}\n", + " {%- if content is string %}\n", + " {{- content }}\n", + " {%- elif content is iterable and content is not mapping %}\n", + " {%- for item in content %}\n", + " {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}\n", + " {%- if is_system_content %}\n", + " {{- raise_exception('System message cannot contain images.') }}\n", + " {%- endif %}\n", + " {%- if do_vision_count %}\n", + " {%- set image_count.value = image_count.value + 1 %}\n", + " {%- endif %}\n", + " {%- if add_vision_id %}\n", + " {{- 'Picture ' ~ image_count.value ~ ': ' }}\n", + " {%- endif %}\n", + " {{- '<|vision_start|><|image_pad|><|vision_end|>' }}\n", + " {%- elif 'video' in item or item.type == 'video' %}\n", + " {%- if is_system_content %}\n", + " {{- raise_exception('System message cannot contain videos.') }}\n", + " {%- endif %}\n", + " {%- if do_vision_count %}\n", + " {%- set video_count.value = video_count.value + 1 %}\n", + " {%- endif %}\n", + " {%- if add_vision_id %}\n", + " {{- 'Video ' ~ video_count.value ~ ': ' }}\n", + " {%- endif %}\n", + " {{- '<|vision_start|><|video_pad|><|vision_end|>' }}\n", + " {%- elif 'text' in item %}\n", + " {{- item.text }}\n", + " {%- else %}\n", + " {{- raise_exception('Unexpected item type in content.') }}\n", + " {%- endif %}\n", + " {%- endfor %}\n", + " {%- elif content is none or content is undefined %}\n", + " {{- '' }}\n", + " {%- else %}\n", + " {{- raise_exception('Unexpected content type.') }}\n", + " {%- endif %}\n", + "{%- endmacro %}\n", + "{%- if not messages %}\n", + " {{- raise_exception('No messages provided.') }}\n", + "{%- endif %}\n", + "{%- if tools and tools is iterable and tools is not mapping %}\n", + " {{- '<|im_start|>system\\n' }}\n", + " {{- \"# Tools\\n\\nYou have access to the following functions:\\n\\n<tools>\" }}\n", + " {%- for tool in tools %}\n", + " {{- \"\\n\" }}\n", + " {{- tool | tojson }}\n", + " {%- endfor %}\n", + " {{- \"\\n</tools>\" }}\n", + " {{- '\\n\\nIf you choose to call a function ONLY reply in the following format with NO suffix:\\n\\n<tool_call>\\n<function=example_function_name>\\n<parameter=example_parameter_1>\\nvalue_1\\n</parameter>\\n<parameter=example_parameter_2>\\nThis is the value for the second parameter\\nthat can span\\nmultiple lines\\n</parameter>\\n</function>\\n</tool_call>\\n\\n<IMPORTANT>\\nReminder:\\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\\n- Required parameters MUST be specified\\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\\n</IMPORTANT>' }}\n", + " {%- if messages[0].role == 'system' %}\n", + " {%- set content = render_content(messages[0].content, false, true)|trim %}\n", + " {%- if content %}\n", + " {{- '\\n\\n' + content }}\n", + " {%- endif %}\n", + " {%- endif %}\n", + " {{- '<|im_end|>\\n' }}\n", + "{%- else %}\n", + " {%- if messages[0].role == 'system' %}\n", + " {%- set content = render_content(messages[0].content, false, true)|trim %}\n", + " {{- '<|im_start|>system\\n' + content + '<|im_end|>\\n' }}\n", + " {%- endif %}\n", + "{%- endif %}\n", + "{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n", + "{%- for message in messages[::-1] %}\n", + " {%- set index = (messages|length - 1) - loop.index0 %}\n", + " {%- if ns.multi_step_tool and message.role == \"user\" %}\n", + " {%- set content = render_content(message.content, false)|trim %}\n", + " {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}\n", + " {%- set ns.multi_step_tool = false %}\n", + " {%- set ns.last_query_index = index %}\n", + " {%- endif %}\n", + " {%- endif %}\n", + "{%- endfor %}\n", + "{%- if ns.multi_step_tool %}\n", + " {{- raise_exception('No user query found in messages.') }}\n", + "{%- endif %}\n", + "{%- for message in messages %}\n", + " {%- set content = render_content(message.content, true)|trim %}\n", + " {%- if message.role == \"system\" %}\n", + " {%- if not loop.first %}\n", + " {{- raise_exception('System message must be at the beginning.') }}\n", + " {%- endif %}\n", + " {%- elif message.role == \"user\" %}\n", + " {{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}\n", + " {%- elif message.role == \"assistant\" %}\n", + " {%- set reasoning_content = '' %}\n", + " {%- if message.reasoning_content is string %}\n", + " {%- set reasoning_content = message.reasoning_content %}\n", + " {%- else %}\n", + " {%- if '</think>' in content %}\n", + " {%- set reasoning_content = content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n", + " {%- set content = content.split('</think>')[-1].lstrip('\\n') %}\n", + " {%- endif %}\n", + " {%- endif %}\n", + " {%- set reasoning_content = reasoning_content|trim %}\n", + " {%- if (preserve_thinking is defined and preserve_thinking is true) or (loop.index0 > ns.last_query_index) %}\n", + " {{- '<|im_start|>' + message.role + '\\n<think>\\n' + reasoning_content + '\\n</think>\\n\\n' + content }}\n", + " {%- else %}\n", + " {{- '<|im_start|>' + message.role + '\\n' + content }}\n", + " {%- endif %}\n", + " {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}\n", + " {%- for tool_call in message.tool_calls %}\n", + " {%- if tool_call.function is defined %}\n", + " {%- set tool_call = tool_call.function %}\n", + " {%- endif %}\n", + " {%- if loop.first %}\n", + " {%- if content|trim %}\n", + " {{- '\\n\\n<tool_call>\\n<function=' + tool_call.name + '>\\n' }}\n", + " {%- else %}\n", + " {{- '<tool_call>\\n<function=' + tool_call.name + '>\\n' }}\n", + " {%- endif %}\n", + " {%- else %}\n", + " {{- '\\n<tool_call>\\n<function=' + tool_call.name + '>\\n' }}\n", + " {%- endif %}\n", + " {%- if tool_call.arguments is defined %}\n", + " {%- set arguments = tool_call.arguments | from_json if tool_call.arguments is string else tool_call.arguments %}\n", + " {%- for args_name, args_value in arguments|items %}\n", + " {{- '<parameter=' + args_name + '>\\n' }}\n", + " {%- set args_value = args_value | string if args_value is string else args_value | tojson | safe %}\n", + " {{- args_value }}\n", + " {{- '\\n</parameter>\\n' }}\n", + " {%- endfor %}\n", + " {%- endif %}\n", + " {{- '</function>\\n</tool_call>' }}\n", + " {%- endfor %}\n", + " {%- endif %}\n", + " {{- '<|im_end|>\\n' }}\n", + " {%- elif message.role == \"tool\" %}\n", + " {%- if loop.previtem and loop.previtem.role != \"tool\" %}\n", + " {{- '<|im_start|>user' }}\n", + " {%- endif %}\n", + " {{- '\\n<tool_response>\\n' }}\n", + " {{- content }}\n", + " {{- '\\n</tool_response>' }}\n", + " {%- if not loop.last and loop.nextitem.role != \"tool\" %}\n", + " {{- '<|im_end|>\\n' }}\n", + " {%- elif loop.last %}\n", + " {{- '<|im_end|>\\n' }}\n", + " {%- endif %}\n", + " {%- else %}\n", + " {{- raise_exception('Unexpected message role.') }}\n", + " {%- endif %}\n", + "{%- endfor %}\n", + "{%- if add_generation_prompt %}\n", + " {{- '<|im_start|>assistant\\n' }}\n", + " {%- if forced_tool_name %}\n", + " {{- '<tool_call>\\n<function=' + forced_tool_name + '>\\n' }}\n", + " {%- elif enable_thinking is defined and enable_thinking is false %}\n", + " {{- '<think>\\n\\n</think>\\n\\n' }}\n", + " {%- else %}\n", + " {{- '<think>\\n' }}\n", + " {%- endif %}\n", + "{%- endif %}" + ] + } +} diff --git a/examples/server/server.py b/examples/server/server.py new file mode 100644 index 0000000000..e8034a2146 --- /dev/null +++ b/examples/server/server.py @@ -0,0 +1,16241 @@ +#!/usr/bin/env -S uv run --script +# /// script +# requires-python = ">=3.8" +# dependencies = [ +# "fastapi", +# "jinja2", +# "llama-cpp-python", +# "numpy", +# "openai", +# "pydantic", +# "safetensors", +# "uvicorn", +# "websockets", +# ] +# /// + +from __future__ import annotations + +import abc +import os +import re +import json +import math +import time +import uuid +import queue +import ctypes +import fnmatch +import base64 +import hashlib +import binascii +import asyncio +import argparse +import threading +import multiprocessing +import copy +import shutil +import inspect +import sys +import tempfile +import urllib.error +import urllib.parse +import urllib.request + +from pathlib import Path +from datetime import datetime +from dataclasses import dataclass, field +from collections import OrderedDict, deque +from openai.types.completion import Completion as OpenAICompletion +from openai.types.completion_choice import ( + CompletionChoice, + Logprobs as CompletionLogprobs, +) +from openai.types.completion_usage import CompletionUsage +from openai.types.chat.chat_completion import ( + ChatCompletion, + Choice as ChatCompletionChoice, + ChoiceLogprobs as ChatCompletionChoiceLogprobs, +) +from openai.types.chat.chat_completion_chunk import ( + ChatCompletionChunk, + Choice as ChatCompletionChunkChoice, + ChoiceDelta, + ChoiceDeltaFunctionCall, + ChoiceDeltaToolCallFunction, + ChoiceDeltaToolCall, + ChoiceLogprobs as ChatCompletionChunkChoiceLogprobs, +) +from openai.types.chat.chat_completion_message import ( + ChatCompletionMessage, + FunctionCall as ChatCompletionMessageFunctionCall, +) +from openai.types.chat.chat_completion_message_tool_call import ( + ChatCompletionMessageToolCall, + Function as ChatCompletionMessageToolCallFunction, +) +from openai.types.chat.chat_completion_token_logprob import ( + ChatCompletionTokenLogprob, + TopLogprob, +) +from typing import ( + Any, + Callable, + Dict, + Generator, + Iterable, + List, + Optional, + Sequence, + Tuple, + Union, + Deque, + Literal, + Iterator, + Protocol, + TypedDict, + cast, + AsyncIterator, +) + +import jinja2 +import uvicorn +import numpy as np + +from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect +from fastapi.middleware.cors import CORSMiddleware +from fastapi.responses import JSONResponse, Response, StreamingResponse + +from jinja2.sandbox import ImmutableSandboxedEnvironment + +from pydantic_core import from_json +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator + +from llama_cpp import llama_cpp # noqa: E402 +from llama_cpp import llama_cpp_ext # noqa: E402 +from llama_cpp import mtmd_cpp # noqa: E402 + + +JSON_GBNF = r""" +root ::= object +value ::= object | array | string | number | ("true" | "false" | "null") ws + +object ::= + "{" ws ( + string ":" ws value + ("," ws string ":" ws value)* + )? "}" ws + +array ::= + "[" ws ( + value + ("," ws value)* + )? "]" ws + +string ::= + "\"" ( + [^"\\\x7F\x00-\x1F] | + "\\" (["\\bfnrt] | "u" [0-9a-fA-F]{4}) + )* "\"" ws + +number ::= ("-"? ([0-9] | [1-9] [0-9]{0,15})) ("." [0-9]+)? ([eE] [-+]? [0-9] [1-9]{0,15})? ws + +ws ::= | " " | "\n" [ \t]{0,20} +""" + + +class JsonSchemaConverter: + @dataclass(frozen=True) + class BuiltinRule: + content: str + deps: Sequence[str] = () + + SPACE_RULE = '" "?' + INVALID_RULE_CHARS_RE = re.compile(r"[^a-zA-Z0-9-]+") + GRAMMAR_LITERAL_ESCAPE_RE = re.compile(r'[\r\n"]') + GRAMMAR_LITERAL_ESCAPES = {"\r": "\\r", "\n": "\\n", '"': '\\"'} + DOTALL = "[\\U00000000-\\U0010FFFF]" + DOT = "[^\\x0A\\x0D]" + NON_LITERAL_SET = set("|.()[]{}*+?") + ESCAPED_IN_REGEXPS_BUT_NOT_IN_LITERALS = set("[]()|{}*+?") + + PRIMITIVE_RULES: Optional[Dict[str, "JsonSchemaConverter.BuiltinRule"]] = None + STRING_FORMAT_RULES: Optional[Dict[str, "JsonSchemaConverter.BuiltinRule"]] = None + RESERVED_NAMES: Optional[set[str]] = None + + @staticmethod + def _build_repetition( + item_rule: str, + min_items: int, + max_items: Optional[int], + separator_rule: Optional[str] = None, + item_rule_is_literal: bool = False, + ) -> str: + if not separator_rule: + if min_items == 0 and max_items == 1: + return f"{item_rule}?" + if min_items == 1 and max_items is None: + return f"{item_rule}+" + + result = "" + + if min_items > 0: + if item_rule_is_literal and separator_rule is None: + result = '"' + (item_rule[1:-1] * min_items) + '"' + else: + result = (f" {separator_rule} " if separator_rule else " ").join( + [item_rule] * min_items + ) + + def opt_repetitions(up_to_n: int, prefix_with_sep: bool = False) -> str: + content = ( + f"{separator_rule} {item_rule}" + if prefix_with_sep and separator_rule + else item_rule + ) + if up_to_n == 0: + return "" + if up_to_n == 1: + return f"({content})?" + if separator_rule and not prefix_with_sep: + return f"({content} {opt_repetitions(up_to_n - 1, prefix_with_sep=True)})?" + return (f"({content} " * up_to_n).rstrip() + (")?" * up_to_n) + + if min_items > 0 and max_items != min_items: + result += " " + + if max_items is not None: + result += opt_repetitions(max_items - min_items, prefix_with_sep=min_items > 0) + else: + item_operator = f"({separator_rule + ' ' if separator_rule else ''}{item_rule})" + if min_items == 0 and separator_rule: + result = f"({item_rule} {item_operator}*)?" + else: + result += f"{item_operator}*" + + return result + + @classmethod + def _primitive_rules(cls) -> Dict[str, "JsonSchemaConverter.BuiltinRule"]: + if cls.PRIMITIVE_RULES is None: + up_to_15_digits = cls._build_repetition("[0-9]", 0, 15) + cls.PRIMITIVE_RULES = { + "boolean": cls.BuiltinRule('("true" | "false") space', []), + "decimal-part": cls.BuiltinRule("[0-9] " + up_to_15_digits, []), + "integral-part": cls.BuiltinRule( + "[0-9] | [1-9] " + up_to_15_digits, + [], + ), + "number": cls.BuiltinRule( + '("-"? integral-part) ("." decimal-part)? ([eE] [-+]? integral-part)? space', + ["integral-part", "decimal-part"], + ), + "integer": cls.BuiltinRule('("-"? integral-part) space', ["integral-part"]), + "value": cls.BuiltinRule( + "object | array | string | number | boolean | null", + ["object", "array", "string", "number", "boolean", "null"], + ), + "object": cls.BuiltinRule( + '"{" space ( string ":" space value ("," space string ":" space value)* )? "}" space', + ["string", "value"], + ), + "array": cls.BuiltinRule( + '"[" space ( value ("," space value)* )? "]" space', + ["value"], + ), + "uuid": cls.BuiltinRule( + r'"\"" ' + + ' "-" '.join("[0-9a-fA-F]" * n for n in [8, 4, 4, 4, 12]) + + r' "\"" space', + [], + ), + "char": cls.BuiltinRule( + r'[^"\\] | "\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F])', + [], + ), + "string": cls.BuiltinRule(r'"\"" char* "\"" space', ["char"]), + "null": cls.BuiltinRule('"null" space', []), + } + return cls.PRIMITIVE_RULES + + @classmethod + def _string_format_rules(cls) -> Dict[str, "JsonSchemaConverter.BuiltinRule"]: + if cls.STRING_FORMAT_RULES is None: + cls.STRING_FORMAT_RULES = { + "date": cls.BuiltinRule( + '[0-9] [0-9] [0-9] [0-9] "-" ( "0" [1-9] | "1" [0-2] ) "-" ( "0" [1-9] | [1-2] [0-9] | "3" [0-1] )', + [], + ), + "time": cls.BuiltinRule( + '([01] [0-9] | "2" [0-3]) ":" [0-5] [0-9] ":" [0-5] [0-9] ( "." [0-9] [0-9] [0-9] )? ( "Z" | ( "+" | "-" ) ( [01] [0-9] | "2" [0-3] ) ":" [0-5] [0-9] )', + [], + ), + "date-time": cls.BuiltinRule('date "T" time', ["date", "time"]), + "date-string": cls.BuiltinRule('"\\"" date "\\"" space', ["date"]), + "time-string": cls.BuiltinRule('"\\"" time "\\"" space', ["time"]), + "date-time-string": cls.BuiltinRule( + '"\\"" date-time "\\"" space', + ["date-time"], + ), + } + return cls.STRING_FORMAT_RULES + + @classmethod + def _reserved_names(cls) -> set[str]: + if cls.RESERVED_NAMES is None: + cls.RESERVED_NAMES = set( + ["root", "dot", *cls._primitive_rules().keys(), *cls._string_format_rules().keys()] + ) + return cls.RESERVED_NAMES + + def __init__( + self, + *, + prop_order: Dict[str, int], + allow_fetch: bool, + dotall: bool, + raw_pattern: bool, + ): + self._prop_order = prop_order + self._allow_fetch = allow_fetch + self._dotall = dotall + self._raw_pattern = raw_pattern + self._rules: Dict[str, str] = {"space": self.SPACE_RULE} + self._refs: Dict[str, Any] = {} + self._refs_being_resolved: set[str] = set() + + def _format_literal(self, literal: str) -> str: + escaped = self.GRAMMAR_LITERAL_ESCAPE_RE.sub( + lambda match: self.GRAMMAR_LITERAL_ESCAPES[match.group(0)], + literal, + ) + return f'"{escaped}"' + + def _add_rule(self, name: str, rule: str) -> str: + escaped_name = self.INVALID_RULE_CHARS_RE.sub("-", name) + if escaped_name not in self._rules or self._rules[escaped_name] == rule: + key = escaped_name + else: + suffix = 0 + while ( + f"{escaped_name}{suffix}" in self._rules + and self._rules[f"{escaped_name}{suffix}"] != rule + ): + suffix += 1 + key = f"{escaped_name}{suffix}" + self._rules[key] = rule + return key + + def resolve_refs(self, schema: Dict[str, Any], url: str) -> Dict[str, Any]: + def visit(node: Any) -> Any: + if isinstance(node, list): + return [visit(child) for child in node] + if isinstance(node, dict): + ref = node.get("$ref") + if ref is not None and ref not in self._refs: + if ref.startswith("https://"): + raise ValueError("remote schema fetch is not allowed") + elif ref.startswith("#/"): + target = schema + ref = f"{url}{ref}" + node["$ref"] = ref + else: + raise ValueError(f"Unsupported ref {ref}") + + for selector in ref.split("#")[-1].split("/")[1:]: + assert target is not None and selector in target, ( + f"Error resolving ref {ref}: {selector} not in {target}" + ) + target = target[selector] + self._refs[ref] = target + else: + for value in node.values(): + visit(value) + return node + + return cast(Dict[str, Any], visit(schema)) + + def _generate_union_rule(self, name: str, alt_schemas: List[Dict[str, Any]]) -> str: + return " | ".join( + self.visit(alt_schema, f"{name}{'-' if name else 'alternative-'}{index}") + for index, alt_schema in enumerate(alt_schemas) + ) + + def _visit_pattern(self, pattern: str, name: str) -> str: + assert pattern.startswith("^") and pattern.endswith("$"), ( + 'Pattern must start with "^" and end with "$"' + ) + pattern = pattern[1:-1] + sub_rule_ids: Dict[str, str] = {} + index = 0 + length = len(pattern) + + def to_rule(item: Tuple[str, bool]) -> str: + text, is_literal = item + return f'"{text}"' if is_literal else text + + def transform() -> Tuple[str, bool]: + nonlocal index + start = index + sequence: List[Tuple[str, bool]] = [] + + def get_dot() -> str: + rule = self.DOTALL if self._dotall else self.DOT + return self._add_rule("dot", rule) + + def join_sequence() -> Tuple[str, bool]: + if len(sequence) == 1: + return sequence[0] + return (" ".join(to_rule(item) for item in sequence), False) + + while index < length: + char = pattern[index] + if char == ".": + sequence.append((get_dot(), False)) + index += 1 + elif char == "(": + index += 1 + if index < length: + assert pattern[index] != "?", ( + f'Unsupported pattern syntax "{pattern[index]}" at index {index} of /{pattern}/' + ) + sequence.append((f"({to_rule(transform())})", False)) + elif char == ")": + index += 1 + assert start > 0 and pattern[start - 1] == "(", ( + f"Unbalanced parentheses; start = {start}, index = {index}, pattern = {pattern}" + ) + return join_sequence() + elif char == "[": + square_brackets = char + index += 1 + while index < length and pattern[index] != "]": + if pattern[index] == "\\": + square_brackets += pattern[index : index + 2] + index += 2 + else: + square_brackets += pattern[index] + index += 1 + assert index < length, ( + f"Unbalanced square brackets; start = {start}, index = {index}, pattern = {pattern}" + ) + square_brackets += "]" + index += 1 + sequence.append((square_brackets, False)) + elif char == "|": + sequence.append(("|", False)) + index += 1 + elif char in ("*", "+", "?"): + sequence[-1] = (to_rule(sequence[-1]) + char, False) + index += 1 + elif char == "{": + curly_brackets = char + index += 1 + while index < length and pattern[index] != "}": + curly_brackets += pattern[index] + index += 1 + assert index < length, ( + f"Unbalanced curly brackets; start = {start}, index = {index}, pattern = {pattern}" + ) + curly_brackets += "}" + index += 1 + numbers = [part.strip() for part in curly_brackets[1:-1].split(",")] + min_times = 0 + max_times: Optional[int] = None + try: + if len(numbers) == 1: + min_times = int(numbers[0]) + max_times = min_times + else: + assert len(numbers) == 2 + min_times = int(numbers[0]) if numbers[0] else 0 + max_times = int(numbers[1]) if numbers[1] else None + except ValueError as exc: + raise ValueError( + f"Invalid quantifier {curly_brackets} in /{pattern}/" + ) from exc + + sub, sub_is_literal = sequence[-1] + if not sub_is_literal: + rule_id = sub_rule_ids.get(sub) + if rule_id is None: + rule_id = self._add_rule(f"{name}-{len(sub_rule_ids) + 1}", sub) + sub_rule_ids[sub] = rule_id + sub = rule_id + + sequence[-1] = ( + self._build_repetition( + f'"{sub}"' if sub_is_literal else sub, + min_times, + max_times, + item_rule_is_literal=sub_is_literal, + ), + False, + ) + else: + literal = "" + while index < length: + if pattern[index] == "\\" and index < length - 1: + next_char = pattern[index + 1] + if next_char in self.ESCAPED_IN_REGEXPS_BUT_NOT_IN_LITERALS: + index += 1 + literal += pattern[index] + index += 1 + else: + literal += pattern[index : index + 2] + index += 2 + elif pattern[index] == '"' and not self._raw_pattern: + literal += '\\"' + index += 1 + elif pattern[index] not in self.NON_LITERAL_SET and ( + index == length - 1 + or literal == "" + or pattern[index + 1] == "." + or pattern[index + 1] not in self.NON_LITERAL_SET + ): + literal += pattern[index] + index += 1 + else: + break + if literal: + sequence.append((literal, True)) + + return join_sequence() + + return self._add_rule( + name, + ( + to_rule(transform()) + if self._raw_pattern + else '"\\"" ' + to_rule(transform()) + ' "\\"" space' + ), + ) + + def _resolve_ref(self, ref: str) -> str: + ref_name = ref.split("/")[-1] + if ref_name not in self._rules and ref not in self._refs_being_resolved: + self._refs_being_resolved.add(ref) + resolved = self._refs[ref] + ref_name = self.visit(resolved, ref_name) + self._refs_being_resolved.remove(ref) + return ref_name + + def _generate_constant_rule(self, value: Any) -> str: + return self._format_literal(json.dumps(value)) + + def visit(self, schema: Dict[str, Any], name: str) -> str: + schema_type = schema.get("type") + schema_format = schema.get("format") + rule_name = name + "-" if name in self._reserved_names() else name or "root" + + ref = schema.get("$ref") + if ref is not None: + return self._add_rule(rule_name, self._resolve_ref(ref)) + + if "oneOf" in schema or "anyOf" in schema: + return self._add_rule( + rule_name, + self._generate_union_rule(name, cast(List[Dict[str, Any]], schema.get("oneOf") or schema["anyOf"])), + ) + + if isinstance(schema_type, list): + return self._add_rule( + rule_name, + self._generate_union_rule(name, [{"type": entry} for entry in schema_type]), + ) + + if "const" in schema: + return self._add_rule(rule_name, self._generate_constant_rule(schema["const"])) + + if "enum" in schema: + rule = " | ".join(self._generate_constant_rule(value) for value in schema["enum"]) + return self._add_rule(rule_name, rule) + + if schema_type in (None, "object") and ( + "properties" in schema + or ("additionalProperties" in schema and schema["additionalProperties"] is not True) + ): + required_props = set(schema.get("required", [])) + property_items = list(cast(Dict[str, Any], schema.get("properties", {})).items()) + return self._add_rule( + rule_name, + self._build_object_rule( + property_items, required_props, name, schema.get("additionalProperties") + ), + ) + + if schema_type in (None, "object") and "allOf" in schema: + allof_required_props: set[str] = set() + allof_property_items: List[Tuple[str, Any]] = [] + + def add_component(component_schema: Dict[str, Any], is_required: bool) -> None: + component_ref = component_schema.get("$ref") + if component_ref is not None: + component_schema = cast(Dict[str, Any], self._refs[component_ref]) + if "properties" in component_schema: + for prop_name, prop_schema in cast(Dict[str, Any], component_schema["properties"]).items(): + allof_property_items.append((prop_name, prop_schema)) + if is_required: + allof_required_props.add(prop_name) + + for entry in cast(List[Dict[str, Any]], schema["allOf"]): + if "anyOf" in entry: + for alt in cast(List[Dict[str, Any]], entry["anyOf"]): + add_component(alt, is_required=False) + else: + add_component(entry, is_required=True) + + return self._add_rule( + rule_name, + self._build_object_rule( + allof_property_items, + allof_required_props, + name, + additional_properties=[], + ), + ) + + if schema_type in (None, "array") and ("items" in schema or "prefixItems" in schema): + items = schema.get("items") or schema["prefixItems"] + if isinstance(items, list): + return self._add_rule( + rule_name, + '"[" space ' + + ' "," space '.join( + self.visit(item, f"{name}{'-' if name else ''}tuple-{index}") + for index, item in enumerate(items) + ) + + ' "]" space', + ) + item_rule_name = self.visit(cast(Dict[str, Any], items), f"{name}{'-' if name else ''}item") + min_items = int(schema.get("minItems", 0)) + max_items = cast(Optional[int], schema.get("maxItems")) + return self._add_rule( + rule_name, + '"[" space ' + + self._build_repetition( + item_rule_name, + min_items, + max_items, + separator_rule='"," space', + ) + + ' "]" space', + ) + + if schema_type in (None, "string") and "pattern" in schema: + return self._visit_pattern(cast(str, schema["pattern"]), rule_name) + + if schema_type in (None, "string") and re.match(r"^uuid[1-5]?$", schema_format or ""): + return self._add_primitive( + "root" if rule_name == "root" else cast(str, schema_format), + self._primitive_rules()["uuid"], + ) + + if schema_type in (None, "string") and f"{schema_format}-string" in self._string_format_rules(): + primitive_name = f"{schema_format}-string" + return self._add_rule( + rule_name, + self._add_primitive(primitive_name, self._string_format_rules()[primitive_name]), + ) + + if schema_type == "string" and ("minLength" in schema or "maxLength" in schema): + char_rule = self._add_primitive("char", self._primitive_rules()["char"]) + min_len = int(schema.get("minLength", 0)) + max_len = cast(Optional[int], schema.get("maxLength")) + return self._add_rule( + rule_name, + r'"\"" ' + + self._build_repetition(char_rule, min_len, max_len) + + r' "\"" space', + ) + + if schema_type == "object" or len(schema) == 0: + return self._add_rule( + rule_name, + self._add_primitive("object", self._primitive_rules()["object"]), + ) + + primitive_rules = self._primitive_rules() + assert schema_type in primitive_rules, f"Unrecognized schema: {schema}" + return self._add_primitive( + "root" if rule_name == "root" else cast(str, schema_type), + primitive_rules[cast(str, schema_type)], + ) + + def _add_primitive(self, name: str, rule: "JsonSchemaConverter.BuiltinRule") -> str: + rule_name = self._add_rule(name, rule.content) + primitive_rules = self._primitive_rules() + string_format_rules = self._string_format_rules() + for dependency in rule.deps: + dependency_rule = primitive_rules.get(dependency) or string_format_rules.get( + dependency + ) + assert dependency_rule is not None, f"Rule {dependency} not known" + if dependency not in self._rules: + self._add_primitive(dependency, dependency_rule) + return rule_name + + def _build_object_rule( + self, + properties: List[Tuple[str, Any]], + required: set[str], + name: str, + additional_properties: Union[bool, Any], + ) -> str: + prop_order = self._prop_order + sorted_props = [ + key + for _, (key, _) in sorted( + enumerate(properties), + key=lambda indexed_key: ( + prop_order.get(indexed_key[1][0], len(prop_order)), + indexed_key[0], + ), + ) + ] + + property_kv_rule_names: Dict[str, str] = {} + for prop_name, prop_schema in properties: + prop_rule_name = self.visit( + cast(Dict[str, Any], prop_schema), + f"{name}{'-' if name else ''}{prop_name}", + ) + property_kv_rule_names[prop_name] = self._add_rule( + f"{name}{'-' if name else ''}{prop_name}-kv", + rf'{self._format_literal(json.dumps(prop_name))} space ":" space {prop_rule_name}', + ) + + required_props = [key for key in sorted_props if key in required] + optional_props = [key for key in sorted_props if key not in required] + + if additional_properties is True or isinstance(additional_properties, dict): + sub_name = f"{name}{'-' if name else ''}additional" + value_rule = self.visit( + {} if additional_properties is True else cast(Dict[str, Any], additional_properties), + f"{sub_name}-value", + ) + property_kv_rule_names["*"] = self._add_rule( + f"{sub_name}-kv", + self._add_primitive("string", self._primitive_rules()["string"]) + + f' ":" space {value_rule}', + ) + optional_props.append("*") + + rule = '"{" space ' + rule += ' "," space '.join(property_kv_rule_names[key] for key in required_props) + + if optional_props: + if required_props: + rule += ' ( "," space ( ' + else: + rule += "( " + + def get_recursive_refs(keys: List[str], first_is_optional: bool) -> str: + head, *rest = keys + kv_rule_name = property_kv_rule_names[head] + result = "" + if head == "*": + if first_is_optional: + result = f"({kv_rule_name})?" + else: + result = kv_rule_name + elif first_is_optional: + result = f'( "," space {kv_rule_name} )?' + else: + result = kv_rule_name + if rest: + result += " " + self._add_rule( + f"{name}{'-' if name else ''}{head}-rest", + get_recursive_refs(rest, first_is_optional=True), + ) + return result + + rule += " | ".join( + get_recursive_refs(optional_props[index:], first_is_optional=False) + for index in range(len(optional_props)) + ) + if required_props: + rule += " )" + rule += " )?" + + rule += ' "}" space' + return rule + + def format_grammar(self) -> str: + return "\n".join( + f"{name} ::= {rule}" + for name, rule in sorted(self._rules.items(), key=lambda item: item[0]) + ) + + @classmethod + def to_gbnf(cls, schema: str, prop_order: Optional[List[str]] = None) -> str: + property_order = prop_order or [] + loaded_schema = json.loads(schema) + order_index = {name: index for index, name in enumerate(property_order)} + converter = cls( + prop_order=order_index, + allow_fetch=False, + dotall=False, + raw_pattern=False, + ) + resolved_schema = converter.resolve_refs(loaded_schema, "stdin") + converter.visit(resolved_schema, "") + return converter.format_grammar() + + +class RadixTrie: + __slots__ = ("root", "sequences", "sequence_lengths") + + @dataclass + class Node: + label: Tuple[int, ...] = () + parent: Optional["RadixTrie.Node"] = None + children: Dict[int, "RadixTrie.Node"] = field(default_factory=dict) + sequences: set[int] = field(default_factory=set) + tails: set[int] = field(default_factory=set) + + def __init__(self) -> None: + self.root = RadixTrie.Node() + self.sequences: Dict[int, RadixTrie.Node] = {} + self.sequence_lengths: Dict[int, int] = {} + + @staticmethod + def _pick_sequence(candidates: set[int], preferred_sequences: Optional[Any]) -> int: + if preferred_sequences is None: + return next(iter(candidates)) + if isinstance(preferred_sequences, OrderedDict): + for sequence_id in reversed(preferred_sequences): + if sequence_id in candidates: + return sequence_id + return next(iter(candidates)) + if isinstance(preferred_sequences, (list, tuple)): + for sequence_id in reversed(preferred_sequences): + if sequence_id in candidates: + return sequence_id + return next(iter(candidates)) + preferred = candidates & preferred_sequences + if preferred: + return next(iter(preferred)) + return next(iter(candidates)) + + @staticmethod + def _common_prefix_len( + label: Sequence[int], + tokens: Sequence[int], + offset: int, + ) -> int: + limit = min(len(label), len(tokens) - offset) + match_len = 0 + while match_len < limit and label[match_len] == tokens[offset + match_len]: + match_len += 1 + return match_len + + def _split_child( + self, + parent: "RadixTrie.Node", + child: "RadixTrie.Node", + prefix_len: int, + ) -> "RadixTrie.Node": + assert 0 < prefix_len < len(child.label) + prefix = child.label[:prefix_len] + suffix = child.label[prefix_len:] + middle = RadixTrie.Node( + label=prefix, + parent=parent, + sequences=set(child.sequences), + ) + parent.children[prefix[0]] = middle + child.label = suffix + child.parent = middle + middle.children[suffix[0]] = child + return middle + + def _locate_prefix_node( + self, + sequence_id: int, + keep_len: int, + ) -> "RadixTrie.Node": + total_len = self.sequence_lengths[sequence_id] + assert 0 <= keep_len <= total_len + if keep_len == 0: + return self.root + node = self.sequences[sequence_id] + drop_len = total_len - keep_len + while node is not self.root: + label_len = len(node.label) + if drop_len == 0: + return node + if drop_len < label_len: + parent = node.parent + assert parent is not None + return self._split_child(parent, node, label_len - drop_len) + drop_len -= label_len + parent = node.parent + assert parent is not None + node = parent + return self.root + + def extend( + self, + sequence_id: int, + tokens: Sequence[int], + ) -> None: + assert sequence_id >= 0 + node = self.sequences.get(sequence_id, self.root) + if tokens: + node.tails.discard(sequence_id) + length = self.sequence_lengths.get(sequence_id, 0) + index = 0 + while index < len(tokens): + token = tokens[index] + child = node.children.get(token) + if child is None: + child = RadixTrie.Node( + label=tuple(tokens[index:]), + parent=node, + sequences={sequence_id}, + ) + node.children[token] = child + node = child + length += len(tokens) - index + index = len(tokens) + break + match_len = self._common_prefix_len(child.label, tokens, index) + if match_len == len(child.label): + child.sequences.add(sequence_id) + node = child + length += match_len + index += match_len + continue + node = self._split_child(node, child, match_len) + node.sequences.add(sequence_id) + length += match_len + index += match_len + if index == len(tokens): + break + suffix = RadixTrie.Node( + label=tuple(tokens[index:]), + parent=node, + sequences={sequence_id}, + ) + node.children[suffix.label[0]] = suffix + node = suffix + length += len(tokens) - index + index = len(tokens) + break + if node is self.root: + self.sequences.pop(sequence_id, None) + self.sequence_lengths.pop(sequence_id, None) + self.root.tails.discard(sequence_id) + else: + self.sequences[sequence_id] = node + self.sequence_lengths[sequence_id] = length + node.tails.add(sequence_id) + + def length(self, sequence_id: int) -> int: + return self.sequence_lengths.get(sequence_id, 0) + + def truncate(self, sequence_id: int, keep_len: int) -> None: + assert sequence_id >= 0 + assert sequence_id in self.sequence_lengths + assert 0 <= keep_len <= self.sequence_lengths[sequence_id] + current_len = self.sequence_lengths[sequence_id] + if keep_len == current_len: + return + boundary = self._locate_prefix_node(sequence_id, keep_len) + node = self.sequences.get(sequence_id, self.root) + if node is not self.root: + node.tails.discard(sequence_id) + while node is not boundary: + node.sequences.remove(sequence_id) + parent = node.parent + assert parent is not None + if not node.sequences: + del parent.children[node.label[0]] + node = parent + if boundary is self.root: + self.sequences.pop(sequence_id, None) + self.sequence_lengths.pop(sequence_id, None) + else: + self.sequences[sequence_id] = boundary + self.sequence_lengths[sequence_id] = keep_len + boundary.tails.add(sequence_id) + + def copy(self, source_sequence_id: int, dest_sequence_id: int, keep_len: int) -> None: + assert source_sequence_id >= 0 + assert dest_sequence_id >= 0 + assert source_sequence_id in self.sequence_lengths + assert dest_sequence_id not in self.sequence_lengths + assert 0 <= keep_len <= self.sequence_lengths[source_sequence_id] + if keep_len == 0: + self.sequences[dest_sequence_id] = self.root + self.sequence_lengths[dest_sequence_id] = 0 + self.root.tails.add(dest_sequence_id) + return + node = self._locate_prefix_node(source_sequence_id, keep_len) + self.sequences[dest_sequence_id] = node + self.sequence_lengths[dest_sequence_id] = keep_len + node.tails.add(dest_sequence_id) + while node is not self.root: + node.sequences.add(dest_sequence_id) + parent = node.parent + assert parent is not None + node = parent + + def tokens(self, sequence_id: int, keep_len: Optional[int] = None) -> List[int]: + length = self.sequence_lengths[sequence_id] + target_len = length if keep_len is None else keep_len + assert 0 <= target_len <= length + node = self.sequences[sequence_id] + labels: List[Tuple[int, ...]] = [] + while node is not self.root: + labels.append(node.label) + parent = node.parent + assert parent is not None + node = parent + values: List[int] = [] + for label in reversed(labels): + values.extend(label) + return values[:target_len] + + def longest_prefix( + self, + tokens: Sequence[int], + preferred_sequences: Optional[Any] = None, + *, + exact_only: bool = False, + ) -> Tuple[int, int]: + node = self.root + longest_sequence_id = -1 + longest_length = 0 + index = 0 + while index < len(tokens): + child = node.children.get(tokens[index]) + if child is None: + break + match_len = self._common_prefix_len(child.label, tokens, index) + if match_len < len(child.label): + break + node = child + index += match_len + candidates = node.tails if exact_only else node.sequences + if candidates: + longest_sequence_id = self._pick_sequence(candidates, preferred_sequences) + longest_length = index + return longest_sequence_id, longest_length + + +class SequenceHistory: + __slots__ = ("_position_lengths", "_root", "_tails", "size") + + @dataclass + class Node: + token: Optional[int] = None + parent: Optional["SequenceHistory.Node"] = None + children: Dict[int, "SequenceHistory.Node"] = field(default_factory=dict) + sequences: set[int] = field(default_factory=set) + position_increment: int = 1 + + def __init__(self) -> None: + self._root = SequenceHistory.Node() + self._tails: Dict[int, SequenceHistory.Node] = {} + self._position_lengths: Dict[int, int] = {} + self.size = 0 + + def extend( + self, + sequence_id: int, + tokens: Sequence[int], + position_increments: Optional[Sequence[int]] = None, + ) -> None: + assert sequence_id >= 0 + if position_increments is None: + position_increments = [1] * len(tokens) + assert len(position_increments) == len(tokens) + node = self._tails.get(sequence_id, self._root) + position_length = self._position_lengths.get(sequence_id, 0) + for token, position_increment in zip(tokens, position_increments): + position_increment = max(0, int(position_increment)) + child = node.children.get(sequence_id) + if child is None: + child = SequenceHistory.Node( + token=token, + parent=node, + position_increment=position_increment, + ) + node.children[sequence_id] = child + self.size += 1 + else: + assert child.parent is node + assert child.token == token + assert child.position_increment == position_increment + child.sequences.add(sequence_id) + position_length += position_increment + node = child + if node is self._root: + self._tails.pop(sequence_id, None) + self._position_lengths.pop(sequence_id, None) + else: + self._tails[sequence_id] = node + self._position_lengths[sequence_id] = position_length + + def position_length(self, sequence_id: int) -> int: + return self._position_lengths.get(sequence_id, 0) + + def position_length_for_prefix(self, sequence_id: int, keep_len: int) -> int: + if keep_len <= 0 or sequence_id not in self._tails: + return 0 + node = self._tails[sequence_id] + increments: List[int] = [] + while node is not self._root: + increments.append(node.position_increment) + parent = node.parent + assert parent is not None + node = parent + increments.reverse() + assert keep_len <= len(increments) + return sum(increments[:keep_len]) + + def copy( + self, + source_sequence_id: int, + dest_sequence_id: int, + source_length: int, + keep_len: int, + ) -> None: + assert source_sequence_id >= 0 + assert dest_sequence_id >= 0 + assert source_sequence_id in self._tails + assert dest_sequence_id not in self._tails + assert 0 <= keep_len <= source_length + node = self._tails[source_sequence_id] + path: List[SequenceHistory.Node] = [] + for _ in range(source_length - keep_len): + parent = node.parent + assert parent is not None + node = parent + while node is not self._root: + path.append(node) + parent = node.parent + assert parent is not None + node = parent + parent = self._root + position_length = 0 + for child in reversed(path): + parent.children[dest_sequence_id] = child + child.sequences.add(dest_sequence_id) + position_length += child.position_increment + parent = child + if keep_len == 0: + self._tails.pop(dest_sequence_id, None) + self._position_lengths.pop(dest_sequence_id, None) + else: + self._tails[dest_sequence_id] = path[0] + self._position_lengths[dest_sequence_id] = position_length + + def truncate( + self, + sequence_id: int, + current_length: int, + keep_len: int, + ) -> None: + assert sequence_id >= 0 + assert sequence_id in self._tails + assert 0 <= keep_len <= current_length + node = self._tails[sequence_id] + drop = current_length - keep_len + position_length = self._position_lengths.get(sequence_id, 0) + while node is not self._root and drop > 0: + node.sequences.remove(sequence_id) + parent = node.parent + assert parent is not None + child = parent.children.get(sequence_id) + if child is node: + del parent.children[sequence_id] + if not node.sequences: + self.size -= 1 + position_length -= node.position_increment + node = parent + drop -= 1 + if node is self._root: + self._tails.pop(sequence_id, None) + self._position_lengths.pop(sequence_id, None) + else: + self._tails[sequence_id] = node + self._position_lengths[sequence_id] = max(0, position_length) + + +class DraftProvider(abc.ABC): + @abc.abstractmethod + def draft( + self, + input_ids: np.ndarray, + /, + *, + seq_id: int, + max_tokens: Optional[int], + ) -> np.ndarray: + raise NotImplementedError() + + def can_draft(self, input_length: int, /, *, seq_id: int) -> bool: + return True + + def process(self, batch: Any, /) -> None: + pass + + def accept(self, seq_id: int, accepted_draft_tokens: int) -> None: + pass + + def truncate(self, seq_id: int, keep_len: int) -> None: + pass + + def copy_sequence( + self, + source_seq_id: int, + dest_seq_id: int, + p0: int, + p1: int, + ) -> None: + pass + + def set_target_processing_enabled(self, enabled: bool) -> None: + pass + + def close(self) -> None: + pass + + +class PromptLookupDecoding(DraftProvider): + def __init__(self, max_ngram_size: int = 2, num_pred_tokens: int = 10) -> None: + self._max_ngram_size = max_ngram_size + self._num_pred_tokens = num_pred_tokens + + def draft( + self, + input_ids: np.ndarray, + /, + *, + seq_id: int, + max_tokens: Optional[int], + ) -> np.ndarray: + input_length = input_ids.shape[0] + if input_length < 2: + return np.array([], dtype=np.intc) + num_pred_tokens = self._num_pred_tokens + if max_tokens is not None: + num_pred_tokens = min(num_pred_tokens, max_tokens) + if num_pred_tokens <= 0: + return np.array([], dtype=np.intc) + max_ngram_size = min(self._max_ngram_size, input_length - 1) + for ngram_size in range(max_ngram_size, 0, -1): + windows = np.lib.stride_tricks.sliding_window_view(input_ids, (ngram_size,)) + ngram = input_ids[-ngram_size:] + matches = np.all(windows == ngram, axis=1) + match_indices = np.nonzero(matches)[0] + for index in match_indices: + start = index + ngram_size + if start >= input_length: + continue + end = min(start + num_pred_tokens, input_length) + if start < end: + return input_ids[start:end].astype(np.intc, copy=False) + return np.array([], dtype=np.intc) + + +class MTPDraftProvider(DraftProvider): + batched_draft = True + sampled_batch_draft = True + + @dataclass + class DraftManyState: + result_index: int + seq_id: int + first_pos: int + keep_len: int + n_predict: int + token: int + drafted: List[int] + embedding: np.ndarray + cache_key: Tuple[Tuple[int, ...], int] + + def __init__( + self, + *, + model: "Model", + draft_model: Any, + context_params: Any, + num_pred_tokens: int, + top_k: int, + p_min: float, + ) -> None: + self.target_ctx = model.ctx + self.model = draft_model + self.n_seq_max = model.n_seq_max + self.n_vocab = model.n_vocab + self.n_embd = int(llama_cpp.llama_model_n_embd_out(self.model)) + if self.n_embd <= 0: + self.n_embd = int(llama_cpp.llama_model_n_embd(self.model)) + if self.n_embd != model.n_embd: + raise RuntimeError( + "MTP draft model output embedding size must match target model " + f"embedding size ({self.n_embd} != {model.n_embd})" + ) + self.num_pred_tokens = max(0, int(num_pred_tokens)) + self.top_k = max(1, int(top_k)) + self.p_min = max(0.0, min(1.0, float(p_min))) + self.ctx = llama_cpp.llama_init_from_model(self.model, context_params) + if self.ctx is None: + raise RuntimeError("failed to create MTP draft context") + ctx_other = llama_cpp_ext.llama_get_ctx_other(self.ctx) + self.is_mem_shared = bool(ctx_other and ctx_other == self.target_ctx) + self.sampled_batch_draft = not self.is_mem_shared + self.n_batch = int(llama_cpp.llama_n_batch(self.ctx)) + mem = llama_cpp.llama_get_memory(self.ctx) + if mem is None: + llama_cpp.llama_free(self.ctx) + raise RuntimeError("failed to access MTP draft memory") + self.mem = mem + self.batch = llama_cpp.llama_batch_init(self.n_batch, self.n_embd, 1) + self._batch_tokens = (llama_cpp.llama_token * self.n_batch)() + self.batch.token = self._batch_tokens + self.batch_embeddings = np.ctypeslib.as_array( + self.batch.embd, + shape=(self.n_batch * self.n_embd,), + ) + self._samplers: List[Any] = [] + self.pending_h = np.zeros((self.n_seq_max, self.n_embd), dtype=np.float32) + self.verify_h: List[np.ndarray] = [ + np.empty((0, self.n_embd), dtype=np.float32) + for _ in range(self.n_seq_max) + ] + self.verify_h_pos: List[List[int]] = [[] for _ in range(self.n_seq_max)] + self.verify_h_rows = [0] * self.n_seq_max + self.ready = [False] * self.n_seq_max + self.ready_pos = [0] * self.n_seq_max + self.context_pos = [0] * self.n_seq_max + self.decode_seconds_total = 0.0 + self.decode_calls_total = 0 + self.decode_tokens_total = 0 + self.decode_failures_total = 0 + self.target_processing_enabled = False + self.set_target_processing_enabled(True) + llama_cpp_ext.llama_set_embeddings_nextn( + self.ctx, + True, + True, + ) + self._init_samplers() + + def _init_samplers(self) -> None: + for seq_id in range(self.n_seq_max): + params = llama_cpp.llama_sampler_chain_default_params() + params.no_perf = True + sampler = llama_cpp.llama_sampler_chain_init(params) + if self.top_k > 1: + llama_cpp.llama_sampler_chain_add( + sampler, + llama_cpp.llama_sampler_init_top_k(self.top_k), + ) + llama_cpp.llama_sampler_chain_add( + sampler, + llama_cpp.llama_sampler_init_greedy(), + ) + self._samplers.append(sampler) + + def _passes_p_min(self, output_index: int) -> bool: + if self.p_min <= 0.0: + return True + logits_ptr = llama_cpp.llama_get_logits_ith(self.ctx, output_index) + if not logits_ptr: + return False + logits = np.ctypeslib.as_array(logits_ptr, shape=(self.n_vocab,)) + n_values = min(self.top_k, self.n_vocab) + if n_values <= 0: + return False + if n_values == self.n_vocab: + values = logits + else: + top_indices = np.argpartition(logits, -n_values)[-n_values:] + values = logits[top_indices] + max_logit = float(np.max(values)) + weights = np.exp(values.astype(np.float64, copy=False) - max_logit) + total = float(np.sum(weights)) + if total <= 0.0 or not math.isfinite(total): + return False + return 1.0 / total >= self.p_min + + def _sample_token(self, output_index: int = 0, *, seq_id: int = 0) -> Optional[int]: + if seq_id < 0 or seq_id >= len(self._samplers): + return None + if not self._passes_p_min(output_index): + return None + sampler = self._samplers[seq_id] + token = int(llama_cpp.llama_sampler_sample(sampler, self.ctx, output_index)) + if token == llama_cpp.LLAMA_TOKEN_NULL: + return None + return token + + def _reset_sampler(self, seq_id: int) -> None: + if 0 <= seq_id < len(self._samplers): + llama_cpp.llama_sampler_reset(self._samplers[seq_id]) + + def close(self) -> None: + self.set_target_processing_enabled(False) + self.batch.token = ctypes.POINTER(llama_cpp.llama_token)() + llama_cpp.llama_batch_free(self.batch) + llama_cpp.llama_free(self.ctx) + for sampler in self._samplers: + llama_cpp.llama_sampler_free(sampler) + self._samplers.clear() + + def set_target_processing_enabled(self, enabled: bool) -> None: + if self.target_processing_enabled == enabled: + return + llama_cpp_ext.llama_set_embeddings_nextn( + self.target_ctx, + enabled, + False, + ) + self.target_processing_enabled = enabled + + def _clear_batch(self) -> None: + self.batch.n_tokens = 0 + + def _batch_embeddings(self) -> np.ndarray: + return self.batch_embeddings + + def _set_batch_embedding_row( + self, + row: int, + embedding: Union[np.ndarray, ctypes.POINTER(ctypes.c_float)], + ) -> None: + row_start = row * self.n_embd + row_end = row_start + self.n_embd + if isinstance(embedding, np.ndarray): + self._batch_embeddings()[row_start:row_end] = embedding + return + self._batch_embeddings()[row_start:row_end] = np.ctypeslib.as_array( + embedding, + shape=(self.n_embd,), + ) + + def _add_batch_token( + self, + *, + token: int, + pos: int, + seq_id: int, + logits: bool, + ) -> None: + slot = int(self.batch.n_tokens) + if slot >= self.n_batch: + raise RuntimeError("MTP draft batch capacity exceeded") + self.batch.token[slot] = int(token) + self.batch.pos[slot] = int(pos) + self.batch.seq_id[slot][0] = int(seq_id) + self.batch.n_seq_id[slot] = 1 + self.batch.logits[slot] = int(logits) + self.batch.n_tokens += 1 + + def _try_decode_batch(self) -> bool: + n_tokens = int(self.batch.n_tokens) + if n_tokens <= 0: + return True + started_at = time.perf_counter() + result = int(llama_cpp.llama_decode(self.ctx, self.batch)) + self.decode_seconds_total += time.perf_counter() - started_at + self.decode_calls_total += 1 + self.decode_tokens_total += n_tokens + if result != 0: + self.decode_failures_total += 1 + return False + return True + + def _decode_batch(self) -> None: + n_tokens = int(self.batch.n_tokens) + if n_tokens <= 0: + return + started_at = time.perf_counter() + result = int(llama_cpp.llama_decode(self.ctx, self.batch)) + self.decode_seconds_total += time.perf_counter() - started_at + self.decode_calls_total += 1 + self.decode_tokens_total += n_tokens + if result != 0: + self.decode_failures_total += 1 + raise RuntimeError(f"MTP draft decode failed with code {result}") + + def metric_definitions( + self, + ) -> List[Tuple[str, str, str, Union[int, float]]]: + decode_tokens_seconds = ( + self.decode_tokens_total / self.decode_seconds_total + if self.decode_seconds_total > 0.0 + else 0.0 + ) + return [ + ( + "counter", + "batch_server:mtp_decode_seconds_total", + "Time spent inside llama_decode() for the MTP draft context.", + self.decode_seconds_total, + ), + ( + "counter", + "batch_server:mtp_decode_calls_total", + "Number of llama_decode() calls made by the MTP draft context.", + self.decode_calls_total, + ), + ( + "counter", + "batch_server:mtp_decode_tokens_total", + "Number of batch rows decoded by the MTP draft context.", + self.decode_tokens_total, + ), + ( + "counter", + "batch_server:mtp_decode_failures_total", + "Number of failed MTP draft context decode calls.", + self.decode_failures_total, + ), + ( + "gauge", + "batch_server:mtp_decode_tokens_seconds", + "Average MTP draft context decode throughput in batch rows/s.", + decode_tokens_seconds, + ), + ] + + def can_draft(self, input_length: int, /, *, seq_id: int) -> bool: + if ( + input_length <= 0 + or seq_id < 0 + or seq_id >= self.n_seq_max + or not self.ready[seq_id] + ): + return False + return self.ready_pos[seq_id] == input_length - 1 + + def process(self, batch: Any, /) -> None: + n_tokens = int(batch.n_tokens) + if ( + n_tokens <= 0 + or not self.target_processing_enabled + or not bool(batch.token) + or bool(batch.embd) + ): + return + + h_tgt = llama_cpp_ext.llama_get_embeddings_nextn(self.target_ctx) + if not h_tgt: + raise RuntimeError("missing target nextn embeddings for MTP") + h_tgt_rows = np.ctypeslib.as_array( + h_tgt, + shape=(n_tokens, self.n_embd), + ) + + previous_row_by_seq: Dict[int, int] = {} + first_pos_by_seq: Dict[int, int] = {} + target_rows_by_seq: Dict[int, List[int]] = {} + aligned_by_seq: Dict[int, bool] = {} + + for start in range(0, n_tokens, self.n_batch): + self._process_rows( + batch, + h_tgt_rows, + start, + min(start + self.n_batch, n_tokens), + previous_row_by_seq, + first_pos_by_seq, + target_rows_by_seq, + aligned_by_seq, + ) + + for seq_id, rows in target_rows_by_seq.items(): + if not aligned_by_seq.get(seq_id, False): + self.ready[seq_id] = False + continue + if rows[-1] - rows[0] + 1 == len(rows): + target_rows = h_tgt_rows[rows[0] : rows[-1] + 1] + else: + target_rows = h_tgt_rows[rows].copy() + target_positions = [int(batch.pos[source_index]) for source_index in rows] + self.verify_h[seq_id] = target_rows + self.verify_h_pos[seq_id] = target_positions + self.verify_h_rows[seq_id] = len(rows) + self.pending_h[seq_id] = target_rows[-1] + self.ready[seq_id] = True + self.ready_pos[seq_id] = target_positions[-1] + 1 + + def _process_rows( + self, + batch: Any, + h_tgt_rows: np.ndarray, + start: int, + end: int, + previous_row_by_seq: Dict[int, int], + first_pos_by_seq: Dict[int, int], + target_rows_by_seq: Dict[int, List[int]], + aligned_by_seq: Dict[int, bool], + ) -> None: + added_pos_by_seq: Dict[int, int] = {} + self._clear_batch() + for index in range(start, end): + if int(batch.n_seq_id[index]) != 1: + raise RuntimeError("MTP requires one sequence id per batch token") + seq_id = int(batch.seq_id[index][0]) + if seq_id < 0 or seq_id >= self.n_seq_max: + raise RuntimeError(f"MTP sequence id out of range: {seq_id}") + pos = int(batch.pos[index]) + first_pos = first_pos_by_seq.setdefault(seq_id, pos) + aligned = first_pos <= 0 or ( + self.ready[seq_id] and self.ready_pos[seq_id] == first_pos + ) + aligned_by_seq.setdefault(seq_id, aligned) + previous_row = previous_row_by_seq.get(seq_id) + if ( + aligned + and not self.is_mem_shared + and pos >= 0 + and pos >= self.context_pos[seq_id] + ): + slot = int(self.batch.n_tokens) + self._add_batch_token( + token=int(batch.token[index]), + pos=pos, + seq_id=seq_id, + logits=False, + ) + if previous_row is None: + self._set_batch_embedding_row(slot, self.pending_h[seq_id]) + else: + self._set_batch_embedding_row(slot, h_tgt_rows[previous_row]) + added_pos_by_seq[seq_id] = pos + previous_row_by_seq[seq_id] = index + target_rows_by_seq.setdefault(seq_id, []).append(index) + + if int(self.batch.n_tokens) > 0: + self._decode_batch() + for seq_id, pos in added_pos_by_seq.items(): + self.context_pos[seq_id] = max(self.context_pos[seq_id], pos + 1) + + def draft( + self, + input_ids: np.ndarray, + /, + *, + seq_id: int, + max_tokens: Optional[int], + ) -> np.ndarray: + if ( + self.num_pred_tokens <= 0 + or input_ids.size == 0 + or seq_id < 0 + or seq_id >= self.n_seq_max + or not self.ready[seq_id] + ): + return np.array([], dtype=np.intc) + n_predict = self.num_pred_tokens + if max_tokens is not None: + n_predict = min(n_predict, max_tokens) + if n_predict <= 0: + return np.array([], dtype=np.intc) + + n_past = int(input_ids.size) - 1 + if self.ready_pos[seq_id] != n_past: + return np.array([], dtype=np.intc) + first_pos = n_past + if first_pos < 0: + return np.array([], dtype=np.intc) + + token = int(input_ids[-1]) + drafted: List[int] = [] + if not self.is_mem_shared and self.context_pos[seq_id] > first_pos: + self.truncate(seq_id, first_pos) + if not self.is_mem_shared and self.context_pos[seq_id] < first_pos: + self.ready[seq_id] = False + return np.array([], dtype=np.intc) + + self._reset_sampler(seq_id) + self._clear_batch() + self._add_batch_token( + token=token, + pos=first_pos, + seq_id=seq_id, + logits=True, + ) + self._set_batch_embedding_row(0, self.pending_h[seq_id]) + if not self._try_decode_batch(): + if not self.is_mem_shared: + self.truncate(seq_id, first_pos) + return np.array([], dtype=np.intc) + if not self.is_mem_shared: + self.context_pos[seq_id] = first_pos + 1 + + while len(drafted) < n_predict: + sampled_token = self._sample_token(seq_id=seq_id) + if sampled_token is None: + break + token = sampled_token + drafted.append(token) + if len(drafted) >= n_predict: + break + h_row = llama_cpp_ext.llama_get_embeddings_nextn_ith(self.ctx, 0) + if not h_row: + break + self._clear_batch() + self._add_batch_token( + token=token, + pos=first_pos if self.is_mem_shared else first_pos + len(drafted), + seq_id=seq_id, + logits=True, + ) + self._set_batch_embedding_row(0, h_row) + if not self._try_decode_batch(): + break + if not self.is_mem_shared: + self.context_pos[seq_id] = first_pos + len(drafted) + 1 + + if not drafted: + if not self.is_mem_shared: + self.truncate(seq_id, n_past) + return np.array([], dtype=np.intc) + if not self.is_mem_shared: + self.truncate(seq_id, n_past) + return np.asarray(drafted, dtype=np.intc) + + def draft_many( + self, + requests: Sequence[Tuple[np.ndarray, int, Optional[int]]], + /, + ) -> List[np.ndarray]: + results = [np.array([], dtype=np.intc) for _ in requests] + active: List["MTPDraftProvider.DraftManyState"] = [] + for result_index, (input_ids, seq_id, max_tokens) in enumerate(requests): + if ( + self.num_pred_tokens <= 0 + or input_ids.size == 0 + or seq_id < 0 + or seq_id >= self.n_seq_max + or not self.ready[seq_id] + ): + continue + n_predict = self.num_pred_tokens + if max_tokens is not None: + n_predict = min(n_predict, max_tokens) + if n_predict <= 0: + continue + if len(active) >= self.n_batch: + break + + n_past = int(input_ids.size) - 1 + if self.ready_pos[seq_id] != n_past: + continue + first_pos = n_past + if first_pos < 0: + continue + if not self.is_mem_shared and self.context_pos[seq_id] > first_pos: + self.truncate(seq_id, first_pos) + if not self.is_mem_shared and self.context_pos[seq_id] < first_pos: + self.ready[seq_id] = False + continue + self._reset_sampler(seq_id) + active.append( + self.DraftManyState( + result_index=result_index, + seq_id=seq_id, + first_pos=first_pos, + keep_len=n_past, + n_predict=n_predict, + token=int(input_ids[-1]), + drafted=[], + embedding=self.pending_h[seq_id], + cache_key=( + tuple(int(token) for token in input_ids.tolist()), + n_predict, + ), + ) + ) + + if not active: + return results + + touched = list(active) + try: + if all(state.n_predict == 1 for state in active): + grouped: Dict[ + Tuple[Tuple[int, ...], int], + List["MTPDraftProvider.DraftManyState"], + ] = {} + for state in active: + grouped.setdefault(state.cache_key, []).append(state) + + representatives = [states[0] for states in grouped.values()] + self._clear_batch() + for row, state in enumerate(representatives): + self._add_batch_token( + token=state.token, + pos=state.first_pos, + seq_id=state.seq_id, + logits=True, + ) + self._set_batch_embedding_row(row, state.embedding) + + if self._try_decode_batch(): + sampled_tokens = [ + self._sample_token(row, seq_id=state.seq_id) + for row, state in enumerate(representatives) + ] + for representative, sampled_token in zip( + representatives, + sampled_tokens, + ): + if sampled_token is None: + continue + if not self.is_mem_shared: + self.context_pos[representative.seq_id] = max( + self.context_pos[representative.seq_id], + representative.first_pos + 1, + ) + for state in grouped[representative.cache_key]: + state.drafted.append(sampled_token) + active = [] + + while active: + self._clear_batch() + for row, state in enumerate(active): + self._add_batch_token( + token=state.token, + pos=( + state.first_pos + if self.is_mem_shared + else state.first_pos + len(state.drafted) + ), + seq_id=state.seq_id, + logits=True, + ) + self._set_batch_embedding_row(row, state.embedding) + + if not self._try_decode_batch(): + break + + next_active: List["MTPDraftProvider.DraftManyState"] = [] + sampled_tokens = [ + self._sample_token(row, seq_id=state.seq_id) + for row, state in enumerate(active) + ] + for row, (state, sampled_token) in enumerate(zip(active, sampled_tokens)): + decoded_pos = ( + state.first_pos + if self.is_mem_shared + else state.first_pos + len(state.drafted) + ) + if not self.is_mem_shared: + self.context_pos[state.seq_id] = max( + self.context_pos[state.seq_id], + decoded_pos + 1, + ) + if sampled_token is None: + continue + h_row_ptr = llama_cpp_ext.llama_get_embeddings_nextn_ith( + self.ctx, row + ) + state.drafted.append(sampled_token) + if len(state.drafted) >= state.n_predict: + continue + if not h_row_ptr: + continue + state.token = sampled_token + state.embedding = np.ctypeslib.as_array( + h_row_ptr, + shape=(self.n_embd,), + ).copy() + next_active.append(state) + active = next_active + finally: + if not self.is_mem_shared: + for state in touched: + self.truncate(state.seq_id, state.keep_len) + + for state in touched: + if state.drafted: + results[state.result_index] = np.asarray( + state.drafted, + dtype=np.intc, + ) + return results + + @dataclass(frozen=True) + class SampledContextRow: + seq_id: int + draft_pos: int + token: int + source_row: Optional[int] + + @dataclass(frozen=True) + class SampledPendingRow: + update_index: int + seq_id: int + draft_pos: int + token: int + source_row: int + + @dataclass(frozen=True) + class SampledOutput: + update_index: int + seq_id: int + output_index: int + keep_len: int + ready_pos: int + + @dataclass(frozen=True) + class SampledBatchUpdate: + seq_id: int + start_pos: int + tokens: List[int] + row_indices: List[int] + target_count: int + sample_index: int + pending_token: Optional[int] + max_tokens: Optional[int] + + @dataclass + class SampledBatchPlan: + context_rows: List["MTPDraftProvider.SampledContextRow"] + pending_rows: List["MTPDraftProvider.SampledPendingRow"] + sample_pending_index_by_update: Dict[int, int] + + @dataclass + class SampledDraftState: + update_index: int + seq_id: int + keep_len: int + pos: int + token: int + drafted: List[int] + n_predict: int + embedding: np.ndarray + + def _build_sampled_batch_plan( + self, + updates: Sequence["MTPDraftProvider.SampledBatchUpdate"], + /, + ) -> "MTPDraftProvider.SampledBatchPlan": + context_rows: List["MTPDraftProvider.SampledContextRow"] = [] + pending_rows: List["MTPDraftProvider.SampledPendingRow"] = [] + sample_pending_index_by_update: Dict[int, int] = {} + + for update_index, update in enumerate(updates): + seq_id = update.seq_id + if seq_id < 0 or seq_id >= self.n_seq_max: + continue + start_pos = update.start_pos + tokens = update.tokens + row_indices = update.row_indices + target_count = update.target_count + sample_index = update.sample_index + pending_token = update.pending_token + if ( + pending_token is None + or start_pos < 0 + or target_count <= 0 + or target_count > len(tokens) + or target_count > len(row_indices) + or sample_index < 0 + or sample_index >= target_count + ): + continue + + for target_index in range(sample_index + 1): + mtp_pos = start_pos + target_index + source_row = ( + None + if target_index == 0 + else row_indices[target_index - 1] + ) + context_rows.append( + self.SampledContextRow( + seq_id=seq_id, + draft_pos=mtp_pos, + token=tokens[target_index], + source_row=source_row, + ) + ) + + actual_pos = start_pos + sample_index + 1 + pending_rows.append( + self.SampledPendingRow( + update_index=update_index, + seq_id=seq_id, + draft_pos=actual_pos, + token=pending_token, + source_row=row_indices[sample_index], + ) + ) + sample_pending_index_by_update[update_index] = len(pending_rows) - 1 + + return self.SampledBatchPlan( + context_rows=context_rows, + pending_rows=pending_rows, + sample_pending_index_by_update=sample_pending_index_by_update, + ) + + def _decode_sampled_context_rows( + self, + context_rows: Sequence["MTPDraftProvider.SampledContextRow"], + h_tgt_rows: np.ndarray, + /, + ) -> None: + self._clear_batch() + decoded_context_rows: List[Tuple[int, int]] = [] + for row in context_rows: + if self.is_mem_shared: + continue + if row.draft_pos < self.context_pos[row.seq_id]: + continue + if row.source_row is None: + embedding = self.pending_h[row.seq_id] + else: + embedding = h_tgt_rows[row.source_row] + self._add_batch_token( + token=row.token, + pos=row.draft_pos, + seq_id=row.seq_id, + logits=False, + ) + self._set_batch_embedding_row(int(self.batch.n_tokens) - 1, embedding) + decoded_context_rows.append((row.seq_id, row.draft_pos)) + + if int(self.batch.n_tokens) > 0: + self._decode_batch() + for seq_id, draft_pos in decoded_context_rows: + self.context_pos[seq_id] = max( + self.context_pos[seq_id], + draft_pos + 1, + ) + + def _decode_sampled_pending_rows( + self, + plan: "MTPDraftProvider.SampledBatchPlan", + h_tgt_rows: np.ndarray, + /, + ) -> List["MTPDraftProvider.SampledOutput"]: + sampled_outputs: List["MTPDraftProvider.SampledOutput"] = [] + pending_rows = plan.pending_rows + + self._clear_batch() + for pending_index, row in enumerate(pending_rows): + if ( + not self.is_mem_shared + and row.draft_pos < self.context_pos[row.seq_id] + ): + continue + is_sample_pending = ( + pending_index + == plan.sample_pending_index_by_update.get(row.update_index) + ) + slot = int(self.batch.n_tokens) + self._add_batch_token( + token=row.token, + pos=row.draft_pos, + seq_id=row.seq_id, + logits=is_sample_pending, + ) + self._set_batch_embedding_row(slot, h_tgt_rows[row.source_row]) + if is_sample_pending: + sampled_outputs.append( + self.SampledOutput( + update_index=row.update_index, + seq_id=row.seq_id, + output_index=slot, + keep_len=row.draft_pos, + ready_pos=row.draft_pos, + ) + ) + + self._decode_batch() + if not self.is_mem_shared: + for row in pending_rows: + self.context_pos[row.seq_id] = max( + self.context_pos[row.seq_id], + row.draft_pos + 1, + ) + + return sampled_outputs + + def _start_sampled_draft_states( + self, + updates: Sequence["MTPDraftProvider.SampledBatchUpdate"], + sampled_outputs: Sequence["MTPDraftProvider.SampledOutput"], + results: List[np.ndarray], + /, + ) -> List["MTPDraftProvider.SampledDraftState"]: + active: List["MTPDraftProvider.SampledDraftState"] = [] + for output in sampled_outputs: + self._reset_sampler(output.seq_id) + sampled_token = self._sample_token( + output.output_index, + seq_id=output.seq_id, + ) + if sampled_token is None: + continue + update = updates[output.update_index] + seq_id = update.seq_id + self.ready[seq_id] = True + n_predict = self.num_pred_tokens + max_tokens = update.max_tokens + if max_tokens is not None: + n_predict = min(n_predict, max_tokens) + if n_predict <= 0: + continue + if n_predict > 1: + h_row = llama_cpp_ext.llama_get_embeddings_nextn_ith( + self.ctx, output.output_index + ) + if h_row: + active.append( + self.SampledDraftState( + update_index=output.update_index, + seq_id=seq_id, + keep_len=output.keep_len, + pos=( + output.keep_len + if self.is_mem_shared + else output.keep_len + 1 + ), + token=sampled_token, + drafted=[sampled_token], + n_predict=n_predict, + embedding=np.ctypeslib.as_array( + h_row, + shape=(self.n_embd,), + ).copy(), + ) + ) + results[output.update_index] = np.asarray([sampled_token], dtype=np.intc) + + return active + + def _extend_sampled_draft_states( + self, + active: List["MTPDraftProvider.SampledDraftState"], + results: List[np.ndarray], + cleanup_keep_len_by_seq: Dict[int, int], + /, + ) -> None: + touched = list(active) + try: + while active: + self._clear_batch() + for batch_row, state in enumerate(active): + self._add_batch_token( + token=state.token, + pos=state.pos, + seq_id=state.seq_id, + logits=True, + ) + self._set_batch_embedding_row(batch_row, state.embedding) + + if not self._try_decode_batch(): + break + + next_active: List["MTPDraftProvider.SampledDraftState"] = [] + sampled_tokens = [ + self._sample_token(batch_row, seq_id=state.seq_id) + for batch_row, state in enumerate(active) + ] + for batch_row, (state, sampled_token) in enumerate( + zip(active, sampled_tokens) + ): + if not self.is_mem_shared: + self.context_pos[state.seq_id] = max( + self.context_pos[state.seq_id], + state.pos + 1, + ) + if sampled_token is None: + continue + state.drafted.append(sampled_token) + if len(state.drafted) >= state.n_predict: + continue + h_row = llama_cpp_ext.llama_get_embeddings_nextn_ith( + self.ctx, batch_row + ) + if not h_row: + continue + state.token = sampled_token + state.embedding = np.ctypeslib.as_array( + h_row, + shape=(self.n_embd,), + ).copy() + if not self.is_mem_shared: + state.pos += 1 + next_active.append(state) + active = next_active + finally: + for state in touched: + cleanup_keep_len_by_seq[state.seq_id] = state.keep_len + if not self.is_mem_shared: + for seq_id, keep_len in cleanup_keep_len_by_seq.items(): + self._truncate_memory(seq_id, keep_len) + + for state in touched: + if state.drafted: + results[state.update_index] = np.asarray( + state.drafted, + dtype=np.intc, + ) + + def process_sampled_batch( + self, + updates: Sequence["MTPDraftProvider.SampledBatchUpdate"], + /, + ) -> List[np.ndarray]: + results = [np.array([], dtype=np.intc) for _ in updates] + if self.num_pred_tokens <= 0 or not updates: + return results + h_tgt = llama_cpp_ext.llama_get_embeddings_nextn(self.target_ctx) + if not h_tgt: + raise RuntimeError("missing target nextn embeddings for MTP") + n_target_rows = max( + ( + max(update.row_indices) + 1 + for update in updates + if update.row_indices + ), + default=0, + ) + if n_target_rows <= 0: + return results + h_tgt_rows = np.ctypeslib.as_array(h_tgt, shape=(n_target_rows, self.n_embd)) + + plan = self._build_sampled_batch_plan(updates) + if not plan.context_rows and not plan.pending_rows: + return results + if ( + len(plan.context_rows) > self.n_batch + or len(plan.pending_rows) > self.n_batch + ): + raise RuntimeError("MTP draft batch capacity exceeded") + + self._decode_sampled_context_rows(plan.context_rows, h_tgt_rows) + sampled_outputs = self._decode_sampled_pending_rows(plan, h_tgt_rows) + cleanup_keep_len_by_seq: Dict[int, int] = {} + for output in sampled_outputs: + update = updates[output.update_index] + sample_index = update.sample_index + sample_source_row = update.row_indices[sample_index] + self.pending_h[output.seq_id] = h_tgt_rows[sample_source_row] + self.ready[output.seq_id] = True + self.ready_pos[output.seq_id] = output.ready_pos + cleanup_keep_len_by_seq[output.seq_id] = output.keep_len + + if not self.is_mem_shared: + for seq_id, keep_len in cleanup_keep_len_by_seq.items(): + self._truncate_memory(seq_id, keep_len) + + if sampled_outputs: + active = self._start_sampled_draft_states( + updates, + sampled_outputs, + results, + ) + self._extend_sampled_draft_states( + active, + results, + cleanup_keep_len_by_seq, + ) + + return results + + def accept(self, seq_id: int, accepted_draft_tokens: int) -> None: + if seq_id < 0 or seq_id >= self.n_seq_max: + return + n_rows = self.verify_h_rows[seq_id] + if n_rows <= 0: + return + row = min(max(accepted_draft_tokens, 0), n_rows - 1) + self.pending_h[seq_id] = self.verify_h[seq_id][row] + self.ready[seq_id] = True + self.ready_pos[seq_id] = self.verify_h_pos[seq_id][row] + 1 + + def _truncate_memory(self, seq_id: int, keep_len: int) -> None: + if seq_id < 0 or seq_id >= self.n_seq_max: + return + if self.is_mem_shared: + self.context_pos[seq_id] = min(self.context_pos[seq_id], keep_len) + return + if not llama_cpp.llama_memory_seq_rm( + self.mem, + seq_id, + keep_len, + -1, + ): + raise RuntimeError( + f"failed to truncate MTP draft sequence {seq_id} at position {keep_len}" + ) + self.context_pos[seq_id] = min(self.context_pos[seq_id], keep_len) + + def truncate(self, seq_id: int, keep_len: int) -> None: + if seq_id < 0 or seq_id >= self.n_seq_max: + return + self._truncate_memory(seq_id, keep_len) + if keep_len <= 0: + self.pending_h[seq_id].fill(0.0) + self.verify_h[seq_id] = np.empty((0, self.n_embd), dtype=np.float32) + self.verify_h_pos[seq_id] = [] + self.verify_h_rows[seq_id] = 0 + self.ready[seq_id] = False + self.ready_pos[seq_id] = 0 + self.context_pos[seq_id] = 0 + return + + if self.ready_pos[seq_id] != keep_len: + self.ready[seq_id] = False + + def copy_sequence( + self, + source_seq_id: int, + dest_seq_id: int, + p0: int, + p1: int, + ) -> None: + if ( + source_seq_id < 0 + or source_seq_id >= self.n_seq_max + or dest_seq_id < 0 + or dest_seq_id >= self.n_seq_max + ): + return + if not self.is_mem_shared: + llama_cpp.llama_memory_seq_cp( + self.mem, + source_seq_id, + dest_seq_id, + p0, + p1, + ) + source_ready_pos = self.ready_pos[source_seq_id] + copied_full_ready_state = p1 < 0 or p1 == source_ready_pos + if self.ready[source_seq_id] and copied_full_ready_state: + self.pending_h[dest_seq_id] = self.pending_h[source_seq_id] + self.verify_h[dest_seq_id] = self.verify_h[source_seq_id].copy() + self.verify_h_pos[dest_seq_id] = list(self.verify_h_pos[source_seq_id]) + self.verify_h_rows[dest_seq_id] = self.verify_h_rows[source_seq_id] + self.ready[dest_seq_id] = True + self.ready_pos[dest_seq_id] = source_ready_pos + self.context_pos[dest_seq_id] = min( + self.context_pos[source_seq_id], + source_ready_pos, + ) + return + + self.pending_h[dest_seq_id].fill(0.0) + self.verify_h[dest_seq_id] = np.empty((0, self.n_embd), dtype=np.float32) + self.verify_h_pos[dest_seq_id] = [] + self.verify_h_rows[dest_seq_id] = 0 + self.ready[dest_seq_id] = False + self.ready_pos[dest_seq_id] = 0 + self.context_pos[dest_seq_id] = 0 + + +class CompletionRequestCancelledError(RuntimeError): + pass + + +class CompletionRequestValidationError(ValueError): + pass + + +class CompletionResponseParsingError(RuntimeError): + pass + + +def omit_additional_properties(schema: Dict[str, Any]) -> None: + schema.pop("additionalProperties", None) + + +class CompletionChunkLogprobs(TypedDict): + tokens: List[str] + text_offset: List[int] + token_logprobs: List[Optional[float]] + top_logprobs: List[Optional[Dict[str, float]]] + + +class CompletionChunkChoice(TypedDict): + text: str + index: int + logprobs: Optional[CompletionChunkLogprobs] + finish_reason: Optional[str] + + +class CompletionChunk(TypedDict): + id: str + object: Literal["text_completion"] + created: int + model: str + choices: List[CompletionChunkChoice] + + +CompletionStream = Generator[CompletionChunk, None, OpenAICompletion] +CompletionPrompt = Union[str, List[int], List[str], List[List[int]]] +EmbeddingInput = Union[str, List[str], List[int], List[List[int]]] + + +class CreateCompletionRequest(BaseModel): + model_config = ConfigDict(extra="ignore") + + prompt: CompletionPrompt = "" + suffix: Optional[str] = None + max_tokens: Optional[int] = Field(default=16, ge=0) + temperature: float = 0.8 + top_p: float = Field(default=0.95, ge=0.0, le=1.0) + echo: bool = False + stop: Optional[Union[str, List[str]]] = None + stream: bool = False + logprobs: Optional[int] = Field(default=None, ge=0) + presence_penalty: Optional[float] = Field(default=0.0, ge=-2.0, le=2.0) + frequency_penalty: Optional[float] = Field(default=0.0, ge=-2.0, le=2.0) + logit_bias: Optional[Dict[str, float]] = None + seed: Optional[int] = None + model: Optional[str] = None + n: int = Field(default=1, ge=1) + best_of: Optional[int] = Field(default=None, ge=1) + user: Optional[str] = None + + @field_validator("logit_bias") + @classmethod + def validate_logit_bias(cls, value: Optional[Dict[str, float]]) -> Optional[Dict[str, float]]: + if value is None: + return None + result: Dict[str, float] = {} + for key, bias in value.items(): + int(key) + result[key] = float(bias) + return result + + @model_validator(mode="after") + def validate_after(self) -> "CreateCompletionRequest": + if self.best_of is None: + self.best_of = self.n + if self.best_of < self.n: + raise ValueError("best_of must be greater than or equal to n") + if self.stream and self.best_of > 1: + raise ValueError("best_of is not supported for streaming completions") + if len(self.normalized_prompt()) > 1 and self.stream: + raise ValueError("streaming does not support multiple prompts") + return self + + def normalized_prompt(self) -> List[Union[str, List[int]]]: + if isinstance(self.prompt, str): + return [self.prompt] + if all(isinstance(token, int) for token in self.prompt): + return [cast(List[int], self.prompt)] + if all(isinstance(prompt, str) for prompt in self.prompt): + return cast(List[Union[str, List[int]]], list(cast(List[str], self.prompt))) + if all( + isinstance(prompt, list) + and all(isinstance(token, int) for token in prompt) + for prompt in self.prompt + ): + return cast( + List[Union[str, List[int]]], + list(cast(List[List[int]], self.prompt)), + ) + raise ValueError("prompt must be a string, token ids, list of strings, or list of token-id lists") + + +class CreateEmbeddingRequest(BaseModel): + model_config = ConfigDict(extra="ignore") + + input: EmbeddingInput + model: str + encoding_format: Literal["float", "base64"] = "float" + dimensions: Optional[int] = Field(default=None, ge=1) + user: Optional[str] = None + + @staticmethod + def _validate_text_input(text: str) -> str: + if text == "": + raise ValueError("embedding input must not contain empty strings") + return text + + @staticmethod + def _validate_token_input(tokens: List[int]) -> List[int]: + if not tokens: + raise ValueError("embedding token input must not be empty") + if len(tokens) > 2048: + raise ValueError("embedding token input must not exceed 2048 tokens") + return tokens + + @model_validator(mode="after") + def validate_after(self) -> "CreateEmbeddingRequest": + self.normalized_input() + return self + + def normalized_input(self) -> List[Union[str, List[int]]]: + if isinstance(self.input, str): + return [self._validate_text_input(self.input)] + if all(isinstance(token, int) for token in self.input): + return [self._validate_token_input(cast(List[int], self.input))] + if all(isinstance(item, str) for item in self.input): + if len(self.input) > 2048: + raise ValueError("embedding input array must not exceed 2048 items") + return [ + self._validate_text_input(item) + for item in cast(List[str], self.input) + ] + if all( + isinstance(item, list) + and all(isinstance(token, int) for token in item) + for item in self.input + ): + if len(self.input) > 2048: + raise ValueError("embedding input array must not exceed 2048 items") + return [ + self._validate_token_input(item) + for item in cast(List[List[int]], self.input) + ] + raise ValueError( + "embedding input must be a string, list of strings, token ids, or list of token-id lists" + ) + + +class EmbeddingDataResponse(BaseModel): + object: Literal["embedding"] = "embedding" + embedding: Union[List[float], str] + index: int + + +class EmbeddingUsageResponse(BaseModel): + prompt_tokens: int + total_tokens: int + + +class CreateEmbeddingResponse(BaseModel): + object: Literal["list"] = "list" + data: List[EmbeddingDataResponse] + model: str + usage: EmbeddingUsageResponse + + @staticmethod + def encode_embedding( + embedding: Sequence[float], + encoding_format: Literal["float", "base64"], + dimensions: Optional[int], + ) -> Union[List[float], str]: + if dimensions is not None: + if dimensions > len(embedding): + raise CompletionRequestValidationError( + f"dimensions ({dimensions}) exceeds embedding size ({len(embedding)})" + ) + embedding = embedding[:dimensions] + if encoding_format == "float": + return [float(value) for value in embedding] + array = np.asarray(embedding, dtype=np.float32) + return base64.b64encode(array.tobytes()).decode("ascii") + + @classmethod + def from_embeddings( + cls, + *, + model: str, + embeddings: Sequence[Sequence[float]], + total_tokens: int, + encoding_format: Literal["float", "base64"], + dimensions: Optional[int], + ) -> "CreateEmbeddingResponse": + return cls( + data=[ + EmbeddingDataResponse( + embedding=cls.encode_embedding( + embedding, + encoding_format, + dimensions, + ), + index=index, + ) + for index, embedding in enumerate(embeddings) + ], + model=model, + usage=EmbeddingUsageResponse( + prompt_tokens=total_tokens, + total_tokens=total_tokens, + ), + ) + + +class ChatCompletionFunctionCall(BaseModel): + model_config = ConfigDict( + extra="allow", + json_schema_extra=omit_additional_properties, + ) + + name: str + arguments: Optional[str] = None + + +class ChatCompletionToolCall(BaseModel): + model_config = ConfigDict(extra="ignore") + + id: Optional[str] = None + type: Literal["function"] = "function" + function: ChatCompletionFunctionCall + + +class ChatCompletionRequestMessage(BaseModel): + model_config = ConfigDict( + extra="allow", + json_schema_extra=omit_additional_properties, + ) + + role: Literal["system", "developer", "user", "assistant", "tool", "function"] = Field( + default="user" + ) + content: Optional[Union[str, List[Dict[str, Any]]]] = Field(default="") + name: Optional[str] = Field(default=None) + tool_call_id: Optional[str] = Field(default=None) + function_call: Optional[ChatCompletionFunctionCall] = Field(default=None) + tool_calls: Optional[List[ChatCompletionToolCall]] = Field(default=None) + + +class ChatTemplateFunctionDefinition(TypedDict, total=False): + name: str + description: Optional[str] + parameters: Optional[Dict[str, Any]] + strict: Optional[bool] + content_type: Optional[str] + + +class ChatTemplateTool(TypedDict, total=False): + type: Literal["function"] + original_type: str + function: ChatTemplateFunctionDefinition + + +class ChatTemplateFunctionCall(TypedDict, total=False): + name: str + + +class ChatTemplateToolChoice(TypedDict, total=False): + type: Literal["function"] + function: ChatTemplateFunctionCall + + +class ChatTemplateResponseFormatJsonSchema(TypedDict, total=False): + name: Optional[str] + description: Optional[str] + schema: Dict[str, Any] + strict: Optional[bool] + + +class ChatTemplateResponseFormat(TypedDict, total=False): + type: Literal["text", "json_object", "json_schema"] + schema: Dict[str, Any] + json_schema: ChatTemplateResponseFormatJsonSchema + + +class ChatCompletionFunctionDefinition(BaseModel): + model_config = ConfigDict(extra="ignore") + + name: str + description: Optional[str] = None + parameters: Optional[Dict[str, Any]] = None + + def to_template_function(self) -> ChatTemplateFunctionDefinition: + function: ChatTemplateFunctionDefinition = {"name": self.name} + if self.description is not None: + function["description"] = self.description + if self.parameters is not None: + function["parameters"] = self.parameters + return function + + +class ChatCompletionToolFunction(BaseModel): + model_config = ConfigDict(extra="ignore") + + name: str + description: Optional[str] = None + parameters: Optional[Dict[str, Any]] = None + strict: Optional[bool] = None + + def to_template_function(self) -> ChatTemplateFunctionDefinition: + function: ChatTemplateFunctionDefinition = {"name": self.name} + if self.description is not None: + function["description"] = self.description + if self.parameters is not None: + function["parameters"] = self.parameters + if self.strict is not None: + function["strict"] = self.strict + return function + + +class ChatCompletionTool(BaseModel): + model_config = ConfigDict(extra="ignore") + + type: Literal["function"] + function: ChatCompletionToolFunction + + def to_template_tool(self) -> ChatTemplateTool: + return { + "type": self.type, + "function": self.function.to_template_function(), + } + + +class ChatCompletionFunctionCallOption(BaseModel): + model_config = ConfigDict(extra="ignore") + + name: str + + def to_template_function_call(self) -> ChatTemplateFunctionCall: + return {"name": self.name} + + +class ChatCompletionToolChoiceObject(BaseModel): + model_config = ConfigDict(extra="ignore") + + type: Literal["function"] + function: ChatCompletionFunctionCallOption + + def to_template_tool_choice(self) -> ChatTemplateToolChoice: + return { + "type": self.type, + "function": self.function.to_template_function_call(), + } + + +class ChatCompletionResponseFormatJsonSchema(BaseModel): + model_config = ConfigDict(extra="ignore", populate_by_name=True) + + name: Optional[str] = None + description: Optional[str] = None + schema_: Optional[Dict[str, Any]] = Field(default=None, alias="schema") + strict: Optional[bool] = None + + def to_template_json_schema(self) -> ChatTemplateResponseFormatJsonSchema: + json_schema: ChatTemplateResponseFormatJsonSchema = {} + if self.name is not None: + json_schema["name"] = self.name + if self.description is not None: + json_schema["description"] = self.description + if self.schema_ is not None: + json_schema["schema"] = self.schema_ + if self.strict is not None: + json_schema["strict"] = self.strict + return json_schema + + +class ChatCompletionResponseFormat(BaseModel): + model_config = ConfigDict(extra="ignore", populate_by_name=True) + + type: Literal["text", "json_object", "json_schema"] + schema_: Optional[Dict[str, Any]] = Field(default=None, alias="schema") + json_schema: Optional[ChatCompletionResponseFormatJsonSchema] = None + + def to_template_response_format(self) -> ChatTemplateResponseFormat: + response_format: ChatTemplateResponseFormat = {"type": self.type} + if self.schema_ is not None: + response_format["schema"] = self.schema_ + if self.json_schema is not None: + response_format["json_schema"] = self.json_schema.to_template_json_schema() + return response_format + + +class CreateChatCompletionRequest(BaseModel): + model_config = ConfigDict(extra="ignore") + + messages: List[ChatCompletionRequestMessage] = Field(default_factory=list) + max_tokens: Optional[int] = Field(default=None, ge=0) + temperature: float = 0.8 + top_p: float = Field(default=0.95, ge=0.0, le=1.0) + stop: Optional[Union[str, List[str]]] = None + stream: bool = False + logprobs: Optional[bool] = Field(default=False) + top_logprobs: Optional[int] = Field(default=None, ge=0) + presence_penalty: Optional[float] = Field(default=0.0, ge=-2.0, le=2.0) + frequency_penalty: Optional[float] = Field(default=0.0, ge=-2.0, le=2.0) + logit_bias: Optional[Dict[str, float]] = None + seed: Optional[int] = None + model: Optional[str] = None + n: int = Field(default=1, ge=1) + user: Optional[str] = None + functions: Optional[List[ChatCompletionFunctionDefinition]] = None + function_call: Optional[ + Union[Literal["none", "auto"], ChatCompletionFunctionCallOption] + ] = None + tools: Optional[List[ChatCompletionTool]] = None + tool_choice: Optional[ + Union[Literal["none", "auto", "required"], ChatCompletionToolChoiceObject] + ] = None + response_format: Optional[ChatCompletionResponseFormat] = None + reasoning_effort: Optional[ + Literal["low", "medium", "high", "minimal", "none"] + ] = None + + @field_validator("logit_bias") + @classmethod + def validate_logit_bias(cls, value: Optional[Dict[str, float]]) -> Optional[Dict[str, float]]: + if value is None: + return None + result: Dict[str, float] = {} + for key, bias in value.items(): + int(key) + result[key] = float(bias) + return result + + @model_validator(mode="after") + def validate_after(self) -> "CreateChatCompletionRequest": + if self.top_logprobs is not None and not self.logprobs: + raise ValueError("top_logprobs requires logprobs=true") + return self + + +class ResponsesFunctionTool(BaseModel): + model_config = ConfigDict(extra="ignore") + + type: Literal["function"] + name: str + description: Optional[str] = None + parameters: Optional[Dict[str, Any]] = None + strict: Optional[bool] = None + + def to_chat_template_tool(self) -> ChatTemplateTool: + return { + "type": "function", + "function": { + "name": self.name, + "description": self.description, + "parameters": self.parameters, + "strict": self.strict, + }, + } + + +class ResponsesCustomToolFormat(BaseModel): + model_config = ConfigDict(extra="ignore") + + type: Optional[str] = None + syntax: Optional[str] = None + definition: Optional[str] = None + + +class ResponsesCustomTool(BaseModel): + model_config = ConfigDict(extra="ignore") + + type: Literal["custom"] + name: str + description: Optional[str] = None + format: Optional[ResponsesCustomToolFormat] = None + strict: Optional[bool] = None + + def to_chat_template_tool(self) -> ChatTemplateTool: + tool_format = self.format + syntax = tool_format.syntax if tool_format is not None else None + definition = tool_format.definition if tool_format is not None else None + description = self.description or "" + text_tool_guidance = ( + "This is a text tool. When calling it, the " + "`function.arguments` field itself must be the raw input string. " + "Do not wrap the input in JSON and do not use an object such as " + "'{\"input\": \"...\"}' or '{\"patch\": \"...\"}'." + ) + if isinstance(syntax, str) and isinstance(definition, str): + if description: + description = ( + f"{description}\n\n{text_tool_guidance}\n\n" + f"{syntax}:\n{definition}" + ) + else: + description = f"{text_tool_guidance}\n\n{syntax}:\n{definition}" + elif description: + description = f"{description}\n\n{text_tool_guidance}" + else: + description = text_tool_guidance + return { + "type": "function", + "original_type": "custom", + "function": { + "name": self.name, + "description": description or None, + "parameters": { + "type": "object", + "properties": { + "input": { + "type": "string", + "description": ( + "Raw input text for this tool. " + "For apply_patch, put the full patch here." + ), + }, + }, + "required": ["input"], + "additionalProperties": False, + }, + "strict": self.strict, + "content_type": "text", + }, + } + + +class ResponsesWebSearchTool(BaseModel): + model_config = ConfigDict(extra="ignore") + + type: Literal["web_search"] + + +class ResponsesNamespaceTool(BaseModel): + model_config = ConfigDict(extra="ignore") + + type: Literal["namespace"] + + +class ResponsesImageGenerationTool(BaseModel): + model_config = ConfigDict(extra="ignore") + + type: Literal["image_generation"] + + +ResponsesToolDefinition = Union[ + ResponsesFunctionTool, + ResponsesCustomTool, + ResponsesWebSearchTool, + ResponsesNamespaceTool, + ResponsesImageGenerationTool, +] + + +class ResponsesToolChoiceObject(BaseModel): + model_config = ConfigDict(extra="ignore") + + type: Literal["function", "custom"] + name: str + + def to_chat_template_tool_choice(self) -> ChatTemplateToolChoice: + return { + "type": "function", + "function": { + "name": self.name, + }, + } + + +ResponsesToolChoice = Union[ + Literal["auto", "none", "required"], + ResponsesToolChoiceObject, +] + + +class ResponsesTextFormatJsonSchema(BaseModel): + model_config = ConfigDict(extra="ignore", populate_by_name=True) + + schema_: Dict[str, Any] = Field(alias="schema") + + def to_template_json_schema(self) -> ChatTemplateResponseFormatJsonSchema: + return {"schema": self.schema_} + + +class ResponsesTextFormat(BaseModel): + model_config = ConfigDict(extra="ignore", populate_by_name=True) + + type: Literal["text", "json_object", "json_schema"] + schema_: Optional[Dict[str, Any]] = Field(default=None, alias="schema") + json_schema: Optional[ResponsesTextFormatJsonSchema] = None + + def to_chat_response_format(self) -> ChatTemplateResponseFormat: + response_format: ChatTemplateResponseFormat = {"type": self.type} + if self.schema_ is not None: + response_format["schema"] = self.schema_ + if self.json_schema is not None: + response_format["json_schema"] = self.json_schema.to_template_json_schema() + return response_format + + +class ResponsesTextOptions(BaseModel): + model_config = ConfigDict(extra="ignore") + + format: Optional[ResponsesTextFormat] = None + + +class ResponsesReasoningOptions(BaseModel): + model_config = ConfigDict(extra="ignore") + + effort: Optional[Literal["low", "medium", "high"]] = None + + +class CreateResponseRequest(BaseModel): + model_config = ConfigDict(extra="ignore") + + input: Any = "" + instructions: Optional[str] = None + max_output_tokens: Optional[int] = Field(default=None, ge=0) + temperature: Optional[float] = Field(default=0.8, ge=0.0, le=2.0) + top_p: Optional[float] = Field(default=0.95, ge=0.0, le=1.0) + stream: bool = False + top_logprobs: Optional[int] = Field(default=None, ge=0) + model: Optional[str] = None + tools: Optional[List["ResponsesToolDefinition"]] = None + tool_choice: Optional["ResponsesToolChoice"] = None + parallel_tool_calls: bool = True + text: Optional["ResponsesTextOptions"] = None + reasoning: Optional["ResponsesReasoningOptions"] = None + metadata: Optional[Dict[str, str]] = None + user: Optional[str] = None + previous_response_id: Optional[str] = None + conversation: Optional[str] = None + store: Optional[bool] = None + truncation: Optional[Literal["auto", "disabled"]] = None + + +class ResponseCreateWebSocketRequest(BaseModel): + model_config = ConfigDict(extra="ignore") + + type: Literal["response.create"] + model: Optional[str] = None + input: Any = "" + instructions: Optional[str] = None + max_output_tokens: Optional[int] = Field(default=None, ge=0) + temperature: Optional[float] = Field(default=0.8, ge=0.0, le=2.0) + top_p: Optional[float] = Field(default=0.95, ge=0.0, le=1.0) + stream: bool = True + top_logprobs: Optional[int] = Field(default=None, ge=0) + tools: Optional[List["ResponsesToolDefinition"]] = None + tool_choice: Optional["ResponsesToolChoice"] = None + parallel_tool_calls: bool = True + text: Optional["ResponsesTextOptions"] = None + reasoning: Optional["ResponsesReasoningOptions"] = None + metadata: Optional[Dict[str, str]] = None + user: Optional[str] = None + previous_response_id: Optional[str] = None + conversation: Optional[str] = None + store: Optional[bool] = None + truncation: Optional[Literal["auto", "disabled"]] = None + generate: Optional[bool] = None + + def to_create_response_request(self) -> CreateResponseRequest: + return CreateResponseRequest.model_validate( + self.model_dump(mode="python", exclude={"type"}) + ) + + +class ModelCardResponse(BaseModel): + id: str + object: Literal["model"] = "model" + created: int + owned_by: str + + +class ModelListResponse(BaseModel): + object: Literal["list"] = "list" + data: List[ModelCardResponse] + + +class HealthzResponse(BaseModel): + status: Literal["ok"] = "ok" + + +@dataclass +class ResponsesWebSocketState: + input_items: List[Any] + output_items: List[Dict[str, Any]] + + +class ConfigFile(BaseModel): + class ServerOptions(BaseModel): + host: str = "127.0.0.1" + port: int = 8000 + + class DiskCacheOptions(BaseModel): + path: str + max_bytes: int = Field(ge=0) + min_tokens: int = Field(default=128, ge=1) + + class FromPretrainedOptions(BaseModel): + repo_id: str + filename: str + additional_files: Optional[List[str]] = None + local_dir: Optional[str] = None + local_dir_use_symlinks: Union[bool, Literal["auto"]] = "auto" + cache_dir: Optional[str] = None + + @staticmethod + def _pattern_has_glob(pattern: str) -> bool: + return any(char in pattern for char in "*?[]") + + @staticmethod + def _subfolder_for_repo_file(repo_file: str) -> Optional[str]: + subfolder = str(Path(repo_file).parent) + return None if subfolder == "." else subfolder + + def _cached_repo_file_list(self) -> List[str]: + from huggingface_hub import scan_cache_dir + + try: + cache_info = scan_cache_dir(self.cache_dir) + except Exception: + return [] + + cached_files: set[str] = set() + for repo in cache_info.repos: + if repo.repo_type != "model" or repo.repo_id != self.repo_id: + continue + for revision in repo.revisions: + for cached_file in revision.files: + cached_files.add(cached_file.file_name) + return sorted(cached_files) + + def _download_repo_file(self, repo_file: str) -> str: + from huggingface_hub import hf_hub_download + + filename = Path(repo_file).name + subfolder = self._subfolder_for_repo_file(repo_file) + + download_kwargs: Dict[str, Any] = dict( + repo_id=self.repo_id, + filename=filename, + subfolder=subfolder, + local_dir=self.local_dir, + cache_dir=self.cache_dir, + ) + hf_hub_download_signature = inspect.signature(hf_hub_download) + supports_local_dir_use_symlinks = ( + "local_dir_use_symlinks" in hf_hub_download_signature.parameters + or any( + parameter.kind is inspect.Parameter.VAR_KEYWORD + for parameter in hf_hub_download_signature.parameters.values() + ) + ) + if supports_local_dir_use_symlinks: + download_kwargs["local_dir_use_symlinks"] = self.local_dir_use_symlinks + + try: + resolved_path = hf_hub_download(**download_kwargs) + except Exception as exc: + try: + cached_path = cast( + str, + hf_hub_download( + repo_id=self.repo_id, + filename=filename, + subfolder=subfolder, + cache_dir=self.cache_dir, + local_files_only=True, + ), + ) + except Exception: + raise exc + + if self.local_dir is None: + return cached_path + + resolved_path = os.path.join(self.local_dir, repo_file) + resolved_parent = os.path.dirname(resolved_path) + if resolved_parent: + os.makedirs(resolved_parent, exist_ok=True) + shutil.copy2(cached_path, resolved_path) + + return cast(str, resolved_path) + + @staticmethod + def _match_single_file(pattern: str, file_list: Sequence[str]) -> str: + matching_files = [file for file in file_list if fnmatch.fnmatch(file, pattern)] + if len(matching_files) == 0: + raise ValueError( + f"No file found matching {pattern}\n\n" + f"Available Files:\n{json.dumps(list(file_list))}" + ) + if len(matching_files) > 1: + raise ValueError( + f"Multiple files found matching {pattern}\n\n" + f"Available Files:\n{json.dumps(list(file_list))}" + ) + return matching_files[0] + + def resolve_model_path(self) -> str: + try: + from huggingface_hub import HfFileSystem + from huggingface_hub.utils import validate_repo_id + except ImportError as exc: + try: + from huggingface_hub import HfFileSystem + from huggingface_hub.utils._validators import validate_repo_id + except ImportError: + raise ImportError( + "model.from_pretrained requires the huggingface-hub package. " + "You can install it with `pip install huggingface-hub`." + ) from exc + + validate_repo_id(self.repo_id) + requested_patterns = [self.filename, *(self.additional_files or [])] + if not any(self._pattern_has_glob(pattern) for pattern in requested_patterns): + model_path = self._download_repo_file(self.filename) + if self.additional_files: + for additional_file_name in self.additional_files: + self._download_repo_file(additional_file_name) + return model_path + + try: + hffs = HfFileSystem() + files = [ + file["name"] if isinstance(file, dict) else file + for file in hffs.ls(self.repo_id, recursive=True) + ] + file_list = [str(Path(file).relative_to(self.repo_id)) for file in files] + except Exception as exc: + file_list = self._cached_repo_file_list() + if not file_list: + raise exc + + matching_file = self._match_single_file(self.filename, file_list) + model_path = self._download_repo_file(matching_file) + + if self.additional_files: + for additional_file_name in self.additional_files: + matching_additional_file = self._match_single_file( + additional_file_name, file_list + ) + self._download_repo_file(matching_additional_file) + + return model_path + + class LoraOptions(BaseModel): + path: Optional[str] = None + from_pretrained: Optional["ConfigFile.FromPretrainedOptions"] = None + scale: float = 1.0 + + @model_validator(mode="after") + def validate_source(self) -> "ConfigFile.LoraOptions": + if (self.path is None) == (self.from_pretrained is None): + raise ValueError("exactly one of lora.path or lora.from_pretrained is required") + return self + + def resolve_path(self) -> str: + if self.from_pretrained is not None: + return self.from_pretrained.resolve_model_path() + assert self.path is not None + return self.path + + class MTMDEmbeddingCacheOptions(BaseModel): + path: str + max_bytes: int = Field(ge=0) + + class MTMDOptions(BaseModel): + mmproj_path: Optional[str] = None + mmproj_from_pretrained: Optional["ConfigFile.FromPretrainedOptions"] = None + embedding_cache: Optional["ConfigFile.MTMDEmbeddingCacheOptions"] = None + allowed_media_domains: Optional[List[str]] = None + allowed_local_media_path: Optional[str] = None + image_max_bytes: int = Field(default=20 * 1024 * 1024, ge=1) + audio_max_bytes: int = Field(default=100 * 1024 * 1024, ge=1) + video_max_bytes: int = Field(default=512 * 1024 * 1024, ge=1) + image_timeout_seconds: float = Field(default=10.0, gt=0.0) + + @model_validator(mode="after") + def validate_source(self) -> "ConfigFile.MTMDOptions": + if (self.mmproj_path is None) == (self.mmproj_from_pretrained is None): + raise ValueError( + "exactly one of model.mtmd.mmproj_path or " + "model.mtmd.mmproj_from_pretrained is required" + ) + return self + + def resolve_mmproj_path(self) -> str: + if self.mmproj_from_pretrained is not None: + return self.mmproj_from_pretrained.resolve_model_path() + assert self.mmproj_path is not None + return self.mmproj_path + + class ModelOptions(BaseModel): + path: Optional[str] = None + alias: Optional[str] = None + chat_template: Optional[str] = None + from_pretrained: Optional["ConfigFile.FromPretrainedOptions"] = None + loras: List["ConfigFile.LoraOptions"] = Field(default_factory=list) + mtmd: Optional["ConfigFile.MTMDOptions"] = None + n_gpu_layers: Optional[int] = None + split_mode: Optional[int] = None + main_gpu: Optional[int] = None + tensor_split: Optional[List[float]] = None + vocab_only: Optional[bool] = None + use_mmap: Optional[bool] = None + use_mlock: Optional[bool] = None + kv_overrides: Optional[Dict[str, Union[bool, int, float, str]]] = None + n_ctx: Optional[int] = None + n_batch: Optional[int] = None + n_ubatch: Optional[int] = None + n_seq_max: Optional[int] = None + threads: Optional[int] = Field( + default_factory=lambda: max(multiprocessing.cpu_count(), 1) + ) + threads_batch: Optional[int] = Field( + default_factory=lambda: max(multiprocessing.cpu_count(), 1) + ) + rope_scaling_type: Optional[int] = None + pooling_type: Optional[int] = None + attention_type: Optional[int] = None + embedding: Optional[bool] = None + rope_freq_base: Optional[float] = None + rope_freq_scale: Optional[float] = None + yarn_ext_factor: Optional[float] = None + yarn_attn_factor: Optional[float] = None + yarn_beta_fast: Optional[float] = None + yarn_beta_slow: Optional[float] = None + yarn_orig_ctx: Optional[int] = None + offload_kqv: Optional[bool] = None + flash_attn: Optional[bool] = None + op_offload: Optional[bool] = None + swa_full: Optional[bool] = None + no_perf: Optional[bool] = None + type_k: Optional[int] = None + type_v: Optional[int] = None + max_seq_len: Optional[int] = None + max_output_tokens: Optional[int] = Field(default=None, ge=0) + kv_unified: bool = True + draft_model: Optional[Literal["prompt-lookup-decoding", "draft-mtp"]] = None + draft_model_path: Optional[str] = None + draft_model_from_pretrained: Optional[ + "ConfigFile.FromPretrainedOptions" + ] = None + draft_model_num_pred_tokens: int = 16 + draft_model_max_ngram_size: int = 2 + draft_model_top_k: int = Field(default=1, ge=1) + draft_model_p_min: float = Field(default=0.0, ge=0.0, le=1.0) + draft_model_max_batch_size: Optional[int] = Field(default=None, ge=1) + draft_model_threads: Optional[int] = Field(default=None, gt=0) + draft_model_threads_batch: Optional[int] = Field(default=None, gt=0) + response_schema: Optional[Dict[str, Any]] = None + store_logits: bool = True + + @model_validator(mode="after") + def validate_source(self) -> "ConfigFile.ModelOptions": + if (self.path is None) == (self.from_pretrained is None): + raise ValueError("exactly one of model.path or model.from_pretrained is required") + if ( + self.draft_model_path is not None + and self.draft_model_from_pretrained is not None + ): + raise ValueError( + "model.draft_model_path and model.draft_model_from_pretrained " + "are mutually exclusive" + ) + return self + + def resolve_draft_model_path(self) -> Optional[str]: + if self.draft_model_from_pretrained is not None: + return self.draft_model_from_pretrained.resolve_model_path() + return self.draft_model_path + + @field_validator("chat_template", mode="before") + @classmethod + def normalize_chat_template(cls, value: Any) -> Any: + if isinstance(value, list): + if not all(isinstance(item, str) for item in value): + raise TypeError("model.chat_template list entries must be strings") + return "".join(value) + return value + + def resolve_model_path(self) -> str: + if self.from_pretrained is not None: + return self.from_pretrained.resolve_model_path() + assert self.path is not None + return self.path + + server: "ConfigFile.ServerOptions" = Field(default_factory=lambda: ConfigFile.ServerOptions()) + model: "ConfigFile.ModelOptions" + disk_cache: Optional["ConfigFile.DiskCacheOptions"] = None + + @classmethod + def load(cls, path: str) -> "ConfigFile": + return cls.model_validate_json(Path(path).read_text()) + + +ConfigFile.model_rebuild() + + +@dataclass(frozen=True) +class MediaInput: + kind: Literal["image", "audio", "video"] + url: Optional[str] = None + data: Optional[str] = None + format: Optional[str] = None + + +class Jinja2ChatFormatter: + def __init__(self, template: str, *, bos_token: str, eos_token: str) -> None: + self._eos_token = eos_token + self._bos_token = bos_token + self._template_text = template + environment = ImmutableSandboxedEnvironment( + loader=jinja2.BaseLoader(), + trim_blocks=True, + lstrip_blocks=True, + ) + environment.filters["from_json"] = self._from_json + self._template = environment.from_string(template) + + @staticmethod + def _from_json(value: Any) -> Any: + if isinstance(value, str): + return json.loads(value) + return value + + @staticmethod + def media_inputs_from_messages( + messages: Sequence[ChatCompletionRequestMessage], + ) -> List[MediaInput]: + media_inputs: List[MediaInput] = [] + for message in messages: + content = message.content + if not isinstance(content, list): + continue + for part in content: + if not isinstance(part, dict): + continue + part_type = part.get("type") + if part_type in {"image_url", "input_image", "image"}: + image_url = part.get("image_url") or part.get("url") + if isinstance(image_url, str): + media_inputs.append(MediaInput(kind="image", url=image_url)) + elif isinstance(image_url, dict) and isinstance(image_url.get("url"), str): + media_inputs.append(MediaInput(kind="image", url=cast(str, image_url["url"]))) + else: + raise ValueError("image_url content part requires a URL string") + continue + if part_type == "audio_url": + audio_url = part.get("audio_url") + if isinstance(audio_url, str): + media_inputs.append(MediaInput(kind="audio", url=audio_url)) + elif isinstance(audio_url, dict) and isinstance(audio_url.get("url"), str): + media_inputs.append(MediaInput(kind="audio", url=cast(str, audio_url["url"]))) + else: + raise ValueError("audio_url content part requires a URL string") + continue + if part_type == "input_audio": + input_audio = part.get("input_audio") + if isinstance(input_audio, dict): + data = input_audio.get("data") + audio_format = input_audio.get("format") + else: + data = part.get("data") + audio_format = part.get("format") + if not isinstance(data, str): + raise ValueError("input_audio content part requires base64 data") + if audio_format is not None and not isinstance(audio_format, str): + raise ValueError("input_audio format must be a string") + media_inputs.append( + MediaInput( + kind="audio", + data=data, + format=cast(Optional[str], audio_format), + ) + ) + continue + if part_type in {"video_url", "input_video", "video"}: + input_video = part.get("input_video") + if isinstance(input_video, dict): + data = input_video.get("data") + video_url = input_video.get("video_url") or input_video.get("url") + video_format = input_video.get("format") + else: + data = part.get("data") + video_url = part.get("video_url") or part.get("url") + video_format = part.get("format") + if isinstance(video_url, dict): + video_url = video_url.get("url") + if isinstance(data, str): + if video_format is not None and not isinstance(video_format, str): + raise ValueError("input_video format must be a string") + media_inputs.append( + MediaInput( + kind="video", + data=data, + format=cast(Optional[str], video_format), + ) + ) + elif isinstance(video_url, str): + media_inputs.append(MediaInput(kind="video", url=video_url)) + else: + raise ValueError("video content part requires base64 data or a URL string") + continue + return media_inputs + + @staticmethod + def _literal_content_for_template( + text: str, + media_marker: Optional[str], + ) -> str: + if media_marker is not None and media_marker in text: + raise ValueError("message content contains reserved multimodal marker text") + return text + + @staticmethod + def _content_part_for_template( + part: Any, + media_marker: Optional[str], + ) -> str: + if isinstance(part, str): + return Jinja2ChatFormatter._literal_content_for_template(part, media_marker) + if not isinstance(part, dict): + raise ValueError("content parts must be strings or objects") + part_type = part.get("type") + if part_type in { + "image_url", + "input_image", + "image", + "audio_url", + "input_audio", + "video_url", + "input_video", + "video", + }: + if media_marker is None: + raise ValueError("multimodal content requires model.mtmd") + return media_marker + if part_type in {"text", "input_text"}: + text = part.get("text") + if not isinstance(text, str): + raise ValueError(f"content part {part_type!r} requires string text") + return Jinja2ChatFormatter._literal_content_for_template(text, media_marker) + raise ValueError(f"unsupported content part type: {part_type!r}") + + @staticmethod + def _messages_for_template( + messages: Sequence[ChatCompletionRequestMessage], + media_marker: Optional[str], + ) -> List[Dict[str, Any]]: + converted: List[Dict[str, Any]] = [] + for message in messages: + data = message.model_dump(exclude_none=True) + content = data.get("content") + if isinstance(content, list): + data["content"] = "".join( + Jinja2ChatFormatter._content_part_for_template( + part, + media_marker, + ) + for part in content + ) + elif isinstance(content, str): + data["content"] = Jinja2ChatFormatter._literal_content_for_template( + content, + media_marker, + ) + converted.append(data) + return converted + + def _render( + self, + *, + messages: List[ChatCompletionRequestMessage], + media_marker: Optional[str] = None, + functions: Optional[List[ChatTemplateFunctionDefinition]] = None, + function_call: Optional[Union[str, ChatTemplateFunctionCall]] = None, + tools: Optional[List[ChatTemplateTool]] = None, + tool_choice: Optional[Union[str, ChatTemplateToolChoice]] = None, + reasoning_effort: Optional[str] = None, + add_generation_prompt: bool, + strftime_now: Callable[[str], str], + ) -> str: + def raise_exception(message: str) -> None: + raise ValueError(message) + + return cast( + str, + self._template.render( + messages=self._messages_for_template(messages, media_marker), + eos_token=self._eos_token, + bos_token=self._bos_token, + raise_exception=raise_exception, + add_generation_prompt=add_generation_prompt, + functions=functions, + function_call=function_call, + tools=tools, + tool_choice=tool_choice, + reasoning_effort=reasoning_effort, + strftime_now=strftime_now, + ), + ) + + def format( + self, + *, + messages: List[ChatCompletionRequestMessage], + media_marker: Optional[str] = None, + functions: Optional[List[ChatTemplateFunctionDefinition]] = None, + function_call: Optional[Union[str, ChatTemplateFunctionCall]] = None, + tools: Optional[List[ChatTemplateTool]] = None, + tool_choice: Optional[Union[str, ChatTemplateToolChoice]] = None, + reasoning_effort: Optional[str] = None, + ) -> Tuple[str, str, List[str]]: + render_time = datetime.now() + + def strftime_now(format_string: str) -> str: + return render_time.strftime(format_string) + + chat_template = self._render( + messages=messages, + media_marker=media_marker, + functions=functions, + function_call=function_call, + tools=tools, + tool_choice=tool_choice, + reasoning_effort=reasoning_effort, + add_generation_prompt=False, + strftime_now=strftime_now, + ) + prompt = self._render( + messages=messages, + media_marker=media_marker, + functions=functions, + function_call=function_call, + tools=tools, + tool_choice=tool_choice, + reasoning_effort=reasoning_effort, + add_generation_prompt=True, + strftime_now=strftime_now, + ) + generation_prompt = prompt[len(chat_template) :] if prompt.startswith(chat_template) else "" + stop = [self._eos_token] if self._eos_token else [] + return prompt, generation_prompt, stop + + +@dataclass +class Token: + token: int + text_bytes: bytes + token_logprob: Optional[float] + top_logprobs: Optional[Dict[str, float]] + + @classmethod + def from_token( + cls, + *, + model: Model, + prev_tokens: Sequence[int], + prev_text_bytes: Optional[Union[bytes, bytearray]] = None, + token: int, + ) -> "Token": + return cls( + token=token, + text_bytes=( + model.token_bytes_with_prev_bytes(prev_tokens, prev_text_bytes, token) + if prev_text_bytes is not None + else model.token_bytes_with_prev(prev_tokens, token) + ), + token_logprob=None, + top_logprobs=None, + ) + + @classmethod + def from_logits( + cls, + *, + model: Model, + formatter: OpenAIFormatter, + prev_tokens: Sequence[int], + prev_text_bytes: Optional[Union[bytes, bytearray]] = None, + token: int, + logits: np.ndarray, + logprobs_count: Optional[int], + need_token_logprob: bool = False, + ) -> "Token": + text_bytes = ( + model.token_bytes_with_prev_bytes(prev_tokens, prev_text_bytes, token) + if prev_text_bytes is not None + else model.token_bytes_with_prev(prev_tokens, token) + ) + if not model.store_logits: + return cls( + token=token, + text_bytes=text_bytes, + token_logprob=None, + top_logprobs=None, + ) + if logprobs_count is None and not need_token_logprob: + return cls( + token=token, + text_bytes=text_bytes, + token_logprob=None, + top_logprobs=None, + ) + logprobs = CompletionScheduler.logits_to_logprobs(logits) + token_logprob = float(logprobs[token]) + top_logprobs: Optional[Dict[str, float]] = None + if logprobs_count is not None: + top_count = min(max(logprobs_count, 0), model.n_vocab) + if top_count > 0: + top_indices = np.argpartition(logprobs, -top_count)[-top_count:] + top_indices = top_indices[np.argsort(logprobs[top_indices])[::-1]] + top_logprobs = { + formatter.decode_text(model.token_bytes(int(index))): float( + logprobs[int(index)] + ) + for index in top_indices + } + else: + top_logprobs = {} + top_logprobs[formatter.decode_text(text_bytes)] = token_logprob + return cls( + token=token, + text_bytes=text_bytes, + token_logprob=token_logprob, + top_logprobs=top_logprobs, + ) + + +@dataclass +class Completion: + request_id: str + index: int + seq_id: int + sampler: "Sampler" + prompt_tokens: List[int] + prompt_length: int + prompt_text: str + multimodal_prompt: bool + max_total_tokens: int + stop_sequences: List[bytes] + logprobs: Optional[int] + completion_tokens: List[int] = field(default_factory=list) + token_records: List[Token] = field(default_factory=list) + rendered_bytes: bytearray = field(default_factory=bytearray) + detokenized_prefix_bytes: bytearray = field(default_factory=bytearray) + pending_input_tokens: List[int] = field(default_factory=list) + draft_tokens: List[int] = field(default_factory=list) + pending_finish_reason: Optional[str] = None + returned_token_count: int = 0 + finished: bool = False + finish_reason: Optional[str] = None + score_sum: float = 0.0 + rank_by_score: bool = False + + @property + def total_tokens(self) -> int: + return self.prompt_length + len(self.completion_tokens) + + @property + def completion_token_count(self) -> int: + return len(self.completion_tokens) + + @property + def needs_token_logprob(self) -> bool: + return self.logprobs is not None or self.rank_by_score + + @property + def max_stop_sequence_length(self) -> int: + return max((len(stop) for stop in self.stop_sequences), default=0) + + +@dataclass +class PromptSegment: + @dataclass + class Media: + embeddings: np.ndarray + positions: np.ndarray + non_causal: bool = False + + kind: Literal["text", "image", "audio", "video"] + start_pos: int + n_pos: int + identity_tokens: List[int] + decode_start_pos: int + decode_n_pos: int + text_tokens: List[int] = field(default_factory=list) + media: Optional["PromptSegment.Media"] = None + + @property + def end_pos(self) -> int: + return self.start_pos + self.n_pos + + @property + def batch_rows(self) -> int: + if self.kind != "text": + if self.media is None: + return 0 + return int(self.media.embeddings.shape[0]) + return len(self.text_tokens) + + @property + def decoder_position_increments(self) -> List[int]: + if not self.identity_tokens: + return [] + if self.kind == "text": + return [1] * len(self.identity_tokens) + return [*([0] * (len(self.identity_tokens) - 1)), self.decode_n_pos] + + def rows_for_capacity(self, offset: int, capacity: int) -> int: + if capacity <= 0: + return 0 + return min(capacity, self.end_pos - self.start_pos - offset) + + def media_slice( + self, + row_offset: int, + row_count: int, + ) -> Tuple[np.ndarray, np.ndarray, List[int]]: + if self.media is None: + raise RuntimeError("media segment is missing embeddings or positions") + row_start = row_offset + row_end = row_offset + row_count + embeddings = self.media.embeddings[row_start:row_end] + if len(self.media.positions) == self.batch_rows: + positions = self.media.positions[row_start:row_end] + else: + positions = ( + self.media.positions.reshape(4, self.batch_rows)[:, row_start:row_end] + .reshape(-1) + ) + return ( + embeddings, + positions, + self.decoder_position_increments[row_start:row_end], + ) + + +@dataclass +class PromptPlan: + text: str + generation_prompt: str + text_tokens: List[int] + identity_tokens: List[int] + segments: List[PromptSegment] + text_token_index_by_pos: Dict[int, int] = field(default_factory=dict) + + @classmethod + def from_tokens( + cls, + text: str, + tokens: List[int], + *, + generation_prompt: str = "", + ) -> "PromptPlan": + segment = PromptSegment( + kind="text", + start_pos=0, + n_pos=len(tokens), + identity_tokens=list(tokens), + decode_start_pos=0, + decode_n_pos=len(tokens), + text_tokens=list(tokens), + ) + return cls( + text=text, + generation_prompt=generation_prompt, + text_tokens=list(tokens), + identity_tokens=list(tokens), + segments=[segment] if tokens else [], + text_token_index_by_pos={pos: pos for pos in range(len(tokens))}, + ) + + @property + def length(self) -> int: + return len(self.identity_tokens) + + @property + def eval_token_count(self) -> int: + return self.length + + def position_increments_up_to(self, pos: int) -> List[int]: + increments: List[int] = [] + for segment in self.segments: + if pos <= segment.start_pos: + break + take = min(pos, segment.end_pos) - segment.start_pos + if take <= 0: + continue + increments.extend(segment.decoder_position_increments[:take]) + if pos <= segment.end_pos: + break + return increments + + def is_boundary(self, pos: int) -> bool: + if pos <= 0 or pos >= self.length: + return True + return any(segment.start_pos == pos or segment.end_pos == pos for segment in self.segments) + + def is_reusable_boundary(self, pos: int) -> bool: + if pos <= 0 or pos >= self.length: + return True + if self.is_boundary(pos): + return True + return self.segment_at(pos).kind == "text" + + def clamp_to_reusable_boundary(self, pos: int) -> int: + if pos <= 0: + return 0 + if pos >= self.length: + return self.length + if self.is_reusable_boundary(pos): + return pos + segment = self.segment_at(pos) + return segment.start_pos + + def decoder_pos_up_to(self, pos: int) -> int: + if not self.is_reusable_boundary(pos): + raise RuntimeError("decoder position requested inside a non-reusable media segment") + if pos <= 0: + return 0 + if pos >= self.length: + if not self.segments: + return 0 + segment = self.segments[-1] + return segment.decode_start_pos + segment.decode_n_pos + segment = self.segment_at(pos) + if segment.kind == "text": + return segment.decode_start_pos + (pos - segment.start_pos) + if pos == segment.start_pos: + return segment.decode_start_pos + if pos == segment.end_pos: + return segment.decode_start_pos + segment.decode_n_pos + raise RuntimeError("decoder position requested inside a media segment") + + def segment_at(self, pos: int) -> PromptSegment: + for segment in self.segments: + if segment.start_pos <= pos < segment.end_pos: + return segment + raise RuntimeError(f"missing prompt segment at position {pos}") + + def has_text_token_at(self, pos: int) -> bool: + return pos in self.text_token_index_by_pos + + def text_token_at(self, pos: int) -> int: + return self.text_tokens[self.text_token_index_by_pos[pos]] + + def prev_text_tokens_at(self, pos: int) -> List[int]: + return self.text_tokens[: self.text_token_index_by_pos[pos]] + + +@dataclass(frozen=True) +class PreparedCompletionParts: + payload: CreateCompletionRequest + prompt_text: str + generation_prompt: str + prompt_plan: PromptPlan + grammar_text: Optional[str] + tool_name: Optional[str] + + +@dataclass(frozen=True) +class ResponsesChatRequestParts: + messages: List[ChatCompletionRequestMessage] + max_tokens: Optional[int] + temperature: float + top_p: float + stream: bool + logprobs: bool + top_logprobs: Optional[int] + model: Optional[str] + user: Optional[str] + tools: Optional[List[ChatTemplateTool]] + tool_choice: Optional[Union[Literal["auto", "none", "required"], ChatTemplateToolChoice]] + response_format: Optional[ChatTemplateResponseFormat] + reasoning_effort: Optional[str] + + +@dataclass +class SchedulerMetrics: + started_at: float = field(default_factory=time.time) + requests_submitted_total: int = 0 + requests_admitted_total: int = 0 + requests_completed_total: int = 0 + requests_cancelled_total: int = 0 + requests_failed_total: int = 0 + prompt_tokens_total: int = 0 + prompt_seconds_total: float = 0.0 + tokens_predicted_total: int = 0 + tokens_predicted_seconds_total: float = 0.0 + scheduler_step_seconds_total: float = 0.0 + process_batch_seconds_total: float = 0.0 + sample_seconds_total: float = 0.0 + draft_seconds_total: float = 0.0 + draft_process_seconds_total: float = 0.0 + draft_generate_seconds_total: float = 0.0 + draft_sampled_batch_seconds_total: float = 0.0 + draft_process_calls_total: int = 0 + draft_generate_calls_total: int = 0 + draft_sampled_batch_calls_total: int = 0 + draft_batches_verified_total: int = 0 + draft_target_tokens_verified_total: int = 0 + draft_target_tokens_wasted_total: int = 0 + draft_tokens_reused_as_pending_total: int = 0 + n_decode_total: int = 0 + scheduler_step_calls_total: int = 0 + process_batch_calls_total: int = 0 + sample_calls_total: int = 0 + n_tokens_max: int = 0 + n_busy_slots_total: int = 0 + checkpoint_hits_total: int = 0 + checkpoint_saves_total: int = 0 + checkpoint_evictions_total: int = 0 + sequence_cache_hits_total: int = 0 + sequence_cache_save_requests_total: int = 0 + sequence_cache_lookup_failures_total: int = 0 + sequence_cache_load_failures_total: int = 0 + sequence_cache_save_failures_total: int = 0 + sequence_cache_tokens_loaded_total: int = 0 + + def observe_decode( + self, + items: Sequence[Any], + elapsed_seconds: float, + ) -> None: + if not items: + return + total_tokens = sum( + item.batch_token_count + for item in items + ) + if total_tokens <= 0: + return + prompt_tokens = sum( + item.batch_token_count + for item in items + if getattr(item, "kind", None) == "prefill" + ) + generation_tokens = total_tokens - prompt_tokens + self.n_decode_total += 1 + self.n_busy_slots_total += len(items) + self.n_tokens_max = max(self.n_tokens_max, total_tokens) + self.prompt_tokens_total += prompt_tokens + if prompt_tokens > 0: + self.prompt_seconds_total += elapsed_seconds * (prompt_tokens / total_tokens) + if generation_tokens > 0: + self.tokens_predicted_seconds_total += elapsed_seconds * ( + generation_tokens / total_tokens + ) + + def observe_predicted_token(self) -> None: + self.tokens_predicted_total += 1 + + +class SequenceCache(Protocol): + """Optional external sequence-state cache; storage and eviction stay outside the scheduler.""" + + @dataclass(frozen=True) + class Match: + tokens: Tuple[int, ...] + has_prompt_logits: bool = False + + @property + def length(self) -> int: + return len(self.tokens) + + @dataclass + class Load: + tokens: List[int] + state_bytes: np.ndarray + prompt_logits: Optional[np.ndarray] = None + + def lookup(self, tokens: Sequence[int]) -> Optional["SequenceCache.Match"]: + ... + + def load( + self, + match: "SequenceCache.Match", + ) -> Optional["SequenceCache.Load"]: + ... + + def save( + self, + tokens: Sequence[int], + state_bytes: np.ndarray, + prompt_logits: Optional[np.ndarray], + ) -> None: + ... + + +@dataclass +class CompletionRequest: + payload: CreateCompletionRequest + prompt_text: str + prompt_tokens: List[int] + prompt_plan: PromptPlan + effective_max_len: int + internal_completion_count: int + prompt_visible_start: int + id: str = field(default_factory=lambda: f"cmpl-{uuid.uuid4().hex}") + created: int = field(default_factory=lambda: int(time.time())) + prompt_cursor: int = 0 + match_sequence_id: int = -1 + match_length: int = 0 + sequence_cache_match: Optional[SequenceCache.Match] = None + sequence_cache_match_length: int = 0 + prompt_logits: Optional[np.ndarray] = None + base_seq_id: Optional[int] = None + sibling_seq_ids: List[int] = field(default_factory=list) + completion_seq_ids: List[int] = field(default_factory=list) + completions: List[Completion] = field(default_factory=list) + admitted: bool = False + prompt_done: bool = False + prompt_checkpoint_saved: bool = False + cancelled: bool = False + prompt_records: List[Token] = field(default_factory=list) + grammar_text: Optional[str] = None + grammar_root: str = "root" + chat_tool_name: Optional[str] = None + on_stream_chunk: Optional[Callable[[CompletionChunk], None]] = None + on_done: Optional[Callable[[OpenAICompletion], None]] = None + on_error: Optional[Callable[[BaseException], None]] = None + + @classmethod + def from_prepared( + cls, + *, + payload: CreateCompletionRequest, + prompt_text: str, + prompt_plan: PromptPlan, + max_seq_len: int, + max_output_tokens: Optional[int], + prompt_visible_start: int, + prompt_records: Optional[List[Token]] = None, + grammar_text: Optional[str] = None, + chat_tool_name: Optional[str] = None, + on_stream_chunk: Optional[Callable[[CompletionChunk], None]] = None, + on_done: Optional[Callable[[OpenAICompletion], None]] = None, + on_error: Optional[Callable[[BaseException], None]] = None, + ) -> "CompletionRequest": + prompt_tokens = list(prompt_plan.identity_tokens) + prompt_eval_tokens = prompt_plan.eval_token_count + if prompt_eval_tokens > max_seq_len: + raise CompletionRequestValidationError("prompt exceeds context window") + ctx_limit = prompt_plan.length + (max_seq_len - prompt_eval_tokens) + if max_output_tokens is not None: + ctx_limit = min(ctx_limit, prompt_plan.length + max_output_tokens) + if payload.max_tokens is None: + effective_max_len = ctx_limit + else: + effective_max_len = min(ctx_limit, prompt_plan.length + payload.max_tokens) + if effective_max_len < prompt_plan.length: + raise CompletionRequestValidationError("prompt exceeds context window") + internal_completion_count = payload.best_of if payload.best_of is not None else payload.n + request = cls( + payload=payload, + prompt_text=prompt_text, + prompt_tokens=prompt_tokens, + prompt_plan=prompt_plan, + effective_max_len=effective_max_len, + internal_completion_count=internal_completion_count, + prompt_visible_start=prompt_visible_start, + prompt_records=list(prompt_records or []), + grammar_text=grammar_text, + chat_tool_name=chat_tool_name, + on_stream_chunk=on_stream_chunk, + on_done=on_done, + on_error=on_error, + ) + return request + + def selected_completions(self) -> List[Completion]: + completions = list(self.completions) + best_of = self.payload.best_of if self.payload.best_of is not None else self.payload.n + if best_of > self.payload.n: + completions.sort( + key=lambda completion: ( + completion.score_sum / max(1, completion.completion_token_count), + completion.score_sum, + ), + reverse=True, + ) + return sorted(completions[: self.payload.n], key=lambda completion: completion.index) + + def capture_prompt_logprobs( + self, + *, + model: Model, + formatter: OpenAIFormatter, + output_indices: Sequence[Optional[int]], + output_positions: Sequence[int], + output_count: int, + output_index_to_logits_index: Callable[[Optional[int], int], Optional[int]], + ) -> None: + if self.payload.logprobs is None or not self.payload.echo: + return + for token_offset, output_index in enumerate(output_indices): + if output_index is None: + continue + next_pos = output_positions[token_offset] + 1 + if not self.prompt_plan.has_text_token_at(next_pos): + continue + text_token_index = self.prompt_plan.text_token_index_by_pos[next_pos] + if text_token_index < self.prompt_visible_start: + continue + logits_index = output_index_to_logits_index(output_index, output_count) + assert logits_index is not None + next_token = self.prompt_plan.text_token_at(next_pos) + record = Token.from_logits( + model=model, + formatter=formatter, + prev_tokens=self.prompt_plan.prev_text_tokens_at(next_pos), + token=next_token, + logits=model.logits(logits_index), + logprobs_count=self.payload.logprobs, + ) + expected_index = text_token_index - self.prompt_visible_start + if expected_index == len(self.prompt_records): + self.prompt_records.append(record) + elif expected_index < len(self.prompt_records): + self.prompt_records[expected_index] = record + else: + raise RuntimeError("prompt logprob order mismatch") + + +class ResponseParser: + @dataclass + class PartialJsonValue: + text: str + schema_type: Optional[str] = None + complete: bool = False + + @dataclass + class PartialJsonObject: + value: "OrderedDict[str, Any]" + complete: bool = False + + @dataclass + class DirectState: + deltas: bool + pending: str + mode: int + done: bool + saw_tool_calls: bool + tool_call_count: int + assistant_prefix: str + leading_capture_field: Optional[str] + leading_capture_start: str + leading_capture_end: str + leading_capture_strip_after: bool + leading_capture_implicit: bool + content_end_markers: Tuple[str, ...] + trim_before_iterator: bool + iterator_start: str + iterator_end: str + stop_markers: Tuple[str, ...] + function_start: str + function_name_end: str + function_end: str + parameter_start: str + parameter_name_end: str + parameter_end: str + + @dataclass + class ItemState: + pending: str + mode: int + tool_call_index: int + tool_name: str + parameter_count: int + visible_tool_call: Optional[Dict[str, Any]] + visible_function: Optional[Dict[str, Any]] + current_parameter: Optional[str] + current_parameter_value: str + current_schema_type: Optional[str] + current_parameter_schema: Dict[str, Any] + + @dataclass + class StreamState: + kind: str + pending: str + mode: str + current_item: Optional[Dict[str, Any]] + current_segment: Optional[Dict[str, Any]] + done: bool + saw_tool_calls: bool + parsed: Optional[Dict[str, Any]] = None + buffer: str = "" + tool_call_count: int = 0 + + _STREAM_PLAN_CACHE: Dict[int, Tuple[Any, Optional[Dict[str, Any]]]] = {} + _TOOL_SCHEMA_CACHE: Dict[int, Tuple[Any, Dict[str, Dict[str, Any]]]] = {} + __slots__ = ( + "_schema", + "_tools", + "_completion_id", + "_choice_index", + "_generation_prompt", + "_tool_schemas", + "_started", + "_text_parts", + "_message", + "_stream_plan", + "_direct", + "_item", + "_stream_state", + "_stream_failed", + ) + DIRECT_MODE_ASSISTANT_PREFIX = 0 + DIRECT_MODE_PRELUDE = 1 + DIRECT_MODE_LEADING_CAPTURE = 2 + DIRECT_MODE_CONTENT = 3 + DIRECT_MODE_TOOL_ITEM = 4 + DIRECT_MODE_AFTER_TOOL_ITEM = 5 + ITEM_MODE_FUNCTION_HEADER = 0 + ITEM_MODE_AFTER_FUNCTION_HEADER = 1 + ITEM_MODE_PARAMETER_NAME = 2 + ITEM_MODE_PARAMETER_VALUE = 3 + ITEM_MODE_AFTER_PARAMETER = 4 + ITEM_MODE_DONE = 5 + + def __init__( + self, + schema: Dict[str, Any], + *, + tools: Optional[List[ChatTemplateTool]] = None, + completion_id: str = "", + choice_index: int = 0, + generation_prompt: str = "", + ) -> None: + self._schema = schema + self._tools = tools + self._completion_id = completion_id + self._choice_index = choice_index + self._generation_prompt = generation_prompt + self._tool_schemas = self._cached_tool_schema_map(tools) + self._started = False + self._text_parts: List[str] = [] + self._message: Dict[str, Any] = {} + self._stream_plan = self._cached_stream_plan(schema) + self._direct = ResponseParser.DirectState( + deltas=bool(self._stream_plan is not None and self._stream_plan.get("direct_deltas")), + pending="", + mode=self.DIRECT_MODE_PRELUDE, + done=False, + saw_tool_calls=False, + tool_call_count=0, + assistant_prefix="", + leading_capture_field=None, + leading_capture_start="", + leading_capture_end="", + leading_capture_strip_after=False, + leading_capture_implicit=False, + content_end_markers=(), + trim_before_iterator=False, + iterator_start="", + iterator_end="", + stop_markers=(), + function_start="", + function_name_end="", + function_end="", + parameter_start="", + parameter_name_end="", + parameter_end="", + ) + self._item = ResponseParser.ItemState( + pending="", + mode=self.ITEM_MODE_FUNCTION_HEADER, + tool_call_index=0, + tool_name="", + parameter_count=0, + visible_tool_call=None, + visible_function=None, + current_parameter=None, + current_parameter_value="", + current_schema_type=None, + current_parameter_schema={}, + ) + if self._direct.deltas and self._stream_plan is not None: + direct_init = self._stream_plan["direct_init"] + ( + self._direct.assistant_prefix, + self._direct.leading_capture_field, + self._direct.leading_capture_start, + self._direct.leading_capture_end, + self._direct.leading_capture_strip_after, + self._direct.leading_capture_implicit, + self._direct.trim_before_iterator, + self._direct.content_end_markers, + self._direct.stop_markers, + self._direct.iterator_start, + self._direct.iterator_end, + self._direct.function_start, + self._direct.function_name_end, + self._direct.function_end, + self._direct.parameter_start, + self._direct.parameter_name_end, + self._direct.parameter_end, + ) = direct_init + self._direct.mode = ( + self.DIRECT_MODE_ASSISTANT_PREFIX + if self._direct.assistant_prefix + else self.DIRECT_MODE_PRELUDE + ) + self._stream_state = None + else: + self._stream_state = ( + self._new_stream_state(self._stream_plan) if self._stream_plan is not None else None + ) + self._stream_failed = False + if self._generation_prompt and self._stream_plan is not None: + success, _ = self._advance_stream_state(self._generation_prompt) + if not success: + self._stream_failed = True + + @property + def started(self) -> bool: + return self._started + + @staticmethod + def _regex_capture(text: str, pattern: str) -> Optional[Any]: + match = re.search(pattern, text, re.S) + if match is None: + return None + group_dict = match.groupdict() + if group_dict: + return {key: value for key, value in group_dict.items() if value is not None} + return match.group(1) if match.lastindex else match.group(0) + + @staticmethod + def _gemma4_tool_call_to_json(text: str) -> str: + strings: List[str] = [] + + def capture(match: re.Match[str]) -> str: + strings.append(match.group(1)) + return f"\x00{len(strings) - 1}\x00" + + stripped = text.strip() + if stripped.startswith("[") and stripped.endswith("]"): + text = "{" + stripped[1:-1] + "}" + text = re.sub(r'<\|"\|>(.*?)<\|"\|>', capture, text, flags=re.S) + text = re.sub(r"(?<=[{,])(\w+):", r'"\1":', text) + for index, value in enumerate(strings): + text = text.replace(f"\x00{index}\x00", json.dumps(value)) + return text + + @staticmethod + def _regex_literal_prefix(pattern: str) -> str: + literal, _ = ResponseParser._regex_literal_prefix_and_remainder(pattern) + return literal + + @staticmethod + def _regex_literal_prefix_and_remainder(pattern: str) -> Tuple[str, str]: + literal: List[str] = [] + index = 0 + while index < len(pattern): + char = pattern[index] + if char == "\\": + index += 1 + if index >= len(pattern): + break + escaped = pattern[index] + if escaped == "n": + literal.append("\n") + elif escaped == "t": + literal.append("\t") + else: + literal.append(escaped) + index += 1 + continue + if char in "[](){}.*+?|^$": + break + literal.append(char) + index += 1 + return "".join(literal), pattern[index:] + + @staticmethod + def _find_regex_group_end(pattern: str, start: int) -> int: + depth = 0 + escaped = False + in_character_class = False + for index in range(start, len(pattern)): + char = pattern[index] + if escaped: + escaped = False + continue + if char == "\\": + escaped = True + continue + if char == "[": + in_character_class = True + continue + if char == "]" and in_character_class: + in_character_class = False + continue + if in_character_class: + continue + if char == "(": + depth += 1 + continue + if char == ")": + depth -= 1 + if depth == 0: + return index + return -1 + + @classmethod + def _consume_optional_literal_prefix( + cls, + pattern: str, + ) -> Optional[Tuple[str, str]]: + if not pattern.startswith("(?:"): + return None + group_end = cls._find_regex_group_end(pattern, 0) + if group_end < 0 or group_end + 1 >= len(pattern) or pattern[group_end + 1] != "?": + return None + literal, remainder = cls._regex_literal_prefix_and_remainder(pattern[3:group_end]) + if not literal or remainder: + return None + return literal, pattern[group_end + 2 :] + + @staticmethod + def _split_regex_alternatives(pattern: str) -> List[str]: + alternatives: List[str] = [] + start = 0 + depth = 0 + escaped = False + in_character_class = False + for index, char in enumerate(pattern): + if escaped: + escaped = False + continue + if char == "\\": + escaped = True + continue + if char == "[": + in_character_class = True + continue + if char == "]" and in_character_class: + in_character_class = False + continue + if in_character_class: + continue + if char == "(": + depth += 1 + continue + if char == ")": + depth -= 1 + continue + if char == "|" and depth == 0: + alternatives.append(pattern[start:index]) + start = index + 1 + alternatives.append(pattern[start:]) + return alternatives + + @classmethod + def _regex_lookahead_literal_specs(cls, pattern: str) -> List[Tuple[str, bool]]: + if not pattern.startswith("(?="): + return [] + group_end = cls._find_regex_group_end(pattern, 0) + if group_end < 0: + return [] + literals: List[Tuple[str, bool]] = [] + for alternative in cls._split_regex_alternatives(pattern[3:group_end]): + strip_leading_whitespace = False + while alternative.startswith(r"\s*"): + strip_leading_whitespace = True + alternative = alternative[3:] + if alternative == "$": + continue + if alternative.endswith("$"): + alternative = alternative[:-1] + literal, _ = cls._regex_literal_prefix_and_remainder(alternative) + if literal: + literals.append((literal, strip_leading_whitespace)) + return literals + + @classmethod + def _regex_capture_parts( + cls, + pattern: str, + ) -> Optional[Tuple[str, str]]: + normalized = pattern.lstrip("^") + captures = [ + (index, token) + for token in ("(.*?)", "(.*)") + if (index := normalized.find(token)) >= 0 + ] + if not captures: + return None + capture_index, capture_token = min(captures, key=lambda item: item[0]) + return normalized[:capture_index], normalized[capture_index + len(capture_token) :] + + @classmethod + def _regex_capture_end_literal_specs(cls, pattern: str) -> List[Tuple[str, bool]]: + capture_parts = cls._regex_capture_parts(pattern) + if capture_parts is None: + return [] + _, suffix_pattern = capture_parts + literal_specs = cls._regex_lookahead_literal_specs(suffix_pattern) + if literal_specs: + return literal_specs + literal, _ = cls._regex_literal_prefix_and_remainder(suffix_pattern) + return [(literal, False)] if literal else [] + + @classmethod + def _regex_capture_end_literals(cls, pattern: str) -> List[str]: + return [literal for literal, _ in cls._regex_capture_end_literal_specs(pattern)] + + @classmethod + def _regex_leading_capture( + cls, + *, + field_name: str, + field_regex: str, + content_regex: Optional[str], + ) -> Optional[Dict[str, Any]]: + capture_parts = cls._regex_capture_parts(field_regex) + if capture_parts is None: + return None + prefix_pattern, _ = capture_parts + prefix_pattern = prefix_pattern.lstrip("^") + optional_prefix = cls._consume_optional_literal_prefix(prefix_pattern) + if optional_prefix is not None: + prefix_pattern = optional_prefix[1] + implicit_at_start = False + optional_capture_start = cls._consume_optional_literal_prefix(prefix_pattern) + if optional_capture_start is not None: + capture_start, prefix_pattern = optional_capture_start + implicit_at_start = True + else: + capture_start, prefix_pattern = cls._regex_literal_prefix_and_remainder(prefix_pattern) + if not capture_start or prefix_pattern: + return None + end_literals = cls._regex_capture_end_literals(field_regex) + if not end_literals: + return None + capture_end = end_literals[0] + strip_after = False + if isinstance(content_regex, str): + escaped_end = re.escape(capture_end) + strip_after = bool(re.search(escaped_end + r"\\s\*", content_regex)) + return { + "field": field_name, + "start": capture_start, + "end": capture_end, + "strip_after": strip_after, + "implicit_at_start": implicit_at_start, + } + + @staticmethod + def _literal_suffix_prefix_length(text: str, literal: str) -> int: + max_length = min(len(text), len(literal) - 1) + if max_length <= 0: + return 0 + tail = text[-max_length:] + if literal[0] not in tail: + return 0 + for prefix_length in range(max_length, 0, -1): + if text.endswith(literal[:prefix_length]): + return prefix_length + return 0 + + @classmethod + def _consume_until_literal( + cls, + text: str, + literal: str, + ) -> Tuple[str, bool, str, str]: + if not literal: + return text, True, "", "" + literal_length = len(literal) + search_from = 0 + first_char = literal[0] + while True: + marker_index = text.find(first_char, search_from) + if marker_index < 0: + return text, False, "", "" + if text.startswith(literal, marker_index): + return text[:marker_index], True, text[marker_index + literal_length :], "" + suffix = text[marker_index:] + if literal.startswith(suffix): + return text[:marker_index], False, "", suffix + search_from = marker_index + 1 + + @staticmethod + def _compile_iterator_pattern(pattern: str) -> Optional[Tuple[str, str]]: + if "(.*?)" not in pattern: + return None + prefix_pattern, suffix_pattern = pattern.split("(.*?)", 1) + prefix = ResponseParser._regex_literal_prefix(prefix_pattern.lstrip("^")) + suffix = ResponseParser._regex_literal_prefix(suffix_pattern) + if not prefix: + return None + return prefix, suffix + + @classmethod + def _compile_iterator_block_pattern( + cls, + pattern: str, + ) -> Optional[Dict[str, Any]]: + first_group = pattern.find("(") + if first_group < 0: + return None + depth = 0 + last_group = -1 + escaped = False + for index in range(first_group, len(pattern)): + char = pattern[index] + if escaped: + escaped = False + continue + if char == "\\": + escaped = True + continue + if char == "(": + depth += 1 + continue + if char == ")": + depth -= 1 + if depth == 0: + last_group = index + break + if last_group < first_group: + return None + start = cls._regex_literal_prefix(pattern[:first_group].lstrip("^")) + if not start: + return None + suffix_pattern = pattern[last_group + 1 :] + if suffix_pattern == "": + return {"start": start, "end": "", "allow_eof": True} + eof_variant = suffix_pattern.endswith("|$)") and suffix_pattern.startswith("(?:") + if eof_variant: + suffix_pattern = suffix_pattern[3:-3] + end = cls._regex_literal_prefix(suffix_pattern) + if not end: + return None + return {"start": start, "end": end, "allow_eof": eof_variant} + + @classmethod + def _compile_capture_pattern( + cls, + pattern: str, + ) -> Optional[Dict[str, Any]]: + normalized = pattern.lstrip("^") + if "(.*?)" in normalized: + prefix_pattern, suffix_pattern = normalized.split("(.*?)", 1) + start = cls._regex_literal_prefix(prefix_pattern) + if not start: + return None + if suffix_pattern == "": + return {"start": start, "end": "", "allow_eof": True} + eof_variant = suffix_pattern.endswith("|$)") and suffix_pattern.startswith("(?:") + if eof_variant: + suffix_pattern = suffix_pattern[3:-3] + end = cls._regex_literal_prefix(suffix_pattern) + if not end: + return None + return {"start": start, "end": end, "allow_eof": eof_variant} + if "(.*)" in normalized: + prefix_pattern, suffix_pattern = normalized.split("(.*)", 1) + start = cls._regex_literal_prefix(prefix_pattern) + if not start: + return None + if suffix_pattern == "": + return {"start": start, "end": "", "allow_eof": True} + eof_variant = suffix_pattern.endswith("|$)") and suffix_pattern.startswith("(?:") + if eof_variant: + suffix_pattern = suffix_pattern[3:-3] + end = cls._regex_literal_prefix(suffix_pattern) + if not end: + return None + return {"start": start, "end": end, "allow_eof": eof_variant} + return None + + @classmethod + def _compile_word_capture_pattern( + cls, + pattern: str, + ) -> Optional[Dict[str, Any]]: + normalized = pattern.lstrip("^") + if r"(\w+)" not in normalized: + return None + prefix_pattern, suffix_pattern = normalized.split(r"(\w+)", 1) + start = cls._regex_literal_prefix(prefix_pattern) + if not start: + return None + end = cls._regex_literal_prefix(suffix_pattern) if suffix_pattern else "" + return { + "kind": "word", + "start": start, + "end": end, + } + + @classmethod + def _consume_until_any_literal( + cls, + text: str, + literals: Sequence[str], + ) -> Tuple[str, Optional[str], str, str]: + for literal in literals: + if literal and len(text) < len(literal) and literal.startswith(text): + return "", None, "", text + literal_first_chars = {literal[0] for literal in literals if literal} + if literal_first_chars and not any(char in text for char in literal_first_chars): + return text, None, "", "" + earliest_index: Optional[int] = None + earliest_literal: Optional[str] = None + for literal in literals: + marker_index = text.find(literal) + if marker_index < 0: + continue + if earliest_index is None or marker_index < earliest_index: + earliest_index = marker_index + earliest_literal = literal + if earliest_index is not None and earliest_literal is not None: + return ( + text[:earliest_index], + earliest_literal, + text[earliest_index + len(earliest_literal) :], + "", + ) + overlap = 0 + for literal in literals: + overlap = max(overlap, cls._literal_suffix_prefix_length(text, literal)) + if overlap: + return text[:-overlap], None, "", text[-overlap:] + return text, None, "", "" + + @classmethod + def _compile_tool_call_item_plan( + cls, + item_schema: Dict[str, Any], + ) -> Optional[Dict[str, Any]]: + item_properties = item_schema.get("properties") + if not isinstance(item_properties, dict): + return None + function_schema = item_properties.get("function") + if not isinstance(function_schema, dict): + return None + function_properties = function_schema.get("properties") + if not isinstance(function_properties, dict): + return None + name_schema = function_properties.get("name") + arguments_schema = function_properties.get("arguments") + if not isinstance(name_schema, dict) or not isinstance(arguments_schema, dict): + return None + name_regex = name_schema.get("x-regex") + argument_regex = arguments_schema.get("x-regex") + argument_key_value = arguments_schema.get("x-regex-key-value") + if ( + name_regex == r"^<function=([^>\n]+)>\n" + and argument_regex == r"^<function=[^>\n]+>\n(.*?)\n</function>$" + and argument_key_value + == r"<parameter=(?P<key>[^>\n]+)>\n(?P<value>.*?)\n</parameter>" + ): + return { + "kind": "tagged-parameters", + "schema": item_schema, + "function_start": "<function=", + "function_name_end": ">\n", + "function_end": "</function>", + "parameter_start": "<parameter=", + "parameter_name_end": ">\n", + "parameter_end": "\n</parameter>", + } + name_capture = ( + cls._compile_word_capture_pattern(name_regex) + if isinstance(name_regex, str) + else None + ) + arguments_capture = ( + cls._compile_capture_pattern(argument_regex) + if isinstance(argument_regex, str) + else None + ) + if ( + isinstance(name_capture, dict) + and isinstance(arguments_capture, dict) + and arguments_schema.get("x-parser") == "json" + ): + arguments_value_schema = { + key: value + for key, value in arguments_schema.items() + if key != "x-regex" + } + return { + "kind": "json-message", + "schema": item_schema, + "name_capture": name_capture, + "arguments_capture": arguments_capture, + "arguments_schema": arguments_value_schema, + } + return { + "kind": "buffered", + "schema": item_schema, + } + + @classmethod + def _compile_segment_message_plan( + cls, + schema: Dict[str, Any], + ) -> Optional[Dict[str, Any]]: + if schema.get("type") != "object": + return None + if isinstance(schema.get("x-regex"), str): + return None + properties = schema.get("properties") + if not isinstance(properties, dict): + return None + segments: List[Dict[str, Any]] = [] + field_segment_count = 0 + for field_name, value_schema in properties.items(): + if not isinstance(value_schema, dict): + continue + if field_name == "tool_calls": + iterator_pattern = value_schema.get("x-regex-iterator") + if not isinstance(iterator_pattern, str): + continue + iterator_capture = cls._compile_iterator_block_pattern(iterator_pattern) + if not isinstance(iterator_capture, dict) or not iterator_capture["start"]: + return None + items_schema = value_schema.get("items") + if not isinstance(items_schema, dict): + return None + item_plan = cls._compile_tool_call_item_plan(items_schema) + if item_plan is None: + return None + segments.append( + { + "kind": "iterator", + "field": field_name, + "start": iterator_capture["start"], + "end": iterator_capture["end"], + "allow_eof": iterator_capture["allow_eof"], + "item": item_plan, + } + ) + continue + field_regex = value_schema.get("x-regex") + if not isinstance(field_regex, str): + continue + capture = cls._compile_capture_pattern(field_regex) + if not isinstance(capture, dict) or not capture["start"]: + return None + segments.append( + { + "kind": "field", + "field": field_name, + "start": capture["start"], + "end": capture["end"], + "allow_eof": capture["allow_eof"], + } + ) + field_segment_count += 1 + if not segments or field_segment_count == 0: + return None + start_literals = tuple(segment["start"] for segment in segments) + if len(start_literals) != len(set(start_literals)): + return None + return { + "kind": "segment-message", + "segments": segments, + "segment_starts": start_literals, + "segments_by_start": {segment["start"]: segment for segment in segments}, + } + + @classmethod + def _compile_tagged_message_plan( + cls, + schema: Dict[str, Any], + ) -> Optional[Dict[str, Any]]: + if schema.get("type") != "object": + return None + properties = schema.get("properties") + if not isinstance(properties, dict): + return None + tool_calls_schema = properties.get("tool_calls") + if not isinstance(tool_calls_schema, dict): + return None + iterator_pattern = tool_calls_schema.get("x-regex-iterator") + if not isinstance(iterator_pattern, str): + return None + iterator = cls._compile_iterator_pattern(iterator_pattern) + if iterator is None: + iterator_capture = cls._compile_iterator_block_pattern(iterator_pattern) + if ( + not isinstance(iterator_capture, dict) + or not iterator_capture["start"] + or iterator_capture["allow_eof"] + ): + return None + iterator = (iterator_capture["start"], iterator_capture["end"]) + items_schema = tool_calls_schema.get("items") + if not isinstance(items_schema, dict): + return None + item_plan = cls._compile_tool_call_item_plan(items_schema) + if item_plan is None: + return None + content_schema = properties.get("content") + content_regex = ( + content_schema.get("x-regex") + if isinstance(content_schema, dict) + else None + ) + assistant_prefix: Optional[str] = None + if isinstance(content_regex, str): + optional_prefix = cls._consume_optional_literal_prefix(content_regex.lstrip("^")) + if optional_prefix is not None: + assistant_prefix = optional_prefix[0] + leading_capture: Optional[Dict[str, Any]] = None + for field_name, value_schema in properties.items(): + if not isinstance(value_schema, dict): + continue + field_regex = value_schema.get("x-regex") + if not isinstance(field_regex, str): + continue + if field_name == "content": + continue + capture = cls._regex_leading_capture( + field_name=field_name, + field_regex=field_regex, + content_regex=content_regex, + ) + if capture is not None: + leading_capture = capture + break + end_markers: List[str] = [] + content_end_marker_specs: List[Tuple[str, bool]] = [] + iterator_start, iterator_end = iterator + if "content" in properties: + end_markers.append(iterator_start) + if isinstance(content_regex, str): + content_end_marker_specs = cls._regex_capture_end_literal_specs(content_regex) + end_markers.extend(literal for literal, _ in content_end_marker_specs) + if not end_markers and iterator_start: + end_markers.append(iterator_start) + deduped_end_markers = tuple(dict.fromkeys(end_markers)) + trim_before_iterator = any( + literal == iterator_start and strip_leading_whitespace + for literal, strip_leading_whitespace in content_end_marker_specs + ) + direct_deltas = item_plan["kind"] == "tagged-parameters" + direct_init = ( + ( + assistant_prefix or "", + leading_capture["field"] if leading_capture is not None else None, + leading_capture["start"] if leading_capture is not None else "", + leading_capture["end"] if leading_capture is not None else "", + bool(leading_capture.get("strip_after")) if leading_capture is not None else False, + bool(leading_capture.get("implicit_at_start")) if leading_capture is not None else False, + trim_before_iterator, + deduped_end_markers, + tuple(marker for marker in deduped_end_markers if marker != iterator_start), + iterator_start, + iterator_end, + item_plan["function_start"], + item_plan["function_name_end"], + item_plan["function_end"], + item_plan["parameter_start"], + item_plan["parameter_name_end"], + item_plan["parameter_end"], + ) + if direct_deltas + else None + ) + return { + "kind": "tagged-message", + "direct_deltas": direct_deltas, + "assistant_prefix": assistant_prefix, + "leading_capture": leading_capture, + "content_field": "content" if "content" in properties else None, + "content_end_markers": deduped_end_markers, + "trim_before_iterator": trim_before_iterator, + "stop_markers": tuple(marker for marker in deduped_end_markers if marker != iterator_start), + "direct_init": direct_init, + "iterator": { + "start": iterator_start, + "end": iterator_end, + "item": item_plan, + }, + } + + @classmethod + def _compile_stream_plan( + cls, + response_schema: Optional[Dict[str, Any]], + ) -> Optional[Dict[str, Any]]: + if not isinstance(response_schema, dict): + return None + if response_schema.get("x-parser") == "json": + return {"kind": "json-root"} + segment_plan = cls._compile_segment_message_plan(response_schema) + if segment_plan is not None: + return segment_plan + return cls._compile_tagged_message_plan(response_schema) + + @classmethod + def _cached_stream_plan( + cls, + response_schema: Optional[Dict[str, Any]], + ) -> Optional[Dict[str, Any]]: + if not isinstance(response_schema, dict): + return None + cache_key = id(response_schema) + cached = cls._STREAM_PLAN_CACHE.get(cache_key) + if cached is not None and cached[0] is response_schema: + return cached[1] + plan = cls._compile_stream_plan(response_schema) + cls._STREAM_PLAN_CACHE[cache_key] = (response_schema, plan) + return plan + + @staticmethod + def _tool_schema_map( + tools: Optional[List[ChatTemplateTool]], + ) -> Dict[str, Dict[str, Any]]: + if tools is None: + return {} + mapping: Dict[str, Dict[str, Any]] = {} + for tool in tools: + if tool.get("type") != "function": + continue + function = tool.get("function", {}) + name = function.get("name") + parameters = function.get("parameters") + if isinstance(name, str) and isinstance(parameters, dict): + mapping[name] = { + "parameters": parameters, + "content_type": function.get("content_type"), + } + return mapping + + @classmethod + def _cached_tool_schema_map( + cls, + tools: Optional[List[ChatTemplateTool]], + ) -> Dict[str, Dict[str, Any]]: + if tools is None: + return {} + cache_key = id(tools) + cached = cls._TOOL_SCHEMA_CACHE.get(cache_key) + if cached is not None and cached[0] is tools: + return cached[1] + mapping = cls._tool_schema_map(tools) + cls._TOOL_SCHEMA_CACHE[cache_key] = (tools, mapping) + return mapping + + def _parameter_schema_for_tool( + self, tool_name: str, parameter_name: str + ) -> Dict[str, Any]: + tool_schema = self._tool_schemas.get(tool_name) + if not isinstance(tool_schema, dict): + return {} + parameters = tool_schema.get("parameters") + if not isinstance(parameters, dict): + return {} + properties = parameters.get("properties") + if not isinstance(properties, dict): + return {} + parameter_schema = properties.get(parameter_name) + if not isinstance(parameter_schema, dict): + return {} + return parameter_schema + + def _tool_content_type(self, tool_name: str) -> Optional[str]: + tool_schema = self._tool_schemas.get(tool_name) + if not isinstance(tool_schema, dict): + return None + content_type = tool_schema.get("content_type") + if isinstance(content_type, str): + return content_type + return None + + def _raw_string_tool_arguments(self, tool_name: str, value: str) -> Optional[Dict[str, str]]: + if self._tools is None: + return None + for tool in self._tools: + if tool.get("type") != "function": + continue + function = tool.get("function", {}) + if function.get("name") != tool_name: + continue + parameters = function.get("parameters") + if not isinstance(parameters, dict): + return None + required = parameters.get("required") + if not isinstance(required, list) or len(required) != 1: + return None + argument_name = required[0] + if not isinstance(argument_name, str): + return None + properties = parameters.get("properties") + if not isinstance(properties, dict): + return None + argument_schema = properties.get(argument_name) + if not isinstance(argument_schema, dict): + return None + argument_type = argument_schema.get("type") + if argument_type == "string" or ( + isinstance(argument_type, list) and "string" in argument_type + ): + return {argument_name: value} + return None + return None + + def _single_string_tool_argument_name(self, tool_name: str) -> Optional[str]: + if self._tools is None: + return None + for tool in self._tools: + if tool.get("type") != "function": + continue + function = tool.get("function", {}) + if function.get("name") != tool_name: + continue + parameters = function.get("parameters") + if not isinstance(parameters, dict): + return None + required = parameters.get("required") + if not isinstance(required, list) or len(required) != 1: + return None + argument_name = required[0] + if not isinstance(argument_name, str): + return None + properties = parameters.get("properties") + if not isinstance(properties, dict): + return None + argument_schema = properties.get(argument_name) + if not isinstance(argument_schema, dict): + return None + argument_type = argument_schema.get("type") + if argument_type == "string" or ( + isinstance(argument_type, list) and "string" in argument_type + ): + return argument_name + return None + return None + + def _text_tool_argument_from_object( + self, + tool_name: str, + arguments: Dict[str, Any], + ) -> Optional[str]: + input_value = arguments.get("input") + if isinstance(input_value, str): + return input_value + argument_name = self._single_string_tool_argument_name(tool_name) + if argument_name is not None: + argument_value = arguments.get(argument_name) + if isinstance(argument_value, str): + return argument_value + if len(arguments) == 1: + argument_value = next(iter(arguments.values())) + if isinstance(argument_value, str): + return argument_value + return None + + def _text_tool_arguments( + self, + tool_name: str, + arguments: Any, + *, + partial: bool, + ) -> Optional[str]: + if isinstance(arguments, str): + parsed_arguments = self._raw_object_tool_arguments(arguments) + if parsed_arguments is not None: + text = self._text_tool_argument_from_object(tool_name, parsed_arguments) + if text is not None: + return text + if partial: + return None + return json.dumps(parsed_arguments, ensure_ascii=False, separators=(",", ":")) + return arguments + if isinstance(arguments, ResponseParser.PartialJsonObject): + arguments = arguments.value + if isinstance(arguments, dict): + text = self._text_tool_argument_from_object(tool_name, arguments) + if text is not None: + return text + if partial: + return None + return json.dumps(arguments, ensure_ascii=False, separators=(",", ":")) + if partial: + return None + return str(arguments) + + @classmethod + def _raw_object_tool_arguments(cls, value: str) -> Optional[Dict[str, Any]]: + candidates = [value] + stripped = value.strip() + if stripped.startswith("{{") and stripped.endswith("}}"): + candidates.append(stripped[1:-1]) + for candidate in candidates: + normalized = cls._gemma4_tool_call_to_json(candidate) + for allow_partial in (False, True): + try: + parsed = from_json(normalized, allow_partial=allow_partial) + except ValueError: + continue + if isinstance(parsed, dict): + return { + key: cls._trim_partial_gemma_quote_marker(value) + if isinstance(value, str) + else value + for key, value in parsed.items() + } + return None + + @staticmethod + def _trim_partial_gemma_quote_marker(value: str) -> str: + quote_marker = '<|"|>' + for prefix_length in range(len(quote_marker) - 1, 0, -1): + if value.endswith(quote_marker[:prefix_length]): + return value[:-prefix_length] + return value + + def _has_text_tools(self) -> bool: + return any( + isinstance(tool_schema, dict) and tool_schema.get("content_type") == "text" + for tool_schema in self._tool_schemas.values() + ) + + @staticmethod + def _append_parsed_text(parsed: Dict[str, Any], key: str, text: str) -> None: + if not text: + return + existing = parsed.get(key) + if isinstance(existing, str): + parsed[key] = existing + text + else: + parsed[key] = text + + def _append_visible_text(self, key: str, text: str) -> None: + if not text: + return + existing = self._message.get(key) + if isinstance(existing, str): + self._message[key] = existing + text + else: + self._message[key] = text + + @staticmethod + def _advance_json_scanner( + item_state: Dict[str, Any], + text: str, + *, + schema_type: Optional[str], + ) -> bool: + started = cast(bool, item_state["json_started"]) + complete = cast(bool, item_state["json_complete"]) + depth = cast(int, item_state["json_depth"]) + in_string = cast(bool, item_state["json_in_string"]) + escaped = cast(bool, item_state["json_escaped"]) + for char in text: + if complete: + if not char.isspace(): + return False + continue + if not started: + if char.isspace(): + continue + if schema_type == "object" and char != "{": + return False + if schema_type == "array" and char != "[": + return False + if char == "{": + started = True + depth = 1 + continue + if char == "[": + started = True + depth = 1 + continue + return False + if in_string: + if escaped: + escaped = False + elif char == "\\": + escaped = True + elif char == '"': + in_string = False + continue + if char == '"': + in_string = True + continue + if char in "{[": + depth += 1 + continue + if char in "}]": + depth -= 1 + if depth < 0: + return False + if depth == 0: + complete = True + continue + item_state["json_started"] = started + item_state["json_complete"] = complete + item_state["json_depth"] = depth + item_state["json_in_string"] = in_string + item_state["json_escaped"] = escaped + return True + + def _new_tool_call_state(self, item_plan: Dict[str, Any]) -> Dict[str, Any]: + if item_plan["kind"] == "buffered": + return {"kind": "buffered", "buffer": ""} + if item_plan["kind"] == "json-message": + return { + "kind": "json-message", + "pending": "", + "mode": "function-name", + "tool_call": { + "type": "function", + "function": { + "name": "", + "arguments": "", + }, + }, + "arguments_text": "", + "json_started": False, + "json_complete": False, + "json_depth": 0, + "json_in_string": False, + "json_escaped": False, + } + return { + "kind": "tagged-parameters", + "pending": "", + "mode": "function-header", + "tool_call": { + "type": "function", + "function": { + "name": "", + "arguments": ResponseParser.PartialJsonObject(OrderedDict(), complete=False), + }, + }, + "current_parameter": None, + "current_schema_type": None, + } + + def _new_stream_state(self, plan: Dict[str, Any]) -> ResponseParser.StreamState: + if plan["kind"] == "json-root": + return ResponseParser.StreamState( + kind="json-root", + pending="", + mode="prelude", + current_item=None, + current_segment=None, + done=False, + saw_tool_calls=False, + parsed={"role": "assistant"}, + buffer="", + ) + if plan.get("direct_deltas"): + return ResponseParser.StreamState( + kind="tagged-message", + pending="", + mode="assistant-prefix" if plan.get("assistant_prefix") else "prelude", + current_item=None, + current_segment=None, + done=False, + saw_tool_calls=False, + tool_call_count=0, + ) + if plan["kind"] == "segment-message": + return ResponseParser.StreamState( + kind="segment-message", + pending="", + mode="segment-start", + current_item=None, + current_segment=None, + done=False, + saw_tool_calls=False, + parsed={"role": "assistant"}, + ) + return ResponseParser.StreamState( + kind="tagged-message", + pending="", + mode="assistant-prefix" if plan.get("assistant_prefix") else "prelude", + current_item=None, + current_segment=None, + done=False, + saw_tool_calls=False, + parsed={"role": "assistant"}, + ) + + def _start_direct_tool_call(self, tool_call_index: int) -> None: + self._item.pending = "" + self._item.mode = self.ITEM_MODE_FUNCTION_HEADER + self._item.tool_call_index = tool_call_index + self._item.tool_name = "" + self._item.parameter_count = 0 + tool_calls = self._message.setdefault("tool_calls", []) + assert isinstance(tool_calls, list) + while len(tool_calls) <= tool_call_index: + tool_calls.append({"function": {"name": "", "arguments": ""}}) + visible_tool_call = tool_calls[tool_call_index] + assert isinstance(visible_tool_call, dict) + self._item.visible_tool_call = visible_tool_call + function = visible_tool_call.setdefault("function", {}) + assert isinstance(function, dict) + self._item.visible_function = function + if tool_call_index == 0: + self._message["function_call"] = self._item.visible_function + self._item.current_parameter = None + self._item.current_parameter_value = "" + self._item.current_schema_type = None + self._item.current_parameter_schema = {} + + def _advance_direct_tool_call_state(self, text: str) -> Tuple[bool, List[Dict[str, Any]]]: + deltas: List[Dict[str, Any]] = [] + mode = self._item.mode + tool_call_index = self._item.tool_call_index + tool_name = self._item.tool_name + parameter_count = self._item.parameter_count + visible_tool_call = self._item.visible_tool_call + visible_function = self._item.visible_function + current_parameter = self._item.current_parameter + current_parameter_value = self._item.current_parameter_value + current_schema_type = self._item.current_schema_type + current_parameter_schema = self._item.current_parameter_schema + buffer = self._item.pending + text + function_start = self._direct.function_start + function_name_end = self._direct.function_name_end + function_end = self._direct.function_end + parameter_start = self._direct.parameter_start + parameter_name_end = self._direct.parameter_name_end + parameter_end = self._direct.parameter_end + + while True: + if mode == self.ITEM_MODE_FUNCTION_HEADER: + if buffer.startswith(function_start): + buffer = buffer[len(function_start) :] + elif function_start.startswith(buffer): + self._item.pending = buffer + self._item.mode = mode + return True, deltas + else: + self._item.pending = "" + return False, deltas + delimiter_index = buffer.find(function_name_end) + if delimiter_index < 0: + self._item.pending = function_start + buffer + self._item.mode = mode + return True, deltas + function_name = buffer[:delimiter_index] + if not function_name: + self._item.pending = "" + return False, deltas + tool_name = function_name + assert visible_tool_call is not None + assert visible_function is not None + visible_function["name"] = function_name + visible_function["arguments"] = "{" + visible_tool_call["id"] = ( + f"call_{self._choice_index}_{function_name}_{self._completion_id}_{tool_call_index}" + ) + visible_tool_call["type"] = "function" + deltas.append( + { + "tool_calls": [ + { + "index": tool_call_index, + "id": ( + f"call_{self._choice_index}_{function_name}_" + f"{self._completion_id}_{tool_call_index}" + ), + "type": "function", + "function": {"name": function_name, "arguments": "{"}, + } + ] + } + ) + buffer = buffer[delimiter_index + len(function_name_end) :] + mode = self.ITEM_MODE_AFTER_FUNCTION_HEADER + continue + if mode == self.ITEM_MODE_AFTER_FUNCTION_HEADER: + if buffer.startswith("\n"): + buffer = buffer[1:] + continue + if buffer.startswith(parameter_start): + mode = self.ITEM_MODE_PARAMETER_NAME + continue + if buffer.startswith(function_end): + buffer = buffer[len(function_end) :] + assert visible_function is not None + visible_function["arguments"] += "}" + deltas.append( + { + "tool_calls": [ + { + "index": tool_call_index, + "function": {"arguments": "}"}, + } + ] + } + ) + mode = self.ITEM_MODE_DONE + continue + if parameter_start.startswith(buffer) or function_end.startswith(buffer): + self._item.pending = buffer + self._item.mode = mode + self._item.tool_name = tool_name + return True, deltas + if not buffer: + self._item.pending = "" + self._item.mode = mode + self._item.tool_name = tool_name + return True, deltas + self._item.pending = "" + return False, deltas + if mode == self.ITEM_MODE_PARAMETER_NAME: + if buffer.startswith(parameter_start): + buffer = buffer[len(parameter_start) :] + elif parameter_start.startswith(buffer): + self._item.pending = buffer + self._item.mode = mode + self._item.tool_name = tool_name + self._item.parameter_count = parameter_count + return True, deltas + else: + self._item.pending = "" + return False, deltas + delimiter_index = buffer.find(parameter_name_end) + if delimiter_index < 0: + self._item.pending = parameter_start + buffer + self._item.mode = mode + self._item.tool_name = tool_name + self._item.parameter_count = parameter_count + return True, deltas + parameter_name = buffer[:delimiter_index] + if not parameter_name: + self._item.pending = "" + return False, deltas + parameter_schema = self._parameter_schema_for_tool(tool_name, parameter_name) + schema_type = parameter_schema.get("type") if isinstance(parameter_schema, dict) else None + current_parameter = parameter_name + current_parameter_value = "" + current_schema_type = schema_type if isinstance(schema_type, str) else None + current_parameter_schema = parameter_schema + key_prefix = json.dumps(parameter_name, ensure_ascii=False, separators=(",", ":")) + ":" + if parameter_count > 0: + key_prefix = "," + key_prefix + if schema_type in {None, "string"}: + key_prefix += '"' + parameter_count += 1 + assert visible_function is not None + visible_function["arguments"] += key_prefix + deltas.append( + { + "tool_calls": [ + { + "index": tool_call_index, + "function": {"arguments": key_prefix}, + } + ] + } + ) + buffer = buffer[delimiter_index + len(parameter_name_end) :] + mode = self.ITEM_MODE_PARAMETER_VALUE + continue + if mode == self.ITEM_MODE_PARAMETER_VALUE: + value_delta, matched, remainder, pending = self._consume_until_literal(buffer, parameter_end) + current_parameter_value += value_delta + if value_delta: + assert visible_function is not None + visible_function["arguments"] += value_delta + deltas.append( + { + "tool_calls": [ + { + "index": tool_call_index, + "function": {"arguments": value_delta}, + } + ] + } + ) + if not matched: + self._item.pending = pending + self._item.mode = mode + self._item.tool_name = tool_name + self._item.parameter_count = parameter_count + self._item.current_parameter = current_parameter + self._item.current_parameter_value = current_parameter_value + self._item.current_schema_type = current_schema_type + self._item.current_parameter_schema = current_parameter_schema + return True, deltas + self._coerce_tool_argument( + cast(str, current_parameter_value), + cast(Dict[str, Any], current_parameter_schema), + tool_name=tool_name, + argument_name=cast(str, current_parameter), + ) + if current_schema_type in {None, "string"}: + assert visible_function is not None + visible_function["arguments"] += '"' + deltas.append( + { + "tool_calls": [ + { + "index": tool_call_index, + "function": {"arguments": '"'}, + } + ] + } + ) + current_parameter = None + current_parameter_value = "" + current_schema_type = None + current_parameter_schema = {} + buffer = remainder + mode = self.ITEM_MODE_AFTER_PARAMETER + continue + if mode == self.ITEM_MODE_AFTER_PARAMETER: + if buffer.startswith("\n"): + buffer = buffer[1:] + continue + if buffer.startswith(parameter_start): + mode = self.ITEM_MODE_PARAMETER_NAME + continue + if buffer.startswith(function_end): + buffer = buffer[len(function_end) :] + assert visible_function is not None + visible_function["arguments"] += "}" + deltas.append( + { + "tool_calls": [ + { + "index": tool_call_index, + "function": {"arguments": "}"}, + } + ] + } + ) + mode = self.ITEM_MODE_DONE + continue + if parameter_start.startswith(buffer) or function_end.startswith(buffer): + self._item.pending = buffer + self._item.mode = mode + self._item.tool_name = tool_name + self._item.parameter_count = parameter_count + return True, deltas + if not buffer: + self._item.pending = "" + self._item.mode = mode + self._item.tool_name = tool_name + self._item.parameter_count = parameter_count + return True, deltas + self._item.pending = "" + return False, deltas + if mode == self.ITEM_MODE_DONE: + if buffer.strip(): + self._item.pending = "" + return False, deltas + self._item.pending = buffer + self._item.mode = mode + self._item.tool_name = tool_name + self._item.parameter_count = parameter_count + return True, deltas + + def _advance_direct_stream_state(self, text: str) -> Tuple[bool, List[Dict[str, Any]]]: + deltas: List[Dict[str, Any]] = [] + mode = self._direct.mode + tool_call_count = self._direct.tool_call_count + saw_tool_calls = self._direct.saw_tool_calls + done = self._direct.done + buffer = self._direct.pending + text + assistant_prefix = self._direct.assistant_prefix + leading_capture_field = self._direct.leading_capture_field + leading_capture_start = self._direct.leading_capture_start + leading_capture_end = self._direct.leading_capture_end + leading_capture_strip_after = self._direct.leading_capture_strip_after + leading_capture_implicit = self._direct.leading_capture_implicit + iterator_start = self._direct.iterator_start + iterator_end = self._direct.iterator_end + content_end_markers = self._direct.content_end_markers + stop_markers = self._direct.stop_markers + + while True: + if mode == self.DIRECT_MODE_ASSISTANT_PREFIX: + if buffer.startswith(assistant_prefix): + buffer = buffer[len(assistant_prefix) :] + mode = self.DIRECT_MODE_PRELUDE + continue + if assistant_prefix.startswith(buffer): + self._direct.pending = buffer + self._direct.mode = mode + return True, deltas + mode = self.DIRECT_MODE_PRELUDE + continue + if mode == self.DIRECT_MODE_PRELUDE: + if leading_capture_field is not None: + if buffer.startswith(leading_capture_start): + buffer = buffer[len(leading_capture_start) :] + mode = self.DIRECT_MODE_LEADING_CAPTURE + continue + if leading_capture_start.startswith(buffer): + self._direct.pending = buffer + self._direct.mode = mode + return True, deltas + if leading_capture_implicit: + tool_index = buffer.find(iterator_start) + end_index = buffer.find(leading_capture_end) + overlap = self._literal_suffix_prefix_length(buffer, leading_capture_end) + if ((end_index >= 0 and (tool_index < 0 or end_index < tool_index)) or overlap): + mode = self.DIRECT_MODE_LEADING_CAPTURE + continue + mode = self.DIRECT_MODE_CONTENT + continue + if mode == self.DIRECT_MODE_LEADING_CAPTURE: + segment, matched, remainder, pending = self._consume_until_literal(buffer, leading_capture_end) + if segment: + assert leading_capture_field is not None + self._append_visible_text(leading_capture_field, segment) + deltas.append({leading_capture_field: segment}) + if not matched: + self._direct.pending = pending + self._direct.mode = mode + self._direct.tool_call_count = tool_call_count + self._direct.saw_tool_calls = saw_tool_calls + self._direct.done = done + return True, deltas + buffer = remainder.lstrip() if leading_capture_strip_after else remainder + mode = self.DIRECT_MODE_CONTENT + continue + if mode == self.DIRECT_MODE_CONTENT: + segment, matched_marker, remainder, pending = self._consume_until_any_literal( + buffer, + content_end_markers, + ) + content_segment = segment.rstrip() if matched_marker == iterator_start and self._direct.trim_before_iterator else segment + if content_segment: + self._append_visible_text("content", content_segment) + deltas.append({"content": content_segment}) + if matched_marker is None: + self._direct.pending = pending + self._direct.mode = mode + self._direct.tool_call_count = tool_call_count + self._direct.saw_tool_calls = saw_tool_calls + self._direct.done = done + return True, deltas + if matched_marker == iterator_start: + saw_tool_calls = True + self._start_direct_tool_call(tool_call_count) + tool_call_count += 1 + mode = self.DIRECT_MODE_TOOL_ITEM + buffer = remainder + continue + done = True + self._direct.pending = "" + self._direct.mode = mode + self._direct.tool_call_count = tool_call_count + self._direct.saw_tool_calls = saw_tool_calls + self._direct.done = done + return True, deltas + if mode == self.DIRECT_MODE_TOOL_ITEM: + item_text, matched, remainder, pending = self._consume_until_literal(buffer, iterator_end) + if item_text: + success, item_deltas = self._advance_direct_tool_call_state(item_text) + if not success: + self._direct.pending = "" + self._direct.mode = mode + self._direct.tool_call_count = tool_call_count + self._direct.saw_tool_calls = saw_tool_calls + self._direct.done = done + return False, deltas + deltas.extend(item_deltas) + if not matched: + self._direct.pending = pending + self._direct.mode = mode + self._direct.tool_call_count = tool_call_count + self._direct.saw_tool_calls = saw_tool_calls + self._direct.done = done + return True, deltas + buffer = remainder.lstrip() + mode = self.DIRECT_MODE_AFTER_TOOL_ITEM + continue + if mode == self.DIRECT_MODE_AFTER_TOOL_ITEM: + if not buffer: + self._direct.pending = "" + self._direct.mode = mode + self._direct.tool_call_count = tool_call_count + self._direct.saw_tool_calls = saw_tool_calls + self._direct.done = done + return True, deltas + if leading_capture_field is not None: + if buffer.startswith(leading_capture_start): + buffer = buffer[len(leading_capture_start) :] + mode = self.DIRECT_MODE_LEADING_CAPTURE + continue + if leading_capture_start.startswith(buffer): + self._direct.pending = buffer + self._direct.mode = mode + self._direct.tool_call_count = tool_call_count + self._direct.saw_tool_calls = saw_tool_calls + self._direct.done = done + return True, deltas + if buffer.startswith(iterator_start): + saw_tool_calls = True + self._start_direct_tool_call(tool_call_count) + tool_call_count += 1 + mode = self.DIRECT_MODE_TOOL_ITEM + buffer = buffer[len(iterator_start) :] + continue + for marker in stop_markers: + if buffer.startswith(marker): + done = True + self._direct.pending = "" + self._direct.mode = mode + self._direct.tool_call_count = tool_call_count + self._direct.saw_tool_calls = saw_tool_calls + self._direct.done = done + return True, deltas + if any(marker.startswith(buffer) for marker in content_end_markers) or buffer.isspace(): + self._direct.pending = buffer + self._direct.mode = mode + self._direct.tool_call_count = tool_call_count + self._direct.saw_tool_calls = saw_tool_calls + self._direct.done = done + return True, deltas + self._direct.pending = "" + self._direct.mode = mode + self._direct.tool_call_count = tool_call_count + self._direct.saw_tool_calls = saw_tool_calls + self._direct.done = done + return False, deltas + + def _advance_tool_call_state( + self, + item_state: Dict[str, Any], + text: str, + item_plan: Dict[str, Any], + ) -> Tuple[bool, List[Dict[str, Any]]]: + deltas: List[Dict[str, Any]] = [] + if item_plan["kind"] == "buffered": + item_state["buffer"] = item_state["buffer"] + text + return True, deltas + if item_plan["kind"] == "json-message": + buffer = item_state["pending"] + text + item_state["pending"] = "" + while True: + mode = item_state["mode"] + if mode == "function-name": + name_capture = item_plan["name_capture"] + name_prefix = name_capture["start"] + if buffer.startswith(name_prefix): + remainder = buffer[len(name_prefix) :] + elif name_prefix.startswith(buffer): + item_state["pending"] = buffer + return True, deltas + else: + return False, deltas + name_end = 0 + while name_end < len(remainder) and ( + remainder[name_end].isalnum() or remainder[name_end] == "_" + ): + name_end += 1 + if name_end == 0: + if not remainder: + item_state["pending"] = buffer + return True, deltas + return False, deltas + if name_end == len(remainder): + item_state["pending"] = buffer + return True, deltas + function_name = remainder[:name_end] + item_state["tool_call"]["function"]["name"] = function_name + tool_call_index = cast(int, item_state["tool_call_index"]) + deltas.append( + { + "tool_calls": [ + { + "index": tool_call_index, + "id": ( + f"call_{self._choice_index}_{function_name}_" + f"{self._completion_id}_{tool_call_index}" + ), + "type": "function", + "function": {"name": function_name}, + } + ] + } + ) + buffer = remainder[name_end:] + item_state["mode"] = "seek-arguments" + continue + if mode == "seek-arguments": + arguments_capture = item_plan["arguments_capture"] + _, matched, remainder, pending = self._consume_until_literal( + buffer, + arguments_capture["start"], + ) + if not matched: + item_state["pending"] = pending + return True, deltas + buffer = remainder + item_state["mode"] = "arguments" + continue + if mode == "arguments": + if buffer: + item_state["arguments_text"] = item_state["arguments_text"] + buffer + deltas.append( + { + "tool_calls": [ + { + "index": cast(int, item_state["tool_call_index"]), + "function": {"arguments": buffer}, + } + ] + } + ) + item_state["pending"] = "" + arguments_schema = cast(Dict[str, Any], item_plan["arguments_schema"]) + schema_type = arguments_schema.get("type") + if not self._advance_json_scanner( + item_state, + buffer, + schema_type=schema_type if isinstance(schema_type, str) else None, + ): + return False, deltas + return True, deltas + + buffer = item_state["pending"] + text + item_state["pending"] = "" + while True: + mode = item_state["mode"] + if mode == "function-header": + function_start = item_plan["function_start"] + if buffer.startswith(function_start): + buffer = buffer[len(function_start) :] + elif function_start.startswith(buffer): + item_state["pending"] = buffer + return True, deltas + else: + return False, deltas + function_name_end = item_plan["function_name_end"] + delimiter_index = buffer.find(function_name_end) + if delimiter_index < 0: + item_state["pending"] = function_start + buffer + return True, deltas + function_name = buffer[:delimiter_index] + if not function_name: + return False, deltas + item_state["tool_call"]["function"]["name"] = function_name + deltas.append( + { + "tool_calls": [ + { + "index": item_state["tool_call_index"], + "id": ( + f"call_{self._choice_index}_{function_name}_" + f"{self._completion_id}_{item_state['tool_call_index']}" + ), + "type": "function", + "function": { + "name": function_name, + "arguments": "{", + }, + } + ] + } + ) + buffer = buffer[delimiter_index + len(function_name_end) :] + item_state["mode"] = "after-function-header" + continue + if mode == "after-function-header": + if buffer.startswith("\n"): + buffer = buffer[1:] + continue + function_end = item_plan["function_end"] + parameter_start = item_plan["parameter_start"] + if buffer.startswith(parameter_start): + item_state["mode"] = "parameter-name" + continue + if buffer.startswith(function_end): + buffer = buffer[len(function_end) :] + item_state["mode"] = "done" + deltas.append( + { + "tool_calls": [ + { + "index": item_state["tool_call_index"], + "function": {"arguments": "}"}, + } + ] + } + ) + continue + if parameter_start.startswith(buffer) or function_end.startswith(buffer): + item_state["pending"] = buffer + return True, deltas + if not buffer: + return True, deltas + return False, deltas + if mode == "parameter-name": + parameter_start = item_plan["parameter_start"] + if buffer.startswith(parameter_start): + buffer = buffer[len(parameter_start) :] + elif parameter_start.startswith(buffer): + item_state["pending"] = buffer + return True, deltas + else: + return False, deltas + parameter_name_end = item_plan["parameter_name_end"] + delimiter_index = buffer.find(parameter_name_end) + if delimiter_index < 0: + item_state["pending"] = parameter_start + buffer + return True, deltas + parameter_name = buffer[:delimiter_index] + if not parameter_name: + return False, deltas + function = item_state["tool_call"]["function"] + arguments = cast(ResponseParser.PartialJsonObject, function["arguments"]) + tool_name = function["name"] + parameter_schema = self._parameter_schema_for_tool(tool_name, parameter_name) + schema_type = ( + parameter_schema.get("type") + if isinstance(parameter_schema, dict) + else None + ) + arguments.value[parameter_name] = ResponseParser.PartialJsonValue( + text="", + schema_type=schema_type if isinstance(schema_type, str) else None, + complete=False, + ) + key_prefix = json.dumps( + parameter_name, + ensure_ascii=False, + separators=(",", ":"), + ) + ":" + if len(arguments.value) > 1: + key_prefix = "," + key_prefix + if schema_type in {None, "string"}: + key_prefix += '"' + deltas.append( + { + "tool_calls": [ + { + "index": item_state["tool_call_index"], + "function": {"arguments": key_prefix}, + } + ] + } + ) + item_state["current_parameter"] = parameter_name + item_state["current_schema_type"] = ( + schema_type if isinstance(schema_type, str) else None + ) + buffer = buffer[delimiter_index + len(parameter_name_end) :] + item_state["mode"] = "parameter-value" + continue + if mode == "parameter-value": + parameter_end = item_plan["parameter_end"] + value_delta, matched, remainder, pending = self._consume_until_literal( + buffer, + parameter_end, + ) + function = item_state["tool_call"]["function"] + arguments = cast(ResponseParser.PartialJsonObject, function["arguments"]) + parameter_name = cast(str, item_state["current_parameter"]) + current_value = arguments.value[parameter_name] + assert isinstance(current_value, ResponseParser.PartialJsonValue) + current_value.text = current_value.text + value_delta + if value_delta: + deltas.append( + { + "tool_calls": [ + { + "index": item_state["tool_call_index"], + "function": {"arguments": value_delta}, + } + ] + } + ) + if not matched: + item_state["pending"] = pending + return True, deltas + tool_name = function["name"] + parameter_schema = self._parameter_schema_for_tool(tool_name, parameter_name) + self._coerce_tool_argument( + current_value.text, + parameter_schema, + tool_name=tool_name, + argument_name=parameter_name, + ) + current_value.complete = True + if current_value.schema_type in {None, "string"}: + deltas.append( + { + "tool_calls": [ + { + "index": item_state["tool_call_index"], + "function": {"arguments": '"'}, + } + ] + } + ) + item_state["current_parameter"] = None + item_state["current_schema_type"] = None + buffer = remainder + item_state["mode"] = "after-parameter" + continue + if mode == "after-parameter": + if buffer.startswith("\n"): + buffer = buffer[1:] + continue + function_end = item_plan["function_end"] + parameter_start = item_plan["parameter_start"] + if buffer.startswith(parameter_start): + item_state["mode"] = "parameter-name" + continue + if buffer.startswith(function_end): + buffer = buffer[len(function_end) :] + cast( + ResponseParser.PartialJsonObject, + item_state["tool_call"]["function"]["arguments"], + ).complete = True + item_state["mode"] = "done" + deltas.append( + { + "tool_calls": [ + { + "index": item_state["tool_call_index"], + "function": {"arguments": "}"}, + } + ] + } + ) + continue + if parameter_start.startswith(buffer) or function_end.startswith(buffer): + item_state["pending"] = buffer + return True, deltas + if not buffer: + return True, deltas + return False, deltas + if mode == "done": + if buffer.strip(): + return False, deltas + item_state["pending"] = buffer + return True, deltas + + def _finish_tool_call_state( + self, + item_state: Dict[str, Any], + item_plan: Dict[str, Any], + ) -> Optional[Dict[str, Any]]: + if item_plan["kind"] == "buffered": + parsed_item = self._parse_response_value( + item_state["buffer"], + item_plan["schema"], + partial=False, + ) + return self._normalize_tool_call_item(parsed_item, partial=False) + if item_plan["kind"] == "json-message": + if item_state["mode"] != "arguments": + return None + if item_state["pending"] or not item_state["json_started"] or not item_state["json_complete"]: + return None + try: + arguments = self._parse_response_value( + item_state["arguments_text"], + item_plan["arguments_schema"], + partial=False, + ) + except CompletionResponseParsingError: + return None + if arguments is None: + return None + item_state["tool_call"]["function"]["arguments"] = arguments + return cast(Dict[str, Any], item_state["tool_call"]) + if item_state["pending"] or item_state["current_parameter"] is not None: + return None + if item_state["mode"] not in {"done", "after-function-header"}: + return None + return cast(Dict[str, Any], item_state["tool_call"]) + + def _tool_call_delta( + self, + *, + tool_call: Dict[str, Any], + tool_call_index: int, + partial: bool, + ) -> Dict[str, Any]: + function = cast(Dict[str, Any], tool_call["function"]) + return { + "tool_calls": [ + { + "index": tool_call_index, + "id": ( + f"call_{self._choice_index}_{function['name']}_" + f"{self._completion_id}_{tool_call_index}" + ), + "type": tool_call.get("type", "function"), + "function": { + "name": function["name"], + "arguments": self._serialize_tool_arguments( + function["arguments"], + partial=partial, + ), + }, + } + ] + } + + def _advance_stream_state(self, text: str) -> Tuple[bool, List[Dict[str, Any]]]: + deltas: List[Dict[str, Any]] = [] + if self._stream_plan is None: + return False, deltas + if self._direct.deltas: + return self._advance_direct_stream_state(text) + if self._stream_state is None: + return False, deltas + if self._stream_plan["kind"] == "segment-message": + state = self._stream_state + parsed = cast(Dict[str, Any], state.parsed) + buffer = state.pending + text + state.pending = "" + while True: + if state.mode == "segment-start": + _, matched_start, remainder, pending = self._consume_until_any_literal( + buffer, + cast(Tuple[str, ...], self._stream_plan["segment_starts"]), + ) + if matched_start is None: + state.pending = buffer if not pending else pending + return True, deltas + state.current_segment = cast( + Dict[str, Any], + self._stream_plan["segments_by_start"][matched_start], + ) + buffer = remainder + state.mode = ( + "segment-tool-item" + if state.current_segment["kind"] == "iterator" + else "segment-field" + ) + if state.current_segment["kind"] == "iterator": + item_state = self._new_tool_call_state(state.current_segment["item"]) + if item_state["kind"] in {"tagged-parameters", "json-message"}: + tool_calls = cast( + List[Dict[str, Any]], + parsed.setdefault("tool_calls", []), + ) + tool_calls.append(item_state["tool_call"]) + item_state["tool_call_index"] = len(tool_calls) - 1 + state.current_item = item_state + state.saw_tool_calls = True + continue + if state.mode == "segment-field": + current_segment = state.current_segment + if not isinstance(current_segment, dict): + return False, deltas + segment_text, matched, remainder, pending = self._consume_until_literal( + buffer, + cast(str, current_segment["end"]), + ) + field_name = cast(str, current_segment["field"]) + self._append_parsed_text(parsed, field_name, segment_text) + if segment_text: + deltas.append({field_name: segment_text}) + if not matched: + state.pending = pending + return True, deltas + buffer = remainder + state.current_segment = None + state.mode = "segment-start" + continue + if state.mode == "segment-tool-item": + current_segment = state.current_segment + if not isinstance(current_segment, dict): + return False, deltas + active_item_state = state.current_item + if not isinstance(active_item_state, dict): + return False, deltas + item_text, matched, remainder, pending = self._consume_until_literal( + buffer, + cast(str, current_segment["end"]), + ) + if item_text: + success, item_deltas = self._advance_tool_call_state( + active_item_state, + item_text, + cast(Dict[str, Any], current_segment["item"]), + ) + if not success: + return False, deltas + deltas.extend(item_deltas) + if not matched: + state.pending = pending + return True, deltas + tool_call = self._finish_tool_call_state( + active_item_state, + cast(Dict[str, Any], current_segment["item"]), + ) + if tool_call is None: + return False, deltas + if active_item_state["kind"] == "buffered": + tool_calls = cast(List[Dict[str, Any]], parsed.setdefault("tool_calls", [])) + tool_call_index = len(tool_calls) + tool_calls.append(tool_call) + deltas.append( + self._tool_call_delta( + tool_call=tool_call, + tool_call_index=tool_call_index, + partial=False, + ) + ) + state.current_item = None + state.current_segment = None + state.mode = "segment-start" + buffer = remainder + continue + return False, deltas + if self._stream_plan["kind"] == "json-root": + buffer = self._stream_state.buffer + text + try: + parsed = self.parse_chat_response(buffer, partial=True) + except CompletionResponseParsingError: + return False, deltas + self._stream_state.buffer = buffer + self._stream_state.parsed = parsed + return True, deltas + + state = self._stream_state + plan = self._stream_plan + buffer = state.pending + text + state.pending = "" + parsed = cast(Dict[str, Any], state.parsed) + iterator_start = plan["iterator"]["start"] + iterator_end = plan["iterator"]["end"] + while True: + mode = state.mode + if mode == "assistant-prefix": + assistant_prefix = plan.get("assistant_prefix") + if not assistant_prefix: + state.mode = "prelude" + continue + if buffer.startswith(assistant_prefix): + buffer = buffer[len(assistant_prefix) :] + state.mode = "prelude" + continue + if assistant_prefix.startswith(buffer): + state.pending = buffer + return True, deltas + state.mode = "prelude" + continue + if mode == "prelude": + leading_capture = plan.get("leading_capture") + if leading_capture is not None: + capture_start = leading_capture["start"] + if buffer.startswith(capture_start): + buffer = buffer[len(capture_start) :] + state.mode = "leading-capture" + continue + if capture_start.startswith(buffer): + state.pending = buffer + return True, deltas + if leading_capture.get("implicit_at_start"): + capture_end = leading_capture["end"] + tool_index = buffer.find(iterator_start) + end_index = buffer.find(capture_end) + overlap = self._literal_suffix_prefix_length(buffer, capture_end) + if ( + end_index >= 0 + and (tool_index < 0 or end_index < tool_index) + ) or overlap: + state.mode = "leading-capture" + continue + state.mode = "content" + continue + if mode == "leading-capture": + leading_capture = cast(Dict[str, Any], plan["leading_capture"]) + segment, matched, remainder, pending = self._consume_until_literal( + buffer, + leading_capture["end"], + ) + self._append_parsed_text(parsed, leading_capture["field"], segment) + if segment: + deltas.append({leading_capture["field"]: segment}) + if not matched: + state.pending = pending + return True, deltas + buffer = remainder.lstrip() if leading_capture.get("strip_after") else remainder + state.mode = "content" + continue + if mode == "content": + content_field = plan.get("content_field") + end_markers = cast(List[str], plan["content_end_markers"]) + segment, matched_marker, remainder, pending = self._consume_until_any_literal( + buffer, + end_markers, + ) + if content_field is not None: + content_segment = ( + segment.rstrip() + if matched_marker == iterator_start and plan.get("trim_before_iterator") + else segment + ) + self._append_parsed_text(parsed, content_field, content_segment) + if content_segment: + deltas.append({content_field: content_segment}) + if matched_marker is None: + state.pending = pending + return True, deltas + if matched_marker == iterator_start: + item_state = self._new_tool_call_state(plan["iterator"]["item"]) + state.saw_tool_calls = True + if item_state["kind"] == "tagged-parameters": + tool_calls = cast(List[Dict[str, Any]], parsed.setdefault("tool_calls", [])) + tool_calls.append(item_state["tool_call"]) + item_state["tool_call_index"] = len(tool_calls) - 1 + state.current_item = item_state + state.mode = "tool-item" + buffer = remainder + continue + state.done = True + state.pending = "" + return True, deltas + if mode == "tool-item": + current_item_state: Optional[Dict[str, Any]] = state.current_item + if not isinstance(current_item_state, dict): + return False, deltas + item_text, matched, remainder, pending = self._consume_until_literal( + buffer, + iterator_end, + ) + if item_text: + success, item_deltas = self._advance_tool_call_state( + current_item_state, + item_text, + plan["iterator"]["item"], + ) + if not success: + return False, deltas + deltas.extend(item_deltas) + if not matched: + state.pending = pending + return True, deltas + tool_call = self._finish_tool_call_state( + current_item_state, + plan["iterator"]["item"], + ) + if tool_call is None: + return False, deltas + if current_item_state["kind"] == "buffered": + parsed.setdefault("tool_calls", []).append(tool_call) + state.current_item = None + buffer = remainder.lstrip() + state.mode = "after-tool-item" + continue + if mode == "after-tool-item": + end_markers = [ + iterator_start, + *[ + marker + for marker in cast(List[str], plan["content_end_markers"]) + if marker != iterator_start + ], + ] + if not buffer: + state.pending = "" + return True, deltas + leading_capture = plan.get("leading_capture") + if leading_capture is not None: + capture_start = leading_capture["start"] + if buffer.startswith(capture_start): + buffer = buffer[len(capture_start) :] + state.mode = "leading-capture" + continue + if capture_start.startswith(buffer): + state.pending = buffer + return True, deltas + if buffer.startswith(iterator_start): + item_state = self._new_tool_call_state(plan["iterator"]["item"]) + state.saw_tool_calls = True + if item_state["kind"] == "tagged-parameters": + tool_calls = cast(List[Dict[str, Any]], parsed.setdefault("tool_calls", [])) + tool_calls.append(item_state["tool_call"]) + item_state["tool_call_index"] = len(tool_calls) - 1 + state.current_item = item_state + state.mode = "tool-item" + buffer = buffer[len(iterator_start) :] + continue + for marker in end_markers[1:]: + if buffer.startswith(marker): + state.done = True + state.pending = "" + return True, deltas + if any(marker.startswith(buffer) for marker in end_markers): + state.pending = buffer + return True, deltas + if buffer.isspace(): + state.pending = buffer + return True, deltas + return False, deltas + + @staticmethod + def _partial_regex_key_value_item( + pattern: str, + text: str, + *, + min_start: int, + ) -> Optional[Tuple[str, str]]: + value_group = "(?P<value>" + value_group_start = pattern.find(value_group) + if value_group_start < 0: + return None + partial_pattern = pattern[:value_group_start] + r"(?P<value>.*)\Z" + partial_match: Optional[re.Match[str]] = None + for match in re.finditer(partial_pattern, text, re.S): + if match.start() < min_start: + continue + partial_match = match + if partial_match is None: + return None + group_dict = partial_match.groupdict() + key = group_dict.get("key") + value = group_dict.get("value") + if key is None or value is None: + return None + value_pattern = "(?P<value>.*?)" + if value_pattern in pattern: + suffix_literal = ResponseParser._regex_literal_prefix( + pattern.split(value_pattern, 1)[1] + ) + for suffix_length in range(len(suffix_literal), 0, -1): + suffix_prefix = suffix_literal[:suffix_length] + if value.endswith(suffix_prefix): + value = value[:-suffix_length] + break + return key, value + + def _trim_partial_tool_call_prefix( + self, + *, + response_text: str, + parsed: Dict[str, Any], + ) -> None: + if not isinstance(parsed.get("content"), str): + return + tool_calls_schema = self._schema.get("properties", {}).get("tool_calls") + if not isinstance(tool_calls_schema, dict): + return + iterator_pattern = tool_calls_schema.get("x-regex-iterator") + if not isinstance(iterator_pattern, str) or "(.*?)" not in iterator_pattern: + return + prefix_pattern = iterator_pattern.split("(.*?)", 1)[0] + literal_prefix = self._regex_literal_prefix(prefix_pattern) + if not literal_prefix: + return + content = cast(str, parsed["content"]) + for prefix_length in range(len(literal_prefix) - 1, 0, -1): + prefix = literal_prefix[:prefix_length] + if content.endswith(prefix) and response_text.endswith(prefix): + trimmed = content[:-prefix_length] + parsed["content"] = trimmed if trimmed else None + break + + def _parse_response_value( + self, + text: Any, + schema: Dict[str, Any], + *, + partial: bool, + ) -> Any: + if "const" in schema: + return schema["const"] + if text is None: + return None + node_content: Any = text + node_regex = schema.get("x-regex") + if node_regex is not None: + if not isinstance(node_content, str): + raise CompletionResponseParsingError( + "response_schema x-regex requires string input" + ) + captured_content = self._regex_capture(node_content, node_regex) + if captured_content is None: + if ( + partial + and schema.get("type") == "object" + and "x-regex-key-value" in schema + ): + captured_content = node_content + else: + return None + node_content = captured_content + node_regex_iterator = schema.get("x-regex-iterator") + if node_regex_iterator is not None: + if schema.get("type") != "array": + raise CompletionResponseParsingError( + "response_schema x-regex-iterator requires array type" + ) + if not isinstance(node_content, str): + raise CompletionResponseParsingError( + "response_schema x-regex-iterator requires string input" + ) + array_values = [] + matches = list(re.finditer(node_regex_iterator, node_content, re.S)) + for match in matches: + item_text = self._regex_capture(match.group(0), node_regex_iterator) + if item_text is not None: + array_values.append(item_text) + if partial and "(.*?)" in node_regex_iterator: + prefix_pattern, suffix_pattern = node_regex_iterator.split("(.*?)", 1) + prefix_matches = list(re.finditer(prefix_pattern, node_content, re.S)) + if prefix_matches: + last_prefix_match = prefix_matches[-1] + if not matches or matches[-1].start() != last_prefix_match.start(): + tail = node_content[last_prefix_match.end() :] + if re.search(suffix_pattern, tail, re.S) is None: + array_values.append(tail) + if not array_values: + return None + node_content = array_values + node_regex_key_value = schema.get("x-regex-key-value") + if node_regex_key_value is not None: + if schema.get("type") != "object": + raise CompletionResponseParsingError( + "response_schema x-regex-key-value requires object type" + ) + if not isinstance(node_content, str): + raise CompletionResponseParsingError( + "response_schema x-regex-key-value requires string input" + ) + key_values: Dict[str, str] = {} + matches = list(re.finditer(node_regex_key_value, node_content, re.S)) + for match in matches: + group_dict = match.groupdict() + if "key" not in group_dict or "value" not in group_dict: + raise CompletionResponseParsingError( + "response_schema x-regex-key-value must define key and value groups" + ) + key = group_dict["key"] + value = group_dict["value"] + if key is None or value is None: + raise CompletionResponseParsingError( + "response_schema x-regex-key-value matched null key or value" + ) + key_values[key] = value + if partial: + min_start = matches[-1].end() if matches else 0 + partial_item = self._partial_regex_key_value_item( + node_regex_key_value, + node_content, + min_start=min_start, + ) + if partial_item is not None: + key_values[partial_item[0]] = partial_item[1] + if not key_values: + return None + node_content = key_values + parser_name = schema.get("x-parser") + if parser_name is not None: + if parser_name != "json": + if parser_name != "gemma4-tool-call": + raise CompletionResponseParsingError( + f"unsupported response_schema x-parser: {parser_name}" + ) + if not isinstance(node_content, str): + raise CompletionResponseParsingError( + "response_schema x-parser='gemma4-tool-call' requires string input" + ) + node_content = self._gemma4_tool_call_to_json(node_content) + parser_name = "json" + if parser_name == "json": + if not isinstance(node_content, str): + raise CompletionResponseParsingError( + "response_schema x-parser='json' requires string input" + ) + try: + parsed = from_json(node_content, allow_partial=partial) + except ValueError as exc: + if ( + self._has_text_tools() + and schema.get("type") == "object" + and schema.get("additionalProperties") is True + and not schema.get("properties") + ): + return node_content + if partial: + return None + raise CompletionResponseParsingError( + "response did not match response_schema JSON parser" + ) from exc + stripped_schema = { + key: value + for key, value in schema.items() + if key + not in { + "x-parser", + "x-parser-args", + "x-regex", + "x-regex-iterator", + "x-regex-key-value", + } + } + return self._parse_response_value( + parsed, stripped_schema, partial=partial + ) + schema_type = schema.get("type") + if schema_type == "string": + return node_content + if schema_type == "array": + if isinstance(node_content, list): + array_values = [] + item_schema = schema.get("items", {}) + for item in node_content: + parsed_item = self._parse_response_value( + item, + item_schema, + partial=partial, + ) + if parsed_item is not None: + array_values.append(parsed_item) + return array_values + return [] + if schema_type == "object": + properties = schema.get("properties", {}) + if isinstance(node_content, dict): + parsed_object: Dict[str, Any] = {} + for key, value_schema in properties.items(): + value = self._parse_response_value( + node_content.get(key), + value_schema, + partial=partial, + ) + if value is None: + continue + if isinstance(value, list) and not value: + continue + parsed_object[key] = value + additional_properties = schema.get("additionalProperties") + if additional_properties is True: + for key, value in node_content.items(): + if key not in parsed_object and key not in properties: + parsed_object[key] = value + elif isinstance(additional_properties, dict): + for key, value in node_content.items(): + if key in parsed_object or key in properties: + continue + parsed_value = self._parse_response_value( + value, + additional_properties, + partial=partial, + ) + if parsed_value is not None: + parsed_object[key] = parsed_value + if not partial: + missing = [ + key + for key in schema.get("required", []) + if key not in parsed_object + ] + if missing: + raise CompletionResponseParsingError( + f"response did not match response_schema: missing {', '.join(missing)}" + ) + return parsed_object + parsed_object_from_text: Dict[str, Any] = {} + for key, value_schema in properties.items(): + value = self._parse_response_value( + node_content, value_schema, partial=partial + ) + if value is None: + continue + if isinstance(value, list) and not value: + continue + parsed_object_from_text[key] = value + if not partial: + missing = [ + key + for key in schema.get("required", []) + if key not in parsed_object_from_text + ] + if missing: + raise CompletionResponseParsingError( + f"response did not match response_schema: missing {', '.join(missing)}" + ) + return parsed_object_from_text + if schema_type == "integer": + if isinstance(node_content, int) and not isinstance(node_content, bool): + return node_content + if partial and isinstance(node_content, str) and not node_content: + return None + try: + return int(node_content) + except (TypeError, ValueError): + return None + if schema_type == "number": + if isinstance(node_content, (int, float)) and not isinstance( + node_content, bool + ): + return node_content + if partial and isinstance(node_content, str) and not node_content: + return None + try: + return float(node_content) + except (TypeError, ValueError): + return None + if schema_type == "boolean": + if isinstance(node_content, bool): + return node_content + if node_content in {"true", "True", 1, "1"}: + return True + if node_content in {"false", "False", 0, "0"}: + return False + return None + one_of = schema.get("oneOf") + if isinstance(one_of, list): + for option in one_of: + value = self._parse_response_value( + node_content, option, partial=partial + ) + if value is not None: + return value + return None + if schema_type is None or schema_type == "any": + return node_content + return None + + def _coerce_tool_argument( + self, + raw_value: str, + schema: Dict[str, Any], + *, + tool_name: str, + argument_name: str, + ) -> Any: + if "oneOf" in schema: + last_error: Optional[BaseException] = None + for variant in schema["oneOf"]: + try: + return self._coerce_tool_argument( + raw_value, + variant, + tool_name=tool_name, + argument_name=argument_name, + ) + except BaseException as exc: + last_error = exc + raise CompletionResponseParsingError( + f"tool '{tool_name}' argument '{argument_name}' did not match any allowed schema" + ) from last_error + schema_type = schema.get("type") + if isinstance(schema_type, list): + last_type_error: Optional[BaseException] = None + for variant_type in schema_type: + try: + return self._coerce_tool_argument( + raw_value, + {**schema, "type": variant_type}, + tool_name=tool_name, + argument_name=argument_name, + ) + except BaseException as exc: + last_type_error = exc + raise CompletionResponseParsingError( + f"tool '{tool_name}' argument '{argument_name}' did not match any allowed type" + ) from last_type_error + if schema_type in {None, "string"}: + return raw_value + try: + decoded = json.loads(raw_value) + except json.JSONDecodeError as exc: + raise CompletionResponseParsingError( + f"tool '{tool_name}' argument '{argument_name}' is not valid JSON for type '{schema_type}'" + ) from exc + if schema_type == "integer": + if isinstance(decoded, bool) or not isinstance(decoded, int): + raise CompletionResponseParsingError( + f"tool '{tool_name}' argument '{argument_name}' must be an integer" + ) + return decoded + if schema_type == "number": + if isinstance(decoded, bool) or not isinstance(decoded, (int, float)): + raise CompletionResponseParsingError( + f"tool '{tool_name}' argument '{argument_name}' must be a number" + ) + return decoded + if schema_type == "boolean": + if not isinstance(decoded, bool): + raise CompletionResponseParsingError( + f"tool '{tool_name}' argument '{argument_name}' must be a boolean" + ) + return decoded + if schema_type == "object": + if not isinstance(decoded, dict): + raise CompletionResponseParsingError( + f"tool '{tool_name}' argument '{argument_name}' must be an object" + ) + return decoded + if schema_type == "array": + if not isinstance(decoded, list): + raise CompletionResponseParsingError( + f"tool '{tool_name}' argument '{argument_name}' must be an array" + ) + return decoded + return decoded + + def _coerce_tool_arguments( + self, + tool_name: str, + arguments: Dict[str, str], + *, + partial: bool, + ) -> Dict[str, Any]: + if self._tool_content_type(tool_name) == "text": + raw_input = arguments.get("input", "") + if not isinstance(raw_input, str): + raw_input = str(raw_input) + return {"input": raw_input} + if self._tools is None: + raise CompletionResponseParsingError( + f"response included tool call '{tool_name}' but the request declared no tools" + ) + tool = next( + ( + candidate + for candidate in self._tools + if candidate.get("type") == "function" + and candidate.get("function", {}).get("name") == tool_name + ), + None, + ) + if tool is None: + raise CompletionResponseParsingError( + f"response included undeclared tool call '{tool_name}'" + ) + parameters = tool.get("function", {}).get("parameters") or { + "type": "object", + "properties": {}, + "required": [], + } + properties = parameters.get("properties", {}) + coerced: Dict[str, Any] = {} + for argument_name, raw_value in arguments.items(): + if argument_name not in properties: + raise CompletionResponseParsingError( + f"tool '{tool_name}' returned unexpected argument '{argument_name}'" + ) + coerced[argument_name] = self._coerce_tool_argument( + raw_value, + properties[argument_name], + tool_name=tool_name, + argument_name=argument_name, + ) + if not partial: + missing = [ + argument_name + for argument_name in parameters.get("required", []) + if argument_name not in coerced + ] + if missing: + raise CompletionResponseParsingError( + f"tool '{tool_name}' response is missing required arguments: {', '.join(missing)}" + ) + return coerced + + def parse_chat_response( + self, + response_text: str, + *, + partial: bool, + ) -> Dict[str, Any]: + full_response_text = self._generation_prompt + response_text + parsed = self._parse_response_value( + full_response_text, + self._schema, + partial=partial, + ) + if not isinstance(parsed, dict): + raise CompletionResponseParsingError("response_schema must produce an object") + if partial: + self._trim_partial_tool_call_prefix( + response_text=full_response_text, + parsed=parsed, + ) + tool_calls = parsed.get("tool_calls") + if isinstance(tool_calls, list): + normalized_tool_calls: List[Dict[str, Any]] = [] + for tool_call in tool_calls: + normalized = self._normalize_tool_call_item(tool_call, partial=partial) + if normalized is None: + continue + normalized_tool_calls.append(normalized) + if normalized_tool_calls: + parsed["tool_calls"] = normalized_tool_calls + else: + parsed.pop("tool_calls", None) + for field in ("reasoning_content", "thinking"): + value = parsed.get(field) + if isinstance(value, str) and not value.strip(): + parsed.pop(field, None) + return parsed + + def _normalize_tool_call_item( + self, + tool_call: Any, + *, + partial: bool, + ) -> Optional[Dict[str, Any]]: + if not isinstance(tool_call, dict): + if partial: + return None + raise CompletionResponseParsingError("tool_calls items must be objects") + function = tool_call.get("function") + if not isinstance(function, dict): + if partial: + return None + raise CompletionResponseParsingError( + "tool_calls items must define a function object" + ) + tool_name = function.get("name") + if not isinstance(tool_name, str) or not tool_name: + if partial: + return None + raise CompletionResponseParsingError( + "tool_calls function name must be a non-empty string" + ) + if self._tool_content_type(tool_name) == "text": + arguments = self._text_tool_arguments( + tool_name, + function.get("arguments", ""), + partial=partial, + ) + if arguments is None: + return None + return { + "type": tool_call.get("type", "function"), + "function": { + "name": tool_name, + "arguments": arguments, + }, + } + arguments = function.get("arguments", {}) + if isinstance(arguments, str): + arguments = self._raw_object_tool_arguments(arguments) or self._raw_string_tool_arguments( + tool_name, arguments + ) + if not isinstance(arguments, (dict, ResponseParser.PartialJsonObject)): + if partial: + return None + raise CompletionResponseParsingError( + "tool_calls function arguments must parse to an object" + ) + argument_values = ( + arguments.value + if isinstance(arguments, ResponseParser.PartialJsonObject) + else arguments + ) + raw_arguments: Dict[str, str] = {} + for argument_name, argument_value in argument_values.items(): + if isinstance(argument_value, ResponseParser.PartialJsonValue): + raw_arguments[argument_name] = argument_value.text + else: + raw_arguments[argument_name] = str(argument_value) + normalized_arguments = ( + arguments + if isinstance(arguments, ResponseParser.PartialJsonObject) + or any( + isinstance(value, ResponseParser.PartialJsonValue) + for value in argument_values.values() + ) + else self._coerce_tool_arguments( + tool_name, + raw_arguments, + partial=partial, + ) + ) + return { + "type": tool_call.get("type", "function"), + "function": { + "name": tool_name, + "arguments": normalized_arguments, + }, + } + + @classmethod + def _serialize_partial_json_prefix(cls, value: Any) -> str: + if isinstance(value, ResponseParser.PartialJsonValue): + if value.complete and value.schema_type in {None, "string"}: + return json.dumps(value.text, ensure_ascii=False, separators=(",", ":")) + if value.schema_type in {None, "string"}: + return json.dumps(value.text, ensure_ascii=False, separators=(",", ":"))[:-1] + return value.text + if isinstance(value, ResponseParser.PartialJsonObject): + return cls._serialize_partial_json_prefix(value.value) + if isinstance(value, dict): + items = list(value.items()) + if not items: + return "{" + rendered = ["{"] + last_index = len(items) - 1 + for index, (key, item_value) in enumerate(items): + if index > 0: + rendered.append(",") + rendered.append(json.dumps(key, ensure_ascii=False, separators=(",", ":"))) + rendered.append(":") + if index == last_index: + rendered.append(cls._serialize_partial_json_prefix(item_value)) + else: + rendered.append( + json.dumps(item_value, ensure_ascii=False, separators=(",", ":")) + ) + return "".join(rendered) + if isinstance(value, list): + if not value: + return "[" + rendered = ["["] + last_index = len(value) - 1 + for index, item_value in enumerate(value): + if index > 0: + rendered.append(",") + if index == last_index: + rendered.append(cls._serialize_partial_json_prefix(item_value)) + else: + rendered.append( + json.dumps(item_value, ensure_ascii=False, separators=(",", ":")) + ) + return "".join(rendered) + if isinstance(value, str): + return json.dumps(value, ensure_ascii=False, separators=(",", ":"))[:-1] + return json.dumps(value, ensure_ascii=False, separators=(",", ":")) + + @classmethod + def _contains_partial_json_value(cls, value: Any) -> bool: + if isinstance(value, ResponseParser.PartialJsonValue): + return True + if isinstance(value, ResponseParser.PartialJsonObject): + return True + if isinstance(value, dict): + return any(cls._contains_partial_json_value(item) for item in value.values()) + if isinstance(value, list): + return any(cls._contains_partial_json_value(item) for item in value) + return False + + @classmethod + def _serialize_partial_json_state(cls, value: Any) -> Tuple[str, bool]: + if isinstance(value, ResponseParser.PartialJsonValue): + if value.schema_type in {None, "string"}: + rendered = json.dumps(value.text, ensure_ascii=False, separators=(",", ":")) + if value.complete: + return rendered, True + return rendered[:-1], False + return value.text, value.complete + if isinstance(value, ResponseParser.PartialJsonObject): + rendered, children_complete = cls._serialize_partial_json_state(value.value) + if value.complete and children_complete: + return rendered, True + if rendered.endswith("}"): + return rendered[:-1], False + return rendered, False + if isinstance(value, dict): + parts = ["{"] + items = list(value.items()) + for index, (key, item_value) in enumerate(items): + if index > 0: + parts.append(",") + parts.append(json.dumps(key, ensure_ascii=False, separators=(",", ":"))) + parts.append(":") + rendered_item, item_complete = cls._serialize_partial_json_state(item_value) + parts.append(rendered_item) + if index == len(items) - 1 and not item_complete: + return "".join(parts), False + parts.append("}") + return "".join(parts), True + if isinstance(value, list): + parts = ["["] + for index, item_value in enumerate(value): + if index > 0: + parts.append(",") + rendered_item, item_complete = cls._serialize_partial_json_state(item_value) + parts.append(rendered_item) + if index == len(value) - 1 and not item_complete: + return "".join(parts), False + parts.append("]") + return "".join(parts), True + return json.dumps(value, ensure_ascii=False, separators=(",", ":")), True + + @classmethod + def _serialize_tool_arguments(cls, arguments: Any, *, partial: bool = False) -> str: + if partial: + if cls._contains_partial_json_value(arguments): + return cls._serialize_partial_json_state(arguments)[0] + return cls._serialize_partial_json_prefix(arguments) + if isinstance(arguments, ResponseParser.PartialJsonObject): + arguments = arguments.value + return json.dumps(arguments, ensure_ascii=False, separators=(",", ":")) + + def _parsed_chat_message( + self, + *, + parsed: Dict[str, Any], + partial: bool = False, + ) -> Dict[str, Any]: + message: Dict[str, Any] = { + "role": parsed.get("role", "assistant"), + } + for key, value in parsed.items(): + if key in {"role", "content", "tool_calls"}: + continue + if value is None or value == "": + continue + message[key] = value + content = parsed.get("content") + tool_calls = parsed.get("tool_calls") + if isinstance(tool_calls, list) and tool_calls: + normalized_tool_calls = [] + for tool_call_index, tool_call in enumerate(tool_calls): + function = tool_call["function"] + if self._tool_content_type(function["name"]) == "text": + arguments = self._text_tool_arguments( + function["name"], + function["arguments"], + partial=partial, + ) + if arguments is None: + continue + else: + arguments = self._serialize_tool_arguments( + function["arguments"], + partial=partial, + ) + normalized_tool_calls.append( + { + "id": f"call_{self._choice_index}_{function['name']}_{self._completion_id}_{tool_call_index}", + "type": tool_call.get("type", "function"), + "function": { + "name": function["name"], + "arguments": arguments, + }, + } + ) + message["content"] = content if content not in {None, ""} else None + message["tool_calls"] = normalized_tool_calls + if len(normalized_tool_calls) == 1: + message["function_call"] = dict(normalized_tool_calls[0]["function"]) + return message + message["content"] = content if content is not None else "" + return message + + def _message_deltas( + self, + previous_message: Dict[str, Any], + message: Dict[str, Any], + ) -> List[Dict[str, Any]]: + deltas: List[Dict[str, Any]] = [] + for key, value in message.items(): + if key in {"role", "content", "tool_calls", "function_call"}: + continue + old_value = previous_message.get(key, "") + if not isinstance(value, str): + if key not in previous_message and value is not None: + deltas.append({key: value}) + continue + if not value: + continue + if isinstance(old_value, str) and value.startswith(old_value): + delta_value = value[len(old_value) :] + else: + delta_value = value + if delta_value: + deltas.append({key: delta_value}) + + new_content = message.get("content") + old_content = previous_message.get("content", "") + if isinstance(new_content, str) and new_content: + if isinstance(old_content, str) and new_content.startswith(old_content): + content_delta = new_content[len(old_content) :] + else: + content_delta = new_content + if content_delta: + deltas.append({"content": content_delta}) + + new_tool_calls = cast(List[Dict[str, Any]], message.get("tool_calls", [])) + old_tool_calls = cast(List[Dict[str, Any]], previous_message.get("tool_calls", [])) + for tool_call_index, tool_call in enumerate(new_tool_calls): + old_tool_call = old_tool_calls[tool_call_index] if tool_call_index < len(old_tool_calls) else None + delta_tool_call: Dict[str, Any] = {"index": tool_call_index} + if old_tool_call is None: + delta_tool_call["id"] = tool_call["id"] + delta_tool_call["type"] = tool_call["type"] + function = cast(Dict[str, Any], tool_call["function"]) + old_function = ( + cast(Dict[str, Any], old_tool_call["function"]) + if old_tool_call is not None + else {} + ) + function_delta: Dict[str, Any] = {} + if function.get("name") and function.get("name") != old_function.get("name"): + function_delta["name"] = function["name"] + arguments = cast(str, function.get("arguments", "")) + old_arguments = cast(str, old_function.get("arguments", "")) + if old_tool_call is None and arguments == "{}": + argument_delta = "" + elif arguments.startswith(old_arguments): + argument_delta = arguments[len(old_arguments) :] + else: + argument_delta = arguments + if argument_delta: + function_delta["arguments"] = argument_delta + if function_delta: + delta_tool_call["function"] = function_delta + if len(delta_tool_call) > 1: + deltas.append({"tool_calls": [delta_tool_call]}) + return deltas + + @staticmethod + def _apply_message_delta(message: Dict[str, Any], delta: Dict[str, Any]) -> None: + if "role" in delta: + message["role"] = delta["role"] + for key, value in delta.items(): + if key in {"role", "tool_calls", "function_call"}: + continue + if isinstance(value, str): + existing = message.get(key) + if isinstance(existing, str): + message[key] = existing + value + else: + message[key] = value + else: + message[key] = value + tool_call_deltas = delta.get("tool_calls") + if not isinstance(tool_call_deltas, list): + return + tool_calls = cast(List[Dict[str, Any]], message.setdefault("tool_calls", [])) + for tool_delta in tool_call_deltas: + if not isinstance(tool_delta, dict): + continue + index = tool_delta.get("index") + if not isinstance(index, int): + continue + while len(tool_calls) <= index: + tool_calls.append({"function": {"name": "", "arguments": ""}}) + tool_call = tool_calls[index] + if "id" in tool_delta: + tool_call["id"] = tool_delta["id"] + if "type" in tool_delta: + tool_call["type"] = tool_delta["type"] + function_delta = tool_delta.get("function") + if not isinstance(function_delta, dict): + continue + function = cast(Dict[str, Any], tool_call.setdefault("function", {})) + name_delta = function_delta.get("name") + if isinstance(name_delta, str): + function["name"] = cast(str, function.get("name", "")) + name_delta + arguments_delta = function_delta.get("arguments") + if isinstance(arguments_delta, str): + function["arguments"] = cast(str, function.get("arguments", "")) + arguments_delta + if tool_calls: + message["function_call"] = dict(cast(Dict[str, Any], tool_calls[0]["function"])) + + def parse_completion_message(self, response_text: str) -> Dict[str, Any]: + parsed = self.parse_chat_response(response_text, partial=False) + return self._parsed_chat_message(parsed=parsed) + + def _stream_state_message(self, *, partial: bool) -> Dict[str, Any]: + assert self._stream_state is not None + if self._stream_plan is not None and self._stream_plan.get("direct_deltas"): + copied = { + key: ( + [ + { + child_key: ( + dict(child_value) + if isinstance(child_value, dict) + else child_value + ) + for child_key, child_value in tool_call.items() + } + for tool_call in cast(List[Dict[str, Any]], value) + ] + if key == "tool_calls" and isinstance(value, list) + else value + ) + for key, value in self._message.items() + } + if copied.get("tool_calls"): + copied["function_call"] = dict( + cast(List[Dict[str, Any]], copied["tool_calls"])[0]["function"] + ) + return copied + parsed = cast(Dict[str, Any], self._stream_state.parsed) + return self._parsed_chat_message(parsed=parsed, partial=partial) + + def _stream_state_complete(self) -> bool: + if self._stream_plan is None: + return False + if self._direct.deltas: + return ( + not self._direct.pending + and self._direct.mode in {self.DIRECT_MODE_CONTENT, self.DIRECT_MODE_AFTER_TOOL_ITEM} + ) + if self._stream_state is None: + return False + if self._stream_plan["kind"] == "json-root": + return True + if self._stream_plan["kind"] == "segment-message": + return ( + not self._stream_state.pending + and self._stream_state.current_item is None + and self._stream_state.current_segment is None + and self._stream_state.mode == "segment-start" + ) + return ( + not self._stream_state.pending + and self._stream_state.current_item is None + and self._stream_state.mode in {"content", "after-tool-item"} + ) + + @staticmethod + def _chat_chunk_payload( + *, + chunk_id: str, + created: int, + model: str, + index: int, + delta: Dict[str, Any], + finish_reason: Optional[str], + logprobs: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + choice: Dict[str, Any] = { + "index": index, + "delta": delta, + "finish_reason": finish_reason, + } + if logprobs is not None: + choice["logprobs"] = logprobs + return { + "id": "chat" + chunk_id, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": [choice], + } + + def _chunk_payloads( + self, + *, + chunk_id: str, + created: int, + model: str, + deltas: List[Dict[str, Any]], + finish_reason: Optional[str], + logprobs: Optional[Dict[str, Any]], + leading_delta: Optional[Dict[str, Any]] = None, + ) -> List[Dict[str, Any]]: + payloads: List[Dict[str, Any]] = [] + logprobs_sent = False + chat_chunk_id = "chat" + chunk_id + index = self._choice_index + if leading_delta is not None: + payloads.append( + { + "id": chat_chunk_id, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": [ + { + "index": index, + "delta": leading_delta, + "finish_reason": None, + } + ], + } + ) + for delta in deltas: + payload_logprobs = None + if ( + logprobs is not None + and delta.get("role") != "assistant" + and not logprobs_sent + ): + payload_logprobs = logprobs + logprobs_sent = True + choice: Dict[str, Any] = { + "index": index, + "delta": delta, + "finish_reason": None, + } + if payload_logprobs is not None: + choice["logprobs"] = payload_logprobs + payloads.append( + { + "id": chat_chunk_id, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": [choice], + } + ) + if finish_reason is not None: + payloads.append( + { + "id": chat_chunk_id, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": [ + { + "index": index, + "delta": {}, + "finish_reason": finish_reason, + } + ], + } + ) + return payloads + + def consume_completion_chunk( + self, + text: str, + *, + chunk_id: str, + created: int, + model: str, + finish_reason: Optional[str], + logprobs: Optional[Dict[str, Any]] = None, + ) -> List[Dict[str, Any]]: + role_delta: Optional[Dict[str, Any]] = None + if not self._started: + self._started = True + self._message["role"] = "assistant" + role_delta = {"role": "assistant"} + self._text_parts.append(text) + + if self._stream_plan is not None and not self._stream_failed: + success, stream_deltas = self._advance_stream_state(text) + if success: + if self._direct.deltas: + if finish_reason is None: + return self._chunk_payloads( + chunk_id=chunk_id, + created=created, + model=model, + deltas=stream_deltas, + finish_reason=None, + logprobs=logprobs, + leading_delta=role_delta, + ) + return self._chunk_payloads( + chunk_id=chunk_id, + created=created, + model=model, + deltas=stream_deltas, + finish_reason=( + "tool_calls" if self._direct.saw_tool_calls else finish_reason + ), + logprobs=logprobs, + leading_delta=role_delta, + ) + elif self._stream_plan["kind"] == "segment-message": + if role_delta is not None: + stream_deltas = [role_delta, *stream_deltas] + for delta in stream_deltas: + self._apply_message_delta(self._message, delta) + if finish_reason is None: + return self._chunk_payloads( + chunk_id=chunk_id, + created=created, + model=model, + deltas=stream_deltas, + finish_reason=None, + logprobs=logprobs, + ) + return self._chunk_payloads( + chunk_id=chunk_id, + created=created, + model=model, + deltas=stream_deltas, + finish_reason=( + "tool_calls" if self._message.get("tool_calls") else finish_reason + ), + logprobs=logprobs, + ) + else: + previous_message = self._message + partial_deltas: List[Dict[str, Any]] = [] + assert self._stream_state is not None + parsed = cast(Dict[str, Any], self._stream_state.parsed) + message = self._parsed_chat_message( + parsed=parsed, + partial=finish_reason is None or not self._stream_state_complete(), + ) + if finish_reason is None: + if role_delta is not None: + partial_deltas.append(role_delta) + partial_deltas.extend(self._message_deltas(previous_message, message)) + self._message = message + return self._chunk_payloads( + chunk_id=chunk_id, + created=created, + model=model, + deltas=partial_deltas, + finish_reason=None, + logprobs=logprobs, + ) + if role_delta is not None: + partial_deltas.append(role_delta) + partial_deltas.extend(self._message_deltas(previous_message, message)) + self._message = message + return self._chunk_payloads( + chunk_id=chunk_id, + created=created, + model=model, + deltas=partial_deltas, + finish_reason=( + "tool_calls" if message.get("tool_calls") else finish_reason + ), + logprobs=logprobs, + ) + else: + self._stream_failed = True + + response_text = "".join(self._text_parts) + parsed = self.parse_chat_response(response_text, partial=finish_reason is None) + message = self._parsed_chat_message(parsed=parsed, partial=finish_reason is None) + previous_message = self._message + deltas: List[Dict[str, Any]] = [] + if role_delta is not None: + deltas.append(role_delta) + deltas.extend(self._message_deltas(previous_message, message)) + self._message = message + return self._chunk_payloads( + chunk_id=chunk_id, + created=created, + model=model, + deltas=deltas, + finish_reason=( + "tool_calls" + if finish_reason is not None and message.get("tool_calls") + else finish_reason + ), + logprobs=logprobs, + ) + + +class OpenAIFormatter: + @dataclass + class ReturnedToken: + index: int + text_bytes: bytes + token: Token + text_offset: int + + @dataclass + class ResponsesOutputItem: + output_index: int + item: Dict[str, Any] + content_index: Optional[int] = None + + @dataclass + class ResponsesStream: + body: CreateResponseRequest + response_id: str + created_at: float + model: str + output: List[Dict[str, Any]] = field(default_factory=list) + sequence_number: int = 0 + started: bool = False + assistant_meta: Dict[str, Any] = field(default_factory=dict) + reasoning_item: Optional["OpenAIFormatter.ResponsesOutputItem"] = None + message_item: Optional["OpenAIFormatter.ResponsesOutputItem"] = None + tool_items: Dict[int, "OpenAIFormatter.ResponsesOutputItem"] = field( + default_factory=dict + ) + final_status: Optional[str] = None + incomplete_details: Optional[Dict[str, Any]] = None + + def __init__(self, model: Model) -> None: + self.model = model + + @staticmethod + def decode_text(data: bytes) -> str: + return data.decode("utf-8", errors="ignore") + + @staticmethod + def encode_sse_payload(payload: BaseModel | Dict[str, Any]) -> bytes: + data = ( + payload.model_dump(mode="json", exclude_none=True) + if isinstance(payload, BaseModel) + else payload + ) + return ( + "data: " + f"{json.dumps(data, ensure_ascii=False, separators=(',', ':'))}\n\n" + ).encode("utf-8") + + @staticmethod + def next_stream_chunk( + stream: Iterator[CompletionChunk], + ) -> Tuple[bool, Optional[CompletionChunk]]: + try: + return False, next(stream) + except StopIteration: + return True, None + + @staticmethod + def next_stream_output( + stream: Iterator[CompletionChunk], + ) -> Tuple[bool, Optional[CompletionChunk], Optional[OpenAICompletion]]: + try: + return False, next(stream), None + except StopIteration as stop: + return True, None, cast(Optional[OpenAICompletion], stop.value) + + @staticmethod + def collect_completion(stream: Iterator[Any]) -> OpenAICompletion: + iterator = iter(stream) + while True: + try: + next(iterator) + except StopIteration as stop: + result = stop.value + if result is None: + raise RuntimeError("missing completion result") + return cast(OpenAICompletion, result) + + @staticmethod + def _tools_for_response_parser( + *, + functions: Optional[List[ChatTemplateFunctionDefinition]] = None, + tools: Optional[List[ChatTemplateTool]] = None, + ) -> Optional[List[ChatTemplateTool]]: + if functions is not None: + return [ + { + "type": "function", + "function": function, + } + for function in functions + ] + return tools + + def _chat_template_text(self) -> str: + if self.model.chat_formatter is None: + return "" + return cast(str, self.model.chat_formatter._template_text) + + def _uses_harmony_channels(self) -> bool: + template = self._chat_template_text() + return "<|channel|>" in template or "<|recipient|>" in template + + def _uses_reasoning_content(self) -> bool: + template = self._chat_template_text() + return "reasoning_content" in template or "<think>" in template + + @staticmethod + def _chat_message(data: Dict[str, Any]) -> ChatCompletionRequestMessage: + return ChatCompletionRequestMessage.model_validate(data) + + def _instructions_role(self) -> Literal["developer", "system"]: + return "developer" if "developer" in self._chat_template_text() else "system" + + @staticmethod + def _response_reasoning_effort( + body: CreateResponseRequest, + ) -> Optional[str]: + reasoning = body.reasoning + if reasoning is None: + return None + effort = reasoning.effort + if effort in {"low", "medium", "high"}: + return cast(str, effort) + return None + + @staticmethod + def _response_text_from_content(content: Any) -> str: + if isinstance(content, str): + return content + if not isinstance(content, list): + raise CompletionRequestValidationError("responses input content must be a string or list") + parts: List[str] = [] + for part in content: + if not isinstance(part, dict): + raise CompletionRequestValidationError("responses content parts must be objects") + part_type = part.get("type") + if part_type in {"input_text", "output_text", "reasoning_text", "summary_text", "text"}: + text = part.get("text") + if not isinstance(text, str): + raise CompletionRequestValidationError( + f"responses content part {part_type!r} requires string text" + ) + parts.append(text) + continue + raise CompletionRequestValidationError( + f"unsupported responses content part type: {part_type!r}" + ) + return "".join(parts) + + @staticmethod + def _response_chat_content_from_content( + content: Any, + ) -> Union[str, List[Dict[str, Any]]]: + if isinstance(content, str): + return content + if not isinstance(content, list): + raise CompletionRequestValidationError("responses input content must be a string or list") + parts: List[Dict[str, Any]] = [] + for part in content: + if not isinstance(part, dict): + raise CompletionRequestValidationError("responses content parts must be objects") + part_type = part.get("type") + if part_type in {"input_text", "output_text", "reasoning_text", "summary_text", "text"}: + text = part.get("text") + if not isinstance(text, str): + raise CompletionRequestValidationError( + f"responses content part {part_type!r} requires string text" + ) + parts.append({"type": "text", "text": text}) + continue + if part_type == "input_image": + image_url = part.get("image_url") + if not isinstance(image_url, str): + raise CompletionRequestValidationError( + "responses input_image content part requires image_url" + ) + parts.append({"type": "image_url", "image_url": {"url": image_url}}) + continue + if part_type == "input_audio": + input_audio = part.get("input_audio") + if isinstance(input_audio, dict): + data = input_audio.get("data") + audio_format = input_audio.get("format") + else: + data = part.get("data") + audio_format = part.get("format") + if not isinstance(data, str): + raise CompletionRequestValidationError( + "responses input_audio content part requires base64 data" + ) + if audio_format is not None and not isinstance(audio_format, str): + raise CompletionRequestValidationError( + "responses input_audio format must be a string" + ) + audio_part: Dict[str, Any] = {"data": data} + if audio_format is not None: + audio_part["format"] = audio_format + parts.append({"type": "input_audio", "input_audio": audio_part}) + continue + if part_type == "audio_url": + audio_url = part.get("audio_url") + if not isinstance(audio_url, str): + raise CompletionRequestValidationError( + "responses audio_url content part requires audio_url" + ) + parts.append({"type": "audio_url", "audio_url": {"url": audio_url}}) + continue + raise CompletionRequestValidationError( + f"unsupported responses content part type: {part_type!r}" + ) + if all(part.get("type") == "text" for part in parts): + return "".join(cast(str, part.get("text", "")) for part in parts) + return parts + + @staticmethod + def _response_reasoning_text(item: Dict[str, Any]) -> str: + content = item.get("content") + if isinstance(content, list) and content: + return OpenAIFormatter._response_text_from_content(content) + summary = item.get("summary") + if isinstance(summary, list) and summary: + return OpenAIFormatter._response_text_from_content(summary) + return "" + + def _response_reasoning_input_message( + self, + *, + text: str, + ) -> ChatCompletionRequestMessage: + if self._uses_harmony_channels(): + return self._chat_message( + { + "role": "assistant", + "content": text, + "channel": "analysis", + } + ) + if self._uses_reasoning_content(): + return self._chat_message( + { + "role": "assistant", + "content": "", + "reasoning_content": text, + } + ) + return self._chat_message( + { + "role": "assistant", + "content": text, + } + ) + + @staticmethod + def _response_tool_call_input_message( + *, + name: str, + arguments: str, + call_id: str, + content_type: Optional[str] = None, + ) -> ChatCompletionRequestMessage: + function: Dict[str, Any] = { + "name": name, + "arguments": arguments, + } + if isinstance(content_type, str) and content_type: + function["content_type"] = content_type + tool_call = { + "id": call_id, + "type": "function", + "function": function, + } + message: Dict[str, Any] = { + "role": "assistant", + "content": None, + "tool_calls": [tool_call], + "function_call": dict(function), + } + return OpenAIFormatter._chat_message(message) + + def _response_function_call_input_message( + self, + item: Dict[str, Any], + ) -> ChatCompletionRequestMessage: + name = item.get("name") + if not isinstance(name, str) or not name: + raise CompletionRequestValidationError("function_call input requires name") + arguments = item.get("arguments", "") + if not isinstance(arguments, str): + raise CompletionRequestValidationError("function_call input requires string arguments") + call_id = item.get("call_id") or item.get("id") or f"call_{uuid.uuid4().hex}" + return self._response_tool_call_input_message( + name=name, + arguments=arguments, + call_id=call_id, + content_type="json", + ) + + def _response_custom_tool_call_input_message( + self, + item: Dict[str, Any], + ) -> ChatCompletionRequestMessage: + name = item.get("name") + if not isinstance(name, str) or not name: + raise CompletionRequestValidationError("custom_tool_call input requires name") + input_text = item.get("input", "") + if not isinstance(input_text, str): + raise CompletionRequestValidationError("custom_tool_call input requires string input") + call_id = item.get("call_id") or item.get("id") or f"call_{uuid.uuid4().hex}" + content_type = item.get("content_type") + if not isinstance(content_type, str) or not content_type: + content_type = "text" + return self._response_tool_call_input_message( + name=name, + arguments=input_text, + call_id=call_id, + content_type=content_type, + ) + + def responses_input_to_chat_messages( + self, + body: CreateResponseRequest, + ) -> List[ChatCompletionRequestMessage]: + if body.conversation is not None: + raise CompletionRequestValidationError( + "conversation is not supported in stateless /v1/responses" + ) + if body.store: + raise CompletionRequestValidationError( + "store=true is not supported in stateless /v1/responses" + ) + if body.truncation not in {None, "disabled"}: + raise CompletionRequestValidationError( + "only truncation='disabled' is supported" + ) + + messages: List[ChatCompletionRequestMessage] = [] + if body.instructions is not None: + messages.append( + self._chat_message( + { + "role": self._instructions_role(), + "content": body.instructions, + } + ) + ) + + if isinstance(body.input, str): + messages.append( + self._chat_message( + { + "role": "user", + "content": body.input, + } + ) + ) + return messages + + items = body.input + if isinstance(items, dict): + items = [items] + if not isinstance(items, list): + raise CompletionRequestValidationError( + "responses input must be a string, object, or list" + ) + + function_names_by_call_id: Dict[str, str] = {} + for item in items: + if not isinstance(item, dict): + raise CompletionRequestValidationError( + "responses input items must be objects" + ) + item_type = item.get("type") + if item_type is None and "role" in item: + item_type = "message" + if item_type == "message": + role = item.get("role", "user") + if role not in { + "user", + "assistant", + "system", + "developer", + "tool", + "function", + }: + raise CompletionRequestValidationError( + f"unsupported responses message role: {role!r}" + ) + if role == "function": + role = "tool" + data: Dict[str, Any] = { + "role": role, + "content": self._response_chat_content_from_content( + item.get("content", "") + ), + } + phase = item.get("phase") + if isinstance(phase, str): + data["phase"] = phase + if role == "assistant" and self._uses_harmony_channels(): + if phase == "commentary": + data["channel"] = "commentary" + elif phase == "final_answer": + data["channel"] = "final" + messages.append(self._chat_message(data)) + continue + if item_type == "reasoning": + text = self._response_reasoning_text(item) + if text: + messages.append(self._response_reasoning_input_message(text=text)) + continue + if item_type == "function_call": + message = self._response_function_call_input_message(item) + call_id = item.get("call_id") or item.get("id") + name = item.get("name") + if isinstance(call_id, str) and isinstance(name, str): + function_names_by_call_id[call_id] = name + messages.append(message) + continue + if item_type == "custom_tool_call": + message = self._response_custom_tool_call_input_message(item) + call_id = item.get("call_id") or item.get("id") + name = item.get("name") + if isinstance(call_id, str) and isinstance(name, str): + function_names_by_call_id[call_id] = name + messages.append(message) + continue + if item_type == "function_call_output": + call_id = item.get("call_id") + if not isinstance(call_id, str) or not call_id: + raise CompletionRequestValidationError( + "function_call_output input requires call_id" + ) + tool_output_data: Dict[str, Any] = { + "role": "tool", + "tool_call_id": call_id, + "content": self._response_text_from_content(item.get("output", "")), + } + name = function_names_by_call_id.get(call_id) + if name is not None: + tool_output_data["name"] = name + messages.append(self._chat_message(tool_output_data)) + continue + if item_type == "custom_tool_call_output": + call_id = item.get("call_id") + if not isinstance(call_id, str) or not call_id: + raise CompletionRequestValidationError( + "custom_tool_call_output input requires call_id" + ) + tool_output_data = { + "role": "tool", + "tool_call_id": call_id, + "content": self._response_text_from_content(item.get("output", "")), + } + name = function_names_by_call_id.get(call_id) + if name is not None: + tool_output_data["name"] = name + messages.append(self._chat_message(tool_output_data)) + continue + raise CompletionRequestValidationError( + f"unsupported responses input item type: {item_type!r}" + ) + return messages + + @staticmethod + def _clone_response_input_items(input_items: Any) -> List[Any]: + if isinstance(input_items, list): + return copy.deepcopy(input_items) + if isinstance(input_items, dict): + return [copy.deepcopy(input_items)] + return [copy.deepcopy(input_items)] + + @staticmethod + def _normalize_response_output_item_for_input( + item: Dict[str, Any], + ) -> Optional[Dict[str, Any]]: + item_type = item.get("type") + if item_type == "message": + role = item.get("role") + if not isinstance(role, str): + return None + normalized: Dict[str, Any] = { + "type": "message", + "role": role, + "content": copy.deepcopy(item.get("content", [])), + } + for key in ("id", "phase", "status"): + value = item.get(key) + if value is not None: + normalized[key] = copy.deepcopy(value) + return normalized + if item_type == "reasoning": + normalized = { + "type": "reasoning", + "content": copy.deepcopy(item.get("content", [])), + } + for key in ("id", "summary", "status"): + value = item.get(key) + if value is not None: + normalized[key] = copy.deepcopy(value) + return normalized + if item_type == "function_call": + normalized = { + "type": "function_call", + "call_id": copy.deepcopy(item.get("call_id")), + "name": copy.deepcopy(item.get("name")), + "arguments": copy.deepcopy(item.get("arguments", "")), + } + namespace = item.get("namespace") + item_id = item.get("id") + if item_id is not None: + normalized["id"] = copy.deepcopy(item_id) + if namespace is not None: + normalized["namespace"] = copy.deepcopy(namespace) + return normalized + if item_type == "custom_tool_call": + normalized = { + "type": "custom_tool_call", + "call_id": copy.deepcopy(item.get("call_id")), + "name": copy.deepcopy(item.get("name")), + "input": copy.deepcopy(item.get("input", "")), + } + namespace = item.get("namespace") + content_type = item.get("content_type") + item_id = item.get("id") + if item_id is not None: + normalized["id"] = copy.deepcopy(item_id) + if namespace is not None: + normalized["namespace"] = copy.deepcopy(namespace) + if content_type is not None: + normalized["content_type"] = copy.deepcopy(content_type) + return normalized + return None + + def _responses_tools_to_chat_tools( + self, + tools: Optional[List[ResponsesToolDefinition]], + ) -> Optional[List[ChatTemplateTool]]: + if tools is None: + return None + chat_tools: List[ChatTemplateTool] = [] + for tool in tools: + if isinstance( + tool, + ( + ResponsesWebSearchTool, + ResponsesNamespaceTool, + ResponsesImageGenerationTool, + ), + ): + continue + if isinstance(tool, ResponsesFunctionTool): + chat_tools.append(tool.to_chat_template_tool()) + continue + if isinstance(tool, ResponsesCustomTool): + chat_tools.append(tool.to_chat_template_tool()) + continue + raise CompletionRequestValidationError( + f"unsupported responses tool type: {tool.type!r}" + ) + return chat_tools + + @staticmethod + def _responses_tool_choice_to_chat_tool_choice( + tool_choice: Optional[ResponsesToolChoice], + ) -> Optional[Union[Literal["auto", "none", "required"], ChatTemplateToolChoice]]: + if tool_choice is None or isinstance(tool_choice, str): + return tool_choice + if tool_choice.type in {"function", "custom"}: + return tool_choice.to_chat_template_tool_choice() + raise CompletionRequestValidationError( + f"unsupported responses tool_choice type: {tool_choice.type!r}" + ) + + @staticmethod + def _response_format_type(response_format: Optional[ChatTemplateResponseFormat]) -> Optional[str]: + if response_format is None: + return None + format_type = response_format.get("type") + if isinstance(format_type, str): + return format_type + return None + + @staticmethod + def _grammar_for_response_format( + response_format: Optional[ChatTemplateResponseFormat], + ) -> Optional[str]: + format_type = OpenAIFormatter._response_format_type(response_format) + if format_type is None or format_type == "text": + return None + if format_type == "json_object": + return JSON_GBNF + if format_type == "json_schema": + assert response_format is not None + schema = response_format.get("schema") + if schema is None and isinstance(response_format.get("json_schema"), dict): + schema = cast(Dict[str, Any], response_format["json_schema"]).get("schema") + if not isinstance(schema, dict): + raise CompletionRequestValidationError( + "json_schema response format requires a schema object" + ) + return JsonSchemaConverter.to_gbnf( + json.dumps(schema, ensure_ascii=False, separators=(",", ":")) + ) + raise CompletionRequestValidationError( + f"unsupported response format type: {format_type!r}" + ) + + def response_request_to_chat_parts( + self, + body: CreateResponseRequest, + ) -> ResponsesChatRequestParts: + chat_tools = self._responses_tools_to_chat_tools(body.tools) + response_format = ( + None + if body.text is None or body.text.format is None + else body.text.format.to_chat_response_format() + ) + return ResponsesChatRequestParts( + messages=self.responses_input_to_chat_messages(body), + max_tokens=body.max_output_tokens, + temperature=0.8 if body.temperature is None else body.temperature, + top_p=0.95 if body.top_p is None else body.top_p, + stream=body.stream, + logprobs=body.top_logprobs is not None, + top_logprobs=body.top_logprobs, + model=body.model, + user=body.user, + tools=chat_tools, + tool_choice=self._responses_tool_choice_to_chat_tool_choice(body.tool_choice), + response_format=response_format, + reasoning_effort=self._response_reasoning_effort(body), + ) + + def _response_parser( + self, + *, + tools: Optional[List[ChatTemplateTool]] = None, + completion_id: str = "", + choice_index: int = 0, + generation_prompt: str = "", + ) -> ResponseParser: + if self.model.response_schema is None: + raise CompletionResponseParsingError("model does not define response_schema") + return ResponseParser( + self.model.response_schema, + tools=tools, + completion_id=completion_id, + choice_index=choice_index, + generation_prompt=generation_prompt, + ) + + def parse_chat_response( + self, + response_text: str, + *, + tools: Optional[List[ChatTemplateTool]] = None, + partial: bool, + generation_prompt: str = "", + ) -> Dict[str, Any]: + return self._response_parser( + tools=tools, + generation_prompt=generation_prompt, + ).parse_chat_response( + response_text, + partial=partial, + ) + + def _chat_tool_name_and_grammar( + self, + *, + tools: Optional[List[ChatTemplateTool]], + function_call: Optional[Union[Literal["none", "auto"], ChatTemplateFunctionCall]], + tool_choice: Optional[Union[Literal["none", "auto", "required"], ChatTemplateToolChoice]], + response_format: Optional[ChatTemplateResponseFormat], + ) -> Tuple[Optional[str], Optional[str]]: + selected_tool_choice: Optional[Union[str, ChatTemplateToolChoice]] = tool_choice + if function_call is not None: + if isinstance(function_call, str): + selected_tool_choice = function_call + else: + selected_tool_choice = { + "type": "function", + "function": { + "name": function_call["name"], + }, + } + grammar_text = self._grammar_for_response_format(response_format) + if not isinstance(selected_tool_choice, dict): + return None, grammar_text + if tools is None: + raise CompletionRequestValidationError("tool choice requires tools") + tool_name = selected_tool_choice["function"]["name"] + tool = next((tool for tool in tools if tool["function"]["name"] == tool_name), None) + if tool is None: + raise CompletionRequestValidationError( + f"Tool choice '{tool_name}' not found in tools." + ) + if self.model.response_schema is None: + return tool_name, JSON_GBNF + return tool_name, grammar_text + + def completion_request_from_chat_request( + self, + body: CreateChatCompletionRequest, + ) -> PreparedCompletionParts: + functions = ( + [function.to_template_function() for function in body.functions] + if body.functions is not None + else None + ) + tools = ( + [tool.to_template_tool() for tool in body.tools] + if body.tools is not None + else None + ) + function_call = ( + body.function_call + if body.function_call is None or isinstance(body.function_call, str) + else body.function_call.to_template_function_call() + ) + tool_choice = ( + body.tool_choice + if body.tool_choice is None or isinstance(body.tool_choice, str) + else body.tool_choice.to_template_tool_choice() + ) + response_format = ( + body.response_format.to_template_response_format() + if body.response_format is not None + else None + ) + parser_tools = self._tools_for_response_parser( + functions=functions, + tools=tools, + ) + try: + prompt_text, generation_prompt, prompt_plan, formatter_stop = self.model.build_chat_prompt( + body.messages, + functions=functions, + function_call=function_call, + tools=tools, + tool_choice=tool_choice, + reasoning_effort=body.reasoning_effort, + ) + tool_name, grammar_text = self._chat_tool_name_and_grammar( + tools=parser_tools, + function_call=function_call, + tool_choice=tool_choice, + response_format=response_format, + ) + except ValueError as exc: + raise CompletionRequestValidationError(str(exc)) from exc + request_stop: List[str] = [] + if body.stop is None: + request_stop = [] + elif isinstance(body.stop, str): + request_stop = [body.stop] + else: + request_stop = list(body.stop) + stop_sequences = [stop for stop in [*request_stop, *formatter_stop] if stop] + deduped_stop: List[str] = [] + seen_stop: set[str] = set() + for stop in stop_sequences: + if stop not in seen_stop: + deduped_stop.append(stop) + seen_stop.add(stop) + payload = CreateCompletionRequest( + prompt=prompt_text, + max_tokens=body.max_tokens, + temperature=body.temperature, + top_p=body.top_p, + echo=False, + stop=deduped_stop or None, + stream=body.stream, + logprobs=( + 0 if body.logprobs and body.top_logprobs is None else body.top_logprobs + ), + presence_penalty=body.presence_penalty, + frequency_penalty=body.frequency_penalty, + logit_bias=body.logit_bias, + seed=body.seed, + model=body.model, + n=body.n, + user=body.user, + ) + return PreparedCompletionParts( + payload=payload, + prompt_text=prompt_text, + generation_prompt=generation_prompt, + prompt_plan=prompt_plan, + grammar_text=grammar_text, + tool_name=tool_name, + ) + + def completion_request_from_response_chat_parts( + self, + body: ResponsesChatRequestParts, + ) -> PreparedCompletionParts: + try: + prompt_text, generation_prompt, prompt_plan, formatter_stop = self.model.build_chat_prompt( + body.messages, + tools=body.tools, + tool_choice=body.tool_choice, + reasoning_effort=body.reasoning_effort, + ) + tool_name, grammar_text = self._chat_tool_name_and_grammar( + tools=body.tools, + function_call=None, + tool_choice=body.tool_choice, + response_format=body.response_format, + ) + except ValueError as exc: + raise CompletionRequestValidationError(str(exc)) from exc + deduped_stop: List[str] = [] + seen_stop: set[str] = set() + for stop in formatter_stop: + if stop and stop not in seen_stop: + deduped_stop.append(stop) + seen_stop.add(stop) + payload = CreateCompletionRequest( + prompt=prompt_text, + max_tokens=body.max_tokens, + temperature=body.temperature, + top_p=body.top_p, + echo=False, + stop=deduped_stop or None, + stream=body.stream, + logprobs=( + 0 if body.logprobs and body.top_logprobs is None else body.top_logprobs + ), + seed=None, + model=body.model, + n=1, + user=body.user, + ) + return PreparedCompletionParts( + payload=payload, + prompt_text=prompt_text, + generation_prompt=generation_prompt, + prompt_plan=prompt_plan, + grammar_text=grammar_text, + tool_name=tool_name, + ) + + @staticmethod + def _response_phase_from_message(message: Dict[str, Any]) -> Optional[str]: + phase = message.get("phase") + if phase in {"commentary", "final_answer"}: + return cast(str, phase) + channel = message.get("channel") + if channel == "commentary": + return "commentary" + if channel == "final": + return "final_answer" + return None + + @staticmethod + def _response_reasoning_text_from_message(message: Dict[str, Any]) -> str: + thinking = message.get("thinking") + if isinstance(thinking, str) and thinking: + return thinking + reasoning_content = message.get("reasoning_content") + if isinstance(reasoning_content, str) and reasoning_content: + return reasoning_content + if message.get("channel") == "analysis": + content = message.get("content") + if isinstance(content, str): + return content + return "" + + @staticmethod + def _response_output_text_from_message(message: Dict[str, Any]) -> str: + if message.get("channel") == "analysis": + return "" + content = message.get("content") + if isinstance(content, str): + return content + return "" + + @staticmethod + def _response_logprobs_from_chat_logprobs( + logprobs: Optional[Dict[str, Any]], + ) -> List[Dict[str, Any]]: + if logprobs is None: + return [] + content = logprobs.get("content") + if not isinstance(content, list): + return [] + response_logprobs: List[Dict[str, Any]] = [] + for entry in content: + if not isinstance(entry, dict): + continue + token = entry.get("token") + logprob = entry.get("logprob") + if not isinstance(token, str) or not isinstance(logprob, (int, float)): + continue + top_logprobs = entry.get("top_logprobs") + converted_top: List[Dict[str, Any]] = [] + if isinstance(top_logprobs, list): + for top in top_logprobs: + if not isinstance(top, dict): + continue + converted_top.append( + { + "token": top.get("token"), + "logprob": top.get("logprob"), + } + ) + response_logprobs.append( + { + "token": token, + "logprob": float(logprob), + "top_logprobs": converted_top or None, + } + ) + return response_logprobs + + def _response_logprobs_from_completion( + self, + logprobs: Optional[CompletionLogprobs], + ) -> List[Dict[str, Any]]: + if logprobs is None: + return [] + response_logprobs: List[Dict[str, Any]] = [] + tokens = logprobs.tokens or [] + token_logprobs_list = logprobs.token_logprobs or [] + top_logprobs_list = logprobs.top_logprobs or [] + for token, token_logprob, top_logprobs in zip( + tokens, + token_logprobs_list, + top_logprobs_list, + ): + if token_logprob is None: + continue + converted_top: List[Dict[str, Any]] = [] + if top_logprobs is not None: + for top_token, top_logprob in top_logprobs.items(): + converted_top.append( + { + "token": top_token, + "logprob": float(top_logprob), + } + ) + response_logprobs.append( + { + "token": token, + "logprob": float(token_logprob), + "top_logprobs": converted_top or None, + } + ) + return response_logprobs + + @staticmethod + def _response_reasoning_item( + *, + item_id: str, + text: str, + status: str, + ) -> Dict[str, Any]: + return { + "id": item_id, + "type": "reasoning", + "summary": [], + "content": [ + { + "type": "reasoning_text", + "text": text, + } + ], + "status": status, + } + + @staticmethod + def _response_message_item( + *, + item_id: str, + text: str, + status: str, + phase: Optional[str], + logprobs: Optional[List[Dict[str, Any]]] = None, + ) -> Dict[str, Any]: + content_item: Dict[str, Any] = { + "type": "output_text", + "text": text, + "annotations": [], + } + if logprobs: + content_item["logprobs"] = logprobs + item: Dict[str, Any] = { + "id": item_id, + "type": "message", + "role": "assistant", + "content": [content_item], + "status": status, + } + if phase is not None: + item["phase"] = phase + return item + + @staticmethod + def _response_function_call_item( + *, + item_id: str, + call_id: str, + name: str, + arguments: str, + status: str, + ) -> Dict[str, Any]: + item: Dict[str, Any] = { + "id": item_id, + "type": "function_call", + "call_id": call_id, + "name": name, + "arguments": arguments, + "status": status, + } + if "." in name: + namespace, bare_name = name.split(".", 1) + item["namespace"] = namespace + item["name"] = bare_name + return item + + @staticmethod + def _response_custom_tool_call_item( + *, + item_id: str, + call_id: str, + name: str, + input_text: str, + ) -> Dict[str, Any]: + item: Dict[str, Any] = { + "id": item_id, + "type": "custom_tool_call", + "call_id": call_id, + "name": name, + "input": input_text, + } + if "." in name: + namespace, bare_name = name.split(".", 1) + item["namespace"] = namespace + item["name"] = bare_name + return item + + @staticmethod + def _responses_tool_type_by_name( + tools: Optional[List[Any]], + ) -> Dict[str, str]: + if tools is None: + return {} + tool_types: Dict[str, str] = {} + for tool in tools: + if isinstance(tool, BaseModel): + tool = tool.model_dump(mode="python", exclude_none=True) + if not isinstance(tool, dict): + continue + tool_type = tool.get("original_type") or tool.get("type") + if not isinstance(tool_type, str): + continue + function = tool.get("function") + if isinstance(function, dict): + name = function.get("name") + else: + name = tool.get("name") + if isinstance(name, str) and name: + tool_types[name] = tool_type + return tool_types + + def _response_output_items_from_message( + self, + *, + response_id: str, + message: Dict[str, Any], + logprobs: Optional[List[Dict[str, Any]]] = None, + tools: Optional[List[ChatTemplateTool]] = None, + ) -> List[Dict[str, Any]]: + items: List[Dict[str, Any]] = [] + tool_types_by_name = self._responses_tool_type_by_name(tools) + reasoning_text = self._response_reasoning_text_from_message(message) + if reasoning_text: + items.append( + self._response_reasoning_item( + item_id=f"rs_{response_id}_0", + text=reasoning_text, + status="completed", + ) + ) + tool_calls = message.get("tool_calls") + if isinstance(tool_calls, list): + for tool_call_index, tool_call in enumerate(tool_calls): + if not isinstance(tool_call, dict): + continue + function = tool_call.get("function") + if not isinstance(function, dict): + continue + name = function.get("name") + arguments = function.get("arguments") + if not isinstance(name, str) or not isinstance(arguments, str): + continue + call_id = tool_call.get("id") + if not isinstance(call_id, str) or not call_id: + call_id = f"call_{response_id}_{tool_call_index}" + tool_type = tool_types_by_name.get(name, "function") + if tool_type == "custom": + items.append( + self._response_custom_tool_call_item( + item_id=f"fc_{response_id}_{tool_call_index}", + call_id=call_id, + name=name, + input_text=arguments, + ) + ) + continue + items.append( + self._response_function_call_item( + item_id=f"fc_{response_id}_{tool_call_index}", + call_id=call_id, + name=name, + arguments=arguments, + status="completed", + ) + ) + output_text = self._response_output_text_from_message(message) + if output_text: + items.append( + self._response_message_item( + item_id=f"msg_{response_id}_0", + text=output_text, + status="completed", + phase=self._response_phase_from_message(message), + logprobs=logprobs, + ) + ) + return items + + def _responses_usage_from_completion( + self, + *, + usage: Optional[CompletionUsage], + output_items: Sequence[Dict[str, Any]], + ) -> Optional[Dict[str, Any]]: + if usage is None: + return None + reasoning_tokens = 0 + for item in output_items: + if item.get("type") != "reasoning": + continue + content = item.get("content") + if not isinstance(content, list): + continue + for part in content: + if isinstance(part, dict) and part.get("type") == "reasoning_text": + text = part.get("text") + if isinstance(text, str) and text: + reasoning_tokens += len( + self.model.tokenize(text, add_bos=False, special=True) + ) + return { + "input_tokens": usage.prompt_tokens, + "input_tokens_details": { + "cached_tokens": 0, + }, + "output_tokens": usage.completion_tokens, + "output_tokens_details": { + "reasoning_tokens": min(reasoning_tokens, usage.completion_tokens), + }, + "total_tokens": usage.total_tokens, + } + + def _response_status_and_incomplete_details( + self, + *, + finish_reason: Optional[str], + ) -> Tuple[str, Optional[Dict[str, str]]]: + if finish_reason == "length": + return "incomplete", {"reason": "max_output_tokens"} + return "completed", None + + def _response_object( + self, + *, + body: CreateResponseRequest, + response_id: str, + created_at: float, + model: str, + output_items: Sequence[Dict[str, Any]], + usage: Optional[Dict[str, Any]], + status: str, + incomplete_details: Optional[Dict[str, str]], + completed_at: Optional[float], + ) -> Dict[str, Any]: + return { + "id": response_id, + "object": "response", + "created_at": created_at, + "completed_at": completed_at, + "error": None, + "incomplete_details": incomplete_details, + "instructions": body.instructions, + "metadata": body.metadata, + "model": model, + "output": list(output_items), + "parallel_tool_calls": body.parallel_tool_calls, + "reasoning": { + "effort": self._response_reasoning_effort(body), + "summary": None, + }, + "store": False, + "temperature": body.temperature, + "tool_choice": ( + body.tool_choice.model_dump(mode="python", exclude_none=True) + if isinstance(body.tool_choice, BaseModel) + else body.tool_choice or "auto" + ), + "tools": ( + [ + tool.model_dump(mode="python", exclude_none=True) + if isinstance(tool, BaseModel) + else tool + for tool in body.tools + ] + if body.tools + else [] + ), + "top_p": body.top_p, + "max_output_tokens": body.max_output_tokens, + "previous_response_id": None, + "status": status, + "text": ( + body.text.model_dump(mode="python", exclude_none=True, by_alias=True) + if body.text is not None + else {"format": {"type": "text"}} + ), + "top_logprobs": body.top_logprobs, + "truncation": body.truncation, + "usage": usage, + "user": body.user, + } + + def _response_from_chat_message( + self, + *, + body: CreateResponseRequest, + response_id: str, + created_at: float, + model: str, + message: Dict[str, Any], + finish_reason: Optional[str], + usage: Optional[CompletionUsage], + logprobs: Optional[List[Dict[str, Any]]] = None, + ) -> Dict[str, Any]: + status, incomplete_details = self._response_status_and_incomplete_details( + finish_reason=finish_reason, + ) + output_items = self._response_output_items_from_message( + response_id=response_id, + message=message, + logprobs=logprobs, + tools=cast(Optional[List[Any]], body.tools), + ) + return self._response_object( + body=body, + response_id=response_id, + created_at=created_at, + model=model, + output_items=output_items, + usage=self._responses_usage_from_completion( + usage=usage, + output_items=output_items, + ), + status=status, + incomplete_details=incomplete_details, + completed_at=time.time() if status in {"completed", "incomplete"} else None, + ) + + def convert_completion_response_to_response( + self, + completion: OpenAICompletion, + body: CreateResponseRequest, + tool_name: Optional[str] = None, + *, + tools: Optional[List[ChatTemplateTool]] = None, + generation_prompt: str = "", + ) -> Dict[str, Any]: + chat_response = self.convert_completion_response_to_chat( + completion, + tool_name, + tools=tools, + generation_prompt=generation_prompt, + ) + if isinstance(chat_response, BaseModel): + chat_payload = chat_response.model_dump(mode="json", exclude_none=True) + else: + chat_payload = chat_response + choice = cast(Dict[str, Any], chat_payload["choices"][0]) + message = cast(Dict[str, Any], choice["message"]) + response_id = "resp_" + completion.id + logprobs = self._response_logprobs_from_completion(completion.choices[0].logprobs) + return self._response_from_chat_message( + body=body, + response_id=response_id, + created_at=float(completion.created), + model=completion.model, + message=message, + finish_reason=cast(Optional[str], choice.get("finish_reason")), + usage=completion.usage, + logprobs=logprobs or None, + ) + + @staticmethod + def _response_event( + state: "OpenAIFormatter.ResponsesStream", + event_type: str, + **payload: Any, + ) -> Dict[str, Any]: + state.sequence_number += 1 + event = { + "type": event_type, + "sequence_number": state.sequence_number, + } + event.update(payload) + return event + + def _response_stream_response( + self, + state: "OpenAIFormatter.ResponsesStream", + *, + status: str, + usage: Optional[Dict[str, Any]], + incomplete_details: Optional[Dict[str, str]], + completed_at: Optional[float], + ) -> Dict[str, Any]: + return self._response_object( + body=state.body, + response_id=state.response_id, + created_at=state.created_at, + model=state.model, + output_items=state.output, + usage=usage, + status=status, + incomplete_details=incomplete_details, + completed_at=completed_at, + ) + + def start_response_stream( + self, + state: "OpenAIFormatter.ResponsesStream", + ) -> List[Dict[str, Any]]: + if state.started: + return [] + state.started = True + response = self._response_stream_response( + state, + status="in_progress", + usage=None, + incomplete_details=None, + completed_at=None, + ) + return [ + self._response_event(state, "response.created", response=response), + self._response_event(state, "response.in_progress", response=response), + ] + + @staticmethod + def _response_item_status( + finish_reason: Optional[str], + ) -> str: + if finish_reason == "length": + return "incomplete" + return "completed" + + def _add_response_stream_item( + self, + state: "OpenAIFormatter.ResponsesStream", + item: Dict[str, Any], + *, + content_index: Optional[int] = None, + ) -> "OpenAIFormatter.ResponsesOutputItem": + item_state = OpenAIFormatter.ResponsesOutputItem( + output_index=len(state.output), + item=item, + content_index=content_index, + ) + state.output.append(item) + return item_state + + def _ensure_reasoning_stream_item( + self, + state: "OpenAIFormatter.ResponsesStream", + ) -> Tuple[List[Dict[str, Any]], "OpenAIFormatter.ResponsesOutputItem"]: + if state.reasoning_item is not None: + return [], state.reasoning_item + item = self._response_reasoning_item( + item_id=f"rs_{state.response_id}_{len(state.output)}", + text="", + status="in_progress", + ) + item_state = self._add_response_stream_item(state, item, content_index=0) + state.reasoning_item = item_state + part = cast(List[Dict[str, Any]], item["content"])[0] + return [ + self._response_event( + state, + "response.output_item.added", + output_index=item_state.output_index, + item=copy.deepcopy(item), + ), + self._response_event( + state, + "response.content_part.added", + item_id=cast(str, item["id"]), + output_index=item_state.output_index, + content_index=0, + part=copy.deepcopy(part), + ), + ], item_state + + def _ensure_message_stream_item( + self, + state: "OpenAIFormatter.ResponsesStream", + ) -> Tuple[List[Dict[str, Any]], "OpenAIFormatter.ResponsesOutputItem"]: + if state.message_item is not None: + return [], state.message_item + item = self._response_message_item( + item_id=f"msg_{state.response_id}_{len(state.output)}", + text="", + status="in_progress", + phase=cast(Optional[str], state.assistant_meta.get("phase")), + ) + item_state = self._add_response_stream_item(state, item, content_index=0) + state.message_item = item_state + part = cast(List[Dict[str, Any]], item["content"])[0] + return [ + self._response_event( + state, + "response.output_item.added", + output_index=item_state.output_index, + item=copy.deepcopy(item), + ), + self._response_event( + state, + "response.content_part.added", + item_id=cast(str, item["id"]), + output_index=item_state.output_index, + content_index=0, + part=copy.deepcopy(part), + ), + ], item_state + + def _ensure_tool_stream_item( + self, + state: "OpenAIFormatter.ResponsesStream", + *, + tool_call_index: int, + call_id: Optional[str], + name: Optional[str], + ) -> Tuple[List[Dict[str, Any]], "OpenAIFormatter.ResponsesOutputItem"]: + existing = state.tool_items.get(tool_call_index) + if existing is not None: + return [], existing + tool_types_by_name = self._responses_tool_type_by_name(state.body.tools) + tool_type = tool_types_by_name.get(name or "", "function") + item_id = f"fc_{state.response_id}_{tool_call_index}" + resolved_call_id = call_id or f"call_{state.response_id}_{tool_call_index}" + resolved_name = name or "" + if tool_type == "custom": + item = self._response_custom_tool_call_item( + item_id=item_id, + call_id=resolved_call_id, + name=resolved_name, + input_text="", + ) + else: + item = self._response_function_call_item( + item_id=item_id, + call_id=resolved_call_id, + name=resolved_name, + arguments="", + status="in_progress", + ) + item_state = self._add_response_stream_item(state, item) + state.tool_items[tool_call_index] = item_state + return [ + self._response_event( + state, + "response.output_item.added", + output_index=item_state.output_index, + item=copy.deepcopy(item), + ) + ], item_state + + def _update_tool_stream_item( + self, + item: Dict[str, Any], + *, + call_id: Optional[str], + name_delta: Optional[str], + arguments_delta: Optional[str], + ) -> None: + if isinstance(call_id, str) and call_id: + item["call_id"] = call_id + if isinstance(name_delta, str) and name_delta: + current_name = cast(str, item.get("name", "")) + if not current_name: + item["name"] = name_delta + elif name_delta == current_name or name_delta.startswith(current_name): + item["name"] = name_delta + elif not current_name.endswith(name_delta): + item["name"] = current_name + name_delta + if isinstance(arguments_delta, str) and arguments_delta: + if item.get("type") == "custom_tool_call": + raw_arguments = cast(str, item.get("_raw_arguments", "")) + arguments_delta + item["_raw_arguments"] = raw_arguments + normalized_input = self._normalize_text_tool_payload(raw_arguments) + if normalized_input is not None: + item["input"] = normalized_input + else: + item["arguments"] = ( + cast(str, item.get("arguments", "")) + arguments_delta + ) + + @staticmethod + def _normalize_text_tool_payload(payload: str) -> Optional[str]: + if payload == "": + return "" + stripped = payload.lstrip() + if not stripped: + return "" + if stripped[0] not in '{["': + return payload + try: + decoded = json.loads(payload) + except Exception: + return None + if isinstance(decoded, str): + return decoded + if isinstance(decoded, dict): + input_value = decoded.get("input") + if isinstance(input_value, str): + return input_value + if len(decoded) == 1: + sole_value = next(iter(decoded.values())) + if isinstance(sole_value, str): + return sole_value + return payload + + def _finalize_response_stream_items( + self, + state: "OpenAIFormatter.ResponsesStream", + *, + finish_reason: Optional[str], + ) -> List[Dict[str, Any]]: + events: List[Dict[str, Any]] = [] + item_status = self._response_item_status(finish_reason) + + if ( + state.reasoning_item is not None + and state.reasoning_item.item["status"] == "in_progress" + ): + item = state.reasoning_item.item + item["status"] = item_status + part = cast(List[Dict[str, Any]], item["content"])[0] + events.append( + self._response_event( + state, + "response.reasoning_text.done", + item_id=cast(str, item["id"]), + output_index=state.reasoning_item.output_index, + content_index=cast(int, state.reasoning_item.content_index), + text=part["text"], + ) + ) + events.append( + self._response_event( + state, + "response.content_part.done", + item_id=cast(str, item["id"]), + output_index=state.reasoning_item.output_index, + content_index=cast(int, state.reasoning_item.content_index), + part=part, + ) + ) + events.append( + self._response_event( + state, + "response.output_item.done", + output_index=state.reasoning_item.output_index, + item=item, + ) + ) + + if ( + state.message_item is not None + and state.message_item.item["status"] == "in_progress" + ): + item = state.message_item.item + item["status"] = item_status + part = cast(List[Dict[str, Any]], item["content"])[0] + events.append( + self._response_event( + state, + "response.output_text.done", + item_id=cast(str, item["id"]), + output_index=state.message_item.output_index, + content_index=cast(int, state.message_item.content_index), + text=part["text"], + logprobs=part.get("logprobs", []), + ) + ) + events.append( + self._response_event( + state, + "response.content_part.done", + item_id=cast(str, item["id"]), + output_index=state.message_item.output_index, + content_index=cast(int, state.message_item.content_index), + part=part, + ) + ) + events.append( + self._response_event( + state, + "response.output_item.done", + output_index=state.message_item.output_index, + item=item, + ) + ) + + for tool_call_index in sorted(state.tool_items): + item_state = state.tool_items[tool_call_index] + item = item_state.item + item.pop("_raw_arguments", None) + if ( + item.get("status") != "in_progress" + and item.get("type") != "custom_tool_call" + ): + continue + if item.get("type") == "custom_tool_call": + events.append( + self._response_event( + state, + "response.custom_tool_call_input.done", + item_id=cast(str, item["id"]), + output_index=item_state.output_index, + input=cast(str, item.get("input", "")), + ) + ) + else: + item["status"] = item_status + events.append( + self._response_event( + state, + "response.function_call_arguments.done", + item_id=cast(str, item["id"]), + output_index=item_state.output_index, + name=cast(str, item["name"]), + arguments=cast(str, item["arguments"]), + ) + ) + events.append( + self._response_event( + state, + "response.output_item.done", + output_index=item_state.output_index, + item=item, + ) + ) + + state.final_status, state.incomplete_details = ( + self._response_status_and_incomplete_details( + finish_reason=finish_reason, + ) + ) + return events + + def convert_chat_chunk_to_response_events( + self, + chunk: ChatCompletionChunk | Dict[str, Any], + state: "OpenAIFormatter.ResponsesStream", + ) -> List[Dict[str, Any]]: + payload = ( + chunk.model_dump(mode="json", exclude_none=True) + if isinstance(chunk, BaseModel) + else chunk + ) + events = self.start_response_stream(state) + choice = cast(Dict[str, Any], payload["choices"][0]) + delta = cast(Dict[str, Any], choice.get("delta", {})) + finish_reason = cast(Optional[str], choice.get("finish_reason")) + logprobs = self._response_logprobs_from_chat_logprobs( + cast(Optional[Dict[str, Any]], choice.get("logprobs")) + ) + + phase = delta.get("phase") + if isinstance(phase, str): + state.assistant_meta["phase"] = phase + if state.message_item is not None: + state.message_item.item["phase"] = phase + + reasoning_delta = delta.get("reasoning_content") + if not isinstance(reasoning_delta, str) or not reasoning_delta: + candidate = delta.get("thinking") + reasoning_delta = candidate if isinstance(candidate, str) else None + if isinstance(reasoning_delta, str) and reasoning_delta: + added, item_state = self._ensure_reasoning_stream_item(state) + events.extend(added) + part = cast(List[Dict[str, Any]], item_state.item["content"])[0] + part["text"] += reasoning_delta + events.append( + self._response_event( + state, + "response.reasoning_text.delta", + item_id=cast(str, item_state.item["id"]), + output_index=item_state.output_index, + content_index=cast(int, item_state.content_index), + delta=reasoning_delta, + ) + ) + + content_delta = delta.get("content") + if isinstance(content_delta, str) and content_delta: + added, item_state = self._ensure_message_stream_item(state) + events.extend(added) + part = cast(List[Dict[str, Any]], item_state.item["content"])[0] + part["text"] += content_delta + if logprobs: + cast(List[Dict[str, Any]], part.setdefault("logprobs", [])).extend( + logprobs + ) + events.append( + self._response_event( + state, + "response.output_text.delta", + item_id=cast(str, item_state.item["id"]), + output_index=item_state.output_index, + content_index=cast(int, item_state.content_index), + delta=content_delta, + logprobs=logprobs, + ) + ) + + tool_calls = delta.get("tool_calls") + if isinstance(tool_calls, list): + for tool_call in tool_calls: + if not isinstance(tool_call, dict): + continue + tool_call_index = tool_call.get("index", 0) + if not isinstance(tool_call_index, int): + continue + function = tool_call.get("function") + if not isinstance(function, dict): + continue + added, item_state = self._ensure_tool_stream_item( + state, + tool_call_index=tool_call_index, + call_id=cast(Optional[str], tool_call.get("id")), + name=cast(Optional[str], function.get("name")), + ) + events.extend(added) + previous_input = cast(str, item_state.item.get("input", "")) + self._update_tool_stream_item( + item_state.item, + call_id=cast(Optional[str], tool_call.get("id")), + name_delta=cast(Optional[str], function.get("name")), + arguments_delta=cast(Optional[str], function.get("arguments")), + ) + arguments_delta = function.get("arguments") + if isinstance(arguments_delta, str) and arguments_delta: + if item_state.item.get("type") == "custom_tool_call": + current_input = cast(str, item_state.item.get("input", "")) + if not current_input or current_input == previous_input: + continue + delta_text = current_input + if current_input.startswith(previous_input): + delta_text = current_input[len(previous_input) :] + events.append( + self._response_event( + state, + "response.custom_tool_call_input.delta", + item_id=cast(str, item_state.item["id"]), + output_index=item_state.output_index, + delta=delta_text, + ) + ) + continue + events.append( + self._response_event( + state, + "response.function_call_arguments.delta", + item_id=cast(str, item_state.item["id"]), + output_index=item_state.output_index, + delta=arguments_delta, + ) + ) + + if finish_reason is not None: + events.extend( + self._finalize_response_stream_items( + state, + finish_reason=finish_reason, + ) + ) + return events + + def response_stream_terminal_events( + self, + state: "OpenAIFormatter.ResponsesStream", + completion: Optional[OpenAICompletion], + ) -> List[Dict[str, Any]]: + if not state.started: + state.started = True + if completion is not None and state.final_status is None: + finish_reason = None + if completion.choices: + finish_reason = completion.choices[0].finish_reason + self._finalize_response_stream_items(state, finish_reason=finish_reason) + status = state.final_status or "completed" + response = self._response_stream_response( + state, + status=status, + usage=( + self._responses_usage_from_completion( + usage=completion.usage if completion is not None else None, + output_items=state.output, + ) + ), + incomplete_details=state.incomplete_details, + completed_at=time.time() if status in {"completed", "incomplete"} else None, + ) + event_type = "response.incomplete" if status == "incomplete" else "response.completed" + return [self._response_event(state, event_type, response=response)] + + def aggregate_completion_results( + self, + results: Sequence[OpenAICompletion], + ) -> OpenAICompletion: + choices: List[CompletionChoice] = [] + prompt_tokens = 0 + completion_tokens = 0 + total_tokens = 0 + for result in results: + for choice in result.choices: + choices.append(choice.model_copy(update={"index": len(choices)})) + assert result.usage is not None + prompt_tokens += result.usage.prompt_tokens + completion_tokens += result.usage.completion_tokens + total_tokens += result.usage.total_tokens + return OpenAICompletion( + id=f"cmpl-{uuid.uuid4().hex}", + object="text_completion", + created=int(time.time()), + model=self.model.model_path, + choices=choices, + usage=CompletionUsage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=total_tokens, + ), + ) + + @staticmethod + def _convert_completion_logprobs_to_chat_choice( + logprobs: Optional[CompletionLogprobs], + ) -> Optional[ChatCompletionChoiceLogprobs]: + if logprobs is None: + return None + tokens = logprobs.tokens or [] + token_logprobs_list = logprobs.token_logprobs or [] + top_logprobs_list = logprobs.top_logprobs or [] + return ChatCompletionChoiceLogprobs( + content=[ + ChatCompletionTokenLogprob( + token=token, + bytes=None, + logprob=token_logprob if token_logprob is not None else 0.0, + top_logprobs=( + [ + TopLogprob( + token=top_token, + logprob=top_logprob, + bytes=None, + ) + for top_token, top_logprob in top_logprobs.items() + ] + if top_logprobs is not None + else [] + ), + ) + for token, token_logprob, top_logprobs in zip( + tokens, + token_logprobs_list, + top_logprobs_list, + ) + ], + refusal=None, + ) + + @staticmethod + def _convert_completion_logprobs_to_chat_chunk( + logprobs: Optional[CompletionLogprobs], + ) -> Optional[ChatCompletionChunkChoiceLogprobs]: + choice_logprobs = OpenAIFormatter._convert_completion_logprobs_to_chat_choice(logprobs) + if choice_logprobs is None: + return None + return ChatCompletionChunkChoiceLogprobs.model_validate( + choice_logprobs.model_dump(mode="python", exclude_none=True) + ) + + @staticmethod + def _chat_message_from_completion_choice( + completion_id: str, + choice: CompletionChoice, + tool_name: Optional[str], + ) -> ChatCompletionMessage: + if tool_name is None: + return ChatCompletionMessage( + role="assistant", + content=choice.text, + ) + tool_id = f"call_{choice.index}_{tool_name}_{completion_id}" + arguments = choice.text + return ChatCompletionMessage( + role="assistant", + content=None, + function_call=ChatCompletionMessageFunctionCall( + name=tool_name, + arguments=arguments, + ), + tool_calls=[ + ChatCompletionMessageToolCall( + id=tool_id, + type="function", + function=ChatCompletionMessageToolCallFunction( + name=tool_name, + arguments=arguments, + ), + ) + ], + ) + + @staticmethod + def _chat_chunk_payload( + *, + chunk_id: str, + created: int, + model: str, + index: int, + delta: Dict[str, Any], + finish_reason: Optional[str], + logprobs: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + choice: Dict[str, Any] = { + "index": index, + "delta": delta, + "finish_reason": finish_reason, + } + if logprobs is not None: + choice["logprobs"] = logprobs + return { + "id": "chat" + chunk_id, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": [choice], + } + + def convert_completion_response_to_chat( + self, + completion: OpenAICompletion, + tool_name: Optional[str] = None, + *, + functions: Optional[List[ChatTemplateFunctionDefinition]] = None, + tools: Optional[List[ChatTemplateTool]] = None, + generation_prompt: str = "", + ) -> ChatCompletion | Dict[str, Any]: + parser_tools = self._tools_for_response_parser( + functions=functions, + tools=tools, + ) + if self.model.response_schema is not None: + choices: List[Dict[str, Any]] = [] + for choice in completion.choices: + parser = self._response_parser( + tools=parser_tools, + completion_id=completion.id, + choice_index=choice.index, + generation_prompt=generation_prompt, + ) + message = parser.parse_completion_message(choice.text) + logprobs = self._convert_completion_logprobs_to_chat_choice(choice.logprobs) + choices.append( + { + "index": choice.index, + "message": message, + "logprobs": ( + logprobs.model_dump(mode="json", exclude_none=True) + if logprobs is not None + else None + ), + "finish_reason": ( + "tool_calls" + if message.get("tool_calls") + else choice.finish_reason + ), + } + ) + return { + "id": "chat" + completion.id, + "object": "chat.completion", + "created": completion.created, + "model": completion.model, + "choices": choices, + "usage": ( + completion.usage.model_dump(mode="json", exclude_none=True) + if completion.usage is not None + else None + ), + } + return ChatCompletion( + id="chat" + completion.id, + object="chat.completion", + created=completion.created, + model=completion.model, + choices=[ + ChatCompletionChoice( + index=choice.index, + message=self._chat_message_from_completion_choice( + completion.id, + choice, + tool_name, + ), + logprobs=self._convert_completion_logprobs_to_chat_choice(choice.logprobs), + finish_reason=cast(Any, ( + "tool_calls" if tool_name is not None else choice.finish_reason + )), + ) + for choice in completion.choices + ], + usage=completion.usage, + ) + + def convert_completion_chunk_to_chat_chunks( + self, + chunk: CompletionChunk, + started_indices: set[int], + tool_name: Optional[str] = None, + *, + functions: Optional[List[ChatTemplateFunctionDefinition]] = None, + tools: Optional[List[ChatTemplateTool]] = None, + parsed_states: Optional[Dict[int, Any]] = None, + generation_prompt: str = "", + ) -> List[ChatCompletionChunk | Dict[str, Any]]: + parser_tools = self._tools_for_response_parser( + functions=functions, + tools=tools, + ) + if self.model.response_schema is not None: + parsed_chunks: List[Dict[str, Any]] = [] + if parsed_states is None: + parsed_states = {} + for choice in chunk["choices"]: + index = choice["index"] + parser = parsed_states.get(index) + if not isinstance(parser, ResponseParser): + parser = self._response_parser( + tools=parser_tools, + completion_id=chunk["id"], + choice_index=index, + generation_prompt=generation_prompt, + ) + parsed_states[index] = parser + logprobs = self._convert_completion_logprobs_to_chat_chunk( + CompletionLogprobs.model_validate(choice["logprobs"]) + if choice["logprobs"] is not None + else None + ) + parsed_chunks.extend( + parser.consume_completion_chunk( + choice["text"], + chunk_id=chunk["id"], + created=chunk["created"], + model=chunk["model"], + finish_reason=choice["finish_reason"], + logprobs=( + logprobs.model_dump(mode="json", exclude_none=True) + if logprobs is not None + else None + ), + ) + ) + if parser.started: + started_indices.add(index) + return cast(List[ChatCompletionChunk | Dict[str, Any]], parsed_chunks) + chat_chunks: List[ChatCompletionChunk] = [] + for choice in chunk["choices"]: + index = choice["index"] + if index not in started_indices: + started_indices.add(index) + chat_chunks.append( + ChatCompletionChunk( + id="chat" + chunk["id"], + object="chat.completion.chunk", + created=chunk["created"], + model=chunk["model"], + choices=[ + ChatCompletionChunkChoice( + index=index, + delta=( + ChoiceDelta( + role="assistant", + content=None, + function_call=None, + tool_calls=None, + ) + if tool_name is not None + else ChoiceDelta(role="assistant") + ), + logprobs=None, + finish_reason=None, + ) + ], + ) + ) + if tool_name is not None: + delta: ChoiceDelta + if choice["finish_reason"] is None: + tool_id = f"call_{index}_{tool_name}_{chunk['id']}" + delta = ChoiceDelta( + role=None, + content=None, + function_call=ChoiceDeltaFunctionCall( + name=tool_name, + arguments=choice["text"], + ), + tool_calls=[ + ChoiceDeltaToolCall( + index=0, + id=tool_id, + type="function", + function=ChoiceDeltaToolCallFunction( + name=tool_name, + arguments=choice["text"], + ), + ) + ], + ) + else: + delta = ChoiceDelta( + role=None, + content=None, + function_call=None, + tool_calls=None, + ) + chat_chunks.append( + ChatCompletionChunk( + id="chat" + chunk["id"], + object="chat.completion.chunk", + created=chunk["created"], + model=chunk["model"], + choices=[ + ChatCompletionChunkChoice( + index=index, + delta=delta, + logprobs=self._convert_completion_logprobs_to_chat_chunk( + CompletionLogprobs.model_validate(choice["logprobs"]) + if choice["logprobs"] is not None + else None + ), + finish_reason=( + "tool_calls" + if choice["finish_reason"] is not None + else None + ), + ) + ], + ) + ) + continue + chat_chunks.append( + ChatCompletionChunk( + id="chat" + chunk["id"], + object="chat.completion.chunk", + created=chunk["created"], + model=chunk["model"], + choices=[ + ChatCompletionChunkChoice( + index=index, + delta=( + ChoiceDelta(content=choice["text"]) + if choice["finish_reason"] is None + else ChoiceDelta() + ), + logprobs=self._convert_completion_logprobs_to_chat_chunk( + CompletionLogprobs.model_validate(choice["logprobs"]) + if choice["logprobs"] is not None + else None + ), + finish_reason=cast(Any, choice["finish_reason"]), + ) + ], + ) + ) + return cast(List[ChatCompletionChunk | Dict[str, Any]], chat_chunks) + + def returned_output_end( + self, + completion: Completion, + finish_reason: Optional[str], + ) -> int: + completion_bytes: bytes | bytearray + if completion.rendered_bytes: + completion_bytes = completion.rendered_bytes + else: + completion_bytes = b"".join(record.text_bytes for record in completion.token_records) + returned_end = len(completion_bytes) + if finish_reason == "stop": + stops = [stop for stop in completion.stop_sequences if stop in completion_bytes] + if stops: + returned_end = min(completion_bytes.index(stop) for stop in stops) + elif finish_reason is None: + holdback = 0 + for stop in completion.stop_sequences: + for size in range(min(len(stop), returned_end), 0, -1): + if completion_bytes.endswith(stop[:size]): + holdback = max(holdback, size) + break + returned_end -= holdback + return returned_end + + def returned_tokens( + self, + completion: Completion, + finish_reason: Optional[str], + *, + start_index: int = 0, + ) -> List["OpenAIFormatter.ReturnedToken"]: + returned_end = self.returned_output_end(completion, finish_reason) + returned_tokens: List[OpenAIFormatter.ReturnedToken] = [] + prefix_bytes = b"" + for index, record in enumerate(completion.token_records): + token_start = len(prefix_bytes) + if token_start >= returned_end: + break + token_end = token_start + len(record.text_bytes) + text_bytes = record.text_bytes + if token_end > returned_end: + if finish_reason is None: + break + text_bytes = text_bytes[: returned_end - token_start] + if index >= start_index: + returned_tokens.append( + OpenAIFormatter.ReturnedToken( + index=index, + text_bytes=text_bytes, + token=record, + text_offset=len(self.decode_text(prefix_bytes)), + ) + ) + prefix_bytes += record.text_bytes + return returned_tokens + + def stream_completion_chunks( + self, + request: CompletionRequest, + completion: Completion, + finish_reason: Optional[str], + ) -> List[CompletionChunk]: + returned_tokens = self.returned_tokens( + completion, + finish_reason, + start_index=completion.returned_token_count, + ) + chunks: List[CompletionChunk] = [] + if completion.logprobs is not None: + for returned_token in returned_tokens: + token = returned_token.token + chunks.append( + { + "id": request.id, + "object": "text_completion", + "created": request.created, + "model": self.model.model_path, + "choices": [ + { + "text": self.decode_text(returned_token.text_bytes), + "index": completion.index, + "logprobs": { + "tokens": [self.decode_text(token.text_bytes)], + "text_offset": [ + len(completion.prompt_text) + returned_token.text_offset + ], + "token_logprobs": [token.token_logprob], + "top_logprobs": [token.top_logprobs], + }, + "finish_reason": None, + } + ], + } + ) + completion.returned_token_count = returned_token.index + 1 + return chunks + + chunk_tokens: List[OpenAIFormatter.ReturnedToken] = [] + for returned_token in returned_tokens: + chunk_tokens.append(returned_token) + chunk_bytes = b"".join(token.text_bytes for token in chunk_tokens) + if returned_token.text_bytes != returned_token.token.text_bytes: + chunks.append( + { + "id": request.id, + "object": "text_completion", + "created": request.created, + "model": self.model.model_path, + "choices": [ + { + "text": self.decode_text(returned_token.text_bytes), + "index": completion.index, + "logprobs": None, + "finish_reason": None, + } + ], + } + ) + completion.returned_token_count = returned_token.index + 1 + chunk_tokens = [] + continue + try: + text = chunk_bytes.decode("utf-8") + except UnicodeError: + continue + chunks.append( + { + "id": request.id, + "object": "text_completion", + "created": request.created, + "model": self.model.model_path, + "choices": [ + { + "text": text, + "index": completion.index, + "logprobs": None, + "finish_reason": None, + } + ], + } + ) + completion.returned_token_count = returned_token.index + 1 + chunk_tokens = [] + return chunks + + def completion_finish_chunk( + self, + request: CompletionRequest, + completion: Completion, + finish_reason: str, + ) -> CompletionChunk: + return { + "id": request.id, + "object": "text_completion", + "created": request.created, + "model": self.model.model_path, + "choices": [ + { + "text": "", + "index": completion.index, + "logprobs": None, + "finish_reason": finish_reason, + } + ], + } + + def _build_completion_choice( + self, + request: CompletionRequest, + completion: Completion, + ) -> CompletionChoice: + returned_tokens = self.returned_tokens(completion, completion.finish_reason) + text_bytes = b"".join(returned_token.text_bytes for returned_token in returned_tokens) + text = self.decode_text(text_bytes) + if request.payload.echo: + text = completion.prompt_text + text + logprobs: Optional[CompletionLogprobs] = None + if completion.logprobs is not None: + offsets: List[int] = [] + token_texts: List[str] = [] + token_logprobs: List[Optional[float]] = [] + top_logprobs: List[Optional[Dict[str, float]]] = [] + text_cursor = request.prompt_text if not request.payload.echo else "" + if request.payload.echo: + for record in request.prompt_records: + offsets.append(len(text_cursor)) + token_texts.append(self.decode_text(record.text_bytes)) + token_logprobs.append(record.token_logprob) + top_logprobs.append(record.top_logprobs) + text_cursor += self.decode_text(record.text_bytes) + text_cursor = request.prompt_text + for returned_token in returned_tokens: + token = returned_token.token + offsets.append(len(completion.prompt_text) + returned_token.text_offset) + if request.payload.echo: + offsets[-1] = len(text_cursor) + token_texts.append(self.decode_text(token.text_bytes)) + token_logprobs.append(token.token_logprob) + top_logprobs.append(token.top_logprobs) + text_cursor += self.decode_text(token.text_bytes) + logprobs = CompletionLogprobs.model_construct( + text_offset=offsets, + token_logprobs=token_logprobs, + tokens=token_texts, + top_logprobs=top_logprobs, + ) + return CompletionChoice( + text=text, + index=completion.index, + logprobs=logprobs, + finish_reason=cast(Any, completion.finish_reason), + ) + + def build_completion_response( + self, + request: CompletionRequest, + completions: Sequence[Completion], + ) -> OpenAICompletion: + completion_tokens = sum(completion.completion_token_count for completion in completions) + prompt_tokens = request.prompt_plan.eval_token_count + return OpenAICompletion( + id=request.id, + object="text_completion", + created=request.created, + model=self.model.model_path, + choices=[ + self._build_completion_choice(request, completion) + for completion in completions + ], + usage=CompletionUsage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ), + ) + + +class Sampler: + TOKEN_DATA_DTYPE = np.dtype( + [("id", np.intc), ("logit", np.single), ("p", np.single)], + align=True, + ) + + def __init__( + self, + *, + seed: int, + vocab: llama_cpp.llama_vocab_p, + n_vocab: int, + top_p: float, + temperature: float, + frequency_penalty: float, + presence_penalty: float, + logit_bias: Optional[Dict[int, float]], + grammar_text: Optional[str] = None, + grammar_root: str = "root", + ) -> None: + params = llama_cpp.llama_sampler_chain_default_params() + self._sampler = llama_cpp.llama_sampler_chain_init(params) + self._closed = False + self._sample_logits_size = 0 + self._sample_logits_token_data: Optional[Any] = None + self._sample_logits_token_array: Optional[Any] = None + self._sample_logits_recarray: Optional[np.recarray] = None + if logit_bias: + bias_array = (llama_cpp.llama_logit_bias * len(logit_bias))() + for index, (token, bias) in enumerate(logit_bias.items()): + bias_array[index].token = ctypes.c_int32(token) + bias_array[index].bias = float(bias) + llama_cpp.llama_sampler_chain_add( + self._sampler, + llama_cpp.llama_sampler_init_logit_bias( + n_vocab, len(logit_bias), bias_array + ), + ) + self.bias_array = bias_array + if frequency_penalty != 0.0 or presence_penalty != 0.0: + llama_cpp.llama_sampler_chain_add( + self._sampler, + llama_cpp.llama_sampler_init_penalties( + 64, + 1.0, + frequency_penalty, + presence_penalty, + ), + ) + if grammar_text is not None: + grammar_sampler = llama_cpp.llama_sampler_init_grammar( + vocab, + grammar_text.encode("utf-8"), + grammar_root.encode("utf-8"), + ) + if grammar_sampler is None: + raise RuntimeError("failed to initialize grammar sampler") + llama_cpp.llama_sampler_chain_add(self._sampler, grammar_sampler) + if temperature < 0.0: + llama_cpp.llama_sampler_chain_add( + self._sampler, llama_cpp.llama_sampler_init_dist(seed) + ) + return + if temperature == 0.0: + llama_cpp.llama_sampler_chain_add( + self._sampler, llama_cpp.llama_sampler_init_greedy() + ) + return + min_keep = 1 + llama_cpp.llama_sampler_chain_add( + self._sampler, llama_cpp.llama_sampler_init_top_p(top_p, min_keep) + ) + llama_cpp.llama_sampler_chain_add( + self._sampler, llama_cpp.llama_sampler_init_temp(temperature) + ) + llama_cpp.llama_sampler_chain_add( + self._sampler, llama_cpp.llama_sampler_init_dist(seed) + ) + + def sample(self, ctx: llama_cpp.llama_context_p, output_index: int) -> int: + return int(llama_cpp.llama_sampler_sample(self._sampler, ctx, output_index)) + + def _ensure_sample_logits_buffer(self, size: int) -> None: + if size == self._sample_logits_size and self._sample_logits_recarray is not None: + return + token_data = (llama_cpp.llama_token_data * size)() + token_data_address = ctypes.addressof(token_data) + recarray = np.recarray( + shape=(size,), + dtype=self.TOKEN_DATA_DTYPE, + buf=cast( + Any, + (llama_cpp.llama_token_data * size).from_address(token_data_address), + ), + ) + recarray.id[:] = np.arange(size, dtype=np.intc) + token_array = llama_cpp.llama_token_data_array( + data=token_data, + size=size, + selected=-1, + sorted=False, + ) + self._sample_logits_size = size + self._sample_logits_token_data = token_data + self._sample_logits_token_array = token_array + self._sample_logits_recarray = recarray + + def sample_logits(self, logits: np.ndarray) -> int: + self._ensure_sample_logits_buffer(len(logits)) + assert self._sample_logits_recarray is not None + assert self._sample_logits_token_array is not None + self._sample_logits_recarray.logit[:] = logits + self._sample_logits_recarray.p.fill(0.0) + self._sample_logits_token_array.selected = -1 + self._sample_logits_token_array.sorted = False + llama_cpp.llama_sampler_apply( + self._sampler, + cast(Any, ctypes.byref(self._sample_logits_token_array)), + ) + token = int(self._sample_logits_recarray.id[self._sample_logits_token_array.selected]) + llama_cpp.llama_sampler_accept(self._sampler, token) + return token + + def close(self) -> None: + if not self._closed: + llama_cpp.llama_sampler_free(self._sampler) + self._closed = True + + +@dataclass +class MTMDEmbedding: + key: str + embeddings: np.ndarray + + +class MTMDEmbeddingCache: + _metadata_version = "1" + + def __init__( + self, + *, + path: str, + max_bytes: int, + model_fingerprint: str, + mmproj_fingerprint: str, + ) -> None: + self.path = Path(path) + self.max_bytes = max_bytes + self.model_fingerprint = model_fingerprint + self.mmproj_fingerprint = mmproj_fingerprint + self.path.mkdir(parents=True, exist_ok=True) + + @staticmethod + def _safe_open(path: Path) -> Any: + try: + from safetensors import safe_open + except ImportError as exc: + raise RuntimeError( + "model.mtmd.embedding_cache requires safetensors. " + "Install it with `pip install safetensors`." + ) from exc + return safe_open(str(path), framework="numpy") + + @staticmethod + def _save_file( + tensors: Dict[str, np.ndarray], + path: Path, + metadata: Dict[str, str], + ) -> None: + from safetensors.numpy import save_file + + save_file(tensors, str(path), metadata=metadata) + + @staticmethod + def fingerprint_file(path: str) -> str: + stat = os.stat(path) + payload = f"{Path(path).resolve()}:{stat.st_size}:{stat.st_mtime_ns}" + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + @staticmethod + def _build_key( + *, + model_fingerprint: str, + mmproj_fingerprint: str, + kind: Literal["image", "audio", "video"], + media_bytes: bytes, + ) -> str: + digest = hashlib.sha256(f"{kind}:".encode("utf-8") + media_bytes).hexdigest() + payload = ":".join( + [ + MTMDEmbeddingCache._metadata_version, + model_fingerprint, + mmproj_fingerprint, + kind, + digest, + ] + ) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + def key_for_media(self, kind: Literal["image", "audio", "video"], media_bytes: bytes) -> str: + return self._build_key( + model_fingerprint=self.model_fingerprint, + mmproj_fingerprint=self.mmproj_fingerprint, + kind=kind, + media_bytes=media_bytes, + ) + + def _path_for_key(self, key: str) -> Path: + return self.path / f"{key}.safetensors" + + def load(self, key: str) -> Optional[MTMDEmbedding]: + entry_path = self._path_for_key(key) + if not entry_path.exists(): + return None + with self._safe_open(entry_path) as tensors: + metadata = tensors.metadata() + if metadata.get("version") != self._metadata_version: + return None + if metadata.get("model") != self.model_fingerprint: + return None + if metadata.get("mmproj") != self.mmproj_fingerprint: + return None + embeddings = np.array(tensors.get_tensor("embeddings"), copy=True) + return MTMDEmbedding(key=key, embeddings=embeddings.astype(np.float32, copy=False)) + + def save(self, key: str, embeddings: np.ndarray) -> None: + if self.max_bytes == 0: + return + tmp_path = self.path / f".{key}.{uuid.uuid4().hex}.tmp" + final_path = self._path_for_key(key) + metadata = { + "version": self._metadata_version, + "model": self.model_fingerprint, + "mmproj": self.mmproj_fingerprint, + "key": key, + } + self._save_file( + {"embeddings": np.ascontiguousarray(embeddings, dtype=np.float32)}, + tmp_path, + metadata, + ) + os.replace(tmp_path, final_path) + self.evict_if_needed() + + def evict_if_needed(self) -> None: + if self.max_bytes <= 0: + return + entries = [path for path in self.path.glob("*.safetensors") if path.is_file()] + total = sum(path.stat().st_size for path in entries) + if total <= self.max_bytes: + return + for path in sorted(entries, key=lambda item: item.stat().st_mtime_ns): + if total <= self.max_bytes: + break + try: + size = path.stat().st_size + path.unlink() + total -= size + except OSError: + continue + + +@dataclass +class MTMDLoadedMedia: + media: MediaInput + media_bytes: bytes + key: str + bitmap: Any + video_ctx: Optional[Any] = None + video_temp_path: Optional[Path] = None + video_callback: Optional[Any] = None + video_frame_count: int = 0 + video_frames_used: int = 0 + + +class MTMDProcessor: + def __init__( + self, + *, + model_path: str, + llama_model: Any, + chat_formatter: Jinja2ChatFormatter, + tokenize: Callable[..., List[int]], + n_embd_inp: int, + n_batch: int, + n_ubatch: int, + n_threads_batch: int, + mmproj_path: str, + embedding_cache: Optional[MTMDEmbeddingCache], + allowed_media_domains: Optional[List[str]], + allowed_local_media_path: Optional[str], + image_max_bytes: int, + audio_max_bytes: int, + video_max_bytes: int, + image_timeout_seconds: float, + ) -> None: + self.chat_formatter = chat_formatter + self.tokenize = tokenize + self.n_embd_inp = n_embd_inp + self.n_batch = n_batch + self.n_ubatch = n_ubatch + self.mmproj_path = mmproj_path + self.embedding_cache = embedding_cache + self.model_fingerprint = MTMDEmbeddingCache.fingerprint_file(model_path) + self.mmproj_fingerprint = MTMDEmbeddingCache.fingerprint_file(mmproj_path) + self.allowed_media_domains = ( + {domain.lower() for domain in allowed_media_domains} + if allowed_media_domains is not None + else set() + ) + self.allowed_local_media_path = ( + Path(allowed_local_media_path).expanduser().resolve() + if allowed_local_media_path is not None + else None + ) + self.image_max_bytes = image_max_bytes + self.audio_max_bytes = audio_max_bytes + self.video_max_bytes = video_max_bytes + self.image_timeout_seconds = image_timeout_seconds + self.lock = threading.Lock() + params = mtmd_cpp.mtmd_context_params_default() + params.n_threads = max(1, n_threads_batch) + self.ctx = mtmd_cpp.mtmd_init_from_file( + mmproj_path.encode("utf-8"), + llama_model, + params, + ) + if self.ctx is None: + raise RuntimeError(f"failed to load MTMD context: {mmproj_path}") + self.supports_vision = bool(mtmd_cpp.mtmd_support_vision(self.ctx)) + self.supports_audio = bool(mtmd_cpp.mtmd_support_audio(self.ctx)) + self.supports_video = self.supports_vision and bool( + mtmd_cpp.mtmd_helper_support_video(self.ctx) + ) + if not self.supports_vision and not self.supports_audio: + mtmd_cpp.mtmd_free(self.ctx) + self.ctx = None + raise RuntimeError(f"MTMD projector does not support image or audio input: {mmproj_path}") + media_marker = mtmd_cpp.mtmd_get_marker(self.ctx) + if media_marker is None: + mtmd_cpp.mtmd_free(self.ctx) + self.ctx = None + raise RuntimeError(f"MTMD projector does not expose a media marker: {mmproj_path}") + self.media_marker = media_marker.decode("utf-8") + + def close(self) -> None: + if self.ctx is not None: + mtmd_cpp.mtmd_free(self.ctx) + self.ctx = None + + def _max_bytes_for_media(self, kind: Literal["image", "audio", "video"]) -> int: + if kind == "image": + return self.image_max_bytes + if kind == "audio": + return self.audio_max_bytes + return self.video_max_bytes + + def _load_media_file(self, kind: Literal["image", "audio", "video"], media_url: str) -> bytes: + if self.allowed_local_media_path is None: + raise CompletionRequestValidationError("local media path is not allowed") + parsed = urllib.parse.urlsplit(media_url) + if parsed.netloc not in {"", "localhost"}: + raise CompletionRequestValidationError("local media path is not allowed") + path = Path(urllib.parse.unquote(parsed.path)).expanduser().resolve() + try: + path.relative_to(self.allowed_local_media_path) + except ValueError as exc: + raise CompletionRequestValidationError("local media path is not allowed") from exc + max_bytes = self._max_bytes_for_media(kind) + try: + if path.stat().st_size > max_bytes: + raise CompletionRequestValidationError(f"{kind} exceeds model.mtmd.{kind}_max_bytes") + data = path.read_bytes() + except OSError as exc: + raise CompletionRequestValidationError(f"failed to read local {kind}: {exc}") from exc + if len(data) > max_bytes: + raise CompletionRequestValidationError(f"{kind} exceeds model.mtmd.{kind}_max_bytes") + return data + + def _validate_remote_media_url(self, media_url: str) -> str: + parsed = urllib.parse.urlsplit(media_url) + if parsed.scheme not in {"http", "https"} or not parsed.hostname: + raise CompletionRequestValidationError( + "only data:, file:, http:, and https: media URLs are supported" + ) + if parsed.username is not None or parsed.password is not None: + raise CompletionRequestValidationError("remote media domain is not allowed") + hostname = parsed.hostname.lower() + if self.allowed_media_domains and hostname not in self.allowed_media_domains: + raise CompletionRequestValidationError("remote media domain is not allowed") + return urllib.parse.urlunsplit(parsed) + + @staticmethod + def _urlopen_without_redirects( + request: urllib.request.Request, + *, + timeout: float, + ) -> Any: + class NoRedirectHandler(urllib.request.HTTPRedirectHandler): + def redirect_request(self, req, fp, code, msg, headers, newurl): # type: ignore[no-untyped-def] + return None + + opener = urllib.request.build_opener(NoRedirectHandler) + return opener.open(request, timeout=timeout) + + def _load_media_url(self, kind: Literal["image", "audio", "video"], media_url: str) -> bytes: + max_bytes = self._max_bytes_for_media(kind) + if media_url.startswith("data:"): + try: + _, encoded = media_url.split(",", 1) + except ValueError as exc: + raise CompletionRequestValidationError(f"invalid data {kind} URL") from exc + data = base64.b64decode(encoded, validate=False) + if len(data) > max_bytes: + raise CompletionRequestValidationError(f"{kind} exceeds model.mtmd.{kind}_max_bytes") + return data + if media_url.startswith("file:"): + return self._load_media_file(kind, media_url) + media_url = self._validate_remote_media_url(media_url) + request = urllib.request.Request( + media_url, + headers={"User-Agent": "llama-cpp-python-batch-mtmd/0"}, + ) + try: + with self._urlopen_without_redirects(request, timeout=self.image_timeout_seconds) as response: + data = response.read(max_bytes + 1) + except (OSError, TimeoutError, urllib.error.URLError) as exc: + raise CompletionRequestValidationError( + f"failed to fetch {kind} URL: {exc}" + ) from exc + if len(data) > max_bytes: + raise CompletionRequestValidationError(f"{kind} exceeds model.mtmd.{kind}_max_bytes") + return data + + def load_media(self, media: MediaInput) -> bytes: + if media.url is not None: + return self._load_media_url(media.kind, media.url) + if media.kind not in {"audio", "video"} or media.data is None: + raise CompletionRequestValidationError(f"{media.kind} input requires a URL") + try: + data = base64.b64decode(media.data, validate=False) + except (ValueError, binascii.Error) as exc: + raise CompletionRequestValidationError(f"input_{media.kind} data must be valid base64") from exc + max_bytes = self._max_bytes_for_media(media.kind) + if len(data) > max_bytes: + raise CompletionRequestValidationError( + f"{media.kind} exceeds model.mtmd.{media.kind}_max_bytes" + ) + return data + + def _create_loaded_media( + self, + media: MediaInput, + media_bytes: bytes, + ) -> MTMDLoadedMedia: + key = ( + self.embedding_cache.key_for_media(media.kind, media_bytes) + if self.embedding_cache is not None + else MTMDEmbeddingCache._build_key( + model_fingerprint=self.model_fingerprint, + mmproj_fingerprint=self.mmproj_fingerprint, + kind=media.kind, + media_bytes=media_bytes, + ) + ) + if media.kind == "video": + return self._create_loaded_video_media(media, media_bytes, key) + buffer = (ctypes.c_uint8 * len(media_bytes)).from_buffer_copy(media_bytes) + wrapper = mtmd_cpp.mtmd_helper_bitmap_init_from_buf_wrapper( + self.ctx, + buffer, + len(media_bytes), + False, + ) + bitmap = wrapper.bitmap + if bitmap is None: + raise CompletionRequestValidationError(f"failed to create MTMD {media.kind} bitmap") + mtmd_cpp.mtmd_bitmap_set_id(bitmap, key.encode("utf-8")) + video_frame_count = 0 + video_ctx = wrapper.video_ctx + if video_ctx: + video_info = mtmd_cpp.mtmd_helper_video_get_info(video_ctx) + video_frame_count = max(0, int(video_info.n_frames)) + return MTMDLoadedMedia( + media=media, + media_bytes=media_bytes, + key=key, + bitmap=bitmap, + video_ctx=video_ctx, + video_frame_count=video_frame_count, + ) + + @staticmethod + def _video_temp_suffix(media: MediaInput) -> str: + extension = (media.format or "mp4").lstrip(".").lower() + if not extension or any(not char.isalnum() for char in extension): + extension = "video" + return f".{extension}" + + def _create_loaded_video_media( + self, + media: MediaInput, + media_bytes: bytes, + key: str, + ) -> MTMDLoadedMedia: + temp_file = tempfile.NamedTemporaryFile( + prefix="llama-cpp-python-mtmd-", + suffix=self._video_temp_suffix(media), + delete=False, + ) + temp_path = Path(temp_file.name) + try: + with temp_file: + temp_file.write(media_bytes) + params = mtmd_cpp.mtmd_helper_video_init_params_default() + video_ctx = mtmd_cpp.mtmd_helper_video_init( + self.ctx, + str(temp_path).encode("utf-8"), + params, + ) + if video_ctx is None: + raise CompletionRequestValidationError("failed to create MTMD video context") + + def read_next( + _chunk_index: int, + _user_data: Any, + out_bitmap: Any, + out_text: Any, + ) -> int: + return int(mtmd_cpp.mtmd_helper_video_read_next(video_ctx, out_bitmap, out_text)) + + callback = mtmd_cpp.mtmd_bitmap_lazy_callback(read_next) + bitmap = mtmd_cpp.mtmd_bitmap_init_lazy( + self.ctx, + key.encode("utf-8"), + ctypes.c_void_p(), + callback, + ) + if bitmap is None: + mtmd_cpp.mtmd_helper_video_free(video_ctx) + raise CompletionRequestValidationError("failed to create MTMD video bitmap") + video_info = mtmd_cpp.mtmd_helper_video_get_info(video_ctx) + return MTMDLoadedMedia( + media=media, + media_bytes=media_bytes, + key=key, + bitmap=bitmap, + video_ctx=video_ctx, + video_temp_path=temp_path, + video_callback=callback, + video_frame_count=max(0, int(video_info.n_frames)), + ) + except Exception: + try: + temp_path.unlink() + except OSError: + pass + raise + + def _media_identity_tokens( + self, + kind: Literal["image", "audio", "video"], + key: str, + n_pos: int, + ) -> List[int]: + tokens: List[int] = [] + for index in range(n_pos): + digest = hashlib.sha256(f"{kind}:{key}:{index}".encode("utf-8")).digest() + tokens.append(-1 - (int.from_bytes(digest[:4], "little") & 0x3FFFFFFF)) + return tokens + + def _encode_media_chunk( + self, + *, + kind: Literal["image", "audio", "video"], + key: str, + chunk: Any, + ) -> np.ndarray: + n_tokens = int(mtmd_cpp.mtmd_input_chunk_get_n_tokens(chunk)) + if self.embedding_cache is not None: + cached = self.embedding_cache.load(key) + if ( + cached is not None + and cached.embeddings.shape == (n_tokens, self.n_embd_inp) + ): + return cached.embeddings + result = int(mtmd_cpp.mtmd_encode_chunk(self.ctx, chunk)) + if result != 0: + raise CompletionRequestValidationError( + f"failed to encode {kind} chunk: error code {result}" + ) + output = mtmd_cpp.mtmd_get_output_embd(self.ctx) + if output is None: + raise CompletionRequestValidationError(f"MTMD {kind} encoder returned no embeddings") + flat = np.ctypeslib.as_array(output, shape=(n_tokens * self.n_embd_inp,)) + embeddings = np.array(flat, dtype=np.float32, copy=True).reshape( + n_tokens, + self.n_embd_inp, + ) + if self.embedding_cache is not None: + self.embedding_cache.save(key, embeddings) + return embeddings + + def _positions_for_chunk(self, chunk: Any, start_pos: int) -> np.ndarray: + n_tokens = int(mtmd_cpp.mtmd_input_chunk_get_n_tokens(chunk)) + if not mtmd_cpp.mtmd_decode_use_mrope(self.ctx): + return np.arange(start_pos, start_pos + n_tokens, dtype=np.int32) + chunk_type = int(mtmd_cpp.mtmd_input_chunk_get_type(chunk)) + if chunk_type == mtmd_cpp.MTMD_INPUT_CHUNK_TYPE_AUDIO: + positions = np.empty((4, n_tokens), dtype=np.int32) + positions[:] = np.arange(start_pos, start_pos + n_tokens, dtype=np.int32) + return positions.reshape(-1) + image_tokens = mtmd_cpp.mtmd_input_chunk_get_tokens_image(chunk) + if image_tokens is None: + raise CompletionRequestValidationError("MTMD image chunk has no image tokens") + positions = np.empty((4, n_tokens), dtype=np.int32) + for index in range(n_tokens): + pos = mtmd_cpp.mtmd_image_tokens_get_decoder_pos( + image_tokens, + llama_cpp.llama_pos(start_pos), + index, + ) + positions[0, index] = int(pos.t) + positions[1, index] = int(pos.y) + positions[2, index] = int(pos.x) + positions[3, index] = int(pos.z) + return positions.reshape(-1) + + def build_prompt_plan( + self, + *, + messages: List[ChatCompletionRequestMessage], + functions: Optional[List[ChatTemplateFunctionDefinition]] = None, + function_call: Optional[Union[str, ChatTemplateFunctionCall]] = None, + tools: Optional[List[ChatTemplateTool]] = None, + tool_choice: Optional[Union[str, ChatTemplateToolChoice]] = None, + reasoning_effort: Optional[str] = None, + ) -> PromptPlan: + media_inputs = Jinja2ChatFormatter.media_inputs_from_messages(messages) + if not media_inputs: + prompt, generation_prompt, _ = self.chat_formatter.format( + messages=messages, + functions=functions, + function_call=function_call, + tools=tools, + tool_choice=tool_choice, + reasoning_effort=reasoning_effort, + ) + tokens = self.tokenize(prompt, add_bos=False, special=True) + return PromptPlan.from_tokens(prompt, tokens, generation_prompt=generation_prompt) + if any(media.kind == "image" for media in media_inputs) and not self.supports_vision: + raise CompletionRequestValidationError("MTMD projector does not support images") + if any(media.kind == "audio" for media in media_inputs) and not self.supports_audio: + raise CompletionRequestValidationError("MTMD projector does not support audio") + if any(media.kind == "video" for media in media_inputs) and not self.supports_video: + raise CompletionRequestValidationError("MTMD projector does not support video") + with self.lock: + return self._build_prompt_plan_locked( + messages=messages, + media_inputs=media_inputs, + functions=functions, + function_call=function_call, + tools=tools, + tool_choice=tool_choice, + reasoning_effort=reasoning_effort, + ) + + def _build_prompt_plan_locked( + self, + *, + messages: List[ChatCompletionRequestMessage], + media_inputs: List[MediaInput], + functions: Optional[List[ChatTemplateFunctionDefinition]], + function_call: Optional[Union[str, ChatTemplateFunctionCall]], + tools: Optional[List[ChatTemplateTool]], + tool_choice: Optional[Union[str, ChatTemplateToolChoice]], + reasoning_effort: Optional[str], + ) -> PromptPlan: + prompt, generation_prompt, _ = self.chat_formatter.format( + messages=messages, + media_marker=self.media_marker, + functions=functions, + function_call=function_call, + tools=tools, + tool_choice=tool_choice, + reasoning_effort=reasoning_effort, + ) + media_bytes_by_index = [self.load_media(media) for media in media_inputs] + loaded_media: List[MTMDLoadedMedia] = [] + chunks: Optional[Any] = None + try: + loaded_media = [ + self._create_loaded_media(media, media_bytes) + for media, media_bytes in zip(media_inputs, media_bytes_by_index) + ] + loaded_media_by_key = {media.key: media for media in loaded_media} + video_media = [media for media in loaded_media if media.media.kind == "video"] + if len(video_media) > 1 and any(media.video_frame_count <= 0 for media in video_media): + raise CompletionRequestValidationError( + "multiple videos require MTMD to report frame counts" + ) + input_text = mtmd_cpp.mtmd_input_text() + input_text.text = prompt.encode("utf-8") + input_text.add_special = False + input_text.parse_special = True + chunks = mtmd_cpp.mtmd_input_chunks_init() + if chunks is None: + raise CompletionRequestValidationError("failed to create MTMD input chunks") + bitmap_array = (mtmd_cpp.mtmd_bitmap_p_ctypes * len(loaded_media))( + *(media.bitmap for media in loaded_media) + ) + result = int( + mtmd_cpp.mtmd_tokenize( + self.ctx, + chunks, + ctypes.byref(input_text), + bitmap_array, + len(loaded_media), + ) + ) + if result != 0: + raise CompletionRequestValidationError( + f"failed to tokenize MTMD prompt: error code {result}" + ) + segments: List[PromptSegment] = [] + identity_tokens: List[int] = [] + text_tokens: List[int] = [] + text_token_index_by_pos: Dict[int, int] = {} + identity_pos = 0 + decode_pos = 0 + video_index = 0 + used_media_keys = set() + n_chunks = int(mtmd_cpp.mtmd_input_chunks_size(chunks)) + for chunk_index in range(n_chunks): + chunk = mtmd_cpp.mtmd_input_chunks_get(chunks, chunk_index) + if chunk is None: + continue + chunk_type = int(mtmd_cpp.mtmd_input_chunk_get_type(chunk)) + if chunk_type == mtmd_cpp.MTMD_INPUT_CHUNK_TYPE_TEXT: + n_tokens_out = ctypes.c_size_t() + tokens_ptr = mtmd_cpp.mtmd_input_chunk_get_tokens_text( + chunk, + ctypes.byref(n_tokens_out), + ) + tokens = ( + [int(tokens_ptr[index]) for index in range(int(n_tokens_out.value))] + if tokens_ptr + else [] + ) + if tokens: + start_pos = identity_pos + segments.append( + PromptSegment( + kind="text", + start_pos=start_pos, + n_pos=len(tokens), + identity_tokens=list(tokens), + decode_start_pos=decode_pos, + decode_n_pos=len(tokens), + text_tokens=list(tokens), + ) + ) + for offset, token in enumerate(tokens): + text_token_index_by_pos[start_pos + offset] = len(text_tokens) + text_tokens.append(token) + identity_tokens.extend(tokens) + identity_pos += len(tokens) + decode_pos += len(tokens) + continue + if chunk_type == mtmd_cpp.MTMD_INPUT_CHUNK_TYPE_IMAGE: + chunk_kind: Literal["image", "audio"] = "image" + if not self.supports_vision: + raise CompletionRequestValidationError("MTMD projector does not support images") + elif chunk_type == mtmd_cpp.MTMD_INPUT_CHUNK_TYPE_AUDIO: + chunk_kind = "audio" + if not self.supports_audio: + raise CompletionRequestValidationError("MTMD projector does not support audio") + else: + raise CompletionRequestValidationError("unsupported MTMD input chunk type") + chunk_id_bytes = mtmd_cpp.mtmd_input_chunk_get_id(chunk) + chunk_id = chunk_id_bytes.decode("utf-8") if chunk_id_bytes else "" + media = loaded_media_by_key.get(chunk_id) + video_frame_index: Optional[int] = None + if media is None and chunk_kind == "image" and video_media: + while ( + video_index < len(video_media) + and video_media[video_index].video_frame_count > 0 + and video_media[video_index].video_frames_used + >= video_media[video_index].video_frame_count + ): + video_index += 1 + if video_index >= len(video_media): + raise CompletionRequestValidationError("MTMD video frame count mismatch") + media = video_media[video_index] + video_frame_index = media.video_frames_used + media.video_frames_used += 1 + if media is None: + raise CompletionRequestValidationError("MTMD media chunk identity mismatch") + if media.media.kind == "video": + if chunk_kind != "image": + raise CompletionRequestValidationError("MTMD video chunk modality mismatch") + kind: Literal["image", "audio", "video"] = "video" + if video_frame_index is None: + video_frame_index = media.video_frames_used + media.video_frames_used += 1 + key = hashlib.sha256( + f"{media.key}:frame:{video_frame_index}".encode("utf-8") + ).hexdigest() + else: + if media.media.kind != chunk_kind: + raise CompletionRequestValidationError("MTMD media chunk modality mismatch") + kind = media.media.kind + key = media.key + used_media_keys.add(media.key) + decode_n_pos = int(mtmd_cpp.mtmd_input_chunk_get_n_pos(chunk)) + if decode_n_pos <= 0: + raise CompletionRequestValidationError("MTMD media chunk has no decoder positions") + embeddings = self._encode_media_chunk(kind=kind, key=key, chunk=chunk) + n_tokens = int(embeddings.shape[0]) + if n_tokens <= 0: + raise CompletionRequestValidationError("MTMD media chunk has no embeddings") + non_causal = bool(mtmd_cpp.mtmd_decode_use_non_causal(self.ctx, chunk)) + segment_identity = self._media_identity_tokens(kind, key, n_tokens) + positions = self._positions_for_chunk(chunk, decode_pos) + segment = PromptSegment( + kind=kind, + start_pos=identity_pos, + n_pos=n_tokens, + identity_tokens=segment_identity, + decode_start_pos=decode_pos, + decode_n_pos=decode_n_pos, + media=PromptSegment.Media( + embeddings=embeddings, + positions=positions, + non_causal=non_causal, + ), + ) + if non_causal and embeddings.shape[0] > min(self.n_batch, self.n_ubatch): + raise CompletionRequestValidationError( + f"non-causal {kind} embedding chunk exceeds model batch limits; " + "increase n_batch and n_ubatch" + ) + segments.append(segment) + identity_tokens.extend(segment_identity) + identity_pos += n_tokens + decode_pos += decode_n_pos + if used_media_keys != {media.key for media in loaded_media}: + raise CompletionRequestValidationError("not all media inputs were consumed by MTMD") + return PromptPlan( + text=prompt, + generation_prompt=generation_prompt, + text_tokens=text_tokens, + identity_tokens=identity_tokens, + segments=segments, + text_token_index_by_pos=text_token_index_by_pos, + ) + finally: + if chunks is not None: + mtmd_cpp.mtmd_input_chunks_free(chunks) + for media in loaded_media: + mtmd_cpp.mtmd_bitmap_free(media.bitmap) + if media.video_ctx: + mtmd_cpp.mtmd_helper_video_free(media.video_ctx) + if media.video_temp_path is not None: + try: + media.video_temp_path.unlink() + except OSError: + pass + + +class Model: + @dataclass(frozen=True) + class LoraAdapter: + path: str + scale: float = 1.0 + + def __init__( + self, + *, + model_path: str, + model_alias: Optional[str] = None, + chat_template: Optional[str] = None, + loras: Optional[List["Model.LoraAdapter"]] = None, + n_gpu_layers: Optional[int] = None, + split_mode: Optional[int] = None, + main_gpu: Optional[int] = None, + tensor_split: Optional[List[float]] = None, + vocab_only: Optional[bool] = None, + use_mmap: Optional[bool] = None, + use_mlock: Optional[bool] = None, + kv_overrides: Optional[Dict[str, Union[bool, int, float, str]]] = None, + n_ctx: Optional[int], + n_batch: Optional[int], + n_ubatch: Optional[int] = None, + n_seq_max: Optional[int], + n_threads: Optional[int], + n_threads_batch: Optional[int], + rope_scaling_type: Optional[int] = None, + pooling_type: Optional[int] = None, + attention_type: Optional[int] = None, + embedding: Optional[bool] = None, + rope_freq_base: Optional[float] = None, + rope_freq_scale: Optional[float] = None, + yarn_ext_factor: Optional[float] = None, + yarn_attn_factor: Optional[float] = None, + yarn_beta_fast: Optional[float] = None, + yarn_beta_slow: Optional[float] = None, + yarn_orig_ctx: Optional[int] = None, + offload_kqv: Optional[bool] = None, + flash_attn: Optional[bool] = None, + op_offload: Optional[bool] = None, + swa_full: Optional[bool] = None, + no_perf: Optional[bool] = None, + type_k: Optional[int] = None, + type_v: Optional[int] = None, + kv_unified: bool = True, + max_seq_len: Optional[int] = None, + max_output_tokens: Optional[int] = None, + draft_model: Optional[str] = None, + draft_model_path: Optional[str] = None, + draft_model_num_pred_tokens: int = 16, + draft_model_max_ngram_size: int = 2, + draft_model_top_k: int = 1, + draft_model_p_min: float = 0.0, + draft_model_max_batch_size: Optional[int] = None, + draft_model_threads: Optional[int] = None, + draft_model_threads_batch: Optional[int] = None, + response_schema: Optional[Dict[str, Any]] = None, + store_logits: bool = True, + ) -> None: + llama_cpp.llama_backend_init() + self.backend_initialized = True + self.model_path = model_path + self.model_alias = model_alias + self.chat_template_override = chat_template + self.response_schema = response_schema + self.store_logits = store_logits + self.max_output_tokens = max_output_tokens + self.draft_model_max_batch_size = draft_model_max_batch_size + self.draft_provider: Optional[DraftProvider] = None + self.loras = list(loras or []) + self._lora_adapters: List[Any] = [] + self._lora_adapter_array: Optional[Any] = None + self._lora_scales_array: Optional[Any] = None + self.draft_llama_model: Optional[Any] = None + model_params, self._c_tensor_split, self._kv_overrides_array = ( + self.build_model_params( + n_gpu_layers=n_gpu_layers, + split_mode=split_mode, + main_gpu=main_gpu, + tensor_split=tensor_split, + vocab_only=vocab_only, + use_mmap=use_mmap, + use_mlock=use_mlock, + kv_overrides=kv_overrides, + ) + ) + llama_model = llama_cpp.llama_model_load_from_file( + model_path.encode("utf-8"), + model_params, + ) + if llama_model is None: + raise RuntimeError(f"failed to load model: {model_path}") + self.llama_model = llama_model + vocab = llama_cpp.llama_model_get_vocab(llama_model) + if vocab is None: + raise RuntimeError("failed to access model vocabulary") + self.vocab = vocab + embedding = self.resolve_embedding_mode(llama_model, embedding) + self.embedding = embedding + self.has_encoder = bool(llama_cpp.llama_model_has_encoder(llama_model)) + self.has_decoder = bool(llama_cpp.llama_model_has_decoder(llama_model)) + if self.has_encoder and not embedding: + raise RuntimeError("encoder models are not supported") + if not self.has_decoder and not (embedding and self.has_encoder): + raise RuntimeError("decoder is required") + if llama_cpp.llama_model_is_recurrent(llama_model): + self.memory_model = "recurrent" + elif llama_cpp.llama_model_is_hybrid(llama_model): + self.memory_model = "hybrid" + else: + self.memory_model = ( + "attention-unified" if kv_unified else "attention-partitioned" + ) + normalized_draft_model = draft_model + required_mtp_batch = max(1, draft_model_num_pred_tokens + 1) + if normalized_draft_model == "draft-mtp": + if n_batch is not None and n_batch < required_mtp_batch: + raise RuntimeError( + "MTP requires n_batch to fit the pending token plus draft tokens " + f"(required {required_mtp_batch}, got {n_batch})" + ) + self.draft_target_batching = normalized_draft_model is not None + if ( + normalized_draft_model is not None + and normalized_draft_model != "draft-mtp" + and not self.memory_model.startswith("attention") + ): + raise RuntimeError( + "speculative decoding is only supported for attention models" + ) + n_ctx_train = int(llama_cpp.llama_model_n_ctx_train(llama_model)) + + context_params = self.build_context_params( + n_ctx=n_ctx if n_ctx is not None else n_ctx_train, + n_batch=n_batch, + n_ubatch=n_ubatch, + n_seq_max=n_seq_max, + n_threads=n_threads, + n_threads_batch=n_threads_batch, + rope_scaling_type=rope_scaling_type, + pooling_type=pooling_type, + attention_type=attention_type, + embedding=embedding, + rope_freq_base=rope_freq_base, + rope_freq_scale=rope_freq_scale, + yarn_ext_factor=yarn_ext_factor, + yarn_attn_factor=yarn_attn_factor, + yarn_beta_fast=yarn_beta_fast, + yarn_beta_slow=yarn_beta_slow, + yarn_orig_ctx=yarn_orig_ctx, + offload_kqv=offload_kqv, + flash_attn=flash_attn, + op_offload=op_offload, + swa_full=swa_full, + no_perf=no_perf, + type_k=type_k, + type_v=type_v, + kv_unified=kv_unified, + n_rs_seq=None, + ctx_type=None, + ) + ctx = llama_cpp.llama_init_from_model(llama_model, context_params) + if ctx is None: + raise RuntimeError("failed to create context") + self.ctx = ctx + mem = llama_cpp.llama_get_memory(ctx) + if mem is None and not embedding: + raise RuntimeError("failed to access model memory") + self.mem = mem + self.n_ctx = int(llama_cpp.llama_n_ctx(ctx)) + self.n_ctx_seq = int(llama_cpp.llama_n_ctx_seq(ctx)) + self.n_seq_max = int(llama_cpp.llama_n_seq_max(ctx)) + self.n_rs_seq = int(llama_cpp.llama_n_rs_seq(ctx)) + self.n_batch = int(llama_cpp.llama_n_batch(ctx)) + self.n_ubatch = int(llama_cpp.llama_n_ubatch(ctx)) + self.n_threads_batch = ( + n_threads_batch + if n_threads_batch is not None + else max(multiprocessing.cpu_count(), 1) + ) + self.mtmd_processor: Optional[MTMDProcessor] = None + self._embedding_batch: Optional[llama_cpp.llama_batch] = None + self._embedding_batch_refs: List[Any] = [] + if normalized_draft_model == "draft-mtp" and self.n_batch < required_mtp_batch: + raise RuntimeError( + "MTP requires runtime n_batch to fit the pending token plus draft tokens " + f"(required {required_mtp_batch}, got {self.n_batch})" + ) + self.n_ctx_train = n_ctx_train + self.n_vocab = int(llama_cpp.llama_vocab_n_tokens(self.vocab)) + self.n_embd = int(llama_cpp.llama_model_n_embd(self.llama_model)) + self.n_embd_inp = int(llama_cpp.llama_model_n_embd_inp(self.llama_model)) + self.n_embd_out = int(llama_cpp.llama_model_n_embd_out(self.llama_model)) + if self.n_embd_out <= 0: + self.n_embd_out = self.n_embd + self.kv_unified = kv_unified + self.max_seq_len_limit = min(self.request_context_limit, self.n_ctx_train) + if max_seq_len is None: + self.max_seq_len = self.max_seq_len_limit + else: + if max_seq_len <= 0: + llama_cpp.llama_free(self.ctx) + llama_cpp.llama_model_free(self.llama_model) + if self.backend_initialized: + llama_cpp.llama_backend_free() + self.backend_initialized = False + raise RuntimeError("max_seq_len must be greater than 0") + if max_seq_len > self.max_seq_len_limit: + llama_cpp.llama_free(self.ctx) + llama_cpp.llama_model_free(self.llama_model) + if self.backend_initialized: + llama_cpp.llama_backend_free() + self.backend_initialized = False + raise RuntimeError( + f"max_seq_len ({max_seq_len}) exceeds runtime limit ({self.max_seq_len_limit})" + ) + self.max_seq_len = max_seq_len + self.batch = llama_cpp.llama_batch_init(self.n_batch, 0, self.n_seq_max) + self.bos_token = int(llama_cpp.llama_vocab_bos(self.vocab)) + self.eos_token = int(llama_cpp.llama_vocab_eos(self.vocab)) + self.cls_token = self.bos_token + self.sep_token = int(llama_cpp.llama_vocab_sep(self.vocab)) + self.fim_pre_token = int(llama_cpp.llama_vocab_fim_pre(self.vocab)) + self.fim_mid_token = int(llama_cpp.llama_vocab_fim_mid(self.vocab)) + self.fim_suf_token = int(llama_cpp.llama_vocab_fim_suf(self.vocab)) + self.add_bos_token = bool(llama_cpp.llama_vocab_get_add_bos(self.vocab)) + self.add_eos_token = bool(llama_cpp.llama_vocab_get_add_eos(self.vocab)) + self.add_space_prefix = ( + self._meta_value("tokenizer.ggml.add_space_prefix") != "false" + ) + if normalized_draft_model is None: + self.draft_provider = None + elif normalized_draft_model == "prompt-lookup-decoding": + self.draft_provider = PromptLookupDecoding( + max_ngram_size=draft_model_max_ngram_size, + num_pred_tokens=draft_model_num_pred_tokens, + ) + elif normalized_draft_model == "draft-mtp": + draft_llama_model = self.llama_model + if draft_model_path is not None: + draft_llama_model = llama_cpp.llama_model_load_from_file( + draft_model_path.encode("utf-8"), + model_params, + ) + if draft_llama_model is None: + llama_cpp.llama_batch_free(self.batch) + llama_cpp.llama_free(self.ctx) + self._free_lora_adapters() + llama_cpp.llama_model_free(self.llama_model) + if self.backend_initialized: + llama_cpp.llama_backend_free() + self.backend_initialized = False + raise RuntimeError(f"failed to load MTP draft model: {draft_model_path}") + self.draft_llama_model = draft_llama_model + if self.n_ubatch < self.n_seq_max: + mtp_n_batch = self.n_batch + else: + mtp_n_batch = min( + self.n_batch, + max(self.n_ubatch, self.n_seq_max, required_mtp_batch), + ) + mtp_context_params = self.build_context_params( + n_ctx=self.n_ctx, + n_batch=mtp_n_batch, + n_ubatch=min(self.n_ubatch, mtp_n_batch), + n_seq_max=self.n_seq_max, + n_threads=( + draft_model_threads + if draft_model_threads is not None + else n_threads + ), + n_threads_batch=( + draft_model_threads_batch + if draft_model_threads_batch is not None + else ( + draft_model_threads + if draft_model_threads is not None + else n_threads_batch + ) + ), + rope_scaling_type=rope_scaling_type, + pooling_type=pooling_type, + attention_type=attention_type, + embedding=embedding, + rope_freq_base=rope_freq_base, + rope_freq_scale=rope_freq_scale, + yarn_ext_factor=yarn_ext_factor, + yarn_attn_factor=yarn_attn_factor, + yarn_beta_fast=yarn_beta_fast, + yarn_beta_slow=yarn_beta_slow, + yarn_orig_ctx=yarn_orig_ctx, + offload_kqv=offload_kqv, + flash_attn=flash_attn, + op_offload=op_offload, + swa_full=swa_full, + no_perf=no_perf, + type_k=type_k, + type_v=type_v, + kv_unified=kv_unified, + n_rs_seq=0, + ctx_type=llama_cpp.LLAMA_CONTEXT_TYPE_MTP, + n_outputs_max=min(mtp_n_batch, self.n_seq_max), + ctx_other=self.ctx, + ) + try: + self.draft_provider = MTPDraftProvider( + model=self, + draft_model=draft_llama_model, + context_params=mtp_context_params, + num_pred_tokens=draft_model_num_pred_tokens, + top_k=draft_model_top_k, + p_min=draft_model_p_min, + ) + except BaseException: + llama_cpp.llama_batch_free(self.batch) + llama_cpp.llama_free(self.ctx) + self._free_lora_adapters() + if self.draft_llama_model is not None: + llama_cpp.llama_model_free(self.draft_llama_model) + self.draft_llama_model = None + llama_cpp.llama_model_free(self.llama_model) + if self.backend_initialized: + llama_cpp.llama_backend_free() + self.backend_initialized = False + raise + else: + raise RuntimeError(f"unsupported draft model: {draft_model}") + try: + self._load_lora_adapters(self.loras) + self._apply_lora_adapters(self.ctx, "target") + if ( + isinstance(self.draft_provider, MTPDraftProvider) + and self.draft_llama_model is None + ): + self._apply_lora_adapters(self.draft_provider.ctx, "MTP draft") + except BaseException: + if self.draft_provider is not None: + self.draft_provider.close() + llama_cpp.llama_batch_free(self.batch) + llama_cpp.llama_free(self.ctx) + self._free_lora_adapters() + if self.draft_llama_model is not None: + llama_cpp.llama_model_free(self.draft_llama_model) + self.draft_llama_model = None + llama_cpp.llama_model_free(self.llama_model) + if self.backend_initialized: + llama_cpp.llama_backend_free() + self.backend_initialized = False + raise + self.chat_formatter = self._build_chat_formatter() + + @staticmethod + def build_model_params( + *, + n_gpu_layers: Optional[int], + split_mode: Optional[int], + main_gpu: Optional[int], + tensor_split: Optional[List[float]], + vocab_only: Optional[bool], + use_mmap: Optional[bool], + use_mlock: Optional[bool], + kv_overrides: Optional[Dict[str, Union[bool, int, float, str]]], + ) -> Tuple[Any, Optional[Any], Optional[Any]]: + model_params = llama_cpp.llama_model_default_params() + if n_gpu_layers is not None: + model_params.n_gpu_layers = 0x7FFFFFFF if n_gpu_layers == -1 else n_gpu_layers + if split_mode is not None: + model_params.split_mode = split_mode + if main_gpu is not None: + model_params.main_gpu = main_gpu + tensor_split_ref = None + if tensor_split is not None: + if len(tensor_split) > llama_cpp.LLAMA_MAX_DEVICES: + raise ValueError( + "tensor_split exceeds maximum supported devices " + f"({llama_cpp.LLAMA_MAX_DEVICES})" + ) + float_array = ctypes.c_float * llama_cpp.LLAMA_MAX_DEVICES + tensor_split_ref = float_array(*tensor_split) + model_params.tensor_split = tensor_split_ref + if vocab_only is not None: + model_params.vocab_only = vocab_only + if use_mmap is not None: + model_params.use_mmap = use_mmap + if use_mlock is not None: + model_params.use_mlock = use_mlock + + kv_overrides_ref = None + if kv_overrides is not None: + kv_overrides_ref = ( + llama_cpp.llama_model_kv_override * (len(kv_overrides) + 1) + )() + for index, (key, value) in enumerate(kv_overrides.items()): + kv_overrides_ref[index].key = key.encode("utf-8") + if isinstance(value, bool): + kv_overrides_ref[index].tag = llama_cpp.LLAMA_KV_OVERRIDE_TYPE_BOOL + kv_overrides_ref[index].value.val_bool = value + elif isinstance(value, int): + kv_overrides_ref[index].tag = llama_cpp.LLAMA_KV_OVERRIDE_TYPE_INT + kv_overrides_ref[index].value.val_i64 = value + elif isinstance(value, float): + kv_overrides_ref[index].tag = llama_cpp.LLAMA_KV_OVERRIDE_TYPE_FLOAT + kv_overrides_ref[index].value.val_f64 = value + elif isinstance(value, str): + encoded = value.encode("utf-8") + if len(encoded) > 128: + raise ValueError(f"kv_overrides value for {key} is too long") + encoded = encoded.ljust(128, b"\0") + kv_overrides_ref[index].tag = llama_cpp.LLAMA_KV_OVERRIDE_TYPE_STR + address = cast( + int, + ctypes.addressof(kv_overrides_ref[index].value) + + cast( + Any, + llama_cpp.llama_model_kv_override_value.val_str, + ).offset, + ) + buffer_start = ctypes.cast(address, ctypes.POINTER(ctypes.c_char)) + ctypes.memmove(buffer_start, encoded, 128) + else: + raise ValueError(f"unsupported kv_override value for {key}: {value!r}") + kv_overrides_ref[-1].key = b"\0" + model_params.kv_overrides = kv_overrides_ref + return model_params, tensor_split_ref, kv_overrides_ref + + @staticmethod + def build_context_params( + *, + n_ctx: Optional[int], + n_batch: Optional[int], + n_ubatch: Optional[int], + n_seq_max: Optional[int], + n_threads: Optional[int], + n_threads_batch: Optional[int], + rope_scaling_type: Optional[int], + pooling_type: Optional[int], + attention_type: Optional[int], + embedding: bool, + rope_freq_base: Optional[float], + rope_freq_scale: Optional[float], + yarn_ext_factor: Optional[float], + yarn_attn_factor: Optional[float], + yarn_beta_fast: Optional[float], + yarn_beta_slow: Optional[float], + yarn_orig_ctx: Optional[int], + offload_kqv: Optional[bool], + flash_attn: Optional[bool], + op_offload: Optional[bool], + swa_full: Optional[bool], + no_perf: Optional[bool], + type_k: Optional[int], + type_v: Optional[int], + kv_unified: bool, + n_rs_seq: Optional[int] = None, + ctx_type: Optional[int] = None, + n_outputs_max: Optional[int] = None, + ctx_other: Optional[Any] = None, + ) -> Any: + context_params = llama_cpp.llama_context_default_params() + if n_ctx is not None: + context_params.n_ctx = n_ctx + if n_batch is not None: + context_params.n_batch = min(int(context_params.n_ctx), n_batch) + if n_ubatch is not None: + context_params.n_ubatch = min(int(context_params.n_batch), n_ubatch) + if n_seq_max is not None: + context_params.n_seq_max = n_seq_max + if n_rs_seq is not None: + context_params.n_rs_seq = n_rs_seq + if n_threads is not None: + context_params.n_threads = n_threads + if n_threads_batch is not None: + context_params.n_threads_batch = n_threads_batch + if ctx_type is not None: + context_params.ctx_type = ctx_type + if n_outputs_max is not None: + context_params.n_outputs_max = n_outputs_max + if ctx_other is not None: + context_params.ctx_other = ctx_other + if rope_scaling_type is not None: + context_params.rope_scaling_type = rope_scaling_type + if pooling_type is not None: + context_params.pooling_type = pooling_type + if attention_type is not None: + context_params.attention_type = attention_type + context_params.embeddings = embedding + if rope_freq_base is not None: + context_params.rope_freq_base = rope_freq_base + if rope_freq_scale is not None: + context_params.rope_freq_scale = rope_freq_scale + if yarn_ext_factor is not None: + context_params.yarn_ext_factor = yarn_ext_factor + if yarn_attn_factor is not None: + context_params.yarn_attn_factor = yarn_attn_factor + if yarn_beta_fast is not None: + context_params.yarn_beta_fast = yarn_beta_fast + if yarn_beta_slow is not None: + context_params.yarn_beta_slow = yarn_beta_slow + if yarn_orig_ctx is not None: + context_params.yarn_orig_ctx = yarn_orig_ctx + if offload_kqv is not None: + context_params.offload_kqv = offload_kqv + if flash_attn is not None: + context_params.flash_attn_type = ( + llama_cpp.LLAMA_FLASH_ATTN_TYPE_ENABLED + if flash_attn + else llama_cpp.LLAMA_FLASH_ATTN_TYPE_DISABLED + ) + if op_offload is not None: + context_params.op_offload = op_offload + if swa_full is not None: + context_params.swa_full = swa_full + if no_perf is not None: + context_params.no_perf = no_perf + if type_k is not None: + context_params.type_k = type_k + if type_v is not None: + context_params.type_v = type_v + context_params.kv_unified = kv_unified + return context_params + + @property + def exact_checkpoints_only(self) -> bool: + return self.memory_model in {"recurrent", "hybrid"} + + @property + def has_attention_budget(self) -> bool: + return self.memory_model != "recurrent" + + @property + def attention_partitioned(self) -> bool: + return self.memory_model == "attention-partitioned" + + @property + def request_context_limit(self) -> int: + if self.attention_partitioned: + return self.n_ctx_seq + return self.n_ctx + + def _load_lora_adapters(self, loras: List["Model.LoraAdapter"]) -> None: + for lora in loras: + adapter = llama_cpp.llama_adapter_lora_init( + self.llama_model, + lora.path.encode("utf-8"), + ) + if adapter is None: + raise RuntimeError(f"failed to load LoRA adapter: {lora.path}") + self._lora_adapters.append(adapter) + + if not self._lora_adapters: + return + + adapter_array_type = llama_cpp.llama_adapter_lora_p_ctypes * len( + self._lora_adapters + ) + scale_array_type = ctypes.c_float * len(self._lora_adapters) + self._lora_adapter_array = adapter_array_type(*self._lora_adapters) + self._lora_scales_array = scale_array_type( + *(float(lora.scale) for lora in loras) + ) + + def _apply_lora_adapters(self, ctx: Any, context_name: str) -> None: + if not self._lora_adapters: + return + if self._lora_adapter_array is None or self._lora_scales_array is None: + raise RuntimeError("LoRA adapter arrays are not initialized") + result = llama_cpp.llama_set_adapters_lora( + ctx, + self._lora_adapter_array, + len(self._lora_adapters), + self._lora_scales_array, + ) + if result: + raise RuntimeError(f"failed to apply LoRA adapters to {context_name} context") + + def _free_lora_adapters(self) -> None: + while self._lora_adapters: + adapter = self._lora_adapters.pop() + llama_cpp.llama_adapter_lora_free(adapter) + self._lora_adapter_array = None + self._lora_scales_array = None + + def close(self) -> None: + if self.mtmd_processor is not None: + self.mtmd_processor.close() + self.mtmd_processor = None + if self.draft_provider is not None: + self.draft_provider.close() + llama_cpp.llama_batch_free(self.batch) + llama_cpp.llama_free(self.ctx) + self._free_lora_adapters() + if self.draft_llama_model is not None: + llama_cpp.llama_model_free(self.draft_llama_model) + self.draft_llama_model = None + llama_cpp.llama_model_free(self.llama_model) + if self.backend_initialized: + llama_cpp.llama_backend_free() + self.backend_initialized = False + + @staticmethod + def _model_meta_key_by_index(llama_model: Any, index: int) -> Optional[str]: + capacity = 256 + while True: + buffer = ctypes.create_string_buffer(capacity) + count = int( + llama_cpp.llama_model_meta_key_by_index( + llama_model, + index, + cast(Any, buffer), + capacity, + ) + ) + if count < 0: + return None + if count < capacity: + return buffer.value.decode("utf-8", errors="ignore") + capacity = count + 1 + + @staticmethod + def _model_meta_value(llama_model: Any, key: str) -> Optional[str]: + encoded = key.encode("utf-8") + capacity = 256 + while True: + buffer = ctypes.create_string_buffer(capacity) + count = int( + llama_cpp.llama_model_meta_val_str( + llama_model, + encoded, + cast(Any, buffer), + capacity, + ) + ) + if count < 0: + return None + if count < capacity: + return buffer.value.decode("utf-8", errors="ignore") + capacity = count + 1 + + @staticmethod + def _parse_pooling_type(value: str) -> Optional[int]: + normalized = value.strip().lower() + try: + return int(normalized) + except ValueError: + return { + "none": llama_cpp.LLAMA_POOLING_TYPE_NONE, + "mean": llama_cpp.LLAMA_POOLING_TYPE_MEAN, + "cls": llama_cpp.LLAMA_POOLING_TYPE_CLS, + "last": llama_cpp.LLAMA_POOLING_TYPE_LAST, + "rank": llama_cpp.LLAMA_POOLING_TYPE_RANK, + }.get(normalized) + + @classmethod + def detect_embedding_model(cls, llama_model: Any) -> bool: + for index in range(int(llama_cpp.llama_model_meta_count(llama_model))): + key = cls._model_meta_key_by_index(llama_model, index) + if key is None or not key.endswith(".pooling_type"): + continue + value = cls._model_meta_value(llama_model, key) + if value is None: + continue + pooling_type = cls._parse_pooling_type(value) + return pooling_type in { + llama_cpp.LLAMA_POOLING_TYPE_MEAN, + llama_cpp.LLAMA_POOLING_TYPE_CLS, + llama_cpp.LLAMA_POOLING_TYPE_LAST, + } + return False + + @classmethod + def resolve_embedding_mode( + cls, + llama_model: Any, + embedding: Optional[bool], + ) -> bool: + if embedding is not None: + return embedding + return cls.detect_embedding_model(llama_model) + + def _meta_value(self, key: str) -> Optional[str]: + return self._model_meta_value(self.llama_model, key) + + def _build_chat_formatter(self) -> Optional[Jinja2ChatFormatter]: + template_text = self.chat_template_override + if template_text is None: + template = llama_cpp.llama_model_chat_template(self.llama_model, None) + if template: + template_text = template.decode("utf-8", errors="ignore") + if not template_text: + return None + bos_token = "" + eos_token = "" + if self.bos_token != -1: + bos_text = llama_cpp.llama_vocab_get_text(self.vocab, self.bos_token) + bos_token = bos_text.decode("utf-8", errors="ignore") if bos_text else "" + if self.eos_token != -1: + eos_text = llama_cpp.llama_vocab_get_text(self.vocab, self.eos_token) + eos_token = eos_text.decode("utf-8", errors="ignore") if eos_text else "" + return Jinja2ChatFormatter( + template=template_text, + bos_token=bos_token, + eos_token=eos_token, + ) + + def tokenize( + self, + text: str, + *, + add_bos: bool = True, + special: bool = True, + ) -> List[int]: + encoded = text.encode("utf-8") + capacity = max(32, len(encoded) + 32) + while True: + tokens = (llama_cpp.llama_token * capacity)() + count = int( + llama_cpp.llama_tokenize( + self.vocab, + encoded, + len(encoded), + tokens, + capacity, + add_bos, + special, + ) + ) + if count >= 0: + return [int(tokens[index]) for index in range(count)] + capacity = max(capacity * 2, -count) + + def build_prompt_tokens(self, prompt: str, suffix: Optional[str]) -> List[int]: + if suffix is None: + return self.tokenize(prompt) + if min(self.fim_pre_token, self.fim_mid_token, self.fim_suf_token) < 0: + raise ValueError("suffix is not supported by this model") + bos_tokens = [self.cls_token if self.cls_token != -1 else self.bos_token] + eos_tokens = [self.sep_token if self.sep_token != -1 else self.eos_token] + if not self.add_bos_token or bos_tokens[:1] == [-1]: + bos_tokens = [] + if not self.add_eos_token and self.sep_token == -1: + eos_tokens = [] + suffix_text = suffix + suffix_space_prefix = 0 + if self.add_space_prefix and suffix_text: + suffix_text = "☺" + suffix_text + suffix_space_prefix = 2 + prefix_tokens = [self.fim_pre_token] + self.tokenize( + prompt, + add_bos=False, + special=False, + ) + suffix_tokens = [self.fim_suf_token] + if suffix_text: + suffix_tokens.extend( + self.tokenize( + suffix_text, + add_bos=False, + special=False, + )[suffix_space_prefix:] + ) + return bos_tokens + prefix_tokens + suffix_tokens + [self.fim_mid_token] + eos_tokens + + def build_chat_prompt( + self, + messages: List[ChatCompletionRequestMessage], + *, + functions: Optional[List[ChatTemplateFunctionDefinition]] = None, + function_call: Optional[Union[str, ChatTemplateFunctionCall]] = None, + tools: Optional[List[ChatTemplateTool]] = None, + tool_choice: Optional[Union[str, ChatTemplateToolChoice]] = None, + reasoning_effort: Optional[str] = None, + ) -> Tuple[str, str, PromptPlan, List[str]]: + if self.chat_formatter is None: + raise ValueError("model does not provide a GGUF chat template") + if self.mtmd_processor is not None: + prompt_plan = self.mtmd_processor.build_prompt_plan( + messages=messages, + functions=functions, + function_call=function_call, + tools=tools, + tool_choice=tool_choice, + reasoning_effort=reasoning_effort, + ) + formatter_stop = [self.chat_formatter._eos_token] if self.chat_formatter._eos_token else [] + return ( + prompt_plan.text, + prompt_plan.generation_prompt, + prompt_plan, + formatter_stop, + ) + prompt, generation_prompt, formatter_stop = self.chat_formatter.format( + messages=messages, + functions=functions, + function_call=function_call, + tools=tools, + tool_choice=tool_choice, + reasoning_effort=reasoning_effort, + ) + prompt_tokens = self.tokenize(prompt, add_bos=False, special=True) + return ( + prompt, + generation_prompt, + PromptPlan.from_tokens( + prompt, + prompt_tokens, + generation_prompt=generation_prompt, + ), + formatter_stop, + ) + + def detokenize(self, tokens: Sequence[int]) -> bytes: + if not tokens: + return b"" + array = (llama_cpp.llama_token * len(tokens))(*tokens) + capacity = max(64, len(tokens) * 16) + while True: + buffer = ctypes.create_string_buffer(capacity) + count = int( + llama_cpp.llama_detokenize( + self.vocab, + array, + len(tokens), + cast(Any, buffer), + capacity, + True, + True, + ) + ) + if count >= 0: + return bytes(buffer.raw[:count]) + capacity = max(capacity * 2, -count) + + def token_bytes(self, token: int) -> bytes: + return self.detokenize([token]) + + def token_bytes_with_prev(self, prev_tokens: Sequence[int], token: int) -> bytes: + current = self.detokenize([*prev_tokens, token]) + previous = self.detokenize(prev_tokens) + return current[len(previous) :] + + def token_bytes_with_prev_bytes( + self, + prev_tokens: Sequence[int], + prev_text_bytes: Union[bytes, bytearray], + token: int, + ) -> bytes: + current = self.detokenize([*prev_tokens, token]) + return current[len(prev_text_bytes) :] + + def clear_batch(self) -> None: + self.batch.n_tokens = 0 + self._embedding_batch = None + self._embedding_batch_refs = [] + + def clear_memory(self) -> None: + if self.mem is not None: + llama_cpp.llama_memory_clear(self.mem, True) + + def add_batch_tokens( + self, + *, + seq_id: int, + start_pos: int, + tokens: Sequence[int], + output_indices: Sequence[Optional[int]], + ) -> None: + if not tokens: + return + for index, token in enumerate(tokens): + slot = self.batch.n_tokens + self.batch.token[slot] = token + self.batch.pos[slot] = start_pos + index + self.batch.seq_id[slot][0] = seq_id + self.batch.n_seq_id[slot] = 1 + self.batch.logits[slot] = int(output_indices[index] is not None) + self.batch.n_tokens += 1 + + def add_batch_embeddings( + self, + *, + seq_id: int, + embeddings: np.ndarray, + positions: np.ndarray, + output_indices: Sequence[Optional[int]], + ) -> None: + if self.batch.n_tokens: + raise RuntimeError("cannot mix token and embedding batches") + if self._embedding_batch is not None: + raise RuntimeError("only one embedding batch is supported per scheduler step") + embeddings = np.ascontiguousarray(embeddings, dtype=np.float32) + positions = np.ascontiguousarray(positions, dtype=np.int32).reshape(-1) + n_tokens = int(embeddings.shape[0]) + if n_tokens == 0: + return + if embeddings.ndim != 2 or embeddings.shape[1] != self.n_embd_inp: + raise RuntimeError("embedding batch shape does not match model input embedding size") + if len(positions) not in {n_tokens, n_tokens * 4}: + raise RuntimeError("embedding position length mismatch") + if len(output_indices) != n_tokens: + raise RuntimeError("embedding output index length mismatch") + pos_array = (llama_cpp.llama_pos * len(positions))( + *[int(pos) for pos in positions] + ) + n_seq_id_array = (ctypes.c_int32 * n_tokens)(*[1] * n_tokens) + seq_id_array = (llama_cpp.llama_seq_id * 1)(llama_cpp.llama_seq_id(seq_id)) + seq_ids_array = (ctypes.POINTER(llama_cpp.llama_seq_id) * (n_tokens + 1))() + for index in range(n_tokens): + seq_ids_array[index] = seq_id_array + logits_array = (ctypes.c_int8 * n_tokens)( + *[int(output_index is not None) for output_index in output_indices] + ) + batch = llama_cpp.llama_batch( + n_tokens=n_tokens, + token=None, + embd=embeddings.ctypes.data_as(ctypes.POINTER(ctypes.c_float)), + pos=pos_array, + n_seq_id=n_seq_id_array, + seq_id=seq_ids_array, + logits=logits_array, + ) + self._embedding_batch = batch + self._embedding_batch_refs = [ + embeddings, + positions, + pos_array, + n_seq_id_array, + seq_id_array, + seq_ids_array, + logits_array, + ] + + def decode(self) -> None: + batch = self._embedding_batch if self._embedding_batch is not None else self.batch + if self.embedding and self.has_encoder: + operation = "llama_encode" + result = int(llama_cpp.llama_encode(self.ctx, batch)) + else: + operation = "llama_decode" + result = int(llama_cpp.llama_decode(self.ctx, batch)) + if result != 0: + raise RuntimeError(f"{operation} failed with code {result}") + + def process_draft_batch(self) -> None: + if self.draft_provider is not None: + self.draft_provider.process(self.batch) + + def set_draft_processing_enabled(self, enabled: bool) -> None: + if self.draft_provider is not None: + self.draft_provider.set_target_processing_enabled(enabled) + + def accept_draft_tokens(self, seq_id: int, accepted_draft_tokens: int) -> None: + if self.draft_provider is not None: + self.draft_provider.accept(seq_id, accepted_draft_tokens) + + def truncate_draft_sequence(self, seq_id: int, keep_len: int) -> None: + if self.draft_provider is not None: + self.draft_provider.truncate(seq_id, keep_len) + + def copy_draft_sequence( + self, + source_seq_id: int, + dest_seq_id: int, + p0: int, + p1: int, + ) -> None: + if self.draft_provider is not None: + self.draft_provider.copy_sequence(source_seq_id, dest_seq_id, p0, p1) + + def logits(self, output_index: int) -> np.ndarray: + ptr = llama_cpp.llama_get_logits_ith(self.ctx, output_index) + if not ptr: + raise RuntimeError(f"missing logits output {output_index}") + return np.ctypeslib.as_array(ptr, shape=(self.n_vocab,)).copy() + + def embed( + self, + inputs: Sequence[Union[str, List[int]]], + ) -> Tuple[List[List[float]], int]: + if not self.embedding: + raise CompletionRequestValidationError( + "model.embedding must be true to use /v1/embeddings" + ) + pooling_type = int(llama_cpp.llama_pooling_type(self.ctx)) + if pooling_type == llama_cpp.LLAMA_POOLING_TYPE_NONE: + raise CompletionRequestValidationError( + "/v1/embeddings requires a pooled embedding model; " + "set model.pooling_type to MEAN, CLS, or LAST" + ) + if pooling_type == llama_cpp.LLAMA_POOLING_TYPE_RANK: + raise CompletionRequestValidationError( + "/v1/embeddings does not support reranking pooling" + ) + if len(inputs) > 2048: + raise CompletionRequestValidationError( + "embedding input batch size exceeds 2048" + ) + + embeddings: List[List[float]] = [] + total_tokens = 0 + batch_sizes: List[int] = [] + batch_token_count = 0 + + def decode_embedding_batch() -> None: + nonlocal batch_token_count + if not batch_sizes: + return + self.clear_memory() + self.decode() + self.clear_batch() + for seq_id in range(len(batch_sizes)): + ptr = llama_cpp.llama_get_embeddings_seq( + self.ctx, + llama_cpp.llama_seq_id(seq_id), + ) + if not ptr: + raise RuntimeError(f"missing embedding output for input {seq_id}") + embeddings.append( + np.ctypeslib.as_array(ptr, shape=(self.n_embd_out,)).astype( + float + ).tolist() + ) + batch_sizes.clear() + batch_token_count = 0 + + try: + self.clear_batch() + self.clear_memory() + for input_item in inputs: + tokens = ( + self.tokenize(input_item) + if isinstance(input_item, str) + else list(input_item) + ) + n_tokens = len(tokens) + if n_tokens == 0: + raise CompletionRequestValidationError( + "embedding input must not be empty" + ) + if n_tokens > self.n_ctx_seq: + raise CompletionRequestValidationError( + f"embedding input has {n_tokens} tokens, exceeding n_ctx_seq ({self.n_ctx_seq})" + ) + if n_tokens > self.n_batch: + raise CompletionRequestValidationError( + f"embedding input has {n_tokens} tokens, exceeding n_batch ({self.n_batch})" + ) + if total_tokens + n_tokens > 300_000: + raise CompletionRequestValidationError( + "embedding request exceeds 300000 total tokens" + ) + if ( + batch_token_count + n_tokens > self.n_batch + or len(batch_sizes) >= self.n_seq_max + ): + decode_embedding_batch() + seq_id = len(batch_sizes) + self.add_batch_tokens( + seq_id=seq_id, + start_pos=0, + tokens=tokens, + output_indices=[0] * n_tokens, + ) + batch_sizes.append(n_tokens) + batch_token_count += n_tokens + total_tokens += n_tokens + decode_embedding_batch() + finally: + self.clear_batch() + self.clear_memory() + return embeddings, total_tokens + + +class SequenceDiskCache(SequenceCache): + """Directory-backed cache for serialized llama.cpp sequence state.""" + + @dataclass + class Entry: + entry_id: int + path: Path + tokens: Tuple[int, ...] + size_bytes: int + has_prompt_logits: bool + last_accessed: float + + @dataclass(frozen=True) + class Header: + tokens: Tuple[int, ...] + has_prompt_logits: bool + + @dataclass + class Payload: + tokens: List[int] + state_bytes: np.ndarray + prompt_logits: Optional[np.ndarray] + + FORMAT_VERSION = "1" + + TENSOR_TOKENS = "tokens" + TENSOR_STATE = "state" + TENSOR_PROMPT_LOGITS = "prompt_logits" + + METADATA_FORMAT = "sequence_cache_format" + METADATA_COMPATIBILITY_KEY = "compatibility_key" + + def __init__( + self, + *, + path: Union[str, Path], + max_bytes: int, + compatibility_key: str, + min_tokens: int = 128, + ) -> None: + self.path = Path(path) + self.path.mkdir(parents=True, exist_ok=True) + self.max_bytes = max(0, int(max_bytes)) + self.min_tokens = max(1, int(min_tokens)) + self.compatibility_key = compatibility_key + self._metadata = { + self.METADATA_FORMAT: self.FORMAT_VERSION, + self.METADATA_COMPATIBILITY_KEY: compatibility_key, + } + self._trie = RadixTrie() + self._entries_by_id: Dict[int, SequenceDiskCache.Entry] = {} + self._entries_by_tokens: Dict[Tuple[int, ...], SequenceDiskCache.Entry] = {} + self._next_entry_id = 0 + self._size_bytes = 0 + self._load_entries() + self._evict_if_needed() + + @staticmethod + def fingerprint_file(path: str) -> str: + stat = os.stat(path) + payload = f"{Path(path).resolve()}:{stat.st_size}:{stat.st_mtime_ns}" + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + @classmethod + def compatibility_key_for_model(cls, model: "Model") -> str: + payload: Dict[str, Any] = { + "model": cls.fingerprint_file(model.model_path), + "memory_model": model.memory_model, + "loras": [ + { + "adapter": cls.fingerprint_file(lora.path), + "scale": float(lora.scale), + } + for lora in model.loras + ], + } + if model.memory_model == "attention-partitioned": + payload["attention_streams"] = model.n_seq_max + encoded = json.dumps( + payload, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + @staticmethod + def _safe_open(path: Path) -> Any: + try: + from safetensors import safe_open + except ImportError as exc: + raise ImportError( + "disk_cache requires safetensors. Install it with `pip install safetensors`." + ) from exc + + return safe_open(str(path), framework="numpy") + + @staticmethod + def _save_file( + tensors: Dict[str, np.ndarray], + path: Path, + metadata: Dict[str, str], + ) -> None: + from safetensors.numpy import save_file + + save_file(tensors, str(path), metadata=metadata) + + def _metadata_compatible(self, metadata: Optional[Dict[str, str]]) -> bool: + if metadata is None: + return False + return all(metadata.get(key) == value for key, value in self._metadata.items()) + + def _read_entry_header(self, path: Path) -> Optional["SequenceDiskCache.Header"]: + with self._safe_open(path) as tensors: + if not self._metadata_compatible(tensors.metadata()): + return None + tokens_array = tensors.get_tensor(self.TENSOR_TOKENS) + has_prompt_logits = self.TENSOR_PROMPT_LOGITS in tensors.keys() + tokens = tuple(int(token) for token in tokens_array.tolist()) + if len(tokens) < self.min_tokens: + return None + return SequenceDiskCache.Header( + tokens=tokens, + has_prompt_logits=has_prompt_logits, + ) + + def _read_entry_payload( + self, + entry: "SequenceDiskCache.Entry", + ) -> "SequenceDiskCache.Payload": + with self._safe_open(entry.path) as tensors: + tokens = [ + int(token) + for token in tensors.get_tensor(self.TENSOR_TOKENS).tolist() + ] + state_bytes = np.ascontiguousarray( + tensors.get_tensor(self.TENSOR_STATE), + dtype=np.uint8, + ).copy() + prompt_logits = ( + tensors.get_tensor(self.TENSOR_PROMPT_LOGITS) + if entry.has_prompt_logits + and self.TENSOR_PROMPT_LOGITS in tensors.keys() + else None + ) + return SequenceDiskCache.Payload( + tokens=tokens, + state_bytes=state_bytes, + prompt_logits=prompt_logits, + ) + + def _write_entry( + self, + path: Path, + tensors: Dict[str, np.ndarray], + ) -> None: + tmp_path = path.with_name(f"{path.name}.tmp") + try: + self._save_file(tensors, tmp_path, self._metadata) + os.replace(tmp_path, path) + except Exception: + try: + tmp_path.unlink() + except FileNotFoundError: + pass + raise + + def _load_entries(self) -> None: + for path in sorted(self.path.glob("*.safetensors")): + try: + header = self._read_entry_header(path) + if header is None: + continue + self._add_entry( + path=path, + tokens=header.tokens, + has_prompt_logits=header.has_prompt_logits, + last_accessed=path.stat().st_mtime, + ) + except Exception: # noqa: BLE001 + continue + + def _add_entry( + self, + *, + path: Path, + tokens: Tuple[int, ...], + has_prompt_logits: bool, + last_accessed: Optional[float] = None, + ) -> None: + existing = self._entries_by_tokens.get(tokens) + if existing is not None: + self._remove_entry(existing) + entry_id = self._next_entry_id + self._next_entry_id += 1 + size_bytes = path.stat().st_size + entry = SequenceDiskCache.Entry( + entry_id=entry_id, + path=path, + tokens=tokens, + size_bytes=size_bytes, + has_prompt_logits=has_prompt_logits, + last_accessed=last_accessed if last_accessed is not None else time.time(), + ) + self._entries_by_id[entry_id] = entry + self._entries_by_tokens[tokens] = entry + self._trie.extend(entry_id, tokens) + self._size_bytes += size_bytes + + def _remove_entry(self, entry: "SequenceDiskCache.Entry") -> None: + if entry.entry_id in self._trie.sequence_lengths: + self._trie.truncate(entry.entry_id, 0) + self._entries_by_id.pop(entry.entry_id, None) + if self._entries_by_tokens.get(entry.tokens) is entry: + self._entries_by_tokens.pop(entry.tokens, None) + self._size_bytes = max(0, self._size_bytes - entry.size_bytes) + try: + entry.path.unlink() + except FileNotFoundError: + pass + + @staticmethod + def _touch(entry: "SequenceDiskCache.Entry") -> None: + entry.last_accessed = time.time() + try: + os.utime(entry.path, None) + except OSError: + pass + + def _oldest_entry(self) -> Optional["SequenceDiskCache.Entry"]: + if not self._entries_by_id: + return None + return min( + self._entries_by_id.values(), + key=lambda item: (item.last_accessed, item.size_bytes), + ) + + def _evict_until(self, target_bytes: int) -> None: + while self._size_bytes > target_bytes: + entry = self._oldest_entry() + if entry is None: + return + self._remove_entry(entry) + + def _evict_if_needed(self) -> None: + if self.max_bytes <= 0: + for entry in list(self._entries_by_id.values()): + self._remove_entry(entry) + return + if self._size_bytes > self.max_bytes: + self._evict_until(self.max_bytes) + self._evict_until(int(self.max_bytes * 0.9)) + + def lookup(self, tokens: Sequence[int]) -> Optional[SequenceCache.Match]: + if not tokens: + return None + entry_id, _ = self._trie.longest_prefix(tokens) + entry = self._entries_by_id.get(entry_id) + if entry is None: + return None + self._touch(entry) + return SequenceCache.Match( + tokens=entry.tokens, + has_prompt_logits=entry.has_prompt_logits, + ) + + def load( + self, + match: SequenceCache.Match, + ) -> Optional[SequenceCache.Load]: + entry = self._entries_by_tokens.get(match.tokens) + if entry is None: + return None + + payload = self._read_entry_payload(entry) + self._touch(entry) + return SequenceCache.Load( + tokens=payload.tokens, + state_bytes=payload.state_bytes, + prompt_logits=( + np.asarray(payload.prompt_logits, dtype=np.float32) + if payload.prompt_logits is not None + else None + ), + ) + + def save( + self, + tokens: Sequence[int], + state_bytes: np.ndarray, + prompt_logits: Optional[np.ndarray], + ) -> None: + if len(tokens) < self.min_tokens or self.max_bytes <= 0: + return + state = np.asarray(state_bytes, dtype=np.uint8) + if state.size <= 0: + return + state = np.ascontiguousarray(state).copy() + entry_tokens = tuple(int(token) for token in tokens) + tensors: Dict[str, np.ndarray] = { + self.TENSOR_TOKENS: np.asarray(entry_tokens, dtype=np.int32), + self.TENSOR_STATE: state, + } + if prompt_logits is not None: + tensors[self.TENSOR_PROMPT_LOGITS] = np.asarray(prompt_logits, dtype=np.float32) + path = self.path / f"{uuid.uuid4().hex}.safetensors" + self._write_entry(path, tensors) + self._add_entry( + path=path, + tokens=entry_tokens, + has_prompt_logits=prompt_logits is not None, + ) + self._evict_if_needed() + + +class MemoryPolicy(abc.ABC): + def __init__(self, scheduler: CompletionScheduler) -> None: + self.scheduler = scheduler + + def reclaim_order(self, best_free: Optional[int]) -> List[int]: + reclaim_order = [seq_id for seq_id in self.scheduler.free_sequences if seq_id != best_free] + if best_free is not None: + reclaim_order.append(best_free) + return reclaim_order + + @staticmethod + def generation_kv_for_request( + request: CompletionRequest, + prompt_length: int, + ) -> int: + return request.internal_completion_count * ( + request.effective_max_len - prompt_length + ) + + @staticmethod + def attention_kv_required( + prompt_kv: int, + reused_kv: int, + generation_kv: int, + ) -> int: + return prompt_kv - reused_kv + generation_kv + + def try_set_sequence_cache_match( + self, + request: CompletionRequest, + resident_reuse_len: int, + required_sequence_ids: int, + required_attn_kv: Optional[int] = None, + skip_attention_budget_when_unbounded: bool = False, + ) -> bool: + cache_match_length = self.scheduler.find_sequence_cache_match( + request, + resident_reuse_len, + ) + if cache_match_length <= resident_reuse_len: + return False + has_sequence_budget = len(self.scheduler.unused_sequences) >= required_sequence_ids + has_attention_budget = required_attn_kv is None or ( + skip_attention_budget_when_unbounded + and not self.scheduler.model.has_attention_budget + ) or self.scheduler.sequence_history.size + required_attn_kv <= self.scheduler.model.n_ctx + if has_sequence_budget and has_attention_budget: + return True + self.scheduler.clear_sequence_cache_match(request) + return False + + @abc.abstractmethod + def can_admit(self, request: CompletionRequest) -> bool: + raise NotImplementedError + + @abc.abstractmethod + def admit_request(self, request: CompletionRequest) -> None: + raise NotImplementedError + + @abc.abstractmethod + def copy_prompt_state( + self, + source_sequence_id: int, + dest_sequence_id: int, + keep_len: int, + ) -> None: + raise NotImplementedError + + +class AttentionMemoryPolicy(MemoryPolicy): + def match_prefix(self, tokens: Sequence[int]) -> Tuple[int, int]: + return self.scheduler.radix_trie.longest_prefix( + tokens, + preferred_sequences=self.scheduler.free_sequences, + ) + + @staticmethod + def reuse_len_for_request(request: CompletionRequest, match_length: int) -> int: + needs_generation = ( + request.payload.max_tokens != 0 + and request.effective_max_len > len(request.prompt_tokens) + ) + reuse_len = match_length + if needs_generation and request.prompt_tokens: + reuse_len = min(reuse_len, len(request.prompt_tokens) - 1) + return request.prompt_plan.clamp_to_reusable_boundary(reuse_len) + + def admit_request(self, request: CompletionRequest) -> None: + match_seq_id = request.match_sequence_id + match_length = request.match_length + reuse_len = self.reuse_len_for_request(request, match_length) + claimable = match_seq_id in self.scheduler.free_sequences + if request.sequence_cache_match is not None: + base_seq_id = self.scheduler.claim_unused_sequence() + reuse_len, request.prompt_logits = self.scheduler.hydrate_sequence_cache_match( + request, + base_seq_id, + ) + elif claimable: + base_seq_id = self.scheduler.claim_free_sequence(match_seq_id) + if self.scheduler.radix_trie.length(base_seq_id) > reuse_len: + self.scheduler.truncate_sequence(base_seq_id, reuse_len) + else: + base_seq_id = self.scheduler.claim_unused_sequence() + if reuse_len > 0 and match_seq_id >= 0: + self.copy_prompt_state(match_seq_id, base_seq_id, reuse_len) + sibling_count = request.internal_completion_count - 1 + sibling_seq_ids: List[int] = [] + for _ in range(sibling_count): + seq_id = self.scheduler.claim_unused_sequence() + sibling_seq_ids.append(seq_id) + self.scheduler.activate_request( + request, + base_seq_id=base_seq_id, + sibling_seq_ids=sibling_seq_ids, + ) + request.prompt_cursor = reuse_len + if request.prompt_cursor == len(request.prompt_tokens): + request.prompt_done = True + self.scheduler.maybe_save_sequence_cache(request) + self.scheduler.start_completions( + request, + prompt_output_index=None, + prompt_logits=request.prompt_logits, + ) + + +class UnifiedAttentionMemoryPolicy(AttentionMemoryPolicy): + def can_admit(self, request: CompletionRequest) -> bool: + match_seq_id, match_length = self.match_prefix(request.prompt_tokens) + match_length = request.prompt_plan.clamp_to_reusable_boundary(match_length) + request.match_sequence_id = match_seq_id + request.match_length = match_length + claimable = match_seq_id in self.scheduler.free_sequences + required_sequence_ids = request.internal_completion_count - int(claimable) + prompt_length = len(request.prompt_tokens) + prompt_kv = request.prompt_plan.eval_token_count + reuse_len = self.reuse_len_for_request(request, match_length) + prefix_credit = match_length if claimable else reuse_len + prefix_credit_kv = max(0, min(prefix_credit, request.prompt_plan.length)) + generation_kv = self.generation_kv_for_request(request, prompt_length) + if self.try_set_sequence_cache_match( + request, + resident_reuse_len=reuse_len, + required_sequence_ids=request.internal_completion_count, + required_attn_kv=self.attention_kv_required( + prompt_kv, + reused_kv=0, + generation_kv=generation_kv, + ), + ): + return True + required_kv = self.attention_kv_required( + prompt_kv, + reused_kv=prefix_credit_kv, + generation_kv=generation_kv, + ) + if ( + len(self.scheduler.unused_sequences) >= required_sequence_ids + and self.scheduler.sequence_history.size + required_kv <= self.scheduler.model.n_ctx + ): + return True + best_free = match_seq_id if claimable else None + for seq_id in self.reclaim_order(best_free): + if len(self.scheduler.unused_sequences) < required_sequence_ids: + self.scheduler.delete_free_sequence(seq_id) + elif self.scheduler.sequence_history.size + required_kv > self.scheduler.model.n_ctx: + if seq_id == best_free and request.match_length > 0: + self.scheduler.truncate_free_sequence(seq_id, request.match_length) + else: + self.scheduler.delete_free_sequence(seq_id) + if ( + len(self.scheduler.unused_sequences) >= required_sequence_ids + and self.scheduler.sequence_history.size + required_kv <= self.scheduler.model.n_ctx + ): + request.match_sequence_id, request.match_length = self.match_prefix( + request.prompt_tokens, + ) + request.match_length = request.prompt_plan.clamp_to_reusable_boundary( + request.match_length + ) + return True + return False + + def copy_prompt_state( + self, + source_sequence_id: int, + dest_sequence_id: int, + keep_len: int, + ) -> None: + self.scheduler.copy_sequence_state( + source_sequence_id, + dest_sequence_id, + keep_len, + ) + + +class PartitionedAttentionMemoryPolicy(AttentionMemoryPolicy): + def can_admit(self, request: CompletionRequest) -> bool: + match_seq_id, match_length = self.match_prefix(request.prompt_tokens) + match_length = request.prompt_plan.clamp_to_reusable_boundary(match_length) + request.match_sequence_id = match_seq_id + request.match_length = match_length + claimable = match_seq_id in self.scheduler.free_sequences + required_sequence_ids = request.internal_completion_count - int(claimable) + reuse_len = self.reuse_len_for_request(request, match_length) + if self.try_set_sequence_cache_match( + request, + resident_reuse_len=reuse_len, + required_sequence_ids=request.internal_completion_count, + ): + return True + if len(self.scheduler.unused_sequences) >= required_sequence_ids: + return True + best_free = match_seq_id if claimable else None + for seq_id in self.reclaim_order(best_free): + self.scheduler.delete_free_sequence(seq_id) + if len(self.scheduler.unused_sequences) >= required_sequence_ids: + request.match_sequence_id, request.match_length = self.match_prefix( + request.prompt_tokens, + ) + request.match_length = request.prompt_plan.clamp_to_reusable_boundary( + request.match_length + ) + return True + return False + + def copy_prompt_state( + self, + source_sequence_id: int, + dest_sequence_id: int, + keep_len: int, + ) -> None: + self.scheduler.copy_sequence_state( + source_sequence_id, + dest_sequence_id, + keep_len, + copy_all_state=True, + ) + + +class CheckpointMemoryPolicy(MemoryPolicy): + def exact_checkpoint_match(self, tokens: Sequence[int]) -> Tuple[int, int]: + match_seq_id, match_length = self.scheduler.radix_trie.longest_prefix( + tokens, + preferred_sequences=self.scheduler.free_sequences, + exact_only=True, + ) + if match_seq_id not in self.scheduler.free_sequences: + return -1, 0 + return match_seq_id, match_length + + def can_admit(self, request: CompletionRequest) -> bool: + match_seq_id, match_length = self.exact_checkpoint_match(request.prompt_tokens) + request.match_sequence_id = match_seq_id + request.match_length = match_length + claimable = match_seq_id in self.scheduler.free_sequences + required_sequence_ids = request.internal_completion_count - int(claimable) + prompt_length = len(request.prompt_tokens) + prompt_kv = request.prompt_plan.eval_token_count + generation_kv = self.generation_kv_for_request(request, prompt_length) + if self.try_set_sequence_cache_match( + request, + resident_reuse_len=match_length, + required_sequence_ids=request.internal_completion_count, + required_attn_kv=self.attention_kv_required( + prompt_kv, + reused_kv=0, + generation_kv=generation_kv, + ), + skip_attention_budget_when_unbounded=True, + ): + return True + required_attn_kv = self.attention_kv_required( + prompt_kv, + reused_kv=max(0, min(match_length, request.prompt_plan.length)), + generation_kv=generation_kv, + ) + if len(self.scheduler.unused_sequences) >= required_sequence_ids and ( + not self.scheduler.model.has_attention_budget + or self.scheduler.sequence_history.size + required_attn_kv <= self.scheduler.model.n_ctx + ): + return True + best_free = match_seq_id if claimable else None + for seq_id in self.reclaim_order(best_free): + self.scheduler.delete_free_sequence(seq_id) + if len(self.scheduler.unused_sequences) >= required_sequence_ids and ( + not self.scheduler.model.has_attention_budget + or self.scheduler.sequence_history.size + required_attn_kv <= self.scheduler.model.n_ctx + ): + request.match_sequence_id, request.match_length = self.exact_checkpoint_match( + request.prompt_tokens, + ) + return True + return False + + def admit_request(self, request: CompletionRequest) -> None: + match_seq_id = request.match_sequence_id + match_length = request.match_length + claimable = match_seq_id in self.scheduler.free_sequences + if request.sequence_cache_match is not None: + base_seq_id = self.scheduler.claim_unused_sequence() + match_length, request.prompt_logits = self.scheduler.hydrate_sequence_cache_match( + request, + base_seq_id, + ) + elif claimable: + base_seq_id = self.scheduler.claim_free_sequence(match_seq_id) + request.prompt_logits = self.scheduler.checkpoint_logits.get(base_seq_id) + self.scheduler.metrics.checkpoint_hits_total += 1 + else: + base_seq_id = self.scheduler.claim_unused_sequence() + request.prompt_logits = None + sibling_count = request.internal_completion_count - 1 + sibling_seq_ids: List[int] = [] + for _ in range(sibling_count): + seq_id = self.scheduler.claim_unused_sequence() + sibling_seq_ids.append(seq_id) + self.scheduler.activate_request( + request, + base_seq_id=base_seq_id, + sibling_seq_ids=sibling_seq_ids, + ) + request.prompt_cursor = match_length + if request.prompt_cursor == len(request.prompt_tokens): + request.prompt_done = True + self.scheduler.maybe_save_prompt_checkpoint(request) + self.scheduler.maybe_save_sequence_cache(request) + self.scheduler.start_completions( + request, + prompt_output_index=None, + prompt_logits=request.prompt_logits, + ) + + def copy_prompt_state( + self, + source_sequence_id: int, + dest_sequence_id: int, + keep_len: int, + ) -> None: + self.scheduler.copy_sequence_state( + source_sequence_id, + dest_sequence_id, + keep_len, + ) + + +class CompletionScheduler: + @dataclass + class BatchItem: + @dataclass + class Prefill: + embeddings: Optional[np.ndarray] = None + positions: Optional[np.ndarray] = None + non_causal: bool = False + prompt_advance_to: Optional[int] = None + + @dataclass + class Decode: + completion_index: int + pending_count: int + accepted_draft_count: int = 0 + sampled_pending_token: Optional[int] = None + rollback_keep_len: Optional[int] = None + rollback_accepted_draft_count: int = 0 + rollback_draft_processed: bool = False + deferred_accept_draft_count: Optional[int] = None + deferred_truncate_draft_len: Optional[int] = None + + kind: Literal["prefill", "decode"] + request_id: str + seq_id: int + start_pos: int + llama_start_pos: int + tokens: List[int] + identity_tokens: List[int] + output_indices: List[Optional[int]] + output_positions: List[int] + position_increments: List[int] + prefill_state: Optional["CompletionScheduler.BatchItem.Prefill"] = None + decode_state: Optional["CompletionScheduler.BatchItem.Decode"] = None + + def require_prefill(self) -> "CompletionScheduler.BatchItem.Prefill": + if self.prefill_state is None: + raise RuntimeError("batch item is not a prefill item") + return self.prefill_state + + def require_decode(self) -> "CompletionScheduler.BatchItem.Decode": + if self.decode_state is None: + raise RuntimeError("batch item is not a decode item") + return self.decode_state + + @property + def batch_token_count(self) -> int: + if ( + self.prefill_state is not None + and self.prefill_state.embeddings is not None + ): + return int(self.prefill_state.embeddings.shape[0]) + return len(self.tokens) + + @classmethod + def prefill( + cls, + *, + request_id: str, + seq_id: int, + start_pos: int, + llama_start_pos: int, + tokens: List[int], + identity_tokens: List[int], + output_indices: List[Optional[int]], + output_positions: List[int], + position_increments: List[int], + embeddings: Optional[np.ndarray] = None, + positions: Optional[np.ndarray] = None, + non_causal: bool = False, + prompt_advance_to: Optional[int] = None, + ) -> "CompletionScheduler.BatchItem": + return cls( + kind="prefill", + request_id=request_id, + seq_id=seq_id, + start_pos=start_pos, + llama_start_pos=llama_start_pos, + tokens=tokens, + identity_tokens=identity_tokens, + output_indices=output_indices, + output_positions=output_positions, + position_increments=position_increments, + prefill_state=cls.Prefill( + embeddings=embeddings, + positions=positions, + non_causal=non_causal, + prompt_advance_to=prompt_advance_to, + ), + ) + + @classmethod + def decode( + cls, + *, + request_id: str, + seq_id: int, + start_pos: int, + llama_start_pos: int, + tokens: List[int], + identity_tokens: List[int], + output_indices: List[Optional[int]], + output_positions: List[int], + position_increments: List[int], + completion_index: int, + pending_count: int, + ) -> "CompletionScheduler.BatchItem": + return cls( + kind="decode", + request_id=request_id, + seq_id=seq_id, + start_pos=start_pos, + llama_start_pos=llama_start_pos, + tokens=tokens, + identity_tokens=identity_tokens, + output_indices=output_indices, + output_positions=output_positions, + position_increments=position_increments, + decode_state=cls.Decode( + completion_index=completion_index, + pending_count=pending_count, + ), + ) + + def __init__( + self, + model: Model, + sequence_cache: Optional[SequenceCache] = None, + ) -> None: + self.model = model + self.sequence_cache = sequence_cache + self.formatter = OpenAIFormatter(model) + self.radix_trie = RadixTrie() + self.sequence_history = SequenceHistory() + self.metrics = SchedulerMetrics() + self.checkpoint_logits: Dict[int, np.ndarray] = {} + self.claimed_sequences: set[int] = set() + self.free_sequences: "OrderedDict[int, None]" = OrderedDict() + self.unused_sequences: List[int] = list(range(self.model.n_seq_max - 1, -1, -1)) + self.requests: Dict[str, CompletionRequest] = {} + self.pending_requests: Deque[CompletionRequest] = deque() + self.active_request_ids: set[str] = set() + self.closed = False + self.sequence_round_robin = 0 + self.speculative_stats: Dict[str, int] = { + "draft_proposals": 0, + "draft_tokens_proposed": 0, + "draft_tokens_accepted": 0, + "draft_tokens_rejected": 0, + } + self.draft_acceptance_length_counts: Dict[int, int] = {} + self.defer_sampled_draft_processing = False + self.memory_policy = self.build_memory_policy() + + def build_memory_policy(self) -> MemoryPolicy: + if self.model.exact_checkpoints_only: + return CheckpointMemoryPolicy(self) + if self.model.attention_partitioned: + return PartitionedAttentionMemoryPolicy(self) + return UnifiedAttentionMemoryPolicy(self) + + def clear_resident_state(self) -> None: + self.model.clear_memory() + self.model.clear_batch() + self.radix_trie = RadixTrie() + self.sequence_history = SequenceHistory() + self.checkpoint_logits.clear() + self.claimed_sequences.clear() + self.free_sequences.clear() + self.unused_sequences = list(range(self.model.n_seq_max - 1, -1, -1)) + for seq_id in range(self.model.n_seq_max): + self.model.truncate_draft_sequence(seq_id, 0) + + def create_embedding( + self, + payload: CreateEmbeddingRequest, + ) -> CreateEmbeddingResponse: + if not self.is_idle(): + raise RuntimeError("embedding requests require an idle scheduler") + self.clear_resident_state() + try: + embeddings, total_tokens = self.model.embed(payload.normalized_input()) + return CreateEmbeddingResponse.from_embeddings( + model=payload.model, + embeddings=embeddings, + total_tokens=total_tokens, + encoding_format=payload.encoding_format, + dimensions=payload.dimensions, + ) + finally: + self.clear_resident_state() + + @staticmethod + def request_needs_prompt_logits(request: CompletionRequest) -> bool: + return request.payload.max_tokens != 0 and request.effective_max_len > len( + request.prompt_tokens + ) + + @staticmethod + def request_needs_uncached_prompt_logprobs(request: CompletionRequest) -> bool: + return request.payload.echo and request.payload.logprobs is not None + + @staticmethod + def clear_sequence_cache_match(request: CompletionRequest) -> None: + request.sequence_cache_match = None + request.sequence_cache_match_length = 0 + + def can_lookup_sequence_cache(self, request: CompletionRequest) -> bool: + return ( + self.sequence_cache is not None + and bool(request.prompt_tokens) + and not self.request_needs_uncached_prompt_logprobs(request) + ) + + def is_sequence_cache_match_usable( + self, + request: CompletionRequest, + match: SequenceCache.Match, + resident_reuse_len: int, + ) -> bool: + match_length = match.length + if ( + match_length <= resident_reuse_len + or match_length > len(request.prompt_tokens) + or tuple(request.prompt_tokens[:match_length]) != match.tokens + or not request.prompt_plan.is_reusable_boundary(match_length) + ): + return False + return not ( + self.request_needs_prompt_logits(request) + and match_length == len(request.prompt_tokens) + and not match.has_prompt_logits + ) + + def find_sequence_cache_match( + self, + request: CompletionRequest, + resident_reuse_len: int, + ) -> int: + self.clear_sequence_cache_match(request) + if not self.can_lookup_sequence_cache(request): + return 0 + assert self.sequence_cache is not None + try: + match = self.sequence_cache.lookup(request.prompt_tokens) + except Exception: # noqa: BLE001 + self.metrics.sequence_cache_lookup_failures_total += 1 + return 0 + if match is None: + return 0 + if not self.is_sequence_cache_match_usable( + request, + match, + resident_reuse_len, + ): + return 0 + request.sequence_cache_match = match + request.sequence_cache_match_length = match.length + return match.length + + def fail_sequence_cache_load( + self, + request: CompletionRequest, + seq_id: int, + ) -> Tuple[int, Optional[np.ndarray]]: + llama_cpp.llama_memory_seq_rm(self.model.mem, seq_id, 0, -1) + self.model.truncate_draft_sequence(seq_id, 0) + self.clear_sequence_cache_match(request) + self.metrics.sequence_cache_load_failures_total += 1 + return 0, None + + def load_sequence_state_bytes( + self, + seq_id: int, + state_bytes: np.ndarray, + ) -> bool: + state = np.ascontiguousarray(state_bytes, dtype=np.uint8) + state_size = int(state.size) + if state_size <= 0: + return False + state_buffer = (ctypes.c_uint8 * state_size).from_buffer(cast(Any, state)) + loaded_bytes = int( + llama_cpp.llama_state_seq_set_data( + self.model.ctx, + state_buffer, + state_size, + cast(llama_cpp.llama_seq_id, seq_id), + ) + ) + return loaded_bytes == state_size + + def save_sequence_state_bytes(self, seq_id: int) -> Optional[np.ndarray]: + c_seq_id = cast(llama_cpp.llama_seq_id, seq_id) + state_size = int(llama_cpp.llama_state_seq_get_size(self.model.ctx, c_seq_id)) + if state_size <= 0: + return None + state_buffer = (ctypes.c_uint8 * state_size)() + state_bytes = int( + llama_cpp.llama_state_seq_get_data( + self.model.ctx, + state_buffer, + state_size, + c_seq_id, + ) + ) + if state_bytes <= 0: + return None + return np.ctypeslib.as_array(state_buffer, shape=(state_bytes,)).copy() + + def hydrate_sequence_cache_match( + self, + request: CompletionRequest, + seq_id: int, + ) -> Tuple[int, Optional[np.ndarray]]: + match = request.sequence_cache_match + if self.sequence_cache is None or match is None: + return 0, None + try: + loaded = self.sequence_cache.load(match) + except Exception: # noqa: BLE001 + return self.fail_sequence_cache_load(request, seq_id) + if loaded is None: + return self.fail_sequence_cache_load(request, seq_id) + tokens = list(loaded.tokens) + expected_length = request.sequence_cache_match_length + if ( + len(tokens) != expected_length + or tokens != request.prompt_tokens[:expected_length] + ): + return self.fail_sequence_cache_load(request, seq_id) + if ( + len(tokens) == len(request.prompt_tokens) + and self.request_needs_prompt_logits(request) + and loaded.prompt_logits is None + ): + return self.fail_sequence_cache_load(request, seq_id) + if not self.load_sequence_state_bytes(seq_id, loaded.state_bytes): + return self.fail_sequence_cache_load(request, seq_id) + self.radix_trie.extend(seq_id, tokens) + self.sequence_history.extend( + seq_id, + tokens, + request.prompt_plan.position_increments_up_to(len(tokens)), + ) + self.model.truncate_draft_sequence(seq_id, 0) + self.metrics.sequence_cache_hits_total += 1 + self.metrics.sequence_cache_tokens_loaded_total += len(tokens) + if len(tokens) == len(request.prompt_tokens): + return len(tokens), loaded.prompt_logits + return len(tokens), None + + def maybe_save_sequence_cache(self, request: CompletionRequest) -> None: + if ( + self.sequence_cache is None + or request.base_seq_id is None + or not request.prompt_tokens + or request.sequence_cache_match_length == len(request.prompt_tokens) + ): + return + state_bytes = self.save_sequence_state_bytes(request.base_seq_id) + if state_bytes is None: + return + try: + self.sequence_cache.save( + request.prompt_tokens, + state_bytes, + request.prompt_logits, + ) + except Exception: # noqa: BLE001 + self.metrics.sequence_cache_save_failures_total += 1 + return + self.metrics.sequence_cache_save_requests_total += 1 + + def close(self) -> None: + self.closed = True + self.model.close() + + def submit_request(self, request: CompletionRequest) -> str: + if self.closed: + raise RuntimeError("scheduler closed") + self.requests[request.id] = request + self.pending_requests.append(request) + self.metrics.requests_submitted_total += 1 + return request.id + + def cancel(self, request_id: str) -> None: + request = self.requests.get(request_id) + if request is not None: + request.cancelled = True + + def is_idle(self) -> bool: + return not self.pending_requests and not self.active_request_ids + + @staticmethod + def logits_to_logprobs(logits: np.ndarray) -> np.ndarray: + logits = logits.astype(np.float32, copy=False) + max_logit = float(np.max(logits)) + shifted = logits - max_logit + return shifted - math.log(float(np.sum(np.exp(shifted, dtype=np.float64)))) + + def step(self) -> bool: + step_started_at = time.perf_counter() + if self.closed: + return False + try: + self.admit_waiting() + batch_items = self.build_batch() + if not batch_items: + if self.maybe_fill_batched_draft_tokens(): + self.finalize_cancelled() + return True + return self.finalize_cancelled() + draft_processing_enabled = self.batch_needs_draft_processing(batch_items) + self.defer_sampled_draft_processing = ( + draft_processing_enabled + and self.supports_sampled_draft_processing + and all(item.kind == "decode" for item in batch_items) + ) + self.model.set_draft_processing_enabled(draft_processing_enabled) + self.model.clear_batch() + for item in batch_items: + prefill = item.prefill_state + if prefill is not None and prefill.embeddings is not None: + if prefill.positions is None: + raise RuntimeError("embedding batch item is missing positions") + self.model.add_batch_embeddings( + seq_id=item.seq_id, + embeddings=prefill.embeddings, + positions=prefill.positions, + output_indices=item.output_indices, + ) + else: + self.model.add_batch_tokens( + seq_id=item.seq_id, + start_pos=item.llama_start_pos, + tokens=item.tokens, + output_indices=item.output_indices, + ) + decode_started_at = time.perf_counter() + try: + non_causal_decode = any( + item.prefill_state is not None and item.prefill_state.non_causal + for item in batch_items + ) + try: + if non_causal_decode: + llama_cpp.llama_set_causal_attn(self.model.ctx, False) + self.model.decode() + finally: + if non_causal_decode: + llama_cpp.llama_set_causal_attn(self.model.ctx, True) + decode_elapsed = time.perf_counter() - decode_started_at + if draft_processing_enabled and not self.defer_sampled_draft_processing: + draft_process_started_at = time.perf_counter() + self.model.process_draft_batch() + draft_process_elapsed = time.perf_counter() - draft_process_started_at + self.metrics.draft_process_seconds_total += draft_process_elapsed + self.metrics.draft_process_calls_total += 1 + self.observe_draft_process(batch_items, draft_process_elapsed) + except BaseException as exc: # noqa: BLE001 + for request_id in list(self.active_request_ids): + self.fail_request(self.requests[request_id], exc) + for request in list(self.pending_requests): + self.fail_request(request, exc) + return True + self.metrics.observe_decode( + batch_items, + decode_elapsed, + ) + process_batch_started_at = time.perf_counter() + self.process_batch(batch_items) + self.metrics.process_batch_seconds_total += ( + time.perf_counter() - process_batch_started_at + ) + self.metrics.process_batch_calls_total += 1 + if self.defer_sampled_draft_processing: + try: + self.process_sampled_draft_batch(batch_items) + finally: + self.apply_deferred_draft_rollbacks(batch_items) + self.defer_sampled_draft_processing = False + self.finalize_cancelled() + return True + finally: + self.metrics.scheduler_step_seconds_total += ( + time.perf_counter() - step_started_at + ) + self.metrics.scheduler_step_calls_total += 1 + + def observe_draft_process( + self, + items: Sequence["CompletionScheduler.BatchItem"], + elapsed_seconds: float, + ) -> None: + if self.model.draft_provider is None or elapsed_seconds <= 0.0: + return + self.metrics.draft_seconds_total += elapsed_seconds + + def item_allows_mtp_processing(self, item: "CompletionScheduler.BatchItem") -> bool: + request = self.requests.get(item.request_id) + if request is None: + return False + if item.kind == "prefill": + return not any(segment.kind != "text" for segment in request.prompt_plan.segments) + decode = item.decode_state + if decode is None: + return False + return not request.completions[decode.completion_index].multimodal_prompt + + def batch_needs_draft_processing( + self, + batch_items: Sequence["CompletionScheduler.BatchItem"], + ) -> bool: + if self.model.draft_provider is None: + return False + if not self.draft_batch_size_allowed(batch_items): + return False + if isinstance(self.model.draft_provider, MTPDraftProvider) and not all( + self.item_allows_mtp_processing(item) for item in batch_items + ): + return False + for item in batch_items: + request = self.requests.get(item.request_id) + if request is None: + continue + if item.kind == "prefill": + if self.request_needs_prompt_logits(request): + return True + continue + decode = item.decode_state + if decode is None: + continue + completion = request.completions[decode.completion_index] + if not completion.finished: + return True + return False + + def draft_batch_size_allowed( + self, + batch_items: Sequence["CompletionScheduler.BatchItem"], + ) -> bool: + max_batch_size = getattr(self.model, "draft_model_max_batch_size", None) + if max_batch_size is None: + return True + draft_batch_size = sum( + 1 + for item in batch_items + if item.kind == "decode" and item.decode_state is not None + ) + return draft_batch_size <= max_batch_size + + def active_draft_batch_size_allowed(self, draft_batch_size: int) -> bool: + max_batch_size = getattr(self.model, "draft_model_max_batch_size", None) + return max_batch_size is None or draft_batch_size <= max_batch_size + + def admit_waiting(self) -> None: + while self.pending_requests: + request = self.pending_requests[0] + if request.cancelled: + self.fail_request(request, CompletionRequestCancelledError("request cancelled")) + continue + if not self.can_admit(request): + break + self.pending_requests.popleft() + self.admit_request(request) + + def can_admit(self, request: CompletionRequest) -> bool: + return self.memory_policy.can_admit(request) + + def admit_request(self, request: CompletionRequest) -> None: + self.memory_policy.admit_request(request) + if request.admitted: + self.metrics.requests_admitted_total += 1 + + def build_batch(self) -> List[CompletionScheduler.BatchItem]: + prompt_requests = [ + self.requests[request_id] + for request_id in self.active_request_ids + if self.requests[request_id].admitted + and not self.requests[request_id].prompt_done + and not self.requests[request_id].cancelled + ] + completions = [ + completion + for request_id in self.active_request_ids + for completion in self.requests[request_id].completions + if (completion.pending_input_tokens or completion.draft_tokens) + and not completion.finished + ] + if not prompt_requests and not completions: + return [] + + for request in prompt_requests: + if request.prompt_cursor < len(request.prompt_tokens): + segment = self._current_prompt_segment(request) + if segment.kind != "text": + token_count = min( + segment.batch_rows + if segment.media is not None and segment.media.non_causal + else self._pending_tokens_length(request), + self.model.n_batch, + ) + if segment.batch_rows > self.model.n_batch: + if segment.media is not None and segment.media.non_causal: + raise RuntimeError( + "non-causal media prompt segment exceeds model.n_batch; " + "increase n_batch" + ) + item, _ = self._build_pending_batch_item( + request, + token_count, + 0, + ) + return [item] + + ordered_sequences = self._ordered_pending_sequences(prompt_requests, completions) + allocations = self._allocate_pending_tokens_for_model(ordered_sequences) + if ordered_sequences: + self.sequence_round_robin += 1 + + items: List[CompletionScheduler.BatchItem] = [] + output_index = 0 + for source in ordered_sequences: + token_count = allocations.get(self._pending_sequence_id(source), 0) + if token_count <= 0: + continue + item, output_index = self._build_pending_batch_item( + source, + token_count, + output_index, + ) + items.append(item) + return self._split_mixed_mtp_batch(items) + + def _split_mixed_mtp_batch( + self, + items: List["CompletionScheduler.BatchItem"], + ) -> List["CompletionScheduler.BatchItem"]: + if not isinstance(self.model.draft_provider, MTPDraftProvider) or len(items) <= 1: + return items + eligibility = [self.item_allows_mtp_processing(item) for item in items] + if all(eligibility) or not any(eligibility): + return items + keep_eligible = eligibility[0] + split_items = [ + item + for item, item_eligible in zip(items, eligibility) + if item_eligible == keep_eligible + ] + self._renumber_output_indices(split_items) + return split_items + + @staticmethod + def _renumber_output_indices( + items: Sequence["CompletionScheduler.BatchItem"], + ) -> None: + next_output_index = 0 + for item in items: + for index, output_index in enumerate(item.output_indices): + if output_index is None: + continue + item.output_indices[index] = next_output_index + next_output_index += 1 + + def _allocate_pending_tokens_for_model( + self, + ordered_sequences: Sequence[Union[CompletionRequest, Completion]], + ) -> Dict[int, int]: + if not self._needs_homogeneous_recurrent_draft_batching(): + return self._allocate_pending_tokens(ordered_sequences, self.model.n_batch) + + active = [ + source + for source in ordered_sequences + if self._pending_tokens_length(source) > 0 + ] + if not active: + return {} + + first_draft_count = self._recurrent_draft_batch_token_count(active[0]) + if first_draft_count > 0: + return self._allocate_homogeneous_recurrent_draft_tokens( + active, + first_draft_count, + ) + + non_draft_sources = [ + source + for source in active + if self._recurrent_draft_batch_token_count(source) == 0 + ] + return self._allocate_pending_tokens(non_draft_sources, self.model.n_batch) + + def _needs_homogeneous_recurrent_draft_batching(self) -> bool: + return bool( + self.model.draft_provider is not None + and getattr(self.model, "draft_target_batching", False) + and self.model.exact_checkpoints_only + and self.model.n_rs_seq > 0 + ) + + def _recurrent_draft_batch_token_count( + self, + source: Union[CompletionRequest, Completion], + ) -> int: + if isinstance(source, CompletionRequest): + return 0 + if ( + source.finished + or source.pending_finish_reason is not None + or not source.pending_input_tokens + or not source.draft_tokens + ): + return 0 + token_count = min( + len(source.pending_input_tokens) + len(source.draft_tokens), + len(source.pending_input_tokens) + self.model.n_rs_seq, + ) + return token_count if token_count > len(source.pending_input_tokens) else 0 + + def _allocate_homogeneous_recurrent_draft_tokens( + self, + sources: Sequence[Union[CompletionRequest, Completion]], + token_count: int, + ) -> Dict[int, int]: + # Recurrent rollback snapshots are only valid if the whole per-sequence + # pending+draft group fits in a single llama.cpp ubatch. + safe_capacity = min(self.model.n_batch, self.model.n_ubatch) + if token_count > safe_capacity: + return self._allocate_recurrent_pending_only_tokens(sources, safe_capacity) + + allocations: Dict[int, int] = {} + remaining_capacity = safe_capacity + for source in sources: + if remaining_capacity < token_count: + break + if self._recurrent_draft_batch_token_count(source) != token_count: + continue + seq_id = self._pending_sequence_id(source) + allocations[seq_id] = token_count + remaining_capacity -= token_count + return allocations + + def _allocate_recurrent_pending_only_tokens( + self, + sources: Sequence[Union[CompletionRequest, Completion]], + capacity: int, + ) -> Dict[int, int]: + allocations: Dict[int, int] = {} + remaining_capacity = capacity + for source in sources: + if remaining_capacity <= 0: + break + if self._recurrent_draft_batch_token_count(source) == 0: + continue + if isinstance(source, CompletionRequest): + continue + pending_count = min(len(source.pending_input_tokens), remaining_capacity) + if pending_count <= 0: + continue + allocations[source.seq_id] = pending_count + remaining_capacity -= pending_count + return allocations + + def _ordered_pending_sequences( + self, + prompt_requests: Sequence[CompletionRequest], + completions: Sequence[Completion], + ) -> List[Union[CompletionRequest, Completion]]: + if not prompt_requests and not completions: + return [] + prompt_by_request_id = {request.id: request for request in prompt_requests} + completions_by_request_id: Dict[str, List[Completion]] = {} + for completion in completions: + completions_by_request_id.setdefault(completion.request_id, []).append(completion) + sequence_list: List[Union[CompletionRequest, Completion]] = [] + for request_id, request in self.requests.items(): + if request_id not in self.active_request_ids or request.cancelled: + continue + prompt_request = prompt_by_request_id.get(request_id) + if prompt_request is not None: + sequence_list.append(prompt_request) + sequence_list.extend(completions_by_request_id.get(request_id, [])) + if not sequence_list: + return [] + start = self.sequence_round_robin % len(sequence_list) + return sequence_list[start:] + sequence_list[:start] + + def _pending_sequence_id( + self, + source: Union[CompletionRequest, Completion], + ) -> int: + if isinstance(source, CompletionRequest): + if source.base_seq_id is None: + raise RuntimeError("prompt sequence is missing base seq id") + return source.base_seq_id + return source.seq_id + + def _current_prompt_segment(self, request: CompletionRequest) -> PromptSegment: + return request.prompt_plan.segment_at(request.prompt_cursor) + + def _pending_tokens_length( + self, + source: Union[CompletionRequest, Completion], + ) -> int: + if isinstance(source, CompletionRequest): + if source.prompt_cursor >= len(source.prompt_tokens): + return 0 + segment = self._current_prompt_segment(source) + if segment.kind != "text": + segment_offset = source.prompt_cursor - segment.start_pos + if segment.media is not None and segment.media.non_causal: + if segment_offset != 0: + raise RuntimeError("non-causal media segment was partially scheduled") + return segment.batch_rows + return segment.end_pos - source.prompt_cursor + return segment.end_pos - source.prompt_cursor + draft_count = ( + len(source.draft_tokens) + if source.pending_finish_reason is None + and getattr(self.model, "draft_target_batching", True) + else 0 + ) + return len(source.pending_input_tokens) + draft_count + + def _pending_allocation_for_capacity( + self, + source: Union[CompletionRequest, Completion], + already_allocated: int, + capacity: int, + ) -> int: + if capacity <= 0: + return 0 + if not isinstance(source, CompletionRequest): + remaining = self._pending_tokens_length(source) - already_allocated + return min(remaining, capacity) + segment = self._current_prompt_segment(source) + if segment.kind == "text": + remaining = self._pending_tokens_length(source) - already_allocated + return min(remaining, capacity) + if segment.media is not None and segment.media.non_causal: + if already_allocated == 0 and segment.batch_rows <= capacity: + return segment.batch_rows + return 0 + segment_offset = source.prompt_cursor - segment.start_pos + already_allocated + return segment.rows_for_capacity(segment_offset, capacity) + + def _allocate_pending_tokens( + self, + sources: Sequence[Union[CompletionRequest, Completion]], + capacity: int, + ) -> Dict[int, int]: + allocations: Dict[int, int] = {} + active = [source for source in sources if self._pending_tokens_length(source) > 0] + remaining_capacity = capacity + + for source in active: + if remaining_capacity <= 0: + break + seq_id = self._pending_sequence_id(source) + allocation = self._pending_allocation_for_capacity( + source, + allocations.get(seq_id, 0), + remaining_capacity, + ) + if allocation <= 0: + continue + allocations[seq_id] = allocations.get(seq_id, 0) + allocation + remaining_capacity -= allocation + + while remaining_capacity > 0 and active: + share = max(1, remaining_capacity // len(active)) + progress = False + next_active: List[Union[CompletionRequest, Completion]] = [] + for source in active: + seq_id = self._pending_sequence_id(source) + remaining_tokens = self._pending_tokens_length(source) - allocations.get(seq_id, 0) + if remaining_tokens <= 0: + continue + allocation = self._pending_allocation_for_capacity( + source, + allocations.get(seq_id, 0), + min(share, remaining_capacity), + ) + if allocation <= 0: + next_active.append(source) + continue + allocations[seq_id] = allocations.get(seq_id, 0) + allocation + remaining_capacity -= allocation + progress = True + if remaining_tokens > allocation and remaining_capacity > 0: + next_active.append(source) + if not progress: + break + active = next_active + return allocations + + def _build_pending_batch_item( + self, + source: Union[CompletionRequest, Completion], + token_count: int, + output_index: int, + ) -> Tuple["CompletionScheduler.BatchItem", int]: + if isinstance(source, CompletionRequest): + if source.base_seq_id is None: + raise RuntimeError("prompt sequence is missing base seq id") + segment = self._current_prompt_segment(source) + if segment.kind != "text": + segment_offset = source.prompt_cursor - segment.start_pos + non_causal = segment.media is not None and segment.media.non_causal + if non_causal: + if segment_offset != 0 or token_count < segment.batch_rows: + raise RuntimeError("non-causal media segment must be scheduled atomically") + row_count = segment.batch_rows + embeddings, positions, position_increments = segment.media_slice( + 0, + row_count, + ) + else: + row_count = segment.rows_for_capacity(segment_offset, token_count) + if row_count <= 0: + raise RuntimeError("media prompt allocation is too small") + embeddings, positions, position_increments = segment.media_slice( + segment_offset, + row_count, + ) + logical_start = source.prompt_cursor + logical_end = logical_start + row_count + output_indices = [None] * row_count + output_positions = [logical_start] * row_count + if output_positions: + output_positions[-1] = logical_end - 1 + needs_last_output = ( + (source.payload.echo and source.payload.logprobs is not None + and source.prompt_plan.has_text_token_at(logical_end)) + or self.model.exact_checkpoints_only + or logical_end == len(source.prompt_tokens) + ) + if needs_last_output and output_indices: + output_indices[-1] = output_index + output_index += 1 + return ( + CompletionScheduler.BatchItem.prefill( + request_id=source.id, + seq_id=source.base_seq_id, + start_pos=logical_start, + llama_start_pos=segment.decode_start_pos + sum( + segment.decoder_position_increments[:segment_offset] + ), + tokens=[], + identity_tokens=list( + segment.identity_tokens[ + segment_offset : segment_offset + row_count + ] + ), + output_indices=output_indices, + output_positions=output_positions, + position_increments=position_increments, + embeddings=embeddings, + positions=positions, + non_causal=non_causal, + prompt_advance_to=logical_end, + ), + output_index, + ) + segment_offset = source.prompt_cursor - segment.start_pos + max_count = min(token_count, segment.end_pos - source.prompt_cursor) + chunk = list(segment.text_tokens[segment_offset : segment_offset + max_count]) + identity_chunk = list( + segment.identity_tokens[segment_offset : segment_offset + max_count] + ) + ends_prompt = source.prompt_cursor + len(identity_chunk) == len(source.prompt_tokens) + output_indices: List[Optional[int]] = [None] * len(chunk) + output_positions = [source.prompt_cursor + index for index in range(len(chunk))] + for index, output_position in enumerate(output_positions): + if ( + source.payload.echo + and source.payload.logprobs is not None + and source.prompt_plan.has_text_token_at(output_position + 1) + ): + output_indices[index] = output_index + output_index += 1 + if self.model.exact_checkpoints_only and chunk and output_indices[-1] is None: + output_indices[-1] = output_index + output_index += 1 + elif ends_prompt and chunk and output_indices[-1] is None: + output_indices[-1] = output_index + output_index += 1 + return ( + CompletionScheduler.BatchItem.prefill( + request_id=source.id, + seq_id=source.base_seq_id, + start_pos=source.prompt_cursor, + llama_start_pos=segment.decode_start_pos + segment_offset, + tokens=chunk, + identity_tokens=identity_chunk, + output_indices=output_indices, + output_positions=output_positions, + position_increments=segment.decoder_position_increments[ + segment_offset : segment_offset + len(identity_chunk) + ], + ), + output_index, + ) + + request = self.requests[source.request_id] + pending_count = min(token_count, len(source.pending_input_tokens)) + draft_count = ( + max(0, token_count - pending_count) + if getattr(self.model, "draft_target_batching", True) + else 0 + ) + scheduled_tokens = [ + *source.pending_input_tokens[:pending_count], + *( + source.draft_tokens[:draft_count] + if source.pending_finish_reason is None + else [] + ), + ] + output_indices = list(range(output_index, output_index + len(scheduled_tokens))) + start_pos = self.radix_trie.length(source.seq_id) + return ( + CompletionScheduler.BatchItem.decode( + request_id=request.id, + seq_id=source.seq_id, + start_pos=start_pos, + llama_start_pos=self.sequence_history.position_length(source.seq_id), + tokens=list(scheduled_tokens), + identity_tokens=list(scheduled_tokens), + output_indices=output_indices, + output_positions=[ + start_pos + index + for index in range(len(scheduled_tokens)) + ], + position_increments=[1] * len(scheduled_tokens), + completion_index=source.index, + pending_count=pending_count, + ), + output_index + len(scheduled_tokens), + ) + + def process_batch(self, items: List[CompletionScheduler.BatchItem]) -> None: + output_count = sum( + output_index is not None + for item in items + for output_index in item.output_indices + ) + for item in items: + request = self.requests[item.request_id] + if request.cancelled: + continue + self.radix_trie.extend(item.seq_id, item.identity_tokens) + self.sequence_history.extend( + item.seq_id, + item.identity_tokens, + item.position_increments, + ) + if item.kind == "prefill": + request.capture_prompt_logprobs( + model=self.model, + formatter=self.formatter, + output_indices=item.output_indices, + output_positions=item.output_positions, + output_count=output_count, + output_index_to_logits_index=self.output_index_to_logits_index, + ) + prompt_output_index = self.output_index_to_logits_index( + self.last_output_index(item.output_indices), + output_count, + ) + prompt_logits = ( + self.model.logits(prompt_output_index) + if prompt_output_index is not None + else None + ) + prefill = item.require_prefill() + if prefill.prompt_advance_to is not None: + request.prompt_cursor = prefill.prompt_advance_to + else: + request.prompt_cursor += len(item.identity_tokens) + if request.prompt_cursor == len(request.prompt_tokens): + request.prompt_done = True + request.prompt_logits = prompt_logits + self.maybe_save_prompt_checkpoint(request) + self.maybe_save_sequence_cache(request) + self.start_completions( + request, + prompt_output_index=prompt_output_index, + prompt_logits=request.prompt_logits, + ) + else: + decode = item.require_decode() + completion = request.completions[decode.completion_index] + self.process_generation_item( + completion, + item, + output_count, + ) + self.finalize_request_if_ready(request) + if not self.defer_sampled_draft_processing: + self.maybe_fill_batched_draft_tokens() + + @property + def should_defer_draft_fill(self) -> bool: + return bool( + self.model.draft_provider is not None + and getattr(self.model.draft_provider, "batched_draft", False) + ) + + @property + def supports_sampled_draft_processing(self) -> bool: + return bool( + self.model.draft_provider is not None + and getattr(self.model.draft_provider, "sampled_batch_draft", False) + ) + + def maybe_fill_draft_tokens(self, completion: Completion) -> None: + if ( + self.model.draft_provider is None + or completion.finished + or completion.pending_finish_reason is not None + or completion.draft_tokens + ): + return + remaining_tokens = completion.max_total_tokens - completion.total_tokens + if not getattr(self.model, "draft_target_batching", True): + remaining_tokens = min(remaining_tokens, 1) + if remaining_tokens <= 0: + return + input_ids = self.draft_input_ids(completion) + if isinstance(self.model.draft_provider, MTPDraftProvider) and completion.multimodal_prompt: + return + can_draft = getattr(self.model.draft_provider, "can_draft", None) + if can_draft is not None and not can_draft( + int(input_ids.shape[0]), + seq_id=completion.seq_id, + ): + return + draft_started_at = time.perf_counter() + try: + proposed = self.model.draft_provider.draft( + input_ids, + seq_id=completion.seq_id, + max_tokens=remaining_tokens, + ) + finally: + draft_elapsed = time.perf_counter() - draft_started_at + self.metrics.draft_seconds_total += draft_elapsed + self.metrics.draft_generate_seconds_total += draft_elapsed + self.metrics.draft_generate_calls_total += 1 + if proposed.size == 0: + return + limited = [int(token) for token in proposed[:remaining_tokens]] + if not limited: + return + completion.draft_tokens = limited + self.speculative_stats["draft_proposals"] += 1 + self.speculative_stats["draft_tokens_proposed"] += len(limited) + + def maybe_fill_batched_draft_tokens(self) -> bool: + if not self.should_defer_draft_fill or self.model.draft_provider is None: + return False + draft_many = getattr(self.model.draft_provider, "draft_many", None) + if draft_many is None: + return False + + completions: List[Completion] = [] + draft_requests: List[Tuple[np.ndarray, int, Optional[int]]] = [] + remaining_by_completion: List[int] = [] + for request_id in self.active_request_ids: + request = self.requests.get(request_id) + if request is None or not request.admitted: + continue + for completion in request.completions: + if ( + completion.finished + or completion.pending_finish_reason is not None + or completion.draft_tokens + ): + continue + remaining_tokens = completion.max_total_tokens - completion.total_tokens + if not getattr(self.model, "draft_target_batching", True): + remaining_tokens = min(remaining_tokens, 1) + if remaining_tokens <= 0: + continue + input_ids = self.draft_input_ids(completion) + if ( + isinstance(self.model.draft_provider, MTPDraftProvider) + and completion.multimodal_prompt + ): + continue + can_draft = getattr(self.model.draft_provider, "can_draft", None) + if can_draft is not None and not can_draft( + int(input_ids.shape[0]), + seq_id=completion.seq_id, + ): + continue + completions.append(completion) + draft_requests.append((input_ids, completion.seq_id, remaining_tokens)) + remaining_by_completion.append(remaining_tokens) + + if not draft_requests: + return False + if not self.active_draft_batch_size_allowed(len(draft_requests)): + return False + + draft_started_at = time.perf_counter() + proposed_many = draft_many(draft_requests) + draft_elapsed = time.perf_counter() - draft_started_at + self.metrics.draft_seconds_total += draft_elapsed + self.metrics.draft_generate_seconds_total += draft_elapsed + self.metrics.draft_generate_calls_total += 1 + made_progress = False + + for completion, proposed, remaining_tokens in zip( + completions, + proposed_many, + remaining_by_completion, + ): + if proposed.size == 0: + continue + limited = [int(token) for token in proposed[:remaining_tokens]] + if not limited: + continue + self.speculative_stats["draft_proposals"] += 1 + self.speculative_stats["draft_tokens_proposed"] += len(limited) + completion.draft_tokens = limited + made_progress = True + return made_progress + + def draft_input_ids(self, completion: Completion) -> np.ndarray: + return np.array( + [*completion.prompt_tokens, *completion.completion_tokens], + dtype=np.intc, + ) + + def process_sampled_draft_batch( + self, + items: Sequence["CompletionScheduler.BatchItem"], + ) -> None: + if self.model.draft_provider is None: + return + process_sampled_batch = getattr( + self.model.draft_provider, + "process_sampled_batch", + None, + ) + if process_sampled_batch is None: + return + if not self.draft_batch_size_allowed(items): + return + + output_count = sum( + output_index is not None + for item in items + for output_index in item.output_indices + ) + completions: List[Completion] = [] + update_items: List["CompletionScheduler.BatchItem"] = [] + updates: List["MTPDraftProvider.SampledBatchUpdate"] = [] + remaining_by_completion: List[int] = [] + batch_row_offset = 0 + for item in items: + item_batch_rows = list( + range(batch_row_offset, batch_row_offset + len(item.tokens)) + ) + batch_row_offset += len(item.tokens) + decode = item.decode_state + if item.kind != "decode" or decode is None: + continue + request = self.requests.get(item.request_id) + if request is None: + continue + completion = request.completions[decode.completion_index] + if completion.finished or completion.pending_finish_reason is not None: + continue + remaining_tokens = completion.max_total_tokens - completion.total_tokens + if remaining_tokens <= 0: + continue + resolved_output_indices = [ + self.output_index_to_logits_index(output_index, output_count) + for output_index in item.output_indices + ] + if any(output_index is None for output_index in resolved_output_indices): + continue + sample_index = decode.pending_count + decode.accepted_draft_count - 1 + target_count = len(item.tokens) + if sample_index < 0 or sample_index >= target_count: + continue + if decode.sampled_pending_token is None: + continue + completions.append(completion) + update_items.append(item) + updates.append( + MTPDraftProvider.SampledBatchUpdate( + seq_id=item.seq_id, + start_pos=item.llama_start_pos, + tokens=item.tokens, + row_indices=item_batch_rows, + target_count=target_count, + sample_index=sample_index, + pending_token=decode.sampled_pending_token, + max_tokens=remaining_tokens, + ) + ) + remaining_by_completion.append(remaining_tokens) + + if not updates: + return + + draft_started_at = time.perf_counter() + proposed_many = process_sampled_batch(updates) + for item in update_items: + item.require_decode().rollback_draft_processed = True + draft_elapsed = time.perf_counter() - draft_started_at + self.metrics.draft_seconds_total += draft_elapsed + self.metrics.draft_sampled_batch_seconds_total += draft_elapsed + self.metrics.draft_sampled_batch_calls_total += 1 + for completion, proposed, remaining_tokens in zip( + completions, + proposed_many, + remaining_by_completion, + ): + if proposed.size == 0: + continue + limited = [int(token) for token in proposed[:remaining_tokens]] + if not limited: + continue + completion.draft_tokens = limited + self.speculative_stats["draft_proposals"] += 1 + self.speculative_stats["draft_tokens_proposed"] += len(limited) + + def rollback_draft_verification( + self, + item: "CompletionScheduler.BatchItem", + keep_len: int, + accepted_draft_count: int, + *, + defer_draft_state: bool, + ) -> None: + decode = item.require_decode() + if defer_draft_state: + decode.rollback_keep_len = keep_len + decode.rollback_accepted_draft_count = accepted_draft_count + decode.rollback_draft_processed = False + return + self.model.accept_draft_tokens(item.seq_id, accepted_draft_count) + self.truncate_sequence(item.seq_id, keep_len) + + def apply_deferred_draft_rollbacks( + self, + items: Sequence["CompletionScheduler.BatchItem"], + ) -> None: + for item in items: + decode = item.require_decode() + rollback_keep_len = decode.rollback_keep_len + if rollback_keep_len is None: + deferred_accept_draft_count = decode.deferred_accept_draft_count + if ( + not decode.rollback_draft_processed + and deferred_accept_draft_count is not None + ): + self.model.accept_draft_tokens( + item.seq_id, + deferred_accept_draft_count, + ) + deferred_truncate_draft_len = decode.deferred_truncate_draft_len + if deferred_truncate_draft_len is not None: + self.model.truncate_draft_sequence( + item.seq_id, + deferred_truncate_draft_len, + ) + decode.deferred_accept_draft_count = None + decode.deferred_truncate_draft_len = None + decode.rollback_draft_processed = False + continue + + draft_state_needs_fallback = not decode.rollback_draft_processed + if draft_state_needs_fallback: + self.model.accept_draft_tokens( + item.seq_id, + decode.rollback_accepted_draft_count, + ) + self.truncate_sequence( + item.seq_id, + rollback_keep_len, + truncate_draft=draft_state_needs_fallback, + ) + decode.rollback_keep_len = None + decode.rollback_accepted_draft_count = 0 + decode.rollback_draft_processed = False + decode.deferred_accept_draft_count = None + decode.deferred_truncate_draft_len = None + + def sample_completion_token(self, completion: Completion, output_index: int) -> int: + sample_started_at = time.perf_counter() + try: + return completion.sampler.sample(self.model.ctx, output_index) + finally: + self.metrics.sample_seconds_total += time.perf_counter() - sample_started_at + self.metrics.sample_calls_total += 1 + + def process_generation_item( + self, + completion: Completion, + item: CompletionScheduler.BatchItem, + output_count: int, + ) -> None: + if completion.finished: + return + decode = item.require_decode() + pending_count = decode.pending_count + resolved_output_indices = [ + self.output_index_to_logits_index(output_index, output_count) + for output_index in item.output_indices + ] + if any(output_index is None for output_index in resolved_output_indices): + raise RuntimeError("generation outputs are required") + logits_indices: List[int] = [] + for output_index in resolved_output_indices: + assert output_index is not None + logits_indices.append(int(output_index)) + if completion.pending_finish_reason is not None: + if self.model.store_logits: + self.checkpoint_logits[completion.seq_id] = self.model.logits( + logits_indices[-1] + ) + completion.pending_input_tokens = completion.pending_input_tokens[pending_count:] + finish_reason: str = completion.pending_finish_reason + completion.pending_finish_reason = None + self.finish_completion(completion, finish_reason) + return + + if pending_count: + completion.pending_input_tokens = completion.pending_input_tokens[pending_count:] + + decoded_draft_tokens = item.tokens[pending_count:] + accepted_draft_count = 0 + defer_draft_state = self.defer_sampled_draft_processing + if decoded_draft_tokens and pending_count <= 0: + raise RuntimeError("draft verification requires at least one pending token") + if decoded_draft_tokens: + self.metrics.draft_batches_verified_total += 1 + self.metrics.draft_target_tokens_verified_total += len(decoded_draft_tokens) + + for draft_index, draft_token in enumerate(decoded_draft_tokens): + verify_output_index = pending_count + draft_index - 1 + if verify_output_index >= len(logits_indices): + raise RuntimeError("missing target output for draft verification") + logits_index = logits_indices[verify_output_index] + sampled_token = self.sample_completion_token(completion, logits_index) + need_logits = self.model.store_logits or completion.needs_token_logprob + logits = self.model.logits(logits_index) if need_logits else None + if self.model.store_logits and logits is not None: + self.checkpoint_logits[completion.seq_id] = logits + if sampled_token != draft_token: + self.record_draft_acceptance_length(accepted_draft_count) + rejected = max(0, len(completion.draft_tokens) - accepted_draft_count) + if rejected > 0: + self.speculative_stats["draft_tokens_rejected"] += rejected + self.metrics.draft_target_tokens_wasted_total += ( + max(0, len(decoded_draft_tokens) - accepted_draft_count - 1) + ) + keep_len = item.start_pos + pending_count + accepted_draft_count + self.rollback_draft_verification( + item, + keep_len, + accepted_draft_count, + defer_draft_state=defer_draft_state, + ) + completion.draft_tokens.clear() + prev_tokens = [*completion.prompt_tokens, *completion.completion_tokens] + if logits is not None: + record = Token.from_logits( + model=self.model, + formatter=self.formatter, + prev_tokens=prev_tokens, + prev_text_bytes=completion.detokenized_prefix_bytes, + token=sampled_token, + logits=logits, + logprobs_count=completion.logprobs, + need_token_logprob=completion.needs_token_logprob, + ) + else: + record = Token.from_token( + model=self.model, + prev_tokens=prev_tokens, + prev_text_bytes=completion.detokenized_prefix_bytes, + token=sampled_token, + ) + mismatch_finish_reason: Optional[str] = self.handle_completion_token( + completion, + sampled_token, + record, + decoded=False, + ) + if mismatch_finish_reason is not None: + completion.pending_finish_reason = mismatch_finish_reason + else: + decode.sampled_pending_token = sampled_token + decode.accepted_draft_count = accepted_draft_count + return + prev_tokens = [*completion.prompt_tokens, *completion.completion_tokens] + if logits is not None: + record = Token.from_logits( + model=self.model, + formatter=self.formatter, + prev_tokens=prev_tokens, + prev_text_bytes=completion.detokenized_prefix_bytes, + token=draft_token, + logits=logits, + logprobs_count=completion.logprobs, + need_token_logprob=completion.needs_token_logprob, + ) + else: + record = Token.from_token( + model=self.model, + prev_tokens=prev_tokens, + prev_text_bytes=completion.detokenized_prefix_bytes, + token=draft_token, + ) + accepted_finish_reason: Optional[str] = self.handle_completion_token( + completion, + draft_token, + record, + decoded=True, + ) + accepted_draft_count += 1 + self.speculative_stats["draft_tokens_accepted"] += 1 + if accepted_finish_reason is not None: + self.record_draft_acceptance_length(accepted_draft_count) + rejected = max(0, len(completion.draft_tokens) - accepted_draft_count) + if rejected > 0: + self.speculative_stats["draft_tokens_rejected"] += rejected + self.metrics.draft_target_tokens_wasted_total += ( + len(decoded_draft_tokens) - accepted_draft_count + ) + keep_len = item.start_pos + pending_count + accepted_draft_count + self.rollback_draft_verification( + item, + keep_len, + accepted_draft_count, + defer_draft_state=defer_draft_state, + ) + completion.draft_tokens.clear() + decode.accepted_draft_count = accepted_draft_count + self.finish_completion(completion, accepted_finish_reason) + return + + decode.accepted_draft_count = accepted_draft_count + if decoded_draft_tokens: + self.record_draft_acceptance_length(accepted_draft_count) + if not defer_draft_state: + self.model.accept_draft_tokens(completion.seq_id, accepted_draft_count) + elif decoded_draft_tokens: + decode.deferred_accept_draft_count = accepted_draft_count + if accepted_draft_count: + completion.draft_tokens = completion.draft_tokens[accepted_draft_count:] + + final_logits_index = logits_indices[-1] + next_token = self.sample_completion_token(completion, final_logits_index) + need_final_logits = self.model.store_logits or completion.needs_token_logprob + final_logits = ( + self.model.logits(final_logits_index) if need_final_logits else None + ) + if self.model.store_logits and final_logits is not None: + self.checkpoint_logits[completion.seq_id] = final_logits + if completion.draft_tokens and next_token != completion.draft_tokens[0]: + self.speculative_stats["draft_tokens_rejected"] += len(completion.draft_tokens) + if not defer_draft_state: + self.model.truncate_draft_sequence( + completion.seq_id, + item.start_pos + len(item.tokens), + ) + else: + decode.deferred_truncate_draft_len = item.start_pos + len(item.tokens) + completion.draft_tokens.clear() + elif completion.draft_tokens and next_token == completion.draft_tokens[0]: + if getattr(self.model, "draft_target_batching", True): + completion.draft_tokens = completion.draft_tokens[1:] + else: + completion.draft_tokens.clear() + self.speculative_stats["draft_tokens_accepted"] += 1 + self.metrics.draft_tokens_reused_as_pending_total += 1 + + prev_tokens = [*completion.prompt_tokens, *completion.completion_tokens] + if final_logits is not None: + record = Token.from_logits( + model=self.model, + formatter=self.formatter, + prev_tokens=prev_tokens, + prev_text_bytes=completion.detokenized_prefix_bytes, + token=next_token, + logits=final_logits, + logprobs_count=completion.logprobs, + need_token_logprob=completion.needs_token_logprob, + ) + else: + record = Token.from_token( + model=self.model, + prev_tokens=prev_tokens, + prev_text_bytes=completion.detokenized_prefix_bytes, + token=next_token, + ) + final_finish_reason: Optional[str] = self.handle_completion_token( + completion, + next_token, + record, + decoded=False, + ) + if final_finish_reason is not None: + completion.pending_finish_reason = final_finish_reason + else: + decode.sampled_pending_token = next_token + + def maybe_save_prompt_checkpoint(self, request: CompletionRequest) -> None: + if ( + not self.model.exact_checkpoints_only + or request.prompt_checkpoint_saved + or request.base_seq_id is None + or not request.prompt_tokens + or request.prompt_logits is None + or not self.unused_sequences + ): + return + checkpoint_seq_id = self.unused_sequences.pop() + self.copy_sequence_state( + request.base_seq_id, + checkpoint_seq_id, + len(request.prompt_tokens), + ) + self.checkpoint_logits[checkpoint_seq_id] = request.prompt_logits + self.free_sequences[checkpoint_seq_id] = None + self.free_sequences.move_to_end(checkpoint_seq_id) + request.prompt_checkpoint_saved = True + self.metrics.checkpoint_saves_total += 1 + + @staticmethod + def output_index_to_logits_index( + output_index: Optional[int], + output_count: int, + ) -> Optional[int]: + if output_index is None: + return None + return output_index - output_count + + @staticmethod + def last_output_index(output_indices: Sequence[Optional[int]]) -> Optional[int]: + for output_index in reversed(output_indices): + if output_index is not None: + return output_index + return None + + def start_completions( + self, + request: CompletionRequest, + prompt_output_index: Optional[int], + prompt_logits: Optional[np.ndarray] = None, + ) -> None: + if request.completions: + return + assert request.base_seq_id is not None + prompt_tokens = list(request.prompt_plan.text_tokens) + prompt_length = len(request.prompt_tokens) + multimodal_prompt = any( + segment.kind != "text" for segment in request.prompt_plan.segments + ) + prompt_text = request.prompt_text + if request.payload.stop is None: + stop_sequences: List[bytes] = [] + elif isinstance(request.payload.stop, str): + stop_sequences = [request.payload.stop.encode("utf-8")] + else: + stop_sequences = [item.encode("utf-8") for item in request.payload.stop] + logit_bias = ( + {int(token): float(bias) for token, bias in request.payload.logit_bias.items()} + if request.payload.logit_bias + else None + ) + prompt_text_bytes = self.model.detokenize(prompt_tokens) if prompt_tokens else b"" + for offset, seq_id in enumerate(request.completion_seq_ids): + if offset > 0: + if prompt_length: + self.memory_policy.copy_prompt_state( + request.base_seq_id, + seq_id, + prompt_length, + ) + sampler = Sampler( + seed=(request.payload.seed or llama_cpp.LLAMA_DEFAULT_SEED) + offset, + vocab=self.model.vocab, + n_vocab=self.model.n_vocab, + top_p=request.payload.top_p, + temperature=request.payload.temperature, + frequency_penalty=request.payload.frequency_penalty or 0.0, + presence_penalty=request.payload.presence_penalty or 0.0, + logit_bias=logit_bias, + grammar_text=request.grammar_text, + grammar_root=request.grammar_root, + ) + request.completions.append( + Completion( + request_id=request.id, + index=offset, + seq_id=seq_id, + sampler=sampler, + prompt_tokens=prompt_tokens, + prompt_length=prompt_length, + prompt_text=prompt_text, + multimodal_prompt=multimodal_prompt, + max_total_tokens=request.effective_max_len, + stop_sequences=stop_sequences, + logprobs=request.payload.logprobs, + detokenized_prefix_bytes=bytearray(prompt_text_bytes), + rank_by_score=( + request.payload.best_of is not None + and request.payload.best_of > request.payload.n + ), + ) + ) + if request.payload.max_tokens == 0 or request.effective_max_len == prompt_length: + for completion in request.completions: + self.finish_completion(completion, "length") + self.finalize_request_if_ready(request) + return + if prompt_output_index is None: + if prompt_logits is not None: + for completion in request.completions: + self.sample_completion_from_logits(completion, prompt_logits) + return + raise RuntimeError("prompt output is required to start generation") + for completion in request.completions: + self.sample_completion(completion, prompt_output_index) + + def sample_completion( + self, + completion: Completion, + output_index: Optional[int], + ) -> None: + if completion.finished: + return + if output_index is None: + raise RuntimeError("missing logits output") + token = self.sample_completion_token(completion, output_index) + prev_tokens = [*completion.prompt_tokens, *completion.completion_tokens] + if self.model.store_logits or completion.needs_token_logprob: + logits = self.model.logits(output_index) + if self.model.store_logits: + self.checkpoint_logits[completion.seq_id] = logits + record = Token.from_logits( + model=self.model, + formatter=self.formatter, + prev_tokens=prev_tokens, + prev_text_bytes=completion.detokenized_prefix_bytes, + token=token, + logits=logits, + logprobs_count=completion.logprobs, + need_token_logprob=completion.needs_token_logprob, + ) + else: + record = Token.from_token( + model=self.model, + prev_tokens=prev_tokens, + prev_text_bytes=completion.detokenized_prefix_bytes, + token=token, + ) + finish_reason = self.handle_completion_token( + completion, + token, + record, + decoded=False, + ) + if finish_reason is not None: + completion.pending_finish_reason = finish_reason + + def sample_completion_from_logits( + self, + completion: Completion, + logits: np.ndarray, + ) -> None: + if completion.finished: + return + token = completion.sampler.sample_logits(logits) + prev_tokens = [*completion.prompt_tokens, *completion.completion_tokens] + if self.model.store_logits or completion.needs_token_logprob: + if self.model.store_logits: + self.checkpoint_logits[completion.seq_id] = logits.copy() + record = Token.from_logits( + model=self.model, + formatter=self.formatter, + prev_tokens=prev_tokens, + prev_text_bytes=completion.detokenized_prefix_bytes, + token=token, + logits=logits, + logprobs_count=completion.logprobs, + need_token_logprob=completion.needs_token_logprob, + ) + else: + record = Token.from_token( + model=self.model, + prev_tokens=prev_tokens, + prev_text_bytes=completion.detokenized_prefix_bytes, + token=token, + ) + finish_reason = self.handle_completion_token( + completion, + token, + record, + decoded=False, + ) + if finish_reason is not None: + completion.pending_finish_reason = finish_reason + + def handle_completion_token( + self, + completion: Completion, + token: int, + record: Token, + *, + decoded: bool, + ) -> Optional[str]: + if record.token_logprob is not None: + completion.score_sum += record.token_logprob + rendered_start = len(completion.rendered_bytes) + completion.completion_tokens.append(token) + self.metrics.observe_predicted_token() + completion.token_records.append(record) + completion.rendered_bytes.extend(record.text_bytes) + completion.detokenized_prefix_bytes.extend(record.text_bytes) + finish_reason: Optional[str] = None + if llama_cpp.llama_vocab_is_eog(self.model.vocab, token): + finish_reason = "stop" + elif completion.total_tokens >= completion.max_total_tokens: + finish_reason = "length" + else: + max_stop_length = completion.max_stop_sequence_length + search_start = max(0, rendered_start - max_stop_length + 1) + if any( + completion.rendered_bytes.find(stop, search_start) != -1 + for stop in completion.stop_sequences + ): + finish_reason = "stop" + if not decoded: + completion.pending_input_tokens.append(token) + if ( + completion.request_id in self.requests + and self.requests[completion.request_id].payload.stream + and finish_reason is None + ): + self.flush_stream_updates(completion, finish_reason=None) + if not decoded and finish_reason is None and not self.should_defer_draft_fill: + self.maybe_fill_draft_tokens(completion) + return finish_reason + + def flush_stream_updates( + self, + completion: Completion, + finish_reason: Optional[str], + ) -> None: + request = self.requests[completion.request_id] + for payload in self.formatter.stream_completion_chunks( + request, + completion, + finish_reason, + ): + if request.on_stream_chunk is not None: + request.on_stream_chunk(payload) + + def finish_completion(self, completion: Completion, finish_reason: str) -> None: + if completion.finished: + return + completion.finished = True + completion.finish_reason = finish_reason + completion.pending_input_tokens.clear() + completion.draft_tokens.clear() + request = self.requests[completion.request_id] + if request.payload.stream: + self.flush_stream_updates(completion, finish_reason=finish_reason) + if request.on_stream_chunk is not None: + request.on_stream_chunk( + self.formatter.completion_finish_chunk( + request, + completion, + finish_reason, + ) + ) + + def finalize_request_if_ready(self, request: CompletionRequest) -> None: + if not request.completions or not all(completion.finished for completion in request.completions): + return + selected = request.selected_completions() + result = self.formatter.build_completion_response(request, selected) + self.metrics.requests_completed_total += 1 + self.release_request(request) + if request.on_done is not None: + request.on_done(result) + + def truncate_sequence( + self, + seq_id: int, + keep_len: int, + *, + truncate_draft: bool = True, + ) -> None: + current_len = self.radix_trie.length(seq_id) + if current_len <= keep_len: + return + keep_pos = self.sequence_history.position_length_for_prefix(seq_id, keep_len) + if not llama_cpp.llama_memory_seq_rm(self.model.mem, seq_id, keep_pos, -1): + raise RuntimeError( + f"failed to truncate model sequence {seq_id} at position {keep_pos}" + ) + if truncate_draft: + self.model.truncate_draft_sequence(seq_id, keep_pos) + self.truncate_sequence_metadata(seq_id, current_len, keep_len) + + def copy_sequence_state( + self, + source_sequence_id: int, + dest_sequence_id: int, + keep_len: int, + *, + copy_all_state: bool = False, + ) -> None: + if keep_len <= 0: + return + source_length = self.radix_trie.length(source_sequence_id) + keep_pos = self.sequence_history.position_length_for_prefix( + source_sequence_id, + keep_len, + ) + copy_p0 = 0 + copy_p1 = keep_pos + if copy_all_state or not self.model.kv_unified: + copy_p0 = -1 + copy_p1 = -1 + llama_cpp.llama_memory_seq_cp( + self.model.mem, + source_sequence_id, + dest_sequence_id, + copy_p0, + copy_p1, + ) + self.model.copy_draft_sequence( + source_sequence_id, + dest_sequence_id, + copy_p0, + copy_p1, + ) + if copy_all_state and source_length > keep_len: + if not llama_cpp.llama_memory_seq_rm( + self.model.mem, + dest_sequence_id, + keep_pos, + -1, + ): + raise RuntimeError( + f"failed to truncate copied model sequence {dest_sequence_id} " + f"at position {keep_pos}" + ) + self.model.truncate_draft_sequence(dest_sequence_id, keep_pos) + self.radix_trie.copy(source_sequence_id, dest_sequence_id, keep_len) + self.sequence_history.copy( + source_sequence_id, + dest_sequence_id, + source_length, + keep_len, + ) + + def truncate_sequence_metadata( + self, + seq_id: int, + current_len: int, + keep_len: int, + ) -> None: + self.radix_trie.truncate(seq_id, keep_len) + self.sequence_history.truncate(seq_id, current_len, keep_len) + self.checkpoint_logits.pop(seq_id, None) + + def truncate_free_sequence(self, seq_id: int, keep_len: int) -> None: + if seq_id not in self.free_sequences: + return + self.truncate_sequence(seq_id, keep_len) + self.free_sequences.move_to_end(seq_id) + + def delete_free_sequence(self, seq_id: int) -> None: + if seq_id not in self.free_sequences: + return + self.truncate_sequence(seq_id, 0) + del self.free_sequences[seq_id] + self.checkpoint_logits.pop(seq_id, None) + self.unused_sequences.append(seq_id) + self.metrics.checkpoint_evictions_total += 1 + + def claim_unused_sequence(self) -> int: + seq_id = self.unused_sequences.pop() + self.claimed_sequences.add(seq_id) + return seq_id + + def claim_free_sequence(self, seq_id: int) -> int: + del self.free_sequences[seq_id] + self.claimed_sequences.add(seq_id) + return seq_id + + def activate_request( + self, + request: CompletionRequest, + *, + base_seq_id: int, + sibling_seq_ids: List[int], + ) -> None: + request.base_seq_id = base_seq_id + request.sibling_seq_ids = sibling_seq_ids + request.completion_seq_ids = [base_seq_id, *sibling_seq_ids] + request.admitted = True + self.active_request_ids.add(request.id) + + def release_request(self, request: CompletionRequest) -> None: + for completion in request.completions: + completion.sampler.close() + for seq_id in request.completion_seq_ids: + if seq_id in self.claimed_sequences: + self.claimed_sequences.remove(seq_id) + self.free_sequences[seq_id] = None + self.free_sequences.move_to_end(seq_id) + self.active_request_ids.discard(request.id) + self.requests.pop(request.id, None) + + def remove_pending_request(self, request: CompletionRequest) -> None: + try: + self.pending_requests.remove(request) + except ValueError: + pass + + def fail_request(self, request: CompletionRequest, exc: BaseException) -> None: + self.remove_pending_request(request) + if request.id in self.active_request_ids or request.admitted: + self.release_request(request) + else: + self.requests.pop(request.id, None) + if isinstance(exc, CompletionRequestCancelledError): + self.metrics.requests_cancelled_total += 1 + else: + self.metrics.requests_failed_total += 1 + if request.on_error is not None: + request.on_error(exc) + + @staticmethod + def _format_prometheus_value(value: Union[int, float]) -> str: + if isinstance(value, int): + return str(value) + if not math.isfinite(value): + raise ValueError(f"invalid Prometheus metric value: {value}") + return format(value, ".16g") + + def sequence_cache_metric_definitions( + self, + ) -> List[Tuple[str, str, str, Union[int, float]]]: + return [ + ( + "counter", + "batch_server:sequence_cache_hits_total", + "Number of external sequence cache entries hydrated.", + self.metrics.sequence_cache_hits_total, + ), + ( + "counter", + "batch_server:sequence_cache_save_requests_total", + "Number of external sequence cache save requests.", + self.metrics.sequence_cache_save_requests_total, + ), + ( + "counter", + "batch_server:sequence_cache_load_failures_total", + "Number of external sequence cache load failures.", + self.metrics.sequence_cache_load_failures_total, + ), + ( + "counter", + "batch_server:sequence_cache_lookup_failures_total", + "Number of external sequence cache lookup failures.", + self.metrics.sequence_cache_lookup_failures_total, + ), + ( + "counter", + "batch_server:sequence_cache_save_failures_total", + "Number of external sequence cache save failures.", + self.metrics.sequence_cache_save_failures_total, + ), + ( + "counter", + "batch_server:sequence_cache_tokens_loaded_total", + "Number of prompt tokens hydrated from external sequence cache.", + self.metrics.sequence_cache_tokens_loaded_total, + ), + ] + + def speculative_metric_definitions( + self, + ) -> List[Tuple[str, str, str, Union[int, float]]]: + proposed = self.speculative_stats["draft_tokens_proposed"] + accepted = self.speculative_stats["draft_tokens_accepted"] + proposals = self.speculative_stats["draft_proposals"] + verified_proposals = sum(self.draft_acceptance_length_counts.values()) + verified_accepted = sum( + accepted_tokens * count + for accepted_tokens, count in self.draft_acceptance_length_counts.items() + ) + acceptance_rate = accepted / proposed if proposed > 0 else 0.0 + average_proposed_tokens = proposed / proposals if proposals > 0 else 0.0 + average_verified_acceptance_length = ( + verified_accepted / verified_proposals if verified_proposals > 0 else 0.0 + ) + return [ + ( + "counter", + "batch_server:draft_proposals_total", + "Number of speculative draft proposal batches generated.", + proposals, + ), + ( + "counter", + "batch_server:draft_tokens_proposed_total", + "Number of speculative draft tokens proposed.", + proposed, + ), + ( + "counter", + "batch_server:draft_tokens_accepted_total", + "Number of speculative draft tokens accepted.", + accepted, + ), + ( + "counter", + "batch_server:draft_tokens_rejected_total", + "Number of speculative draft tokens rejected.", + self.speculative_stats["draft_tokens_rejected"], + ), + ( + "counter", + "batch_server:draft_batches_verified_total", + "Number of completion items that target-verified draft tokens.", + self.metrics.draft_batches_verified_total, + ), + ( + "counter", + "batch_server:draft_target_tokens_verified_total", + "Number of draft tokens decoded by the target for verification.", + self.metrics.draft_target_tokens_verified_total, + ), + ( + "counter", + "batch_server:draft_target_tokens_wasted_total", + "Number of target-verified draft tokens rejected before emission.", + self.metrics.draft_target_tokens_wasted_total, + ), + ( + "counter", + "batch_server:draft_tokens_reused_as_pending_total", + "Number of draft tokens accepted by the next pending-token sample.", + self.metrics.draft_tokens_reused_as_pending_total, + ), + ( + "gauge", + "batch_server:draft_acceptance_rate", + "Fraction of proposed draft tokens accepted, including reused pending matches.", + acceptance_rate, + ), + ( + "gauge", + "batch_server:draft_average_proposed_tokens", + "Average number of tokens in generated draft proposals.", + average_proposed_tokens, + ), + ( + "gauge", + "batch_server:draft_average_verified_acceptance_length", + "Average accepted draft-token prefix length for target-verified proposals only.", + average_verified_acceptance_length, + ), + ] + + def record_draft_acceptance_length(self, accepted_tokens: int) -> None: + accepted_tokens = max(0, int(accepted_tokens)) + self.draft_acceptance_length_counts[accepted_tokens] = ( + self.draft_acceptance_length_counts.get(accepted_tokens, 0) + 1 + ) + + def draft_acceptance_length_metric_lines(self) -> List[str]: + lines = [ + "# HELP batch_server:draft_acceptance_length_total Number of verified draft proposals by accepted-token prefix length.", + "# TYPE batch_server:draft_acceptance_length_total counter", + ] + for accepted_tokens, count in sorted(self.draft_acceptance_length_counts.items()): + lines.append( + f'batch_server:draft_acceptance_length_total{{accepted_tokens="{accepted_tokens}"}} {count}' + ) + return lines + + def render_prometheus_metrics(self) -> str: + active_completions = sum( + 1 + for request_id in self.active_request_ids + for completion in self.requests[request_id].completions + if not completion.finished + ) + checkpoint_entries = sum( + 1 for seq_id in self.free_sequences if seq_id in self.checkpoint_logits + ) + prompt_tokens_seconds = ( + self.metrics.prompt_tokens_total / self.metrics.prompt_seconds_total + if self.metrics.prompt_seconds_total > 0 + else 0.0 + ) + predicted_tokens_seconds = ( + self.metrics.tokens_predicted_total + / self.metrics.tokens_predicted_seconds_total + if self.metrics.tokens_predicted_seconds_total > 0 + else 0.0 + ) + n_busy_slots_per_decode = ( + self.metrics.n_busy_slots_total / self.metrics.n_decode_total + if self.metrics.n_decode_total > 0 + else 0.0 + ) + draft_provider_metrics: List[Tuple[str, str, str, Union[int, float]]] = [] + if self.model.draft_provider is not None: + draft_metric_definitions = getattr( + self.model.draft_provider, + "metric_definitions", + None, + ) + if draft_metric_definitions is not None: + draft_provider_metrics = cast( + List[Tuple[str, str, str, Union[int, float]]], + draft_metric_definitions(), + ) + metrics_def: List[Tuple[str, str, str, Union[int, float]]] = [ + ( + "counter", + "llamacpp:prompt_tokens_total", + "Number of prompt tokens processed.", + self.metrics.prompt_tokens_total, + ), + ( + "counter", + "llamacpp:prompt_seconds_total", + "Estimated prompt processing time in seconds.", + self.metrics.prompt_seconds_total, + ), + ( + "counter", + "llamacpp:tokens_predicted_total", + "Number of generated tokens processed.", + self.metrics.tokens_predicted_total, + ), + ( + "counter", + "llamacpp:tokens_predicted_seconds_total", + "Estimated generation processing time in seconds.", + self.metrics.tokens_predicted_seconds_total, + ), + ( + "counter", + "batch_server:scheduler_step_seconds_total", + "Total scheduler step wall time.", + self.metrics.scheduler_step_seconds_total, + ), + ( + "counter", + "batch_server:process_batch_seconds_total", + "Time spent processing decoded batches after llama_decode().", + self.metrics.process_batch_seconds_total, + ), + ( + "counter", + "batch_server:sample_seconds_total", + "Time spent sampling target logits.", + self.metrics.sample_seconds_total, + ), + ( + "counter", + "batch_server:draft_seconds_total", + "Speculative draft processing time in seconds.", + self.metrics.draft_seconds_total, + ), + ( + "counter", + "batch_server:draft_process_seconds_total", + "Time spent feeding target batches into the draft context.", + self.metrics.draft_process_seconds_total, + ), + ( + "counter", + "batch_server:draft_generate_seconds_total", + "Time spent generating speculative draft tokens.", + self.metrics.draft_generate_seconds_total, + ), + ( + "counter", + "batch_server:draft_sampled_batch_seconds_total", + "Time spent generating speculative tokens from sampled target batches.", + self.metrics.draft_sampled_batch_seconds_total, + ), + ( + "counter", + "batch_server:draft_process_calls_total", + "Number of target-to-draft context processing phases.", + self.metrics.draft_process_calls_total, + ), + ( + "counter", + "batch_server:draft_generate_calls_total", + "Number of speculative draft generation phases.", + self.metrics.draft_generate_calls_total, + ), + ( + "counter", + "batch_server:draft_sampled_batch_calls_total", + "Number of sampled target batch draft phases.", + self.metrics.draft_sampled_batch_calls_total, + ), + ( + "counter", + "llamacpp:n_decode_total", + "Total number of llama_decode() calls.", + self.metrics.n_decode_total, + ), + ( + "counter", + "batch_server:scheduler_step_calls_total", + "Number of scheduler steps.", + self.metrics.scheduler_step_calls_total, + ), + ( + "counter", + "batch_server:process_batch_calls_total", + "Number of decoded batch processing phases.", + self.metrics.process_batch_calls_total, + ), + ( + "counter", + "batch_server:sample_calls_total", + "Number of target logit sampling calls.", + self.metrics.sample_calls_total, + ), + ( + "counter", + "llamacpp:n_tokens_max", + "Largest observed n_tokens.", + self.metrics.n_tokens_max, + ), + ( + "gauge", + "llamacpp:n_busy_slots_per_decode", + "Average number of busy sequences per llama_decode() call.", + n_busy_slots_per_decode, + ), + ( + "gauge", + "llamacpp:prompt_tokens_seconds", + "Average prompt throughput in tokens/s.", + prompt_tokens_seconds, + ), + ( + "gauge", + "llamacpp:predicted_tokens_seconds", + "Average generation throughput in tokens/s.", + predicted_tokens_seconds, + ), + ( + "gauge", + "llamacpp:requests_processing", + "Number of requests processing.", + len(self.active_request_ids), + ), + ( + "gauge", + "llamacpp:requests_deferred", + "Number of requests deferred.", + len(self.pending_requests), + ), + ( + "counter", + "batch_server:requests_submitted_total", + "Number of requests submitted.", + self.metrics.requests_submitted_total, + ), + ( + "counter", + "batch_server:requests_admitted_total", + "Number of requests admitted.", + self.metrics.requests_admitted_total, + ), + ( + "counter", + "batch_server:requests_completed_total", + "Number of requests completed successfully.", + self.metrics.requests_completed_total, + ), + ( + "counter", + "batch_server:requests_cancelled_total", + "Number of requests cancelled.", + self.metrics.requests_cancelled_total, + ), + ( + "counter", + "batch_server:requests_failed_total", + "Number of requests failed.", + self.metrics.requests_failed_total, + ), + ( + "gauge", + "batch_server:active_completions", + "Number of active unfinished completions.", + active_completions, + ), + ( + "gauge", + "batch_server:claimed_sequences", + "Number of claimed sequence ids.", + len(self.claimed_sequences), + ), + ( + "gauge", + "batch_server:free_sequences", + "Number of reusable free sequence ids.", + len(self.free_sequences), + ), + ( + "gauge", + "batch_server:unused_sequences", + "Number of unused sequence ids.", + len(self.unused_sequences), + ), + ( + "gauge", + "batch_server:checkpoint_entries", + "Number of free sequence checkpoints with stored logits.", + checkpoint_entries, + ), + ( + "counter", + "batch_server:checkpoint_hits_total", + "Number of exact checkpoint hits.", + self.metrics.checkpoint_hits_total, + ), + ( + "counter", + "batch_server:checkpoint_saves_total", + "Number of prompt checkpoints saved.", + self.metrics.checkpoint_saves_total, + ), + ( + "counter", + "batch_server:checkpoint_evictions_total", + "Number of free checkpoints evicted.", + self.metrics.checkpoint_evictions_total, + ), + *self.sequence_cache_metric_definitions(), + *self.speculative_metric_definitions(), + *draft_provider_metrics, + ( + "gauge", + "batch_server:radix_trie_sequences", + "Number of sequences tracked in the radix trie.", + len(self.radix_trie.sequence_lengths), + ), + ( + "gauge", + "batch_server:radix_trie_tokens", + "Total tokens tracked in the radix trie.", + self.sequence_history.size, + ), + ] + lines: List[str] = [] + for metric_type, name, help_text, value in metrics_def: + lines.append(f"# HELP {name} {help_text}") + lines.append(f"# TYPE {name} {metric_type}") + lines.append(f"{name} {self._format_prometheus_value(value)}") + lines.extend(self.draft_acceptance_length_metric_lines()) + lines.append("") + return "\n".join(lines) + + def finalize_cancelled(self) -> bool: + finalized = False + for request in list(self.pending_requests): + if request.cancelled: + self.fail_request(request, CompletionRequestCancelledError("request cancelled")) + finalized = True + for request_id in list(self.active_request_ids): + request = self.requests[request_id] + if request.cancelled: + self.fail_request(request, CompletionRequestCancelledError("request cancelled")) + finalized = True + return finalized + + +class CompletionService: + def __init__(self, scheduler: CompletionScheduler) -> None: + self.scheduler = scheduler + self.formatter = scheduler.formatter + self.condition = threading.Condition() + self.commands: Deque[Callable[[], None]] = deque() + self.closed = False + self.worker = threading.Thread( + target=self.run_loop, + name="completion-service", + daemon=True, + ) + self.worker.start() + + def close(self) -> None: + with self.condition: + self.closed = True + self.condition.notify_all() + self.worker.join() + self.scheduler.close() + + def enqueue(self, command: Callable[[], None]) -> None: + with self.condition: + if self.closed: + raise RuntimeError("completion service closed") + self.commands.append(command) + self.condition.notify_all() + + def call_on_scheduler(self, callback: Callable[[], Any]) -> Any: + result_box: Dict[str, Any] = {} + error_box: Dict[str, BaseException] = {} + done = threading.Event() + + def run_callback() -> None: + try: + result_box["result"] = callback() + except BaseException as exc: # noqa: BLE001 + error_box["error"] = exc + finally: + done.set() + + self.enqueue(run_callback) + done.wait() + if "error" in error_box: + raise error_box["error"] + return result_box.get("result") + + def call_on_idle_scheduler(self, callback: Callable[[], Any]) -> Any: + result_box: Dict[str, Any] = {} + error_box: Dict[str, BaseException] = {} + done = threading.Event() + + def run_callback() -> None: + if not self.scheduler.is_idle(): + with self.condition: + self.commands.appendleft(run_callback) + self.condition.notify_all() + return + try: + result_box["result"] = callback() + except BaseException as exc: # noqa: BLE001 + error_box["error"] = exc + finally: + done.set() + + self.enqueue(run_callback) + done.wait() + if "error" in error_box: + raise error_box["error"] + return result_box.get("result") + + def create_embedding( + self, + payload: CreateEmbeddingRequest, + ) -> CreateEmbeddingResponse: + embedding = self.call_on_idle_scheduler( + lambda: self.scheduler.create_embedding(payload) + ) + return cast(CreateEmbeddingResponse, embedding) + + def render_prometheus_metrics(self) -> str: + metrics = self.call_on_scheduler(self.scheduler.render_prometheus_metrics) + return cast(str, metrics) + + def run_loop(self) -> None: + while True: + with self.condition: + while not self.closed and not self.commands and self.scheduler.is_idle(): + self.condition.wait() + if self.closed and not self.commands and self.scheduler.is_idle(): + return + commands = list(self.commands) + self.commands.clear() + for command in commands: + command() + progressed = self.scheduler.step() + if progressed: + continue + with self.condition: + if self.closed and not self.commands and self.scheduler.is_idle(): + return + if not self.commands: + self.condition.wait(timeout=0.01) + + def submit_request( + self, + request: CompletionRequest, + ) -> Tuple[CompletionStream, Callable[[], None]]: + mailbox: "queue.Queue[object]" = queue.Queue() + done = threading.Event() + cancelled = threading.Event() + sentinel = object() + result_box: Dict[str, OpenAICompletion] = {} + error_box: Dict[str, BaseException] = {} + + def on_stream_chunk(chunk: CompletionChunk) -> None: + mailbox.put(chunk) + + def on_done(result: OpenAICompletion) -> None: + result_box["result"] = result + done.set() + mailbox.put(sentinel) + + def on_error(exc: BaseException) -> None: + error_box["error"] = exc + done.set() + mailbox.put(sentinel) + + request.on_stream_chunk = on_stream_chunk + request.on_done = on_done + request.on_error = on_error + + def cancel() -> None: + if cancelled.is_set(): + return + cancelled.set() + try: + def cancel_request() -> None: + self.scheduler.cancel(request.id) + + self.enqueue(cancel_request) + except RuntimeError: + pass + + def stream() -> CompletionStream: + try: + while True: + item = mailbox.get() + if item is sentinel: + break + yield cast(CompletionChunk, item) + finally: + if not done.is_set(): + cancel() + if "error" in error_box: + raise error_box["error"] + result = result_box.get("result") + if result is None: + raise RuntimeError("missing completion result") + return result + + def submit_request() -> None: + self.scheduler.submit_request(request) + + self.enqueue(submit_request) + return stream(), cancel + + def submit( + self, + payload: CreateCompletionRequest, + ) -> Tuple[CompletionStream, Callable[[], None]]: + request = self.request_from_payload(payload) + return self.submit_request(request) + + def request_from_payload( + self, + payload: CreateCompletionRequest, + ) -> CompletionRequest: + prompt_text, prompt_plan = self.prepare_completion_prompt(payload) + return self.request_from_prepared( + payload=payload, + prompt_text=prompt_text, + prompt_plan=prompt_plan, + ) + + def request_from_prepared( + self, + *, + payload: CreateCompletionRequest, + prompt_text: str, + prompt_plan: PromptPlan, + grammar_text: Optional[str] = None, + chat_tool_name: Optional[str] = None, + on_stream_chunk: Optional[Callable[[CompletionChunk], None]] = None, + on_done: Optional[Callable[[OpenAICompletion], None]] = None, + on_error: Optional[Callable[[BaseException], None]] = None, + ) -> CompletionRequest: + model = self.scheduler.model + prompt_visible_start = self.prompt_visible_start(prompt_plan) + return CompletionRequest.from_prepared( + payload=payload, + prompt_text=prompt_text, + prompt_plan=prompt_plan, + max_seq_len=model.max_seq_len, + max_output_tokens=model.max_output_tokens, + prompt_visible_start=prompt_visible_start, + prompt_records=self.initial_prompt_records( + payload, + prompt_plan, + prompt_visible_start, + ), + grammar_text=grammar_text, + chat_tool_name=chat_tool_name, + on_stream_chunk=on_stream_chunk, + on_done=on_done, + on_error=on_error, + ) + + def prepare_completion_prompt( + self, + payload: CreateCompletionRequest, + ) -> Tuple[str, PromptPlan]: + model = self.scheduler.model + prompts = payload.normalized_prompt() + if len(prompts) != 1: + raise CompletionRequestValidationError("multiple prompts are not supported") + prompt_item = prompts[0] + if isinstance(prompt_item, str): + prompt_text = prompt_item + try: + prompt_tokens = model.build_prompt_tokens(prompt_text, payload.suffix) + except ValueError as exc: + raise CompletionRequestValidationError(str(exc)) from exc + else: + if payload.suffix is not None: + raise CompletionRequestValidationError( + "suffix is not supported with token id prompts" + ) + prompt_tokens = list(prompt_item) + prompt_text = model.detokenize(prompt_tokens).decode("utf-8", errors="ignore") + return prompt_text, PromptPlan.from_tokens(prompt_text, prompt_tokens) + + def prompt_visible_start(self, prompt_plan: PromptPlan) -> int: + model = self.scheduler.model + if prompt_plan.text_tokens and prompt_plan.text_tokens[0] == model.bos_token: + return 1 + return 0 + + def initial_prompt_records( + self, + payload: CreateCompletionRequest, + prompt_plan: PromptPlan, + prompt_visible_start: int, + ) -> List[Token]: + if not (payload.echo and payload.logprobs is not None): + return [] + if prompt_visible_start >= len(prompt_plan.text_tokens): + return [] + model = self.scheduler.model + first_pos = prompt_visible_start + first_token = prompt_plan.text_tokens[first_pos] + return [ + Token( + token=first_token, + text_bytes=model.token_bytes_with_prev( + prompt_plan.text_tokens[:first_pos], + first_token, + ), + token_logprob=None, + top_logprobs=None, + ) + ] + + +def create_app() -> FastAPI: + app = FastAPI() + app.state.service = None + app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], + ) + + async def watch_http_disconnect( + http_request: Request, + cancel: Callable[[], None], + ) -> None: + while True: + message = await http_request.receive() + if message["type"] == "http.disconnect": + cancel() + return + + async def disconnected_cancelled_response_or_raise( + http_request: Request, + exc: BaseException, + ) -> Response: + if isinstance(exc, CompletionRequestCancelledError): + if await http_request.is_disconnected(): + return Response(status_code=204) + raise exc + + def bad_request(exc: CompletionRequestValidationError) -> HTTPException: + return HTTPException(status_code=400, detail=str(exc)) + + async def collect_completion_result( + formatter: OpenAIFormatter, + http_request: Request, + stream: CompletionStream, + cancel: Callable[[], None], + ) -> OpenAICompletion | Response: + disconnect_task = asyncio.create_task( + watch_http_disconnect(http_request, cancel) + ) + try: + return await asyncio.to_thread(formatter.collect_completion, stream) + except asyncio.CancelledError: + cancel() + raise + except BaseException as exc: + return await disconnected_cancelled_response_or_raise(http_request, exc) + finally: + disconnect_task.cancel() + stream.close() + + async def collect_completion_results( + formatter: OpenAIFormatter, + http_request: Request, + submissions: Sequence[ + Tuple[ + CompletionStream, + Callable[[], None], + ] + ], + ) -> List[OpenAICompletion] | Response: + streams = [stream for stream, _ in submissions] + cancel_all = [cancel for _, cancel in submissions] + + def cancel_all_requests() -> None: + for cancel in cancel_all: + cancel() + + disconnect_task = asyncio.create_task( + watch_http_disconnect( + http_request, + cancel_all_requests, + ) + ) + try: + return await asyncio.gather( + *( + asyncio.to_thread(formatter.collect_completion, stream) + for stream in streams + ) + ) + except asyncio.CancelledError: + cancel_all_requests() + raise + except BaseException as exc: + return await disconnected_cancelled_response_or_raise(http_request, exc) + finally: + disconnect_task.cancel() + for stream in streams: + stream.close() + + async def stream_sse_chunks( + formatter: OpenAIFormatter, + http_request: Request, + stream: CompletionStream, + cancel: Callable[[], None], + chunk_payloads: Callable[[CompletionChunk], Iterable[Any]], + ) -> AsyncIterator[bytes]: + disconnect_task = asyncio.create_task( + watch_http_disconnect(http_request, cancel) + ) + try: + while True: + done, chunk = await asyncio.to_thread( + formatter.next_stream_chunk, + stream, + ) + if done: + break + assert chunk is not None + for payload in chunk_payloads(chunk): + yield formatter.encode_sse_payload(payload) + yield b"data: [DONE]\n\n" + except asyncio.CancelledError: + cancel() + raise + except BaseException as exc: + cancel() + if ( + isinstance(exc, CompletionRequestCancelledError) + and await http_request.is_disconnected() + ): + return + raise + finally: + disconnect_task.cancel() + + async def stream_sse_outputs( + formatter: OpenAIFormatter, + http_request: Request, + stream: CompletionStream, + cancel: Callable[[], None], + *, + initial_payloads: Iterable[BaseModel | Dict[str, Any]], + chunk_payloads: Callable[[CompletionChunk], Iterable[Any]], + done_payloads: Callable[ + [Optional[OpenAICompletion]], Iterable[BaseModel | Dict[str, Any]] + ], + ) -> AsyncIterator[bytes]: + disconnect_task = asyncio.create_task( + watch_http_disconnect(http_request, cancel) + ) + try: + for payload in initial_payloads: + yield formatter.encode_sse_payload(payload) + while True: + done, chunk, result = await asyncio.to_thread( + formatter.next_stream_output, + stream, + ) + if done: + for payload in done_payloads(result): + yield formatter.encode_sse_payload(payload) + break + assert chunk is not None + for payload in chunk_payloads(chunk): + yield formatter.encode_sse_payload(payload) + yield b"data: [DONE]\n\n" + except asyncio.CancelledError: + cancel() + raise + except BaseException as exc: + cancel() + if ( + isinstance(exc, CompletionRequestCancelledError) + and await http_request.is_disconnected() + ): + return + raise + finally: + disconnect_task.cancel() + + async def stream_websocket_responses( + websocket: WebSocket, + formatter: OpenAIFormatter, + stream: CompletionStream, + _: Callable[[], None], + *, + initial_payloads: Iterable[BaseModel | Dict[str, Any]], + chunk_payloads: Callable[[CompletionChunk], Iterable[Any]], + done_payloads: Callable[ + [Optional[OpenAICompletion]], Iterable[BaseModel | Dict[str, Any]] + ], + ) -> None: + try: + for payload in initial_payloads: + await websocket.send_json(payload) + while True: + done, chunk, result = await asyncio.to_thread( + formatter.next_stream_output, + stream, + ) + if done: + for payload in done_payloads(result): + await websocket.send_json(payload) + break + assert chunk is not None + for payload in chunk_payloads(chunk): + await websocket.send_json(payload) + finally: + stream.close() + + @app.post("/v1/completions") + async def create_completion( # pyright: ignore[reportUnusedFunction] + http_request: Request, + body: CreateCompletionRequest, + ): + service: CompletionService = app.state.service + formatter = service.formatter + prompts = body.normalized_prompt() + if len(prompts) > 1: + if body.stream: + raise HTTPException( + status_code=400, + detail="streaming does not support multiple prompts", + ) + try: + submissions = [ + service.submit(body.model_copy(update={"prompt": prompt})) + for prompt in prompts + ] + except CompletionRequestValidationError as exc: + raise bad_request(exc) from exc + results = await collect_completion_results( + formatter, http_request, submissions + ) + if isinstance(results, Response): + return results + return JSONResponse( + formatter.aggregate_completion_results(results).model_dump( + mode="json", + exclude_none=True, + ) + ) + try: + stream, cancel = service.submit( + body.model_copy(update={"prompt": prompts[0]}) + ) + except CompletionRequestValidationError as exc: + raise bad_request(exc) from exc + if body.stream: + return StreamingResponse( + stream_sse_chunks( + formatter, + http_request, + stream, + cancel, + lambda chunk: [chunk], + ), + media_type="text/event-stream", + ) + result = await collect_completion_result( + formatter, http_request, stream, cancel + ) + if isinstance(result, Response): + return result + return JSONResponse(result.model_dump(mode="json", exclude_none=True)) + + @app.post("/v1/embeddings", response_model=CreateEmbeddingResponse) + async def create_embedding( # pyright: ignore[reportUnusedFunction] + body: CreateEmbeddingRequest, + ) -> JSONResponse: + service: CompletionService = app.state.service + try: + embedding = await asyncio.to_thread(service.create_embedding, body) + except CompletionRequestValidationError as exc: + raise bad_request(exc) from exc + return JSONResponse(embedding.model_dump(mode="json", exclude_none=True)) + + @app.post("/v1/chat/completions") + async def create_chat_completion( # pyright: ignore[reportUnusedFunction] + http_request: Request, body: CreateChatCompletionRequest + ): + service: CompletionService = app.state.service + formatter = service.formatter + try: + parts = formatter.completion_request_from_chat_request( + body, + ) + except CompletionRequestValidationError as exc: + raise bad_request(exc) from exc + template_functions = ( + [function.to_template_function() for function in body.functions] + if body.functions is not None + else None + ) + template_tools = ( + [tool.to_template_tool() for tool in body.tools] + if body.tools is not None + else None + ) + try: + request = service.request_from_prepared( + payload=parts.payload, + prompt_text=parts.prompt_text, + prompt_plan=parts.prompt_plan, + grammar_text=parts.grammar_text, + chat_tool_name=parts.tool_name, + ) + stream, cancel = service.submit_request(request) + except CompletionRequestValidationError as exc: + raise bad_request(exc) from exc + if body.stream: + started_indices: set[int] = set() + + def chat_chunk_payloads( + completion_chunk: CompletionChunk, + ) -> Iterable[BaseModel | Dict[str, Any]]: + return formatter.convert_completion_chunk_to_chat_chunks( + completion_chunk, + started_indices, + parts.tool_name, + functions=template_functions, + tools=template_tools, + parsed_states=parsed_states, + generation_prompt=parts.generation_prompt, + ) + + parsed_states: Dict[int, Dict[str, Any]] = {} + return StreamingResponse( + stream_sse_chunks( + formatter, + http_request, + stream, + cancel, + chat_chunk_payloads, + ), + media_type="text/event-stream", + ) + completion = await collect_completion_result( + formatter, http_request, stream, cancel + ) + if isinstance(completion, Response): + return completion + chat_response = formatter.convert_completion_response_to_chat( + completion, + parts.tool_name, + functions=template_functions, + tools=template_tools, + generation_prompt=parts.generation_prompt, + ) + if isinstance(chat_response, BaseModel): + return JSONResponse( + chat_response.model_dump(mode="json", exclude_none=True) + ) + return JSONResponse(chat_response) + + @app.post("/v1/responses") + async def create_response( # pyright: ignore[reportUnusedFunction] + http_request: Request, + body: CreateResponseRequest, + ): + service: CompletionService = app.state.service + formatter = service.formatter + try: + chat_parts = formatter.response_request_to_chat_parts(body) + parts = formatter.completion_request_from_response_chat_parts(chat_parts) + except CompletionRequestValidationError as exc: + raise bad_request(exc) from exc + response_tools = chat_parts.tools + try: + request = service.request_from_prepared( + payload=parts.payload, + prompt_text=parts.prompt_text, + prompt_plan=parts.prompt_plan, + grammar_text=parts.grammar_text, + chat_tool_name=parts.tool_name, + ) + stream, cancel = service.submit_request(request) + except CompletionRequestValidationError as exc: + raise bad_request(exc) from exc + if body.stream: + started_indices: set[int] = set() + parsed_states: Dict[int, ResponseParser] = {} + stream_state = OpenAIFormatter.ResponsesStream( + body=body, + response_id="resp_" + request.id, + created_at=float(request.created), + model=service.scheduler.model.model_path, + ) + + def response_chunk_payloads( + completion_chunk: CompletionChunk, + ) -> Iterable[BaseModel | Dict[str, Any]]: + chat_chunks = formatter.convert_completion_chunk_to_chat_chunks( + completion_chunk, + started_indices, + parts.tool_name, + tools=response_tools, + parsed_states=parsed_states, + generation_prompt=parts.generation_prompt, + ) + payloads: List[BaseModel | Dict[str, Any]] = [] + for chat_chunk in chat_chunks: + payloads.extend( + formatter.convert_chat_chunk_to_response_events( + chat_chunk, + stream_state, + ) + ) + return payloads + + return StreamingResponse( + stream_sse_outputs( + formatter, + http_request, + stream, + cancel, + initial_payloads=formatter.start_response_stream(stream_state), + chunk_payloads=response_chunk_payloads, + done_payloads=lambda completion: ( + formatter.response_stream_terminal_events( + stream_state, + completion, + ) + ), + ), + media_type="text/event-stream", + ) + completion = await collect_completion_result( + formatter, http_request, stream, cancel + ) + if isinstance(completion, Response): + return completion + return JSONResponse( + formatter.convert_completion_response_to_response( + completion, + body, + parts.tool_name, + tools=response_tools, + generation_prompt=parts.generation_prompt, + ) + ) + + @app.websocket("/v1/responses") + async def responses_websocket( # pyright: ignore[reportUnusedFunction] + websocket: WebSocket, + ): + await websocket.accept() + service: CompletionService = app.state.service + formatter = service.formatter + # HTTP /v1/responses remains stateless. The websocket transport keeps + # per-connection response history so Codex can use previous_response_id + # within a single live session. + websocket_response_history: Dict[str, ResponsesWebSocketState] = {} + + def websocket_request_with_ephemeral_history( + ws_body: ResponseCreateWebSocketRequest, + ) -> CreateResponseRequest: + body = ws_body.to_create_response_request() + previous_response_id = body.previous_response_id + if previous_response_id is None: + return body + prior_state = websocket_response_history.get(previous_response_id) + if prior_state is None: + raise CompletionRequestValidationError( + f"unknown previous_response_id: {previous_response_id}" + ) + current_items = formatter._clone_response_input_items(body.input) + replay_items = formatter._clone_response_input_items( + prior_state.input_items + ) + for item in prior_state.output_items: + normalized = formatter._normalize_response_output_item_for_input(item) + if normalized is not None: + replay_items.append(normalized) + replay_items.extend(current_items) + return body.model_copy( + update={ + "input": replay_items, + "previous_response_id": None, + } + ) + + async def send_error(message: str) -> None: + await websocket.send_json( + { + "type": "error", + "error": { + "message": message, + }, + } + ) + + try: + while True: + try: + payload = await websocket.receive_json() + except WebSocketDisconnect: + break + except Exception: + await send_error("invalid websocket request payload") + continue + + try: + ws_body = ResponseCreateWebSocketRequest.model_validate(payload) + body = websocket_request_with_ephemeral_history(ws_body) + if ws_body.generate is False: + body = body.model_copy(update={"max_output_tokens": 0}) + chat_parts = formatter.response_request_to_chat_parts(body) + parts = formatter.completion_request_from_response_chat_parts(chat_parts) + response_tools = chat_parts.tools + request = service.request_from_prepared( + payload=parts.payload, + prompt_text=parts.prompt_text, + prompt_plan=parts.prompt_plan, + grammar_text=parts.grammar_text, + chat_tool_name=parts.tool_name, + ) + stream, cancel = service.submit_request(request) + except CompletionRequestValidationError as exc: + await send_error(str(exc)) + continue + + started_indices: set[int] = set() + parsed_states: Dict[int, ResponseParser] = {} + stream_state = OpenAIFormatter.ResponsesStream( + body=body, + response_id="resp_" + request.id, + created_at=float(request.created), + model=service.scheduler.model.model_path, + ) + + def response_chunk_payloads( + completion_chunk: CompletionChunk, + ) -> Iterable[BaseModel | Dict[str, Any]]: + chat_chunks = formatter.convert_completion_chunk_to_chat_chunks( + completion_chunk, + started_indices, + parts.tool_name, + tools=response_tools, + parsed_states=parsed_states, + generation_prompt=parts.generation_prompt, + ) + payloads: List[BaseModel | Dict[str, Any]] = [] + for chat_chunk in chat_chunks: + payloads.extend( + formatter.convert_chat_chunk_to_response_events( + chat_chunk, + stream_state, + ) + ) + return payloads + + try: + await stream_websocket_responses( + websocket, + formatter, + stream, + cancel, + initial_payloads=formatter.start_response_stream(stream_state), + chunk_payloads=response_chunk_payloads, + done_payloads=lambda completion: ( + formatter.response_stream_terminal_events( + stream_state, + completion, + ) + ), + ) + websocket_response_history[stream_state.response_id] = ( + ResponsesWebSocketState( + input_items=formatter._clone_response_input_items( + body.input + ), + output_items=copy.deepcopy(stream_state.output), + ) + ) + except WebSocketDisconnect: + cancel() + break + except BaseException as exc: + cancel() + await send_error(str(exc)) + except WebSocketDisconnect: + pass + + @app.get("/v1/models", response_model=ModelListResponse) + async def list_models() -> ModelListResponse: # pyright: ignore[reportUnusedFunction] + service = app.state.service + model = service.scheduler.model + model_id = getattr(model, "model_alias", None) or model.model_path + created = int(time.time()) + return ModelListResponse( + data=[ + ModelCardResponse( + id=model_id, + created=created, + owned_by="llama-cpp-python", + ) + ] + ) + + @app.get("/healthz", response_model=HealthzResponse) + async def healthz() -> HealthzResponse: # pyright: ignore[reportUnusedFunction] + return HealthzResponse() + + @app.get("/metrics") + async def metrics() -> Response: # pyright: ignore[reportUnusedFunction] + service = cast(Optional[CompletionService], getattr(app.state, "service", None)) + if service is None: + raise HTTPException(status_code=503, detail="completion service unavailable") + payload = await asyncio.to_thread(service.render_prometheus_metrics) + return Response( + payload, + media_type="text/plain; version=0.0.4; charset=utf-8", + ) + + return app + + +APP = create_app() + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("-C", "--config-file", required=True) + args = parser.parse_args() + config = ConfigFile.load(args.config_file) + model_path = config.model.resolve_model_path() + loras = [ + Model.LoraAdapter(path=lora.resolve_path(), scale=lora.scale) + for lora in config.model.loras + ] + model = Model( + model_path=model_path, + model_alias=config.model.alias, + chat_template=config.model.chat_template, + loras=loras, + n_gpu_layers=config.model.n_gpu_layers, + split_mode=config.model.split_mode, + main_gpu=config.model.main_gpu, + tensor_split=config.model.tensor_split, + vocab_only=config.model.vocab_only, + use_mmap=config.model.use_mmap, + use_mlock=config.model.use_mlock, + kv_overrides=config.model.kv_overrides, + n_ctx=config.model.n_ctx, + n_batch=config.model.n_batch, + n_ubatch=config.model.n_ubatch, + n_seq_max=config.model.n_seq_max, + n_threads=config.model.threads, + n_threads_batch=config.model.threads_batch, + rope_scaling_type=config.model.rope_scaling_type, + pooling_type=config.model.pooling_type, + attention_type=config.model.attention_type, + embedding=config.model.embedding, + rope_freq_base=config.model.rope_freq_base, + rope_freq_scale=config.model.rope_freq_scale, + yarn_ext_factor=config.model.yarn_ext_factor, + yarn_attn_factor=config.model.yarn_attn_factor, + yarn_beta_fast=config.model.yarn_beta_fast, + yarn_beta_slow=config.model.yarn_beta_slow, + yarn_orig_ctx=config.model.yarn_orig_ctx, + offload_kqv=config.model.offload_kqv, + flash_attn=config.model.flash_attn, + op_offload=config.model.op_offload, + swa_full=config.model.swa_full, + no_perf=config.model.no_perf, + type_k=config.model.type_k, + type_v=config.model.type_v, + kv_unified=config.model.kv_unified, + max_seq_len=config.model.max_seq_len, + max_output_tokens=config.model.max_output_tokens, + draft_model=config.model.draft_model, + draft_model_path=config.model.resolve_draft_model_path(), + draft_model_num_pred_tokens=config.model.draft_model_num_pred_tokens, + draft_model_max_ngram_size=config.model.draft_model_max_ngram_size, + draft_model_top_k=config.model.draft_model_top_k, + draft_model_p_min=config.model.draft_model_p_min, + draft_model_max_batch_size=config.model.draft_model_max_batch_size, + draft_model_threads=config.model.draft_model_threads, + draft_model_threads_batch=config.model.draft_model_threads_batch, + response_schema=config.model.response_schema, + store_logits=config.model.store_logits, + ) + if config.model.mtmd is not None: + if model.chat_formatter is None: + raise RuntimeError("MTMD requires a GGUF chat template") + mmproj_path = config.model.mtmd.resolve_mmproj_path() + embedding_cache: Optional[MTMDEmbeddingCache] = None + if config.model.mtmd.embedding_cache is not None: + embedding_cache = MTMDEmbeddingCache( + path=config.model.mtmd.embedding_cache.path, + max_bytes=config.model.mtmd.embedding_cache.max_bytes, + model_fingerprint=MTMDEmbeddingCache.fingerprint_file(model_path), + mmproj_fingerprint=MTMDEmbeddingCache.fingerprint_file(mmproj_path), + ) + model.mtmd_processor = MTMDProcessor( + model_path=model.model_path, + llama_model=model.llama_model, + chat_formatter=model.chat_formatter, + tokenize=model.tokenize, + n_embd_inp=model.n_embd_inp, + n_batch=model.n_batch, + n_ubatch=model.n_ubatch, + n_threads_batch=model.n_threads_batch, + mmproj_path=mmproj_path, + embedding_cache=embedding_cache, + allowed_media_domains=config.model.mtmd.allowed_media_domains, + allowed_local_media_path=config.model.mtmd.allowed_local_media_path, + image_max_bytes=config.model.mtmd.image_max_bytes, + audio_max_bytes=config.model.mtmd.audio_max_bytes, + video_max_bytes=config.model.mtmd.video_max_bytes, + image_timeout_seconds=config.model.mtmd.image_timeout_seconds, + ) + sequence_cache: Optional[SequenceCache] = None + if config.disk_cache is not None: + sequence_cache = SequenceDiskCache( + path=config.disk_cache.path, + max_bytes=config.disk_cache.max_bytes, + min_tokens=config.disk_cache.min_tokens, + compatibility_key=SequenceDiskCache.compatibility_key_for_model(model), + ) + scheduler = CompletionScheduler(model, sequence_cache=sequence_cache) + APP.state.service = CompletionService(scheduler) + try: + uvicorn.run( + APP, host=config.server.host, port=config.server.port, log_level="info" + ) + finally: + APP.state.service.close() + + +if __name__ == "__main__": + main() diff --git a/llama_cpp/__init__.py b/llama_cpp/__init__.py index dce1764f65..13668893fe 100644 --- a/llama_cpp/__init__.py +++ b/llama_cpp/__init__.py @@ -1,2 +1,4 @@ from .llama_cpp import * from .llama import * + +__version__ = "0.3.28" diff --git a/llama_cpp/_ctypes_extensions.py b/llama_cpp/_ctypes_extensions.py new file mode 100644 index 0000000000..e88ed387df --- /dev/null +++ b/llama_cpp/_ctypes_extensions.py @@ -0,0 +1,131 @@ +from __future__ import annotations + +import sys +import os +import ctypes +import functools +import pathlib + +from typing import ( + Any, + Callable, + List, + Union, + Optional, + TYPE_CHECKING, + TypeVar, + Generic, +) +from typing_extensions import TypeAlias + + +# Load the library +def load_shared_library(lib_base_name: str, base_path: pathlib.Path): + """Platform independent shared library loader""" + # Searching for the library in the current directory under the name "libllama" (default name + # for llamacpp) and "llama" (default name for this repo) + lib_paths: List[pathlib.Path] = [] + # Determine the file extension based on the platform + if sys.platform.startswith("linux") or sys.platform.startswith("freebsd"): + lib_paths += [ + base_path / f"lib{lib_base_name}.so", + ] + elif sys.platform == "darwin": + lib_paths += [ + base_path / f"lib{lib_base_name}.so", + base_path / f"lib{lib_base_name}.dylib", + ] + elif sys.platform == "win32": + lib_paths += [ + base_path / f"{lib_base_name}.dll", + base_path / f"lib{lib_base_name}.dll", + ] + else: + raise RuntimeError("Unsupported platform") + + cdll_args = dict() # type: ignore + + # Add the library directory to the DLL search path on Windows (if needed) + if sys.platform == "win32": + os.add_dll_directory(str(base_path)) + os.environ["PATH"] = str(base_path) + os.pathsep + os.environ["PATH"] + + if sys.platform == "win32" and sys.version_info >= (3, 8): + os.add_dll_directory(str(base_path)) + if "CUDA_PATH" in os.environ: + os.add_dll_directory(os.path.join(os.environ["CUDA_PATH"], "bin")) + os.add_dll_directory(os.path.join(os.environ["CUDA_PATH"], "lib")) + if "HIP_PATH" in os.environ: + os.add_dll_directory(os.path.join(os.environ["HIP_PATH"], "bin")) + os.add_dll_directory(os.path.join(os.environ["HIP_PATH"], "lib")) + cdll_args["winmode"] = ctypes.RTLD_GLOBAL + + # Try to load the shared library, handling potential errors + for lib_path in lib_paths: + if lib_path.exists(): + try: + return ctypes.CDLL(str(lib_path), **cdll_args) # type: ignore + except Exception as e: + raise RuntimeError(f"Failed to load shared library '{lib_path}': {e}") + + raise FileNotFoundError( + f"Shared library with base name '{lib_base_name}' not found" + ) + + +# ctypes sane type hint helpers +# +# - Generic Pointer and Array types +# - PointerOrRef type with a type hinted byref function +# +# NOTE: Only use these for static type checking not for runtime checks +# no good will come of that + +if TYPE_CHECKING: + CtypesCData = TypeVar("CtypesCData", bound=ctypes._CData) # type: ignore + + CtypesArray: TypeAlias = ctypes.Array[CtypesCData] # type: ignore + + CtypesPointer: TypeAlias = ctypes._Pointer[CtypesCData] # type: ignore + + CtypesVoidPointer: TypeAlias = ctypes.c_void_p + + class CtypesRef(Generic[CtypesCData]): + pass + + CtypesPointerOrRef: TypeAlias = Union[ + CtypesPointer[CtypesCData], CtypesRef[CtypesCData] + ] + + CtypesFuncPointer: TypeAlias = ctypes._FuncPointer # type: ignore + +F = TypeVar("F", bound=Callable[..., Any]) + + +def ctypes_function_for_shared_library(lib: ctypes.CDLL): + """Decorator for defining ctypes functions with type hints""" + + def ctypes_function( + name: str, argtypes: List[Any], restype: Any, enabled: bool = True + ): + def decorator(f: F) -> F: + if enabled: + func = getattr(lib, name) + func.argtypes = argtypes + func.restype = restype + functools.wraps(f)(func) + return func + else: + return f + + return decorator + + return ctypes_function + + +def _byref(obj: CtypesCData, offset: Optional[int] = None) -> CtypesRef[CtypesCData]: + """Type-annotated version of ctypes.byref""" + ... + + +byref = _byref if TYPE_CHECKING else ctypes.byref diff --git a/llama_cpp/_ggml.py b/llama_cpp/_ggml.py new file mode 100644 index 0000000000..5ece01e032 --- /dev/null +++ b/llama_cpp/_ggml.py @@ -0,0 +1,12 @@ +"""Internal module use at your own risk + +This module provides a minimal interface for working with ggml tensors from llama-cpp-python +""" + +import os +import pathlib + +import llama_cpp._ctypes_extensions as ctypes_ext + +libggml_base_path = pathlib.Path(os.path.abspath(os.path.dirname(__file__))) / "lib" +libggml = ctypes_ext.load_shared_library("ggml", libggml_base_path) diff --git a/llama_cpp/_internals.py b/llama_cpp/_internals.py new file mode 100644 index 0000000000..b0fe94d01f --- /dev/null +++ b/llama_cpp/_internals.py @@ -0,0 +1,886 @@ +from __future__ import annotations + +import os +import ctypes +import warnings + +from typing import ( + Dict, + List, + Tuple, + Optional, + Sequence, + Callable, + Union, +) +from dataclasses import dataclass, field +from contextlib import ExitStack + +import numpy as np +import numpy.typing as npt + +from .llama_types import * +from .llama_grammar import LlamaGrammar +from ._utils import suppress_stdout_stderr + +import llama_cpp.llama_cpp as llama_cpp + + +# Python wrappers over llama.h structs + + +class LlamaModel: + """Intermediate Python wrapper for a llama.cpp llama_model. + NOTE: For stability it's recommended you use the Llama class instead.""" + + def __init__( + self, + *, + path_model: str, + params: llama_cpp.llama_model_params, + verbose: bool = True, + ): + self.path_model = path_model + self.params = params + self.verbose = verbose + self._exit_stack = ExitStack() + # LlamaModel does not use samplers, but close() can run after partial init. + self.sampler = None + self.custom_samplers = [] + + model = None + + if not os.path.exists(path_model): + raise ValueError(f"Model path does not exist: {path_model}") + + with suppress_stdout_stderr(disable=verbose): + model = llama_cpp.llama_model_load_from_file( + self.path_model.encode("utf-8"), self.params + ) + + if model is None: + raise ValueError(f"Failed to load model from file: {path_model}") + + vocab = llama_cpp.llama_model_get_vocab(model) + + if vocab is None: + raise ValueError(f"Failed to get vocab from model: {path_model}") + + self.model = model + self.vocab = vocab + + def free_model(): + if self.model is None: + return + llama_cpp.llama_model_free(self.model) + self.model = None + + self._exit_stack.callback(free_model) + + def close(self): + if self.sampler is not None: + # NOTE: Must remove custom samplers before free or llama.cpp will try to free them + for i, _ in reversed(self.custom_samplers): + llama_cpp.llama_sampler_chain_remove(self.sampler, i) + self.custom_samplers.clear() + self._exit_stack.close() + + def __del__(self): + self.close() + + def vocab_type(self) -> int: + return llama_cpp.llama_vocab_type(self.vocab) + + def n_vocab(self) -> int: + return llama_cpp.llama_vocab_n_tokens(self.vocab) + + def n_ctx_train(self) -> int: + return llama_cpp.llama_model_n_ctx_train(self.model) + + def n_embd(self) -> int: + return llama_cpp.llama_model_n_embd(self.model) + + def rope_freq_scale_train(self) -> float: + return llama_cpp.llama_model_rope_freq_scale_train(self.model) + + def desc(self) -> str: + buf = ctypes.create_string_buffer(1024) + llama_cpp.llama_model_desc(self.model, buf, 1024) + return buf.value.decode("utf-8") + + def size(self) -> int: + return llama_cpp.llama_model_size(self.model) + + def n_params(self) -> int: + return llama_cpp.llama_model_n_params(self.model) + + def get_tensor(self, name: str) -> ctypes.c_void_p: + raise NotImplementedError("get_tensor is not implemented in llama.cpp") + + # Vocab + + def token_get_text(self, token: int) -> str: + return llama_cpp.llama_vocab_get_text(self.vocab, token).decode("utf-8") + + def token_get_score(self, token: int) -> float: + return llama_cpp.llama_vocab_get_score(self.vocab, token) + + def token_get_attr(self, token: int) -> int: + return llama_cpp.llama_vocab_get_attr(self.vocab, token) + + # Special tokens + + def token_bos(self) -> int: + return llama_cpp.llama_vocab_bos(self.vocab) + + def token_eos(self) -> int: + return llama_cpp.llama_vocab_eos(self.vocab) + + def token_cls(self) -> int: + return llama_cpp.llama_vocab_bos(self.vocab) + + def token_sep(self) -> int: + return llama_cpp.llama_vocab_sep(self.vocab) + + def token_nl(self) -> int: + return llama_cpp.llama_vocab_nl(self.vocab) + + def token_prefix(self) -> int: + return llama_cpp.llama_vocab_fim_pre(self.vocab) + + def token_middle(self) -> int: + return llama_cpp.llama_vocab_fim_mid(self.vocab) + + def token_suffix(self) -> int: + return llama_cpp.llama_vocab_fim_suf(self.vocab) + + def token_eot(self) -> int: + return llama_cpp.llama_vocab_eot(self.vocab) + + def add_bos_token(self) -> bool: + return llama_cpp.llama_vocab_get_add_bos(self.vocab) + + def add_eos_token(self) -> bool: + return llama_cpp.llama_vocab_get_add_eos(self.vocab) + + # Tokenization + + def tokenize(self, text: bytes, add_bos: bool, special: bool): + n_ctx = self.n_ctx_train() + tokens = (llama_cpp.llama_token * n_ctx)() + n_tokens = llama_cpp.llama_tokenize( + self.vocab, text, len(text), tokens, n_ctx, add_bos, special + ) + if n_tokens < 0: + n_tokens = abs(n_tokens) + tokens = (llama_cpp.llama_token * n_tokens)() + n_tokens = llama_cpp.llama_tokenize( + self.vocab, text, len(text), tokens, n_tokens, add_bos, special + ) + if n_tokens < 0: + raise RuntimeError( + f'Failed to tokenize: text="{text}" n_tokens={n_tokens}' + ) + return list(tokens[:n_tokens]) + + def token_to_piece(self, token: int, special: bool = False) -> bytes: + buf = ctypes.create_string_buffer(32) + llama_cpp.llama_token_to_piece(self.vocab, token, buf, 32, 0, special) + return bytes(buf) + + def detokenize(self, tokens: List[int], special: bool = False) -> bytes: + output = b"" + size = 32 + buffer = (ctypes.c_char * size)() + for token in tokens: + n = llama_cpp.llama_token_to_piece( + self.vocab, llama_cpp.llama_token(token), buffer, size, 0, special + ) + assert n <= size + output += bytes(buffer[:n]) + # NOTE: Llama1 models automatically added a space at the start of the prompt + # this line removes a leading space if the first token is a beginning of sentence token + return ( + output[1:] + if len(tokens) > 0 and tokens[0] == self.token_bos() and output[0:1] == b" " + else output + ) + + # Extra + def metadata(self) -> Dict[str, str]: + metadata: Dict[str, str] = {} + buffer_size = 1024 + buffer = ctypes.create_string_buffer(buffer_size) + # zero the buffer + buffer.value = b"\0" * buffer_size + # iterate over model keys + for i in range(llama_cpp.llama_model_meta_count(self.model)): + nbytes = llama_cpp.llama_model_meta_key_by_index( + self.model, i, buffer, buffer_size + ) + if nbytes > buffer_size: + buffer_size = nbytes + 1 + buffer = ctypes.create_string_buffer(buffer_size) + nbytes = llama_cpp.llama_model_meta_key_by_index( + self.model, i, buffer, buffer_size + ) + key = buffer.value.decode("utf-8") + nbytes = llama_cpp.llama_model_meta_val_str_by_index( + self.model, i, buffer, buffer_size + ) + if nbytes > buffer_size: + buffer_size = nbytes + 1 + buffer = ctypes.create_string_buffer(buffer_size) + nbytes = llama_cpp.llama_model_meta_val_str_by_index( + self.model, i, buffer, buffer_size + ) + value = buffer.value.decode("utf-8") + metadata[key] = value + return metadata + + @staticmethod + def default_params(): + """Get the default llama_model_params.""" + return llama_cpp.llama_model_default_params() + + +class LlamaContext: + """Intermediate Python wrapper for a llama.cpp llama_context. + NOTE: For stability it's recommended you use the Llama class instead.""" + + def __init__( + self, + *, + model: LlamaModel, + params: llama_cpp.llama_context_params, + verbose: bool = True, + ): + self.model = model + self.params = params + self.verbose = verbose + self._exit_stack = ExitStack() + + ctx = llama_cpp.llama_init_from_model(self.model.model, self.params) + + if ctx is None: + raise ValueError("Failed to create llama_context") + + self.ctx = ctx + self.memory = llama_cpp.llama_get_memory(self.ctx) + self.sampler = None # LlamaContext doesn't manage samplers directly, but some cleanup code expects this attribute + + def free_ctx(): + if self.ctx is None: + return + llama_cpp.llama_free(self.ctx) + self.ctx = None + + self._exit_stack.callback(free_ctx) + + def close(self): + self._exit_stack.close() + + def __del__(self): + self.close() + + def n_ctx(self) -> int: + return llama_cpp.llama_n_ctx(self.ctx) + + def pooling_type(self) -> int: + return llama_cpp.llama_pooling_type(self.ctx) + + def kv_cache_clear(self): + # Embedding models with non-causal attention may not allocate memory. + if self.memory is None: + return + llama_cpp.llama_memory_clear(self.memory, True) + + def kv_cache_seq_rm(self, seq_id: int, p0: int, p1: int) -> bool: + assert self.memory is not None, "Memory is not initialized" + seq_id = seq_id if seq_id >= 0 else 0 + return llama_cpp.llama_memory_seq_rm(self.memory, seq_id, p0, p1) + + def kv_cache_seq_cp(self, seq_id_src: int, seq_id_dst: int, p0: int, p1: int): + assert self.memory is not None, "Memory is not initialized" + llama_cpp.llama_memory_seq_cp(self.memory, seq_id_src, seq_id_dst, p0, p1) + + def kv_cache_seq_keep(self, seq_id: int): + assert self.memory is not None, "Memory is not initialized" + llama_cpp.llama_memory_seq_keep(self.memory, seq_id) + + def kv_cache_seq_shift(self, seq_id: int, p0: int, p1: int, shift: int): + assert self.memory is not None, "Memory is not initialized" + llama_cpp.llama_memory_seq_add(self.memory, seq_id, p0, p1, shift) + + def get_state_size(self) -> int: + return llama_cpp.llama_state_get_size(self.ctx) + + # TODO: copy_state_data + + # TODO: set_state_data + + # TODO: llama_state_load_file + + # TODO: llama_state_save_file + + def decode(self, batch: LlamaBatch): + return_code = llama_cpp.llama_decode( + self.ctx, + batch.batch, + ) + if return_code != 0: + raise RuntimeError(f"llama_decode returned {return_code}") + + def encode(self, batch: LlamaBatch): + return_code = llama_cpp.llama_encode( + self.ctx, + batch.batch, + ) + if return_code != 0: + raise RuntimeError(f"llama_encode returned {return_code}") + + def set_n_threads(self, n_threads: int, n_threads_batch: int): + llama_cpp.llama_set_n_threads(self.ctx, n_threads, n_threads_batch) + + def get_logits(self): + return llama_cpp.llama_get_logits(self.ctx) + + def get_logits_ith(self, i: int): + return llama_cpp.llama_get_logits_ith(self.ctx, i) + + def get_embeddings(self): + return llama_cpp.llama_get_embeddings(self.ctx) + + def get_embeddings_ith(self, i: int): + return llama_cpp.llama_get_embeddings_ith(self.ctx, i) + + def get_embeddings_seq(self, seq_id: int): + return llama_cpp.llama_get_embeddings_seq(self.ctx, seq_id) + + # Sampling functions - deprecated, use LlamaSampler instead + + def set_rng_seed(self, seed: int): + raise NotImplementedError( + "set_rng_seed is deprecated, use LlamaSampler instead" + ) + + def sample_repetition_penalties( + self, + candidates: "_LlamaTokenDataArray", + last_tokens_data: "llama_cpp.Array[llama_cpp.llama_token]", + penalty_last_n: int, + penalty_repeat: float, + penalty_freq: float, + penalty_present: float, + ): + raise NotImplementedError( + "sample_repetition_penalties is deprecated, use LlamaSampler instead" + ) + + def sample_softmax(self, candidates: "_LlamaTokenDataArray"): + raise NotImplementedError( + "sample_softmax is deprecated, use LlamaSampler instead" + ) + + def sample_top_k(self, candidates: "_LlamaTokenDataArray", k: int, min_keep: int): + raise NotImplementedError( + "sample_top_k is deprecated, use LlamaSampler instead" + ) + + def sample_top_p(self, candidates: "_LlamaTokenDataArray", p: float, min_keep: int): + raise NotImplementedError( + "sample_top_p is deprecated, use LlamaSampler instead" + ) + + def sample_min_p(self, candidates: "_LlamaTokenDataArray", p: float, min_keep: int): + raise NotImplementedError( + "sample_min_p is deprecated, use LlamaSampler instead" + ) + + def sample_typical( + self, candidates: "_LlamaTokenDataArray", p: float, min_keep: int + ): + raise NotImplementedError( + "sample_typical is deprecated, use LlamaSampler instead" + ) + + def sample_temp(self, candidates: "_LlamaTokenDataArray", temp: float): + raise NotImplementedError("sample_temp is deprecated, use LlamaSampler instead") + + def sample_grammar(self, candidates: "_LlamaTokenDataArray", grammar: LlamaGrammar): + raise NotImplementedError( + "sample_grammar is deprecated, use LlamaSampler instead" + ) + + def sample_token_mirostat( + self, + candidates: "_LlamaTokenDataArray", + tau: float, + eta: float, + m: int, + mu: llama_cpp.CtypesPointerOrRef[ctypes.c_float], + ) -> int: + raise NotImplementedError( + "sample_token_mirostat is deprecated, use LlamaSampler instead" + ) + + def sample_token_mirostat_v2( + self, + candidates: "_LlamaTokenDataArray", + tau: float, + eta: float, + mu: llama_cpp.CtypesPointerOrRef[ctypes.c_float], + ) -> int: + raise NotImplementedError( + "sample_token_mirostat_v2 is deprecated, use LlamaSampler instead" + ) + + def sample_token_greedy(self, candidates: "_LlamaTokenDataArray") -> int: + raise NotImplementedError( + "sample_token_greedy is deprecated, use LlamaSampler instead" + ) + + def sample_token(self, candidates: "_LlamaTokenDataArray") -> int: + raise NotImplementedError( + "sample_token is deprecated, use LlamaSampler instead" + ) + + # Grammar + def grammar_accept_token(self, grammar: LlamaGrammar, token: int): + raise NotImplementedError( + "grammar_accept_token is deprecated, use LlamaSampler instead" + ) + + def reset_timings(self): + llama_cpp.llama_perf_context_reset(self.ctx) + + def print_timings(self): + llama_cpp.llama_perf_context_print(self.ctx) + + # Utility functions + @staticmethod + def default_params(): + """Get the default llama_context_params.""" + return llama_cpp.llama_context_default_params() + + +class LlamaBatch: + def __init__( + self, *, n_tokens: int, embd: int, n_seq_max: int, verbose: bool = True + ): + self._n_tokens = n_tokens + self.embd = embd + self.n_seq_max = n_seq_max + self.verbose = verbose + self._exit_stack = ExitStack() + + batch = llama_cpp.llama_batch_init(self._n_tokens, self.embd, self.n_seq_max) + + if batch is None: + raise ValueError("Failed to create llama_batch") + + self.batch = batch + self.sampler = None # LlamaBatch doesn't use samplers, but some cleanup code expects this attribute + + def free_batch(): + if self.batch is None: + return + llama_cpp.llama_batch_free(self.batch) + self.batch = None + + self._exit_stack.callback(free_batch) + + def close(self): + self._exit_stack.close() + + def __del__(self): + self.close() + + def n_tokens(self) -> int: + return self.batch.n_tokens + + def reset(self): + self.batch.n_tokens = 0 + + def set_batch(self, batch: Sequence[int], n_past: int, logits_all: bool): + n_tokens = len(batch) + self.batch.n_tokens = n_tokens + for i in range(n_tokens): + self.batch.token[i] = batch[i] + self.batch.pos[i] = n_past + i + self.batch.seq_id[i][0] = 0 + self.batch.n_seq_id[i] = 1 + self.batch.logits[i] = logits_all + self.batch.logits[n_tokens - 1] = True + + def add_sequence(self, batch: Sequence[int], seq_id: int, logits_all: bool): + n_tokens = len(batch) + n_tokens0 = self.batch.n_tokens + self.batch.n_tokens += n_tokens + for i in range(n_tokens): + j = n_tokens0 + i + self.batch.token[j] = batch[i] + self.batch.pos[j] = i + self.batch.seq_id[j][0] = seq_id + self.batch.n_seq_id[j] = 1 + self.batch.logits[j] = logits_all + self.batch.logits[n_tokens0 + n_tokens - 1] = True + + +class LlamaTokenDataArray: + def __init__(self, *, n_vocab: int): + self.n_vocab = n_vocab + self.candidates_data = np.recarray( + (self.n_vocab,), + dtype=np.dtype( + [("id", np.intc), ("logit", np.single), ("p", np.single)], align=True + ), + ) + self.candidates = llama_cpp.llama_token_data_array( + data=self.candidates_data.ctypes.data_as(llama_cpp.llama_token_data_p), + size=self.n_vocab, + sorted=False, + ) + self.default_candidates_data_id = np.arange(self.n_vocab, dtype=np.intc) # type: ignore + self.default_candidates_data_p = np.zeros(self.n_vocab, dtype=np.single) + self.sampler = None # LlamaTokenDataArray doesn't use samplers, but some cleanup code expects this attribute + + def copy_logits(self, logits: npt.NDArray[np.single]): + self.candidates_data.id[:] = self.default_candidates_data_id + self.candidates_data.logit[:] = logits + self.candidates_data.p[:] = self.default_candidates_data_p + self.candidates.sorted = False + self.candidates.size = self.n_vocab + + +# Embedding functions + + +def normalize_embedding(embedding): + norm = float(np.linalg.norm(embedding)) + if norm == 0.0: + return embedding + return [v / norm for v in embedding] + + +# Python wrappers over common/sampling structs + + +@dataclass +class LlamaSamplingParams: + n_prev: int = 64 + n_probs: int = 0 + top_k: int = 40 + top_p: float = 0.95 + min_p: float = 0.05 + tfs_z: float = 1.00 + typical_p: float = 1.00 + temp: float = 0.80 + penalty_last_n: int = 64 + penalty_repeat: float = 1.0 + penalty_freq: float = 0.00 + penalty_present: float = 0.00 + mirostat: int = 0 + mirostat_tau: float = 5.00 + mirostat_eta: float = 0.10 + penalize_nl: bool = True + + grammar: str = "" + + cfg_negative_prompt: str = "" + cfg_scale: float = 1.00 + + logit_bias: dict[int, float] = field(default_factory=dict) + + +@dataclass +class LlamaSamplingContext: + params: LlamaSamplingParams = field(default_factory=LlamaSamplingParams) + mirostat_mu: ctypes.c_float = field(default_factory=ctypes.c_float) + grammar: Optional[LlamaGrammar] = None + # NOTE: Missing parsed_grammar + prev: list[int] = field(default_factory=list) + cur: list[llama_cpp.llama_token_data] = field(default_factory=list) + + def reset(self): + self.prev = [] + self.cur = [] + if self.grammar is not None: + self.grammar.reset() + + def cp(self): + return LlamaSamplingContext( + params=self.params, + mirostat_mu=self.mirostat_mu, + grammar=self.grammar, + prev=self.prev.copy(), + cur=self.cur.copy(), + ) + + def last(self) -> Optional[int]: + if len(self.prev) > 0: + return self.prev[-1] + else: + return None + + def prev_str(self, ctx_main: LlamaContext, n: int) -> str: + return ctx_main.model.detokenize(self.prev[-n:]).decode("utf-8") + + def sample( + self, + ctx_main: LlamaContext, + idx: int = 0, + logits_array: Optional[npt.NDArray[np.single]] = None, + ): + # This method is deprecated in favor of using LlamaSampler directly + raise NotImplementedError( + "LlamaSamplingContext.sample is deprecated, use LlamaSampler instead" + ) + + def accept(self, ctx_main: LlamaContext, id: int, apply_grammar: bool): + self.prev.append(id) + + +class CustomSampler: + def __init__(self, apply_func: Callable[[llama_cpp.llama_token_data_array], None]): + self.apply_func = apply_func + + def apply_wrapper( + sampler: llama_cpp.llama_sampler_p, + cur_p: llama_cpp.llama_token_data_array_p, + ): + self.apply_func(cur_p) + + def free_wrapper(sampler: llama_cpp.llama_sampler_p): + pass + + sampler_i = llama_cpp.llama_sampler_i() + sampler_i.apply = llama_cpp.llama_sampler_i_apply(apply_wrapper) + self._apply_wrapper_ref = apply_wrapper + + sampler_i.name = llama_cpp.llama_sampler_i_name(0) + sampler_i.accept = llama_cpp.llama_sampler_i_accept(0) + sampler_i.reset = llama_cpp.llama_sampler_i_reset(0) + sampler_i.clone = llama_cpp.llama_sampler_i_clone(0) + sampler_i.free = llama_cpp.llama_sampler_i_free(0) + + self.sampler = llama_cpp.llama_sampler() + self.sampler.iface = ctypes.pointer(sampler_i) + self.sampler.ctx = None + + def get_sampler(self) -> llama_cpp.llama_sampler_p: + return ctypes.pointer(self.sampler) + + +class LlamaSampler: + def __init__(self): + params = llama_cpp.llama_sampler_chain_default_params() + self.sampler = llama_cpp.llama_sampler_chain_init(params) + self.custom_samplers: List[Tuple[int, CustomSampler]] = [] + self._exit_stack = ExitStack() + + def free_sampler(): + if self.sampler is not None: + # NOTE: Must remove custom samplers before free or llama.cpp will try to free them + for i, _ in reversed(self.custom_samplers): + llama_cpp.llama_sampler_chain_remove(self.sampler, i) + llama_cpp.llama_sampler_free(self.sampler) + self.sampler = None + + self._exit_stack.callback(free_sampler) + + def close(self): + self._exit_stack.close() + + def __del__(self): + self.close() + + def add_greedy(self): + sampler = llama_cpp.llama_sampler_init_greedy() + llama_cpp.llama_sampler_chain_add(self.sampler, sampler) + + def add_dist(self, seed: int): + sampler = llama_cpp.llama_sampler_init_dist(seed) + llama_cpp.llama_sampler_chain_add(self.sampler, sampler) + + def add_softmax(self): + warnings.warn( + "add_softmax is deprecated; llama_sampler_init_dist now samples directly from logits", + DeprecationWarning, + stacklevel=2, + ) + + def add_top_k(self, k: int): + sampler = llama_cpp.llama_sampler_init_top_k(k) + llama_cpp.llama_sampler_chain_add(self.sampler, sampler) + + def add_top_p(self, p: float, min_keep: int = 1): + sampler = llama_cpp.llama_sampler_init_top_p(p, min_keep) + llama_cpp.llama_sampler_chain_add(self.sampler, sampler) + + def add_min_p(self, p: float, min_keep: int = 1): + sampler = llama_cpp.llama_sampler_init_min_p(p, min_keep) + llama_cpp.llama_sampler_chain_add(self.sampler, sampler) + + def add_typical(self, p: float, min_keep: int = 1): + sampler = llama_cpp.llama_sampler_init_typical(p, min_keep) + llama_cpp.llama_sampler_chain_add(self.sampler, sampler) + + def add_temp(self, temp: float): + sampler = llama_cpp.llama_sampler_init_temp(temp) + llama_cpp.llama_sampler_chain_add(self.sampler, sampler) + + def add_temp_ext(self, t: float, delta: float, exponent: float): + sampler = llama_cpp.llama_sampler_init_temp_ext(t, delta, exponent) + llama_cpp.llama_sampler_chain_add(self.sampler, sampler) + + def add_xtc(self, p: float, t: float, min_keep: int, seed: int): + sampler = llama_cpp.llama_sampler_init_xtc(p, t, min_keep, seed) + llama_cpp.llama_sampler_chain_add(self.sampler, sampler) + + def add_top_n_sigma(self, n: float): + sampler = llama_cpp.llama_sampler_init_top_n_sigma(n) + llama_cpp.llama_sampler_chain_add(self.sampler, sampler) + + def add_mirostat(self, n_vocab: int, seed: int, tau: float, eta: float, m: int): + sampler = llama_cpp.llama_sampler_init_mirostat(n_vocab, seed, tau, eta, m) + llama_cpp.llama_sampler_chain_add(self.sampler, sampler) + + def add_mirostat_v2(self, seed: int, tau: float, eta: float): + sampler = llama_cpp.llama_sampler_init_mirostat_v2(seed, tau, eta) + llama_cpp.llama_sampler_chain_add(self.sampler, sampler) + + def add_grammar(self, model: LlamaModel, grammar: LlamaGrammar): + sampler = llama_cpp.llama_sampler_init_grammar( + model.vocab, grammar._grammar.encode("utf-8"), grammar._root.encode("utf-8") + ) + llama_cpp.llama_sampler_chain_add(self.sampler, sampler) + + def add_grammar_lazy_patterns( + self, + model: LlamaModel, + grammar: LlamaGrammar, + trigger_patterns: List[str], + trigger_tokens: List[int], + ): + # Convert patterns to C array + pattern_ptrs = (ctypes.c_char_p * len(trigger_patterns))() + for i, pattern in enumerate(trigger_patterns): + pattern_ptrs[i] = pattern.encode("utf-8") + + # Convert tokens to C array + token_array = (llama_cpp.llama_token * len(trigger_tokens))(*trigger_tokens) + + sampler = llama_cpp.llama_sampler_init_grammar_lazy_patterns( + model.vocab, + grammar._grammar.encode("utf-8"), + grammar._root.encode("utf-8"), + pattern_ptrs, + len(trigger_patterns), + token_array, + len(trigger_tokens), + ) + llama_cpp.llama_sampler_chain_add(self.sampler, sampler) + + def add_penalties( + self, + penalty_last_n: int, + penalty_repeat: float, + penalty_freq: float, + penalty_present: float, + ): + sampler = llama_cpp.llama_sampler_init_penalties( + penalty_last_n, + penalty_repeat, + penalty_freq, + penalty_present, + ) + llama_cpp.llama_sampler_chain_add(self.sampler, sampler) + + def add_dry( + self, + model: LlamaModel, + n_ctx_train: int, + dry_multiplier: float, + dry_base: float, + dry_allowed_length: int, + dry_penalty_last_n: int, + seq_breakers: List[str], + ): + # Convert seq_breakers to C array + breaker_ptrs = (ctypes.c_char_p * len(seq_breakers))() + for i, breaker in enumerate(seq_breakers): + breaker_ptrs[i] = breaker.encode("utf-8") + + sampler = llama_cpp.llama_sampler_init_dry( + model.vocab, + n_ctx_train, + dry_multiplier, + dry_base, + dry_allowed_length, + dry_penalty_last_n, + breaker_ptrs, + len(seq_breakers), + ) + llama_cpp.llama_sampler_chain_add(self.sampler, sampler) + + def add_logit_bias(self, n_vocab: int, logit_bias: Dict[int, float]): + # Convert logit_bias dict to C array + bias_array = (llama_cpp.llama_logit_bias * len(logit_bias))() + for i, (token, bias) in enumerate(logit_bias.items()): + bias_array[i].token = token + bias_array[i].bias = bias + + sampler = llama_cpp.llama_sampler_init_logit_bias( + n_vocab, len(logit_bias), bias_array + ) + llama_cpp.llama_sampler_chain_add(self.sampler, sampler) + + def add_infill(self, model: LlamaModel): + sampler = llama_cpp.llama_sampler_init_infill(model.vocab) + llama_cpp.llama_sampler_chain_add(self.sampler, sampler) + + def add_custom( + self, apply_func: Callable[[llama_cpp.llama_token_data_array], None] + ): + custom_sampler = CustomSampler(apply_func) + sampler = custom_sampler.get_sampler() + llama_cpp.llama_sampler_chain_add(self.sampler, sampler) + # NOTE: Must remove custom samplers before free or llama.cpp will try to free them + self.custom_samplers.append( + (llama_cpp.llama_sampler_chain_n(self.sampler) - 1, custom_sampler) + ) + + def get_seed(self) -> int: + return llama_cpp.llama_sampler_get_seed(self.sampler) + + def sample(self, ctx: LlamaContext, idx: int = -1) -> int: + return llama_cpp.llama_sampler_sample(self.sampler, ctx.ctx, idx) + + def accept(self, token: int): + llama_cpp.llama_sampler_accept(self.sampler, token) + + def reset(self): + llama_cpp.llama_sampler_reset(self.sampler) + + def clone(self): + # NOTE: Custom samplers cannot be cloned due to Python callback limitations + if self.custom_samplers: + raise NotImplementedError( + "Cannot clone LlamaSampler that contains custom samplers" + ) + + cloned_sampler = llama_cpp.llama_sampler_clone(self.sampler) + # Create a new wrapper around the cloned sampler + new_sampler = LlamaSampler.__new__(LlamaSampler) + new_sampler.sampler = cloned_sampler + new_sampler.custom_samplers = [] + new_sampler._exit_stack = ExitStack() + + def free_sampler(): + if new_sampler.sampler is not None: + llama_cpp.llama_sampler_free(new_sampler.sampler) + new_sampler.sampler = None + + new_sampler._exit_stack.callback(free_sampler) + return new_sampler diff --git a/llama_cpp/_logger.py b/llama_cpp/_logger.py new file mode 100644 index 0000000000..31d89d0995 --- /dev/null +++ b/llama_cpp/_logger.py @@ -0,0 +1,50 @@ +import sys +import ctypes +import logging + +import llama_cpp + +# enum ggml_log_level { +# GGML_LOG_LEVEL_NONE = 0, +# GGML_LOG_LEVEL_INFO = 1, +# GGML_LOG_LEVEL_WARN = 2, +# GGML_LOG_LEVEL_ERROR = 3, +# GGML_LOG_LEVEL_DEBUG = 4, +# GGML_LOG_LEVEL_CONT = 5, // continue previous log +# }; +GGML_LOG_LEVEL_TO_LOGGING_LEVEL = { + 0: logging.CRITICAL, + 1: logging.INFO, + 2: logging.WARNING, + 3: logging.ERROR, + 4: logging.DEBUG, + 5: logging.DEBUG, +} + +logger = logging.getLogger("llama-cpp-python") + +_last_log_level = GGML_LOG_LEVEL_TO_LOGGING_LEVEL[0] + + +# typedef void (*ggml_log_callback)(enum ggml_log_level level, const char * text, void * user_data); +@llama_cpp.llama_log_callback +def llama_log_callback( + level: int, + text: bytes, + user_data: ctypes.c_void_p, +): + # TODO: Correctly implement continue previous log + global _last_log_level + log_level = ( + GGML_LOG_LEVEL_TO_LOGGING_LEVEL[level] if level != 5 else _last_log_level + ) + if logger.level <= GGML_LOG_LEVEL_TO_LOGGING_LEVEL[level]: + print(text.decode("utf-8"), end="", flush=True, file=sys.stderr) + _last_log_level = log_level + + +llama_cpp.llama_log_set(llama_log_callback, ctypes.c_void_p(0)) + + +def set_verbose(verbose: bool): + logger.setLevel(logging.DEBUG if verbose else logging.ERROR) diff --git a/llama_cpp/_utils.py b/llama_cpp/_utils.py new file mode 100644 index 0000000000..08ede64205 --- /dev/null +++ b/llama_cpp/_utils.py @@ -0,0 +1,99 @@ +import os +import sys + +from typing import Any, Dict + +# Avoid "LookupError: unknown encoding: ascii" when open() called in a destructor +outnull_file = open(os.devnull, "w") +errnull_file = open(os.devnull, "w") + +STDOUT_FILENO = 1 +STDERR_FILENO = 2 + + +class suppress_stdout_stderr(object): + # NOTE: these must be "saved" here to avoid exceptions when using + # this context manager inside of a __del__ method + sys = sys + os = os + + def __init__(self, disable: bool = True): + self.disable = disable + + # Oddly enough this works better than the contextlib version + def __enter__(self): + if self.disable: + return self + + self.old_stdout_fileno_undup = STDOUT_FILENO + self.old_stderr_fileno_undup = STDERR_FILENO + + self.old_stdout_fileno = self.os.dup(self.old_stdout_fileno_undup) + self.old_stderr_fileno = self.os.dup(self.old_stderr_fileno_undup) + + self.old_stdout = self.sys.stdout + self.old_stderr = self.sys.stderr + + # In Jupyter notebooks, ipykernel replaces sys.stdout/stderr with + # OutStream objects that hold their own copy of the original fd in + # _original_stdstream_copy. This bypasses our dup2 redirect, so we + # need to point that copy at the real fd temporarily. + # https://github.com/ipython/ipykernel/blob/912923542d55e6c80c0e7c1f94c648b52a225011/ipykernel/iostream.py#L618-L622 + self._saved_stdout_copy = getattr( + self.sys.stdout, "_original_stdstream_copy", None + ) + self._saved_stderr_copy = getattr( + self.sys.stderr, "_original_stdstream_copy", None + ) + if self._saved_stdout_copy is not None: + self.sys.stdout._original_stdstream_copy = self.old_stdout_fileno_undup + if self._saved_stderr_copy is not None: + self.sys.stderr._original_stdstream_copy = self.old_stderr_fileno_undup + + self.os.dup2(outnull_file.fileno(), self.old_stdout_fileno_undup) + self.os.dup2(errnull_file.fileno(), self.old_stderr_fileno_undup) + + self.sys.stdout = outnull_file + self.sys.stderr = errnull_file + return self + + def __exit__(self, *_): + if self.disable: + return + + self.sys.stdout = self.old_stdout + self.sys.stderr = self.old_stderr + + self.os.dup2(self.old_stdout_fileno, self.old_stdout_fileno_undup) + self.os.dup2(self.old_stderr_fileno, self.old_stderr_fileno_undup) + + self.os.close(self.old_stdout_fileno) + self.os.close(self.old_stderr_fileno) + + # Restore ipykernel's OutStream fd copies + if self._saved_stdout_copy is not None: + self.sys.stdout._original_stdstream_copy = self._saved_stdout_copy + if self._saved_stderr_copy is not None: + self.sys.stderr._original_stdstream_copy = self._saved_stderr_copy + + +class MetaSingleton(type): + """ + Metaclass for implementing the Singleton pattern. + """ + + _instances: Dict[type, Any] = {} + + def __call__(cls, *args: Any, **kwargs: Any) -> Any: + if cls not in cls._instances: + cls._instances[cls] = super(MetaSingleton, cls).__call__(*args, **kwargs) + return cls._instances[cls] + + +class Singleton(object, metaclass=MetaSingleton): + """ + Base class for implementing the Singleton pattern. + """ + + def __init__(self): + super(Singleton, self).__init__() diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py index edb68c9e58..4a09b55ee5 100644 --- a/llama_cpp/llama.py +++ b/llama_cpp/llama.py @@ -1,246 +1,196 @@ +from __future__ import annotations + import os import sys import uuid import time -import math +import json +import ctypes +import typing +import random +import fnmatch +import warnings +import contextlib import multiprocessing -from abc import ABC, abstractmethod + from typing import ( + Any, List, + Literal, Optional, Union, Generator, Sequence, Iterator, Deque, - Tuple, Callable, + Dict, ) -from collections import deque, OrderedDict +from collections import deque +from pathlib import Path -import diskcache -from . import llama_cpp from .llama_types import * +from .llama_grammar import LlamaGrammar +from .llama_cache import ( + BaseLlamaCache, + LlamaCache, # type: ignore + LlamaDiskCache, # type: ignore + LlamaRAMCache, # type: ignore +) +from .llama_tokenizer import BaseLlamaTokenizer, LlamaTokenizer +import llama_cpp.llama_cpp as llama_cpp +import llama_cpp.llama_chat_format as llama_chat_format + +from llama_cpp.llama_speculative import LlamaDraftModel import numpy as np import numpy.typing as npt - -class BaseLlamaCache(ABC): - """Base cache class for a llama.cpp model.""" - - def __init__(self, capacity_bytes: int = (2 << 30)): - self.capacity_bytes = capacity_bytes - - @property - @abstractmethod - def cache_size(self) -> int: - raise NotImplementedError - - def _find_longest_prefix_key( - self, - key: Tuple[int, ...], - ) -> Optional[Tuple[int, ...]]: - pass - - @abstractmethod - def __getitem__(self, key: Sequence[int]) -> "LlamaState": - raise NotImplementedError - - @abstractmethod - def __contains__(self, key: Sequence[int]) -> bool: - raise NotImplementedError - - @abstractmethod - def __setitem__(self, key: Sequence[int], value: "LlamaState") -> None: - raise NotImplementedError - - -class LlamaRAMCache(BaseLlamaCache): - """Cache for a llama.cpp model using RAM.""" - - def __init__(self, capacity_bytes: int = (2 << 30)): - super().__init__(capacity_bytes) - self.capacity_bytes = capacity_bytes - self.cache_state: OrderedDict[Tuple[int, ...], "LlamaState"] = OrderedDict() - - @property - def cache_size(self): - return sum([state.llama_state_size for state in self.cache_state.values()]) - - def _find_longest_prefix_key( - self, - key: Tuple[int, ...], - ) -> Optional[Tuple[int, ...]]: - min_len = 0 - min_key = None - keys = ( - (k, Llama.longest_token_prefix(k, key)) for k in self.cache_state.keys() - ) - for k, prefix_len in keys: - if prefix_len > min_len: - min_len = prefix_len - min_key = k - return min_key - - def __getitem__(self, key: Sequence[int]) -> "LlamaState": - key = tuple(key) - _key = self._find_longest_prefix_key(key) - if _key is None: - raise KeyError("Key not found") - value = self.cache_state[_key] - self.cache_state.move_to_end(_key) - return value - - def __contains__(self, key: Sequence[int]) -> bool: - return self._find_longest_prefix_key(tuple(key)) is not None - - def __setitem__(self, key: Sequence[int], value: "LlamaState"): - key = tuple(key) - if key in self.cache_state: - del self.cache_state[key] - self.cache_state[key] = value - while self.cache_size > self.capacity_bytes and len(self.cache_state) > 0: - self.cache_state.popitem(last=False) - - -# Alias for backwards compatibility -LlamaCache = LlamaRAMCache - - -class LlamaDiskCache(BaseLlamaCache): - """Cache for a llama.cpp model using disk.""" - - def __init__( - self, cache_dir: str = ".cache/llama_cache", capacity_bytes: int = (2 << 30) - ): - super().__init__(capacity_bytes) - self.cache = diskcache.Cache(cache_dir) - - @property - def cache_size(self): - return int(self.cache.volume()) # type: ignore - - def _find_longest_prefix_key( - self, - key: Tuple[int, ...], - ) -> Optional[Tuple[int, ...]]: - min_len = 0 - min_key: Optional[Tuple[int, ...]] = None - for k in self.cache.iterkeys(): # type: ignore - prefix_len = Llama.longest_token_prefix(k, key) - if prefix_len > min_len: - min_len = prefix_len - min_key = k # type: ignore - return min_key - - def __getitem__(self, key: Sequence[int]) -> "LlamaState": - key = tuple(key) - _key = self._find_longest_prefix_key(key) - if _key is None: - raise KeyError("Key not found") - value: "LlamaState" = self.cache.pop(_key) # type: ignore - # NOTE: This puts an integer as key in cache, which breaks, - # Llama.longest_token_prefix(k, key) above since k is not a tuple of ints/tokens - # self.cache.push(_key, side="front") # type: ignore - return value - - def __contains__(self, key: Sequence[int]) -> bool: - return self._find_longest_prefix_key(tuple(key)) is not None - - def __setitem__(self, key: Sequence[int], value: "LlamaState"): - print("LlamaDiskCache.__setitem__: called", file=sys.stderr) - key = tuple(key) - if key in self.cache: - print("LlamaDiskCache.__setitem__: delete", file=sys.stderr) - del self.cache[key] - self.cache[key] = value - print("LlamaDiskCache.__setitem__: set", file=sys.stderr) - while self.cache_size > self.capacity_bytes and len(self.cache) > 0: - key_to_remove = next(iter(self.cache)) - del self.cache[key_to_remove] - print("LlamaDiskCache.__setitem__: trim", file=sys.stderr) - - -class LlamaState: - def __init__( - self, - input_ids: npt.NDArray[np.intc], - scores: npt.NDArray[np.single], - n_tokens: int, - llama_state: bytes, - llama_state_size: int, - ): - self.input_ids = input_ids - self.scores = scores - self.n_tokens = n_tokens - self.llama_state = llama_state - self.llama_state_size = llama_state_size - - -LogitsProcessor = Callable[[List[int], List[float]], List[float]] - - -class LogitsProcessorList(List[LogitsProcessor]): - def __call__(self, input_ids: List[int], scores: List[float]) -> List[float]: - for processor in self: - scores = processor(input_ids, scores) - return scores - - -StoppingCriteria = Callable[[List[int], List[float]], bool] - - -class StoppingCriteriaList(List[StoppingCriteria]): - def __call__(self, input_ids: List[int], logits: List[float]) -> bool: - return any([stopping_criteria(input_ids, logits) for stopping_criteria in self]) +import llama_cpp._internals as internals +from ._logger import set_verbose +from ._utils import suppress_stdout_stderr class Llama: """High-level Python wrapper for a llama.cpp model.""" + __backend_initialized = False + def __init__( self, model_path: str, - # NOTE: These parameters are likely to change in the future. - n_ctx: int = 512, - n_parts: int = -1, + *, + # Model Params n_gpu_layers: int = 0, - seed: int = 1337, - f16_kv: bool = True, - logits_all: bool = False, + split_mode: int = llama_cpp.LLAMA_SPLIT_MODE_LAYER, + main_gpu: int = 0, + tensor_split: Optional[List[float]] = None, vocab_only: bool = False, use_mmap: bool = True, use_mlock: bool = False, - embedding: bool = False, - n_threads: Optional[int] = None, + kv_overrides: Optional[Dict[str, Union[bool, int, float, str]]] = None, + # Context Params + seed: int = llama_cpp.LLAMA_DEFAULT_SEED, + n_ctx: int = 512, n_batch: int = 512, + n_ubatch: int = 512, + n_threads: Optional[int] = None, + n_threads_batch: Optional[int] = None, + rope_scaling_type: Optional[ + int + ] = llama_cpp.LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED, + pooling_type: int = llama_cpp.LLAMA_POOLING_TYPE_UNSPECIFIED, + attention_type: int = llama_cpp.LLAMA_ATTENTION_TYPE_UNSPECIFIED, + rope_freq_base: float = 0.0, + rope_freq_scale: float = 0.0, + yarn_ext_factor: float = -1.0, + yarn_attn_factor: float = 1.0, + yarn_beta_fast: float = 32.0, + yarn_beta_slow: float = 1.0, + yarn_orig_ctx: int = 0, + logits_all: bool = False, + embedding: bool = False, + offload_kqv: bool = True, + flash_attn: bool = False, + op_offload: Optional[bool] = None, + swa_full: Optional[bool] = None, + # Sampling Params + no_perf: bool = False, last_n_tokens_size: int = 64, + # LoRA Params lora_base: Optional[str] = None, + lora_scale: float = 1.0, lora_path: Optional[str] = None, - low_vram: bool = False, + # Backend Params + numa: Union[bool, int] = False, + # Chat Format Params + chat_format: Optional[str] = None, + chat_handler: Optional[llama_chat_format.LlamaChatCompletionHandler] = None, + # Speculative Decoding + draft_model: Optional[LlamaDraftModel] = None, + # Tokenizer Override + tokenizer: Optional[BaseLlamaTokenizer] = None, + # KV cache quantization + type_k: Optional[int] = None, + type_v: Optional[int] = None, + # Misc + spm_infill: bool = False, verbose: bool = True, + # Extra Params + **kwargs, # type: ignore ): """Load a llama.cpp model from `model_path`. + Examples: + Basic usage + + >>> import llama_cpp + >>> model = llama_cpp.Llama( + ... model_path="path/to/model", + ... ) + >>> print(model("The quick brown fox jumps ", stop=["."])["choices"][0]["text"]) + the lazy dog + + Loading a chat model + + >>> import llama_cpp + >>> model = llama_cpp.Llama( + ... model_path="path/to/model", + ... chat_format="llama-2", + ... ) + >>> print(model.create_chat_completion( + ... messages=[{ + ... "role": "user", + ... "content": "what is the meaning of life?" + ... }] + ... )) + Args: model_path: Path to the model. - n_ctx: Maximum context size. - n_parts: Number of parts to split the model into. If -1, the number of parts is automatically determined. - seed: Random seed. -1 for random. - f16_kv: Use half-precision for key/value cache. - logits_all: Return logits for all tokens, not just the last token. + n_gpu_layers: Number of layers to offload to GPU (-ngl). If -1, all layers are offloaded. + split_mode: How to split the model across GPUs. See llama_cpp.LLAMA_SPLIT_* for options. + main_gpu: main_gpu interpretation depends on split_mode: LLAMA_SPLIT_MODE_NONE: the GPU that is used for the entire model. LLAMA_SPLIT_MODE_ROW: the GPU that is used for small tensors and intermediate results. LLAMA_SPLIT_MODE_LAYER: ignored + tensor_split: How split tensors should be distributed across GPUs. If None, the model is not split. vocab_only: Only load the vocabulary no weights. use_mmap: Use mmap if possible. use_mlock: Force the system to keep the model in RAM. + kv_overrides: Key-value overrides for the model. + seed: RNG seed, -1 for random + n_ctx: Text context, 0 = from model + n_batch: Prompt processing maximum batch size + n_ubatch: Physical batch size + n_threads: Number of threads to use for generation + n_threads_batch: Number of threads to use for batch processing + rope_scaling_type: RoPE scaling type, from `enum llama_rope_scaling_type`. ref: https://github.com/ggerganov/llama.cpp/pull/2054 + pooling_type: Pooling type, from `enum llama_pooling_type`. + attention_type: Attention type, from `enum llama_attention_type`. + rope_freq_base: RoPE base frequency, 0 = from model + rope_freq_scale: RoPE frequency scaling factor, 0 = from model + yarn_ext_factor: YaRN extrapolation mix factor, negative = from model + yarn_attn_factor: YaRN magnitude scaling factor + yarn_beta_fast: YaRN low correction dim + yarn_beta_slow: YaRN high correction dim + yarn_orig_ctx: YaRN original context size + logits_all: Return logits for all tokens, not just the last token. Must be True for completion to return logprobs. embedding: Embedding mode only. - n_threads: Number of threads to use. If None, the number of threads is automatically determined. - n_batch: Maximum number of prompt tokens to batch together when calling llama_eval. + offload_kqv: Offload K, Q, V to GPU. + flash_attn: Use flash attention. + op_offload: offload host tensor operations to device + swa_full: use full-size SWA cache (https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055) + no_perf: Measure performance timings. last_n_tokens_size: Maximum number of tokens to keep in the last_n_tokens deque. lora_base: Optional path to base model, useful if using a quantized base model and you want to apply LoRA to an f16 model. lora_path: Path to a LoRA file to apply to the model. + numa: numa policy + chat_format: String specifying the chat format to use when calling create_chat_completion. + chat_handler: Optional chat handler to use when calling create_chat_completion. + draft_model: Optional draft model to use for speculative decoding. + tokenizer: Optional tokenizer to override the default tokenizer from llama.cpp. verbose: Print verbose output to stderr. + type_k: KV cache data type for K (default: f16) + type_v: KV cache data type for V (default: f16) + spm_infill: Use Suffix/Prefix/Middle pattern for infill (instead of Prefix/Suffix/Middle) as some models prefer this. Raises: ValueError: If the model path does not exist. @@ -249,90 +199,378 @@ def __init__( A Llama instance. """ self.verbose = verbose + self._stack = contextlib.ExitStack() + + set_verbose(verbose) + + if not Llama.__backend_initialized: + with suppress_stdout_stderr(disable=verbose): + llama_cpp.llama_backend_init() + Llama.__backend_initialized = True + + if isinstance(numa, bool): + self.numa = ( + llama_cpp.GGML_NUMA_STRATEGY_DISTRIBUTE + if numa + else llama_cpp.GGML_NUMA_STRATEGY_DISABLED + ) + else: + self.numa = numa + + if self.numa != llama_cpp.GGML_NUMA_STRATEGY_DISABLED: + with suppress_stdout_stderr(disable=verbose): + llama_cpp.llama_numa_init(self.numa) + self.model_path = model_path - self.params = llama_cpp.llama_context_default_params() - self.params.n_ctx = n_ctx - self.params.n_gpu_layers = n_gpu_layers - self.params.seed = seed - self.params.f16_kv = f16_kv - self.params.logits_all = logits_all - self.params.vocab_only = vocab_only - self.params.use_mmap = use_mmap if lora_path is None else False - self.params.use_mlock = use_mlock - self.params.embedding = embedding - self.params.low_vram = low_vram + # Model Params + self.model_params = llama_cpp.llama_model_default_params() + self.model_params.n_gpu_layers = ( + 0x7FFFFFFF if n_gpu_layers == -1 else n_gpu_layers + ) # 0x7FFFFFFF is INT32 max, will be auto set to all layers + self.model_params.split_mode = split_mode + self.model_params.main_gpu = main_gpu + self.tensor_split = tensor_split + self._c_tensor_split = None + if self.tensor_split is not None: + if len(self.tensor_split) > llama_cpp.LLAMA_MAX_DEVICES: + raise ValueError( + f"Attempt to split tensors that exceed maximum supported devices. Current LLAMA_MAX_DEVICES={llama_cpp.LLAMA_MAX_DEVICES}" + ) + # Type conversion and expand the list to the length of LLAMA_MAX_DEVICES + FloatArray = ctypes.c_float * llama_cpp.LLAMA_MAX_DEVICES + self._c_tensor_split = FloatArray( + *tensor_split # type: ignore + ) # keep a reference to the array so it is not gc'd + self.model_params.tensor_split = self._c_tensor_split + self.model_params.vocab_only = vocab_only + self.model_params.use_mmap = use_mmap if lora_path is None else False + self.model_params.use_mlock = use_mlock + + # kv_overrides is the original python dict + self.kv_overrides = kv_overrides + if kv_overrides is not None: + # _kv_overrides_array is a ctypes.Array of llama_model_kv_override Structs + kvo_array_len = len(kv_overrides) + 1 # for sentinel element + self._kv_overrides_array = ( + llama_cpp.llama_model_kv_override * kvo_array_len + )() + + for i, (k, v) in enumerate(kv_overrides.items()): + self._kv_overrides_array[i].key = k.encode("utf-8") + if isinstance(v, bool): + self._kv_overrides_array[ + i + ].tag = llama_cpp.LLAMA_KV_OVERRIDE_TYPE_BOOL + self._kv_overrides_array[i].value.val_bool = v + elif isinstance(v, int): + self._kv_overrides_array[ + i + ].tag = llama_cpp.LLAMA_KV_OVERRIDE_TYPE_INT + self._kv_overrides_array[i].value.val_i64 = v + elif isinstance(v, float): + self._kv_overrides_array[ + i + ].tag = llama_cpp.LLAMA_KV_OVERRIDE_TYPE_FLOAT + self._kv_overrides_array[i].value.val_f64 = v + elif isinstance(v, str): # type: ignore + v_bytes = v.encode("utf-8") + if len(v_bytes) > 128: # TODO: Make this a constant + raise ValueError(f"Value for {k} is too long: {v}") + v_bytes = v_bytes.ljust(128, b"\0") + self._kv_overrides_array[ + i + ].tag = llama_cpp.LLAMA_KV_OVERRIDE_TYPE_STR + # copy min(v_bytes, 128) to str_value + address = typing.cast( + int, + ctypes.addressof(self._kv_overrides_array[i].value) + + llama_cpp.llama_model_kv_override_value.val_str.offset, + ) + buffer_start = ctypes.cast(address, ctypes.POINTER(ctypes.c_char)) + ctypes.memmove( + buffer_start, + v_bytes, + 128, + ) + else: + raise ValueError(f"Unknown value type for {k}: {v}") + + self._kv_overrides_array[ + -1 + ].key = b"\0" # ensure sentinel element is zeroed + self.model_params.kv_overrides = self._kv_overrides_array + self.n_batch = min(n_ctx, n_batch) # ??? + self.n_threads = n_threads or max(multiprocessing.cpu_count() // 2, 1) + self.n_threads_batch = n_threads_batch or multiprocessing.cpu_count() + + # Used by the sampler + self._seed = seed or llama_cpp.LLAMA_DEFAULT_SEED + + # Context Params + self.context_params = llama_cpp.llama_context_default_params() + self.context_params.n_ctx = n_ctx + self.context_params.n_batch = self.n_batch + self.context_params.n_ubatch = min(self.n_batch, n_ubatch) + self.context_params.n_threads = self.n_threads + self.context_params.n_threads_batch = self.n_threads_batch + self.context_params.rope_scaling_type = ( + rope_scaling_type + if rope_scaling_type is not None + else llama_cpp.LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED + ) + self.context_params.pooling_type = pooling_type + self.context_params.attention_type = attention_type + self.context_params.rope_freq_base = ( + rope_freq_base if rope_freq_base != 0.0 else 0 + ) + self.context_params.rope_freq_scale = ( + rope_freq_scale if rope_freq_scale != 0.0 else 0 + ) + self.context_params.yarn_ext_factor = ( + yarn_ext_factor if yarn_ext_factor != 0.0 else 0 + ) + self.context_params.yarn_attn_factor = ( + yarn_attn_factor if yarn_attn_factor != 0.0 else 0 + ) + self.context_params.yarn_beta_fast = ( + yarn_beta_fast if yarn_beta_fast != 0.0 else 0 + ) + self.context_params.yarn_beta_slow = ( + yarn_beta_slow if yarn_beta_slow != 0.0 else 0 + ) + self.context_params.yarn_orig_ctx = yarn_orig_ctx if yarn_orig_ctx != 0 else 0 + self._logits_all = logits_all if draft_model is None else True + self.context_params.embeddings = embedding # TODO: Rename to embeddings + self.context_params.offload_kqv = offload_kqv + self.context_params.flash_attn_type = ( + llama_cpp.LLAMA_FLASH_ATTN_TYPE_ENABLED + if flash_attn + else llama_cpp.LLAMA_FLASH_ATTN_TYPE_DISABLED + ) + + if op_offload is not None: + self.context_params.op_offload = op_offload + + if swa_full is not None: + self.context_params.swa_full = swa_full + + # KV cache quantization + if type_k is not None: + self.context_params.type_k = type_k + if type_v is not None: + self.context_params.type_v = type_v + # Sampling Params + self.context_params.no_perf = no_perf self.last_n_tokens_size = last_n_tokens_size - self.n_batch = min(n_ctx, n_batch) self.cache: Optional[BaseLlamaCache] = None - self.n_threads = n_threads or max(multiprocessing.cpu_count() // 2, 1) - self.lora_base = lora_base + self.lora_scale = lora_scale self.lora_path = lora_path - ### DEPRECATED ### - self.n_parts = n_parts - ### DEPRECATED ### + self.spm_infill = spm_infill if not os.path.exists(model_path): raise ValueError(f"Model path does not exist: {model_path}") - self.model = llama_cpp.llama_load_model_from_file( - self.model_path.encode("utf-8"), self.params + self._model = self._stack.enter_context( + contextlib.closing( + internals.LlamaModel( + path_model=self.model_path, + params=self.model_params, + verbose=self.verbose, + ) + ) ) - assert self.model is not None - self.ctx = llama_cpp.llama_new_context_with_model(self.model, self.params) + # Override tokenizer + self.tokenizer_ = tokenizer or LlamaTokenizer(self) + + # Set the default value for the context and correct the batch + if n_ctx == 0: + n_ctx = self._model.n_ctx_train() + self.n_batch = min(n_ctx, n_batch) + self.context_params.n_ctx = self._model.n_ctx_train() + self.context_params.n_batch = self.n_batch + self.context_params.n_ubatch = min(self.n_batch, n_ubatch) + + if embedding: + self.context_params.n_seq_max = min( + self.n_batch, + llama_cpp.llama_max_parallel_sequences(), + ) + self.context_params.kv_unified = True + self._ctx = self._stack.enter_context( + contextlib.closing( + internals.LlamaContext( + model=self._model, + params=self.context_params, + verbose=self.verbose, + ) + ) + ) - assert self.ctx is not None + self._batch = self._stack.enter_context( + contextlib.closing( + internals.LlamaBatch( + n_tokens=self.n_batch, + embd=0, + n_seq_max=self.context_params.n_ctx, + verbose=self.verbose, + ) + ) + ) + + self._lora_adapter: Optional[llama_cpp.llama_adapter_lora_p] = None if self.lora_path: - if llama_cpp.llama_model_apply_lora_from_file( - self.model, - llama_cpp.c_char_p(self.lora_path.encode("utf-8")), - llama_cpp.c_char_p(self.lora_base.encode("utf-8")) - if self.lora_base is not None - else llama_cpp.c_char_p(0), - llama_cpp.c_int(self.n_threads), - ): + self._lora_adapter = llama_cpp.llama_adapter_lora_init( + self._model.model, + self.lora_path.encode("utf-8"), + ) + if self._lora_adapter is None: + raise RuntimeError( + f"Failed to initialize LoRA adapter from lora path: {self.lora_path}" + ) + + def free_lora_adapter(): + if self._lora_adapter is None: + return + llama_cpp.llama_adapter_lora_free(self._lora_adapter) + self._lora_adapter = None + + self._stack.callback(free_lora_adapter) + + adapters = (llama_cpp.llama_adapter_lora_p_ctypes * 1)(self._lora_adapter) + scales = (ctypes.c_float * 1)(self.lora_scale) + if llama_cpp.llama_set_adapters_lora(self._ctx.ctx, adapters, 1, scales): raise RuntimeError( - f"Failed to apply LoRA from lora path: {self.lora_path} to base path: {self.lora_base}" + f"Failed to set LoRA adapter from lora path: {self.lora_path}" ) if self.verbose: print(llama_cpp.llama_print_system_info().decode("utf-8"), file=sys.stderr) + self.chat_format = chat_format + self.chat_handler = chat_handler + self._chat_handlers: Dict[ + str, llama_chat_format.LlamaChatCompletionHandler + ] = {} + + self.draft_model = draft_model + self._n_vocab = self.n_vocab() self._n_ctx = self.n_ctx() - size = llama_cpp.c_size_t(self._n_vocab) - sorted = llama_cpp.c_bool(False) - self._candidates_data = np.array( - [], - dtype=np.dtype( - [("id", np.intc), ("logit", np.single), ("p", np.single)], align=True - ), - ) - self._candidates_data.resize(3, self._n_vocab, refcheck=False) - candidates = llama_cpp.llama_token_data_array( - data=self._candidates_data.ctypes.data_as(llama_cpp.llama_token_data_p), - size=size, - sorted=sorted, - ) - self._candidates = candidates - self._token_nl = Llama.token_nl() - self._token_eos = Llama.token_eos() - self._candidates_data_id = np.arange(self._n_vocab, dtype=np.intc) # type: ignore - self._candidates_data_p = np.zeros(self._n_vocab, dtype=np.single) + + self._token_nl = self.token_nl() + self._token_eos = self.token_eos() + + self._candidates = internals.LlamaTokenDataArray(n_vocab=self._n_vocab) self.n_tokens = 0 self.input_ids: npt.NDArray[np.intc] = np.ndarray((n_ctx,), dtype=np.intc) self.scores: npt.NDArray[np.single] = np.ndarray( - (n_ctx, self._n_vocab), dtype=np.single + (n_ctx if logits_all == True else n_batch, self._n_vocab), dtype=np.single + ) + + self._mirostat_mu = ctypes.c_float( + 2.0 * 5.0 + ) # TODO: Move this to sampling context + + try: + self.metadata = self._model.metadata() + except Exception as e: + self.metadata = {} + if self.verbose: + print(f"Failed to load metadata: {e}", file=sys.stderr) + + if self.verbose: + print(f"Model metadata: {self.metadata}", file=sys.stderr) + + eos_token_id = self.token_eos() + bos_token_id = self.token_bos() + + eos_token = ( + self._model.token_get_text(eos_token_id) if eos_token_id != -1 else "" + ) + bos_token = ( + self._model.token_get_text(bos_token_id) if bos_token_id != -1 else "" ) + # Unfortunately the llama.cpp API does not return metadata arrays, so we can't get template names from tokenizer.chat_templates + template_choices = dict( + (name[10:], template) + for name, template in self.metadata.items() + if name.startswith("tokenizer.chat_template.") + ) + + if "tokenizer.chat_template" in self.metadata: + template_choices["chat_template.default"] = self.metadata[ + "tokenizer.chat_template" + ] + + if self.verbose and template_choices: + print( + f"Available chat formats from metadata: {', '.join(template_choices.keys())}", + file=sys.stderr, + ) + + for name, template in template_choices.items(): + self._chat_handlers[name] = llama_chat_format.Jinja2ChatFormatter( + template=template, + eos_token=eos_token, + bos_token=bos_token, + stop_token_ids=[eos_token_id], + ).to_chat_handler() + + if ( + self.chat_format is None + and self.chat_handler is None + and "chat_template.default" in template_choices + ): + chat_format = llama_chat_format.guess_chat_format_from_gguf_metadata( + self.metadata + ) + + if chat_format is not None: + self.chat_format = chat_format + if self.verbose: + print(f"Guessed chat format: {chat_format}", file=sys.stderr) + else: + if self.verbose: + print( + f"Using gguf chat template: {template_choices['chat_template.default']}", + file=sys.stderr, + ) + print(f"Using chat eos_token: {eos_token}", file=sys.stderr) + print(f"Using chat bos_token: {bos_token}", file=sys.stderr) + + self.chat_format = "chat_template.default" + + if self.chat_format is None and self.chat_handler is None: + self.chat_format = "llama-2" + if self.verbose: + print( + f"Using fallback chat format: {self.chat_format}", file=sys.stderr + ) + + self._sampler = None + + # Cache recurrent/hybrid model detection to avoid repeated FFI calls + self._is_recurrent = llama_cpp.llama_model_is_recurrent(self._model.model) + self._is_hybrid = llama_cpp.llama_model_is_hybrid(self._model.model) + + @property + def ctx(self) -> llama_cpp.llama_context_p: + return self._ctx.ctx + + @property + def model(self) -> llama_cpp.llama_model_p: + return self._model.model + @property def _input_ids(self) -> npt.NDArray[np.intc]: return self.input_ids[: self.n_tokens] @@ -349,14 +587,18 @@ def eval_tokens(self) -> Deque[int]: def eval_logits(self) -> Deque[List[float]]: return deque( self.scores[: self.n_tokens, :].tolist(), - maxlen=self._n_ctx if self.params.logits_all else 1, + maxlen=self._n_ctx if self._logits_all else 1, ) - def tokenize(self, text: bytes, add_bos: bool = True) -> List[int]: + def tokenize( + self, text: bytes, add_bos: bool = True, special: bool = False + ) -> List[int]: """Tokenize a string. Args: text: The utf-8 encoded string to tokenize. + add_bos: Whether to add a beginning of sequence token. + special: Whether to tokenize special tokens. Raises: RuntimeError: If the tokenization failed. @@ -364,48 +606,27 @@ def tokenize(self, text: bytes, add_bos: bool = True) -> List[int]: Returns: A list of tokens. """ - assert self.ctx is not None - n_ctx = self._n_ctx - tokens = (llama_cpp.llama_token * n_ctx)() - n_tokens = llama_cpp.llama_tokenize( - self.ctx, - text, - tokens, - llama_cpp.c_int(n_ctx), - llama_cpp.c_bool(add_bos), - ) - if n_tokens < 0: - n_tokens = abs(n_tokens) - tokens = (llama_cpp.llama_token * n_tokens)() - n_tokens = llama_cpp.llama_tokenize( - self.ctx, - text, - tokens, - llama_cpp.c_int(n_tokens), - llama_cpp.c_bool(add_bos), - ) - if n_tokens < 0: - raise RuntimeError( - f'Failed to tokenize: text="{text}" n_tokens={n_tokens}' - ) - return list(tokens[:n_tokens]) + return self.tokenizer_.tokenize(text, add_bos, special) - def detokenize(self, tokens: List[int]) -> bytes: + def detokenize( + self, + tokens: List[int], + prev_tokens: Optional[List[int]] = None, + special: bool = False, + ) -> bytes: """Detokenize a list of tokens. Args: tokens: The list of tokens to detokenize. + prev_tokens: The list of previous tokens. Offset mapping will be performed if provided. + special: Whether to detokenize special tokens. Returns: The detokenized string. """ - assert self.ctx is not None - output = b"" - for token in tokens: - output += llama_cpp.llama_token_to_str( - self.ctx, llama_cpp.llama_token(token) - ) - return output + return self.tokenizer_.detokenize( + tokens, prev_tokens=prev_tokens, special=special + ) def set_cache(self, cache: Optional[BaseLlamaCache]): """Set the cache. @@ -415,183 +636,156 @@ def set_cache(self, cache: Optional[BaseLlamaCache]): """ self.cache = cache + def set_seed(self, seed: int): + """Set the random seed. + + Args: + seed: The random seed. + """ + self._seed = seed + def reset(self): """Reset the model state.""" self.n_tokens = 0 + if self._is_recurrent or self._is_hybrid: + mem = llama_cpp.llama_get_memory(self._ctx.ctx) + if mem is not None: + llama_cpp.llama_memory_clear(mem, True) + def eval(self, tokens: Sequence[int]): """Evaluate a list of tokens. Args: tokens: The list of tokens to evaluate. """ - assert self.ctx is not None - n_ctx = self._n_ctx + self._ctx.kv_cache_seq_rm(-1, self.n_tokens, -1) for i in range(0, len(tokens), self.n_batch): batch = tokens[i : min(len(tokens), i + self.n_batch)] - n_past = min(n_ctx - len(batch), len(self._input_ids)) + n_past = self.n_tokens n_tokens = len(batch) - return_code = llama_cpp.llama_eval( - ctx=self.ctx, - tokens=(llama_cpp.llama_token * len(batch))(*batch), - n_tokens=llama_cpp.c_int(n_tokens), - n_past=llama_cpp.c_int(n_past), - n_threads=llama_cpp.c_int(self.n_threads), + self._batch.set_batch( + batch=batch, n_past=n_past, logits_all=self._logits_all ) - if return_code != 0: - raise RuntimeError(f"llama_eval returned {return_code}") + self._ctx.decode(self._batch) # Save tokens - self.input_ids[self.n_tokens : self.n_tokens + n_tokens] = batch + self.input_ids[n_past : n_past + n_tokens] = batch # Save logits - rows = n_tokens if self.params.logits_all else 1 - cols = self._n_vocab - offset = ( - 0 if self.params.logits_all else n_tokens - 1 - ) # NOTE: Only save the last token logits if logits_all is False - self.scores[self.n_tokens + offset : self.n_tokens + n_tokens, :].reshape( - -1 - )[:] = llama_cpp.llama_get_logits(self.ctx)[: rows * cols] + if self._logits_all: + rows = n_tokens + cols = self._n_vocab + logits = np.ctypeslib.as_array( + self._ctx.get_logits(), shape=(rows * cols,) + ) + self.scores[n_past : n_past + n_tokens, :].reshape(-1)[::] = logits + else: + # rows = 1 + # cols = self._n_vocab + # logits = np.ctypeslib.as_array( + # self._ctx.get_logits(), shape=(rows * cols,) + # ) + # self.scores[n_past + n_tokens - 1, :].reshape(-1)[::] = logits + # NOTE: Now that sampling is done inside the sampler, logits are only needed for logprobs which requires logits_all + pass # Update n_tokens self.n_tokens += n_tokens - def _sample( + def _init_sampler( self, - last_n_tokens_data, # type: llama_cpp.Array[llama_cpp.llama_token] - last_n_tokens_size: llama_cpp.c_int, - top_k: llama_cpp.c_int, - top_p: llama_cpp.c_float, - temp: llama_cpp.c_float, - tfs_z: llama_cpp.c_float, - repeat_penalty: llama_cpp.c_float, - frequency_penalty: llama_cpp.c_float, - presence_penalty: llama_cpp.c_float, - mirostat_mode: llama_cpp.c_int, - mirostat_tau: llama_cpp.c_float, - mirostat_eta: llama_cpp.c_float, + top_k: int = 40, + top_p: float = 0.95, + min_p: float = 0.05, + typical_p: float = 1.0, + temp: float = 0.80, + repeat_penalty: float = 1.0, + frequency_penalty: float = 0.0, + presence_penalty: float = 0.0, + tfs_z: float = 1.0, + mirostat_mode: int = 0, + mirostat_eta: float = 0.1, + mirostat_tau: float = 5.0, penalize_nl: bool = True, logits_processor: Optional[LogitsProcessorList] = None, + grammar: Optional[LlamaGrammar] = None, ): - assert self.ctx is not None - assert self.n_tokens > 0 - n_vocab = self._n_vocab - n_ctx = self._n_ctx - top_k = llama_cpp.c_int(n_vocab) if top_k.value <= 0 else top_k - last_n_tokens_size = ( - llama_cpp.c_int(n_ctx) - if last_n_tokens_size.value < 0 - else last_n_tokens_size - ) - logits: npt.NDArray[np.single] = self._scores[-1, :] + sampler = internals.LlamaSampler() if logits_processor is not None: - logits = np.array( - logits_processor(self._input_ids.tolist(), logits.tolist()), - dtype=np.single, - ) - self._scores[-1, :] = logits - - nl_logit = logits[self._token_nl] - candidates = self._candidates - candidates_data = self._candidates_data - candidates_data["id"][:] = self._candidates_data_id # type: ignore - candidates_data["logit"][:] = logits - candidates_data["p"][:] = self._candidates_data_p # type: ignore - candidates.data = candidates_data.ctypes.data_as(llama_cpp.llama_token_data_p) - candidates.sorted = llama_cpp.c_bool(False) - candidates.size = llama_cpp.c_size_t(n_vocab) - llama_cpp.llama_sample_repetition_penalty( - ctx=self.ctx, - last_tokens_data=last_n_tokens_data, - last_tokens_size=last_n_tokens_size, - candidates=llama_cpp.ctypes.byref(candidates), # type: ignore - penalty=repeat_penalty, - ) - llama_cpp.llama_sample_frequency_and_presence_penalties( - ctx=self.ctx, - candidates=llama_cpp.ctypes.byref(candidates), # type: ignore - last_tokens_data=last_n_tokens_data, - last_tokens_size=last_n_tokens_size, - alpha_frequency=frequency_penalty, - alpha_presence=presence_penalty, + # Create and add a custom sampler + def apply_func(token_data_array: llama_cpp.llama_token_data_array_p): + size = token_data_array.contents.size + data_soa = token_data_array.contents.data + data_soa_address = ctypes.addressof(data_soa.contents) + # NOTE: This is probably broken + recarray = np.recarray( + shape=(size,), + dtype=np.dtype( + [("id", np.intc), ("logit", np.single), ("p", np.single)], + align=True, + ), + buf=(llama_cpp.llama_token_data * size).from_address( + data_soa_address + ), + ) + for logit_processor in logits_processor: + recarray.logit[:] = logit_processor(self._input_ids, recarray.logit) + + sampler.add_custom(apply_func) + + sampler.add_penalties( + # n_vocab=self._n_vocab, + # special_eos_id=self._token_eos, + # linefeed_id=self._token_nl, + penalty_last_n=self.last_n_tokens_size, + penalty_repeat=repeat_penalty, + penalty_freq=frequency_penalty, + penalty_present=presence_penalty, + # penalize_nl=penalize_nl, + # ignore_eos=False, ) - if not penalize_nl: - candidates.data[self._token_nl].logit = llama_cpp.c_float(nl_logit) - if temp.value == 0.0: - return llama_cpp.llama_sample_token_greedy( - ctx=self.ctx, - candidates=llama_cpp.ctypes.byref(candidates), # type: ignore - ) - elif mirostat_mode.value == 1: - mirostat_mu = llama_cpp.c_float(2.0 * mirostat_tau.value) - mirostat_m = llama_cpp.c_int(100) - llama_cpp.llama_sample_temperature( - ctx=self.ctx, - candidates=llama_cpp.ctypes.byref(candidates), # type: ignore - temp=temp, - ) - return llama_cpp.llama_sample_token_mirostat( - ctx=self.ctx, - candidates=llama_cpp.ctypes.byref(candidates), # type: ignore - tau=mirostat_tau, - eta=mirostat_eta, - mu=llama_cpp.ctypes.byref(mirostat_mu), # type: ignore - m=mirostat_m, - ) - elif mirostat_mode.value == 2: - mirostat_mu = llama_cpp.c_float(2.0 * mirostat_tau.value) - llama_cpp.llama_sample_temperature( - ctx=self.ctx, - candidates=llama_cpp.ctypes.byref(candidates), # type: ignore - temp=temp, - ) - return llama_cpp.llama_sample_token_mirostat_v2( - ctx=self.ctx, - candidates=llama_cpp.ctypes.byref(candidates), # type: ignore - tau=mirostat_tau, - eta=mirostat_eta, - mu=llama_cpp.ctypes.byref(mirostat_mu), # type: ignore - ) + + if grammar is not None: + sampler.add_grammar(self._model, grammar) + + if temp < 0.0: + sampler.add_dist(self._seed) + elif temp == 0.0: + sampler.add_greedy() else: - llama_cpp.llama_sample_top_k( - ctx=self.ctx, - candidates=llama_cpp.ctypes.byref(candidates), # type: ignore - k=top_k, - min_keep=llama_cpp.c_size_t(1), - ) - llama_cpp.llama_sample_tail_free( - ctx=self.ctx, - candidates=llama_cpp.ctypes.byref(candidates), # type: ignore - z=tfs_z, - min_keep=llama_cpp.c_size_t(1), - ) - llama_cpp.llama_sample_typical( - ctx=self.ctx, - candidates=llama_cpp.ctypes.byref(candidates), # type: ignore - p=llama_cpp.c_float(1.0), - min_keep=llama_cpp.c_size_t(1), - ) - llama_cpp.llama_sample_top_p( - ctx=self.ctx, - candidates=llama_cpp.ctypes.byref(candidates), # type: ignore - p=top_p, - min_keep=llama_cpp.c_size_t(1), - ) - llama_cpp.llama_sample_temperature( - ctx=self.ctx, - candidates=llama_cpp.ctypes.byref(candidates), # type: ignore - temp=temp, - ) - return llama_cpp.llama_sample_token( - ctx=self.ctx, - candidates=llama_cpp.ctypes.byref(candidates), # type: ignore - ) + if mirostat_mode == 1: + mirostat_m = 100 + sampler.add_mirostat( + self._n_vocab, + self._seed, + mirostat_tau, + mirostat_eta, + mirostat_m, + ) + elif mirostat_mode == 2: + sampler.add_mirostat_v2( + self._seed, + mirostat_tau, + mirostat_eta, + ) + else: + n_probs = 0 + min_keep = max(1, n_probs) + sampler.add_top_k(top_k) + sampler.add_typical(typical_p, min_keep) + sampler.add_top_p(top_p, min_keep) + sampler.add_min_p(min_p, min_keep) + sampler.add_temp(temp) + sampler.add_dist(self._seed) + return sampler def sample( self, top_k: int = 40, top_p: float = 0.95, + min_p: float = 0.05, + typical_p: float = 1.0, temp: float = 0.80, - repeat_penalty: float = 1.1, + repeat_penalty: float = 1.0, frequency_penalty: float = 0.0, presence_penalty: float = 0.0, tfs_z: float = 1.0, @@ -600,6 +794,8 @@ def sample( mirostat_tau: float = 5.0, penalize_nl: bool = True, logits_processor: Optional[LogitsProcessorList] = None, + grammar: Optional[LlamaGrammar] = None, + idx: Optional[int] = None, ): """Sample a token from the model. @@ -612,36 +808,47 @@ def sample( Returns: The sampled token. """ + assert self.n_tokens > 0 + + tmp_sampler = False + + if self._sampler is None: + tmp_sampler = True + self._sampler = self._init_sampler( + top_k=top_k, + top_p=top_p, + min_p=min_p, + typical_p=typical_p, + temp=temp, + repeat_penalty=repeat_penalty, + frequency_penalty=frequency_penalty, + presence_penalty=presence_penalty, + tfs_z=tfs_z, + mirostat_mode=mirostat_mode, + mirostat_tau=mirostat_tau, + mirostat_eta=mirostat_eta, + penalize_nl=penalize_nl, + logits_processor=logits_processor, + grammar=grammar, + ) + + ridx = idx - self.n_tokens if idx is not None else -1 + assert self.ctx is not None - last_n_tokens_data = [llama_cpp.llama_token(0)] * max( - 0, self.last_n_tokens_size - len(self._input_ids) - ) + self._input_ids[-self.last_n_tokens_size :].tolist() - return self._sample( - last_n_tokens_data=(llama_cpp.llama_token * self.last_n_tokens_size)( - *last_n_tokens_data - ), - last_n_tokens_size=llama_cpp.c_int(self.last_n_tokens_size), - top_k=llama_cpp.c_int(top_k), - top_p=llama_cpp.c_float(top_p), - temp=llama_cpp.c_float(temp), - tfs_z=llama_cpp.c_float(tfs_z), - repeat_penalty=llama_cpp.c_float(repeat_penalty), - frequency_penalty=llama_cpp.c_float(frequency_penalty), - presence_penalty=llama_cpp.c_float(presence_penalty), - mirostat_mode=llama_cpp.c_int(mirostat_mode), - mirostat_tau=llama_cpp.c_float(mirostat_tau), - mirostat_eta=llama_cpp.c_float(mirostat_eta), - penalize_nl=penalize_nl, - logits_processor=logits_processor, - ) + token = self._sampler.sample(self._ctx, ridx) + if tmp_sampler: + self._sampler = None + return token def generate( self, tokens: Sequence[int], top_k: int = 40, top_p: float = 0.95, + min_p: float = 0.05, + typical_p: float = 1.0, temp: float = 0.80, - repeat_penalty: float = 1.1, + repeat_penalty: float = 1.0, reset: bool = True, frequency_penalty: float = 0.0, presence_penalty: float = 0.0, @@ -649,15 +856,17 @@ def generate( mirostat_mode: int = 0, mirostat_tau: float = 5.0, mirostat_eta: float = 0.1, + penalize_nl: bool = True, logits_processor: Optional[LogitsProcessorList] = None, stopping_criteria: Optional[StoppingCriteriaList] = None, + grammar: Optional[LlamaGrammar] = None, ) -> Generator[int, Optional[Sequence[int]], None]: """Create a generator of tokens from a prompt. Examples: >>> llama = Llama("models/ggml-7b.bin") >>> tokens = llama.tokenize(b"Hello, world!") - >>> for token in llama.generate(tokens, top_k=40, top_p=0.95, temp=1.0, repeat_penalty=1.1): + >>> for token in llama.generate(tokens, top_k=40, top_p=0.95, temp=1.0, repeat_penalty=1.0): ... print(llama.detokenize([token])) Args: @@ -671,52 +880,130 @@ def generate( Yields: The generated tokens. """ - assert self.ctx is not None + # Reset mirostat sampling + self._mirostat_mu = ctypes.c_float(2.0 * mirostat_tau) + self._sampler = self._init_sampler( + top_k=top_k, + top_p=top_p, + min_p=min_p, + typical_p=typical_p, + temp=temp, + repeat_penalty=repeat_penalty, + frequency_penalty=frequency_penalty, + presence_penalty=presence_penalty, + tfs_z=tfs_z, + mirostat_mode=mirostat_mode, + mirostat_tau=mirostat_tau, + mirostat_eta=mirostat_eta, + penalize_nl=penalize_nl, + logits_processor=logits_processor, + grammar=grammar, + ) - if reset and len(self._input_ids) > 0: + # Check for kv cache prefix match + if reset and self.n_tokens > 0: longest_prefix = 0 for a, b in zip(self._input_ids, tokens[:-1]): if a == b: longest_prefix += 1 else: break - if longest_prefix > 0: + + # Recurrent and hybrid models cannot rewind state; reset if needed + if ( + self._is_recurrent or self._is_hybrid + ) and longest_prefix < self.n_tokens: + longest_prefix = 0 + reset = True if self.verbose: - print("Llama.generate: prefix-match hit", file=sys.stderr) - reset = False - tokens = tokens[longest_prefix:] - self.n_tokens = longest_prefix + print( + "Llama.generate: recurrent/hybrid model requires full state reset", + file=sys.stderr, + ) + if longest_prefix > 0: + if self._ctx.kv_cache_seq_rm(-1, longest_prefix, -1): + reset = False + tokens = tokens[longest_prefix:] + self.n_tokens = longest_prefix + if self.verbose: + print( + f"Llama.generate: {longest_prefix} prefix-match hit, " + f"remaining {len(tokens)} prompt tokens to eval", + file=sys.stderr, + ) + elif self.verbose: + print( + f"Llama.generate: {longest_prefix} prefix-match found " + f"but partial kv removal not supported, re-evaluating full prompt", + file=sys.stderr, + ) + + # Reset the model state if reset: self.reset() + # # Reset the grammar + # if grammar is not None: + # grammar.reset() + + sample_idx = self.n_tokens + len(tokens) - 1 + tokens = list(tokens) + + # Eval and sample while True: self.eval(tokens) - token = self.sample( - top_k=top_k, - top_p=top_p, - temp=temp, - repeat_penalty=repeat_penalty, - frequency_penalty=frequency_penalty, - presence_penalty=presence_penalty, - tfs_z=tfs_z, - mirostat_mode=mirostat_mode, - mirostat_tau=mirostat_tau, - mirostat_eta=mirostat_eta, - logits_processor=logits_processor, - ) - if stopping_criteria is not None and stopping_criteria( - self._input_ids.tolist(), self._scores[-1, :].tolist() - ): - return - tokens_or_none = yield token - tokens = [token] - if tokens_or_none is not None: - tokens.extend(tokens_or_none) + while sample_idx < self.n_tokens: + token = self.sample( + top_k=top_k, + top_p=top_p, + min_p=min_p, + typical_p=typical_p, + temp=temp, + repeat_penalty=repeat_penalty, + frequency_penalty=frequency_penalty, + presence_penalty=presence_penalty, + tfs_z=tfs_z, + mirostat_mode=mirostat_mode, + mirostat_tau=mirostat_tau, + mirostat_eta=mirostat_eta, + logits_processor=logits_processor, + grammar=grammar, + penalize_nl=penalize_nl, + idx=sample_idx, + ) + + sample_idx += 1 + if stopping_criteria is not None and stopping_criteria( + self._input_ids[:sample_idx], + self._scores[sample_idx - self.n_tokens, :], + ): + return + tokens_or_none = yield token + tokens.clear() + tokens.append(token) + if tokens_or_none is not None: + tokens.extend(tokens_or_none) + + if sample_idx < self.n_tokens and token != self._input_ids[sample_idx]: + self.n_tokens = sample_idx + self._ctx.kv_cache_seq_rm(-1, self.n_tokens, -1) + break + + if self.draft_model is not None: + self.input_ids[self.n_tokens : self.n_tokens + len(tokens)] = tokens + draft_tokens = self.draft_model( + self.input_ids[: self.n_tokens + len(tokens)] + ) + tokens.extend( + draft_tokens.astype(int)[ + : self._n_ctx - self.n_tokens - len(tokens) + ] + ) def create_embedding( self, input: Union[str, List[str]], model: Optional[str] = None - ) -> Embedding: + ) -> CreateEmbeddingResponse: """Embed a string. Args: @@ -725,80 +1012,181 @@ def create_embedding( Returns: An embedding object. """ - assert self.ctx is not None model_name: str = model if model is not None else self.model_path - if self.params.embedding == False: + input = input if isinstance(input, list) else [input] + + # get numeric embeddings + embeds: Union[List[List[float]], List[List[List[float]]]] + total_tokens: int + embeds, total_tokens = self.embed(input, return_count=True) # type: ignore + + # convert to CreateEmbeddingResponse + data: List[Embedding] = [ + { + "object": "embedding", + "embedding": emb, + "index": idx, + } + for idx, emb in enumerate(embeds) + ] + + return { + "object": "list", + "data": data, + "model": model_name, + "usage": { + "prompt_tokens": total_tokens, + "total_tokens": total_tokens, + }, + } + + def embed( + self, + input: Union[str, List[str]], + normalize: bool = False, + truncate: bool = True, + return_count: bool = False, + ): + """Embed a string. + + Args: + input: The utf-8 encoded string to embed. + + Returns: + A list of embeddings + """ + n_embd = self.n_embd() + n_batch = self.n_batch + n_seq_max = self.context_params.n_seq_max + + # get pooling information + pooling_type = self.pooling_type() + # In embedding mode every input token must be marked as an output, regardless of + # pooling type. llama.cpp would otherwise override per-token `logits[i]` and emit + # "embeddings required but some input tokens were not marked as outputs -> + # overriding" once per input. Pooling NONE vs MEAN/CLS only changes how the + # per-token outputs are read back (see decode_batch below), not whether they are + # produced. See abetlen/llama-cpp-python#2208. + logits_all = True + + if self.context_params.embeddings is False: raise RuntimeError( "Llama model must be created with embedding=True to call this method" ) if self.verbose: - llama_cpp.llama_reset_timings(self.ctx) + llama_cpp.llama_perf_context_reset(self._ctx.ctx) if isinstance(input, str): inputs = [input] else: inputs = input - data: List[EmbeddingData] = [] + # reset batch + self._batch.reset() + + # decode and fetch embeddings + data: Union[List[List[float]], List[List[List[float]]]] = [] + + def decode_batch(seq_sizes: List[int]): + self._ctx.kv_cache_clear() + self._ctx.decode(self._batch) + self._batch.reset() + + # store embeddings + if pooling_type == llama_cpp.LLAMA_POOLING_TYPE_NONE: + pos: int = 0 + for i, size in enumerate(seq_sizes): + ptr = llama_cpp.llama_get_embeddings(self._ctx.ctx) + embedding: List[List[float]] = [ + ptr[pos + j * n_embd : pos + (j + 1) * n_embd] + for j in range(size) + ] + if normalize: + embedding = [ + internals.normalize_embedding(e) for e in embedding + ] + data.append(embedding) + pos += size + else: + for i in range(len(seq_sizes)): + ptr = llama_cpp.llama_get_embeddings_seq(self._ctx.ctx, i) + embedding: List[float] = ptr[:n_embd] + if normalize: + embedding = internals.normalize_embedding(embedding) + data.append(embedding) + + # init state total_tokens = 0 - for index, input in enumerate(inputs): - tokens = self.tokenize(input.encode("utf-8")) - self.reset() - self.eval(tokens) + s_batch = [] + t_batch = 0 + p_batch = 0 + + # accumulate batches and encode + for text in inputs: + tokens = self.tokenize(text.encode("utf-8")) + if truncate: + tokens = tokens[:n_batch] + n_tokens = len(tokens) total_tokens += n_tokens - embedding = llama_cpp.llama_get_embeddings(self.ctx)[ - : llama_cpp.llama_n_embd(self.ctx) - ] - data.append( - { - "object": "embedding", - "embedding": embedding, - "index": index, - } - ) - if self.verbose: - llama_cpp.llama_print_timings(self.ctx) + # check for overrun + if n_tokens > n_batch: + raise ValueError( + f"Requested tokens ({n_tokens}) exceed batch size of {n_batch}" + ) + + # time to eval batch + if t_batch + n_tokens > n_batch or p_batch >= n_seq_max: + decode_batch(s_batch) + s_batch = [] + t_batch = 0 + p_batch = 0 - return { - "object": "list", - "data": data, - "model": model_name, - "usage": { - "prompt_tokens": total_tokens, - "total_tokens": total_tokens, - }, - } + # add to batch + self._batch.add_sequence(tokens, p_batch, logits_all) - def embed(self, input: str) -> List[float]: - """Embed a string. + # update batch stats + s_batch.append(n_tokens) + t_batch += n_tokens + p_batch += 1 - Args: - input: The utf-8 encoded string to embed. + # hanlde last batch + decode_batch(s_batch) - Returns: - A list of embeddings - """ - return list(map(float, self.create_embedding(input)["data"][0]["embedding"])) + if self.verbose: + llama_cpp.llama_perf_context_print(self._ctx.ctx) + + output = data[0] if isinstance(input, str) else data + + self._ctx.kv_cache_clear() + self.reset() + + if return_count: + return output, total_tokens + else: + return output def _create_completion( self, - prompt: str, + prompt: Union[str, List[int]], suffix: Optional[str] = None, - max_tokens: int = 16, + max_tokens: Optional[int] = 16, temperature: float = 0.8, top_p: float = 0.95, + min_p: float = 0.05, + typical_p: float = 1.0, logprobs: Optional[int] = None, echo: bool = False, stop: Optional[Union[str, List[str]]] = [], frequency_penalty: float = 0.0, presence_penalty: float = 0.0, - repeat_penalty: float = 1.1, + repeat_penalty: float = 1.0, top_k: int = 40, stream: bool = False, + seed: Optional[int] = None, tfs_z: float = 1.0, mirostat_mode: int = 0, mirostat_tau: float = 5.0, @@ -806,14 +1194,92 @@ def _create_completion( model: Optional[str] = None, stopping_criteria: Optional[StoppingCriteriaList] = None, logits_processor: Optional[LogitsProcessorList] = None, - ) -> Union[Iterator[Completion], Iterator[CompletionChunk]]: - assert self.ctx is not None + grammar: Optional[LlamaGrammar] = None, + logit_bias: Optional[Dict[int, float]] = None, + ) -> Union[ + Iterator[CreateCompletionResponse], Iterator[CreateCompletionStreamResponse] + ]: + assert suffix is None or suffix.__class__ is str completion_id: str = f"cmpl-{str(uuid.uuid4())}" created: int = int(time.time()) - completion_tokens: List[int] = [] + bos_token_id: int = self.token_bos() + cls_token_id: int = self._model.token_cls() + sep_token_id: int = self._model.token_sep() + prefix_token_id: int = 0 # self._model.token_prefix() # TODO: Fix + middle_token_id: int = 0 # self._model.token_middle() # TODO: Fix + suffix_token_id: int = 0 # self._model.token_suffix() # TODO: Fix + add_space_prefix: bool = ( + self.metadata.get("tokenizer.ggml.add_space_prefix", "true") == "true" + ) + bos_tokens: List[int] = [cls_token_id if cls_token_id != -1 else bos_token_id] + eos_tokens: List[int] = [ + sep_token_id if sep_token_id != -1 else self.token_eos() + ] + + if ( + (isinstance(prompt, list) and suffix is None) + or not self._model.add_bos_token() + or bos_tokens[:1] == [-1] + ): + bos_tokens = [] + + if (isinstance(prompt, list) and suffix is None) or ( + not self._model.add_eos_token() and sep_token_id == -1 + ): + eos_tokens = [] + + suffix_space_prefix: int = 0 + # Tokenizer hack to remove leading space + if add_space_prefix and suffix_token_id >= 0 and suffix: + suffix = "☺" + suffix + suffix_space_prefix = 2 + + # If prompt is empty, initialize completion with BOS token to avoid + # detokenization including a space at the beginning of the completion + completion_tokens: List[int] = [] if len(prompt) > 0 else [bos_token_id] # Add blank space to start of prompt to match OG llama tokenizer - prompt_tokens: List[int] = self.tokenize(b" " + prompt.encode("utf-8")) + prefix_tokens: List[int] = ( + [prefix_token_id] if prefix_token_id >= 0 and suffix is not None else [] + ) + ( + ( + self.tokenize( + prompt.encode("utf-8"), + add_bos=False, + special=(prefix_token_id < 0 or suffix is None), + ) + if prompt != "" + else [] + ) + if isinstance(prompt, str) + else prompt + ) + suffix_tokens: List[int] = ( + ( + [suffix_token_id] + + ( + self.tokenize(suffix.encode("utf-8"), add_bos=False, special=False)[ + suffix_space_prefix: + ] + if suffix + else [] + ) + ) + if suffix_token_id >= 0 and suffix is not None + else [] + ) + middle_tokens: List[int] = ( + [middle_token_id] if middle_token_id >= 0 and suffix is not None else [] + ) + prompt_tokens: List[int] = ( + bos_tokens + + ( + (suffix_tokens + prefix_tokens + middle_tokens) + if self.spm_infill + else (prefix_tokens + suffix_tokens + middle_tokens) + ) + + eos_tokens + ) text: bytes = b"" returned_tokens: int = 0 stop = ( @@ -821,17 +1287,45 @@ def _create_completion( ) model_name: str = model if model is not None else self.model_path + if prompt_tokens[:2] == [self.token_bos()] * 2: + warnings.warn( + f'Detected duplicate leading "{self._model.token_get_text(self.token_bos())}" in prompt, this will likely reduce response quality, consider removing it...', + RuntimeWarning, + ) + + # NOTE: This likely doesn't work correctly for the first token in the prompt + # because of the extra space added to the start of the prompt_tokens + if logit_bias is not None: + logit_bias_map = {int(k): float(v) for k, v in logit_bias.items()} + + def logit_bias_processor( + input_ids: npt.NDArray[np.intc], + scores: npt.NDArray[np.single], + ) -> npt.NDArray[np.single]: + new_scores = np.copy( + scores + ) # Does it make sense to copy the whole array or can we just overwrite the original one? + for input_id, score in logit_bias_map.items(): + new_scores[input_id] = score + scores[input_id] + return new_scores + + _logit_bias_processor = LogitsProcessorList([logit_bias_processor]) + if logits_processor is None: + logits_processor = _logit_bias_processor + else: + logits_processor = logits_processor.extend(_logit_bias_processor) + if self.verbose: - llama_cpp.llama_reset_timings(self.ctx) + self._ctx.reset_timings() - if len(prompt_tokens) >= llama_cpp.llama_n_ctx(self.ctx): + if len(prompt_tokens) >= self._n_ctx: raise ValueError( - f"Requested tokens exceed context window of {llama_cpp.llama_n_ctx(self.ctx)}" + f"Requested tokens ({len(prompt_tokens)}) exceed context window of {llama_cpp.llama_n_ctx(self.ctx)}" ) - if max_tokens <= 0: + if max_tokens is None or max_tokens <= 0: # Unlimited, depending on n_ctx. - max_tokens = llama_cpp.llama_n_ctx(self.ctx) - len(prompt_tokens) + max_tokens = self._n_ctx - len(prompt_tokens) # Truncate max_tokens if requested tokens would exceed the context window max_tokens = ( @@ -845,7 +1339,7 @@ def _create_completion( else: stop_sequences = [] - if logprobs is not None and self.params.logits_all is False: + if logprobs is not None and self._logits_all is False: raise ValueError( "logprobs is not supported for models created with logits_all=False" ) @@ -867,12 +1361,19 @@ def _create_completion( if self.verbose: print("Llama._create_completion: cache miss", file=sys.stderr) + if seed is not None: + self.set_seed(seed) + else: + self.set_seed(random.Random(self._seed).randint(0, 2**32)) + finish_reason = "length" multibyte_fix = 0 for token in self.generate( prompt_tokens, top_k=top_k, top_p=top_p, + min_p=min_p, + typical_p=typical_p, temp=temperature, tfs_z=tfs_z, mirostat_mode=mirostat_mode, @@ -883,15 +1384,16 @@ def _create_completion( repeat_penalty=repeat_penalty, stopping_criteria=stopping_criteria, logits_processor=logits_processor, + grammar=grammar, ): - if token == self._token_eos: - text = self.detokenize(completion_tokens) + if llama_cpp.llama_vocab_is_eog(self._model.vocab, token): + text = self.detokenize(completion_tokens, prev_tokens=prompt_tokens) finish_reason = "stop" break completion_tokens.append(token) - all_text = self.detokenize(completion_tokens) + all_text = self.detokenize(completion_tokens, prev_tokens=prompt_tokens) # Contains multi-byte UTF8 for k, char in enumerate(all_text[-3:]): @@ -915,7 +1417,10 @@ def _create_completion( if stream: remaining_tokens = completion_tokens[returned_tokens:] - remaining_text = self.detokenize(remaining_tokens) + remaining_text = self.detokenize( + remaining_tokens, + prev_tokens=prompt_tokens + completion_tokens[:returned_tokens], + ) remaining_length = len(remaining_text) # We want to avoid yielding any characters from @@ -930,24 +1435,40 @@ def _create_completion( break token_end_position = 0 - for token in remaining_tokens: - token_end_position += len(self.detokenize([token])) - # Check if stop sequence is in the token - if token_end_position >= ( - remaining_length - first_stop_position - 1 - ): - break - logprobs_or_none: Optional[CompletionLogprobs] = None - if logprobs is not None: - token_str = self.detokenize([token]).decode( - "utf-8", errors="ignore" + + if logprobs is not None: + # not sure how to handle this branch when dealing + # with CJK output, so keep it unchanged + for token in remaining_tokens: + if token == bos_token_id: + continue + token_end_position += len( + self.detokenize( + [token], + prev_tokens=prompt_tokens + + completion_tokens[:returned_tokens], + ) ) + # Check if stop sequence is in the token + if token_end_position > ( + remaining_length - first_stop_position + ): + break + token_str = self.detokenize( + [token], + prev_tokens=prompt_tokens + + completion_tokens[:returned_tokens], + ).decode("utf-8", errors="ignore") text_offset = len(prompt) + len( - self.detokenize(completion_tokens[:returned_tokens]) + self.detokenize( + completion_tokens[:returned_tokens], + prev_tokens=prompt_tokens + + completion_tokens[:returned_tokens], + ).decode("utf-8", errors="ignore") ) token_offset = len(prompt_tokens) + returned_tokens - logits = self._scores[token_offset - 1, :].tolist() - current_logprobs = Llama.logits_to_logprobs(logits) + logits = self._scores[token_offset - 1, :] + current_logprobs = Llama.logits_to_logprobs(logits).tolist() sorted_logprobs = list( sorted( zip(current_logprobs, range(len(current_logprobs))), @@ -963,70 +1484,130 @@ def _create_completion( top_logprob.update({token_str: current_logprobs[int(token)]}) logprobs_or_none = { "tokens": [ - self.detokenize([token]).decode( - "utf-8", errors="ignore" - ) + self.detokenize( + [token], + prev_tokens=prompt_tokens + + completion_tokens[:returned_tokens], + ).decode("utf-8", errors="ignore") ], "text_offset": [text_offset], "token_logprobs": [current_logprobs[int(token)]], "top_logprobs": [top_logprob], } - returned_tokens += 1 - yield { - "id": completion_id, - "object": "text_completion", - "created": created, - "model": model_name, - "choices": [ - { - "text": self.detokenize([token]).decode( - "utf-8", errors="ignore" - ), - "index": 0, - "logprobs": logprobs_or_none, - "finish_reason": None, - } - ], - } + returned_tokens += 1 + yield { + "id": completion_id, + "object": "text_completion", + "created": created, + "model": model_name, + "choices": [ + { + "text": self.detokenize( + [token], + prev_tokens=prompt_tokens + + completion_tokens[:returned_tokens], + ).decode("utf-8", errors="ignore"), + "index": 0, + "logprobs": logprobs_or_none, + "finish_reason": None, + } + ], + } + else: + while len(remaining_tokens) > 0: + decode_success = False + for i in range(1, len(remaining_tokens) + 1): + try: + bs = self.detokenize( + remaining_tokens[:i], + prev_tokens=prompt_tokens + + completion_tokens[:returned_tokens], + ) + ts = bs.decode("utf-8") + decode_success = True + break + except UnicodeError: + pass + else: + break + if not decode_success: + # all remaining tokens cannot be decoded to a UTF-8 character + break + token_end_position += len(bs) + if token_end_position > ( + remaining_length - first_stop_position + ): + break + remaining_tokens = remaining_tokens[i:] + returned_tokens += i + + yield { + "id": completion_id, + "object": "text_completion", + "created": created, + "model": model_name, + "choices": [ + { + "text": ts, + "index": 0, + "logprobs": None, + "finish_reason": None, + } + ], + } if len(completion_tokens) >= max_tokens: - text = self.detokenize(completion_tokens) + text = self.detokenize(completion_tokens, prev_tokens=prompt_tokens) finish_reason = "length" break if stopping_criteria is not None and stopping_criteria( - self._input_ids.tolist(), self._scores[-1, :].tolist() + self._input_ids, self._scores[-1, :] ): - text = self.detokenize(completion_tokens) + text = self.detokenize(completion_tokens, prev_tokens=prompt_tokens) finish_reason = "stop" if self.verbose: - llama_cpp.llama_print_timings(self.ctx) + self._ctx.print_timings() if stream: remaining_tokens = completion_tokens[returned_tokens:] - all_text = self.detokenize(remaining_tokens) - any_stop = [s for s in stop_sequences if s in all_text] + remaining_text = self.detokenize( + remaining_tokens, + prev_tokens=prompt_tokens + completion_tokens[:returned_tokens], + ) + any_stop = [s for s in stop_sequences if s in remaining_text] if len(any_stop) > 0: - end = min(all_text.index(stop) for stop in any_stop) + end = min(remaining_text.index(stop) for stop in any_stop) else: - end = len(all_text) + end = len(remaining_text) token_end_position = 0 for token in remaining_tokens: - token_end_position += len(self.detokenize([token])) + token_end_position += len( + self.detokenize( + [token], + prev_tokens=prompt_tokens + completion_tokens[:returned_tokens], + ) + ) logprobs_or_none: Optional[CompletionLogprobs] = None if logprobs is not None: + if token == bos_token_id: + continue token_str = self.detokenize([token]).decode( "utf-8", errors="ignore" ) text_offset = len(prompt) + len( - self.detokenize(completion_tokens[:returned_tokens]) + self.detokenize( + completion_tokens[:returned_tokens], + prev_tokens=prompt_tokens + + completion_tokens[:returned_tokens], + ) ) token_offset = len(prompt_tokens) + returned_tokens - 1 - logits = self._scores[token_offset, :].tolist() - current_logprobs = Llama.logits_to_logprobs(logits) + logits = self._scores[token_offset, :] + current_logprobs = Llama.logits_to_logprobs(logits).tolist() sorted_logprobs = list( sorted( zip(current_logprobs, range(len(current_logprobs))), @@ -1068,20 +1649,6 @@ def _create_completion( } ], } - yield { - "id": completion_id, - "object": "text_completion", - "created": created, - "model": model_name, - "choices": [ - { - "text": "", - "index": 0, - "logprobs": None, - "finish_reason": finish_reason, - } - ], - } break returned_tokens += 1 yield { @@ -1100,25 +1667,26 @@ def _create_completion( } ], } - yield { - "id": completion_id, - "object": "text_completion", - "created": created, - "model": model_name, - "choices": [ - { - "text": "", - "index": 0, - "logprobs": None, - "finish_reason": finish_reason, - } - ], - } + yield { + "id": completion_id, + "object": "text_completion", + "created": created, + "model": model_name, + "choices": [ + { + "text": "", + "index": 0, + "logprobs": None, + "finish_reason": finish_reason, + } + ], + } if self.cache: if self.verbose: print("Llama._create_completion: cache save", file=sys.stderr) self.cache[prompt_tokens + completion_tokens] = self.save_state() - print("Llama._create_completion: cache saved", file=sys.stderr) + if self.verbose: + print("Llama._create_completion: cache saved", file=sys.stderr) return if self.cache: @@ -1131,7 +1699,7 @@ def _create_completion( if echo: text_str = prompt + text_str - if suffix is not None: + if suffix_token_id < 0 and suffix is not None: text_str = text_str + suffix logprobs_or_none: Optional[CompletionLogprobs] = None @@ -1144,23 +1712,35 @@ def _create_completion( top_logprobs: List[Optional[Dict[str, float]]] = [] if echo: - # Remove leading BOS token - all_tokens = prompt_tokens[1:] + completion_tokens + # Remove leading BOS token if exists + all_tokens = ( + prompt_tokens[1 if prompt_tokens[0] == self.token_bos() else 0 :] + + completion_tokens + ) else: all_tokens = completion_tokens all_token_strs = [ - self.detokenize([token]).decode("utf-8", errors="ignore") - for token in all_tokens + self.detokenize([token], prev_tokens=all_tokens[:i]).decode( + "utf-8", errors="ignore" + ) + for i, token in enumerate(all_tokens) ] - all_logprobs = [ - Llama.logits_to_logprobs(row.tolist()) for row in self._scores - ][token_offset:] - for token, token_str, logprobs_token in zip( - all_tokens, all_token_strs, all_logprobs + all_logprobs = Llama.logits_to_logprobs(self._scores)[token_offset:] + # TODO: may be able to change this loop to use np.take_along_dim + for idx, (token, token_str, logprobs_token) in enumerate( + zip(all_tokens, all_token_strs, all_logprobs) ): - text_offsets.append(text_offset) - text_offset += len(token_str) + if token == bos_token_id: + continue + text_offsets.append( + text_offset + + len( + self.detokenize(all_tokens[:idx]).decode( + "utf-8", errors="ignore" + ) + ) + ) tokens.append(token_str) sorted_logprobs = list( sorted( @@ -1169,7 +1749,9 @@ def _create_completion( ) token_logprobs.append(logprobs_token[int(token)]) top_logprob: Optional[Dict[str, float]] = { - self.detokenize([i]).decode("utf-8", errors="ignore"): logprob + self.detokenize([i], prev_tokens=all_tokens[:idx]).decode( + "utf-8", errors="ignore" + ): logprob for logprob, i in sorted_logprobs[:logprobs] } top_logprob.update({token_str: logprobs_token[int(token)]}) @@ -1209,19 +1791,22 @@ def _create_completion( def create_completion( self, - prompt: str, + prompt: Union[str, List[int]], suffix: Optional[str] = None, - max_tokens: int = 128, + max_tokens: Optional[int] = 16, temperature: float = 0.8, top_p: float = 0.95, + min_p: float = 0.05, + typical_p: float = 1.0, logprobs: Optional[int] = None, echo: bool = False, stop: Optional[Union[str, List[str]]] = [], frequency_penalty: float = 0.0, presence_penalty: float = 0.0, - repeat_penalty: float = 1.1, + repeat_penalty: float = 1.0, top_k: int = 40, stream: bool = False, + seed: Optional[int] = None, tfs_z: float = 1.0, mirostat_mode: int = 0, mirostat_tau: float = 5.0, @@ -1229,21 +1814,37 @@ def create_completion( model: Optional[str] = None, stopping_criteria: Optional[StoppingCriteriaList] = None, logits_processor: Optional[LogitsProcessorList] = None, - ) -> Union[Completion, Iterator[CompletionChunk]]: + grammar: Optional[LlamaGrammar] = None, + logit_bias: Optional[Dict[int, float]] = None, + ) -> Union[CreateCompletionResponse, Iterator[CreateCompletionStreamResponse]]: """Generate text from a prompt. Args: prompt: The prompt to generate text from. suffix: A suffix to append to the generated text. If None, no suffix is appended. - max_tokens: The maximum number of tokens to generate. If max_tokens <= 0, the maximum number of tokens to generate is unlimited and depends on n_ctx. + max_tokens: The maximum number of tokens to generate. If max_tokens <= 0 or None, the maximum number of tokens to generate is unlimited and depends on n_ctx. temperature: The temperature to use for sampling. - top_p: The top-p value to use for sampling. + top_p: The top-p value to use for nucleus sampling. Nucleus sampling described in academic paper "The Curious Case of Neural Text Degeneration" https://arxiv.org/abs/1904.09751 + min_p: The min-p value to use for minimum p sampling. Minimum P sampling as described in https://github.com/ggerganov/llama.cpp/pull/3841 + typical_p: The typical-p value to use for sampling. Locally Typical Sampling implementation described in the paper https://arxiv.org/abs/2202.00666. logprobs: The number of logprobs to return. If None, no logprobs are returned. echo: Whether to echo the prompt. stop: A list of strings to stop generation when encountered. + frequency_penalty: The penalty to apply to tokens based on their frequency in the prompt. + presence_penalty: The penalty to apply to tokens based on their presence in the prompt. repeat_penalty: The penalty to apply to repeated tokens. - top_k: The top-k value to use for sampling. + top_k: The top-k value to use for sampling. Top-K sampling described in academic paper "The Curious Case of Neural Text Degeneration" https://arxiv.org/abs/1904.09751 stream: Whether to stream the results. + seed: The seed to use for sampling. + tfs_z: The tail-free sampling parameter. Tail Free Sampling described in https://www.trentonbricken.com/Tail-Free-Sampling/. + mirostat_mode: The mirostat sampling mode. + mirostat_tau: The target cross-entropy (or surprise) value you want to achieve for the generated text. A higher value corresponds to more surprising or less predictable text, while a lower value corresponds to less surprising or more predictable text. + mirostat_eta: The learning rate used to update `mu` based on the error between the target and observed surprisal of the sampled word. A larger learning rate will cause `mu` to be updated more quickly, while a smaller learning rate will result in slower updates. + model: The name to use for the model in the completion object. + stopping_criteria: A list of stopping criteria to use. + logits_processor: A list of logits processors to use. + grammar: A grammar to use for constrained sampling. + logit_bias: A logit bias to use. Raises: ValueError: If the requested tokens exceed the context window. @@ -1255,9 +1856,11 @@ def create_completion( completion_or_chunks = self._create_completion( prompt=prompt, suffix=suffix, - max_tokens=max_tokens, + max_tokens=-1 if max_tokens is None else max_tokens, temperature=temperature, top_p=top_p, + min_p=min_p, + typical_p=typical_p, logprobs=logprobs, echo=echo, stop=stop, @@ -1266,6 +1869,7 @@ def create_completion( repeat_penalty=repeat_penalty, top_k=top_k, stream=stream, + seed=seed, tfs_z=tfs_z, mirostat_mode=mirostat_mode, mirostat_tau=mirostat_tau, @@ -1273,9 +1877,11 @@ def create_completion( model=model, stopping_criteria=stopping_criteria, logits_processor=logits_processor, + grammar=grammar, + logit_bias=logit_bias, ) if stream: - chunks: Iterator[CompletionChunk] = completion_or_chunks + chunks: Iterator[CreateCompletionStreamResponse] = completion_or_chunks return chunks completion: Completion = next(completion_or_chunks) # type: ignore return completion @@ -1284,17 +1890,20 @@ def __call__( self, prompt: str, suffix: Optional[str] = None, - max_tokens: int = 128, + max_tokens: Optional[int] = 16, temperature: float = 0.8, top_p: float = 0.95, + min_p: float = 0.05, + typical_p: float = 1.0, logprobs: Optional[int] = None, echo: bool = False, stop: Optional[Union[str, List[str]]] = [], frequency_penalty: float = 0.0, presence_penalty: float = 0.0, - repeat_penalty: float = 1.1, + repeat_penalty: float = 1.0, top_k: int = 40, stream: bool = False, + seed: Optional[int] = None, tfs_z: float = 1.0, mirostat_mode: int = 0, mirostat_tau: float = 5.0, @@ -1302,21 +1911,37 @@ def __call__( model: Optional[str] = None, stopping_criteria: Optional[StoppingCriteriaList] = None, logits_processor: Optional[LogitsProcessorList] = None, - ) -> Union[Completion, Iterator[CompletionChunk]]: + grammar: Optional[LlamaGrammar] = None, + logit_bias: Optional[Dict[int, float]] = None, + ) -> Union[CreateCompletionResponse, Iterator[CreateCompletionStreamResponse]]: """Generate text from a prompt. Args: prompt: The prompt to generate text from. suffix: A suffix to append to the generated text. If None, no suffix is appended. - max_tokens: The maximum number of tokens to generate. If max_tokens <= 0, the maximum number of tokens to generate is unlimited and depends on n_ctx. + max_tokens: The maximum number of tokens to generate. If max_tokens <= 0 or None, the maximum number of tokens to generate is unlimited and depends on n_ctx. temperature: The temperature to use for sampling. - top_p: The top-p value to use for sampling. + top_p: The top-p value to use for nucleus sampling. Nucleus sampling described in academic paper "The Curious Case of Neural Text Degeneration" https://arxiv.org/abs/1904.09751 + min_p: The min-p value to use for minimum p sampling. Minimum P sampling as described in https://github.com/ggerganov/llama.cpp/pull/3841 + typical_p: The typical-p value to use for sampling. Locally Typical Sampling implementation described in the paper https://arxiv.org/abs/2202.00666. logprobs: The number of logprobs to return. If None, no logprobs are returned. echo: Whether to echo the prompt. stop: A list of strings to stop generation when encountered. + frequency_penalty: The penalty to apply to tokens based on their frequency in the prompt. + presence_penalty: The penalty to apply to tokens based on their presence in the prompt. repeat_penalty: The penalty to apply to repeated tokens. - top_k: The top-k value to use for sampling. + top_k: The top-k value to use for sampling. Top-K sampling described in academic paper "The Curious Case of Neural Text Degeneration" https://arxiv.org/abs/1904.09751 stream: Whether to stream the results. + seed: The seed to use for sampling. + tfs_z: The tail-free sampling parameter. Tail Free Sampling described in https://www.trentonbricken.com/Tail-Free-Sampling/. + mirostat_mode: The mirostat sampling mode. + mirostat_tau: The target cross-entropy (or surprise) value you want to achieve for the generated text. A higher value corresponds to more surprising or less predictable text, while a lower value corresponds to less surprising or more predictable text. + mirostat_eta: The learning rate used to update `mu` based on the error between the target and observed surprisal of the sampled word. A larger learning rate will cause `mu` to be updated more quickly, while a smaller learning rate will result in slower updates. + model: The name to use for the model in the completion object. + stopping_criteria: A list of stopping criteria to use. + logits_processor: A list of logits processors to use. + grammar: A grammar to use for constrained sampling. + logit_bias: A logit bias to use. Raises: ValueError: If the requested tokens exceed the context window. @@ -1331,6 +1956,8 @@ def __call__( max_tokens=max_tokens, temperature=temperature, top_p=top_p, + min_p=min_p, + typical_p=typical_p, logprobs=logprobs, echo=echo, stop=stop, @@ -1339,6 +1966,7 @@ def __call__( repeat_penalty=repeat_penalty, top_k=top_k, stream=stream, + seed=seed, tfs_z=tfs_z, mirostat_mode=mirostat_mode, mirostat_tau=mirostat_tau, @@ -1346,206 +1974,227 @@ def __call__( model=model, stopping_criteria=stopping_criteria, logits_processor=logits_processor, + grammar=grammar, + logit_bias=logit_bias, ) - def _convert_text_completion_to_chat( - self, completion: Completion - ) -> ChatCompletion: - return { - "id": "chat" + completion["id"], - "object": "chat.completion", - "created": completion["created"], - "model": completion["model"], - "choices": [ - { - "index": 0, - "message": { - "role": "assistant", - "content": completion["choices"][0]["text"], - }, - "finish_reason": completion["choices"][0]["finish_reason"], - } - ], - "usage": completion["usage"], - } - - def _convert_text_completion_chunks_to_chat( - self, - chunks: Iterator[CompletionChunk], - ) -> Iterator[ChatCompletionChunk]: - for i, chunk in enumerate(chunks): - if i == 0: - yield { - "id": "chat" + chunk["id"], - "model": chunk["model"], - "created": chunk["created"], - "object": "chat.completion.chunk", - "choices": [ - { - "index": 0, - "delta": { - "role": "assistant", - }, - "finish_reason": None, - } - ], - } - yield { - "id": "chat" + chunk["id"], - "model": chunk["model"], - "created": chunk["created"], - "object": "chat.completion.chunk", - "choices": [ - { - "index": 0, - "delta": { - "content": chunk["choices"][0]["text"], - } - if chunk["choices"][0]["finish_reason"] is None - else {}, - "finish_reason": chunk["choices"][0]["finish_reason"], - } - ], - } - def create_chat_completion( self, - messages: List[ChatCompletionMessage], + messages: List[ChatCompletionRequestMessage], + functions: Optional[List[ChatCompletionFunction]] = None, + function_call: Optional[ChatCompletionRequestFunctionCall] = None, + tools: Optional[List[ChatCompletionTool]] = None, + tool_choice: Optional[ChatCompletionToolChoiceOption] = None, temperature: float = 0.2, top_p: float = 0.95, top_k: int = 40, + min_p: float = 0.05, + typical_p: float = 1.0, stream: bool = False, stop: Optional[Union[str, List[str]]] = [], - max_tokens: int = 256, + seed: Optional[int] = None, + response_format: Optional[ChatCompletionRequestResponseFormat] = None, + max_tokens: Optional[int] = None, presence_penalty: float = 0.0, frequency_penalty: float = 0.0, - repeat_penalty: float = 1.1, + repeat_penalty: float = 1.0, tfs_z: float = 1.0, mirostat_mode: int = 0, mirostat_tau: float = 5.0, mirostat_eta: float = 0.1, model: Optional[str] = None, logits_processor: Optional[LogitsProcessorList] = None, - ) -> Union[ChatCompletion, Iterator[ChatCompletionChunk]]: + grammar: Optional[LlamaGrammar] = None, + logit_bias: Optional[Dict[int, float]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + ) -> Union[ + CreateChatCompletionResponse, Iterator[CreateChatCompletionStreamResponse] + ]: """Generate a chat completion from a list of messages. Args: messages: A list of messages to generate a response for. + functions: A list of functions to use for the chat completion. + function_call: A function call to use for the chat completion. + tools: A list of tools to use for the chat completion. + tool_choice: A tool choice to use for the chat completion. temperature: The temperature to use for sampling. - top_p: The top-p value to use for sampling. - top_k: The top-k value to use for sampling. + top_p: The top-p value to use for nucleus sampling. Nucleus sampling described in academic paper "The Curious Case of Neural Text Degeneration" https://arxiv.org/abs/1904.09751 + top_k: The top-k value to use for sampling. Top-K sampling described in academic paper "The Curious Case of Neural Text Degeneration" https://arxiv.org/abs/1904.09751 + min_p: The min-p value to use for minimum p sampling. Minimum P sampling as described in https://github.com/ggerganov/llama.cpp/pull/3841 + typical_p: The typical-p value to use for sampling. Locally Typical Sampling implementation described in the paper https://arxiv.org/abs/2202.00666. stream: Whether to stream the results. stop: A list of strings to stop generation when encountered. - max_tokens: The maximum number of tokens to generate. If max_tokens <= 0, the maximum number of tokens to generate is unlimited and depends on n_ctx. + seed: The seed to use for sampling. + response_format: The response format to use for the chat completion. Use { "type": "json_object" } to contstrain output to only valid json. + max_tokens: The maximum number of tokens to generate. If max_tokens <= 0 or None, the maximum number of tokens to generate is unlimited and depends on n_ctx. + presence_penalty: The penalty to apply to tokens based on their presence in the prompt. + frequency_penalty: The penalty to apply to tokens based on their frequency in the prompt. repeat_penalty: The penalty to apply to repeated tokens. + tfs_z: The tail-free sampling parameter. + mirostat_mode: The mirostat sampling mode. + mirostat_tau: The mirostat sampling tau parameter. + mirostat_eta: The mirostat sampling eta parameter. + model: The name to use for the model in the completion object. + logits_processor: A list of logits processors to use. + grammar: A grammar to use. + logit_bias: A logit bias to use. Returns: Generated chat completion or a stream of chat completion chunks. """ - stop = ( - stop if isinstance(stop, list) else [stop] if isinstance(stop, str) else [] - ) - chat_history = "".join( - f'### {"Human" if message["role"] == "user" else "Assistant"}:{message["content"]}' - for message in messages + handler = ( + self.chat_handler + or self._chat_handlers.get(self.chat_format) + or llama_chat_format.get_chat_completion_handler(self.chat_format) ) - PROMPT = chat_history + "### Assistant:" - PROMPT_STOP = ["### Assistant:", "### Human:"] - completion_or_chunks = self( - prompt=PROMPT, - stop=PROMPT_STOP + stop, + return handler( + llama=self, + messages=messages, + functions=functions, + function_call=function_call, + tools=tools, + tool_choice=tool_choice, temperature=temperature, top_p=top_p, top_k=top_k, + min_p=min_p, + typical_p=typical_p, + logprobs=logprobs, + top_logprobs=top_logprobs, stream=stream, + stop=stop, + seed=seed, + response_format=response_format, max_tokens=max_tokens, - repeat_penalty=repeat_penalty, presence_penalty=presence_penalty, frequency_penalty=frequency_penalty, + repeat_penalty=repeat_penalty, tfs_z=tfs_z, mirostat_mode=mirostat_mode, mirostat_tau=mirostat_tau, mirostat_eta=mirostat_eta, model=model, logits_processor=logits_processor, + grammar=grammar, + logit_bias=logit_bias, ) - if stream: - chunks: Iterator[CompletionChunk] = completion_or_chunks # type: ignore - return self._convert_text_completion_chunks_to_chat(chunks) - else: - completion: Completion = completion_or_chunks # type: ignore - return self._convert_text_completion_to_chat(completion) - def __del__(self): - if self.model is not None: - llama_cpp.llama_free_model(self.model) - self.model = None - if self.ctx is not None: - llama_cpp.llama_free(self.ctx) - self.ctx = None + def create_chat_completion_openai_v1( + self, + *args: Any, + **kwargs: Any, + ): + """Generate a chat completion with return type based on the the OpenAI v1 API. + + OpenAI python package is required to use this method. + + You can install it with `pip install openai`. + + Args: + *args: Positional arguments to pass to create_chat_completion. + **kwargs: Keyword arguments to pass to create_chat_completion. + + Returns: + Generated chat completion or a stream of chat completion chunks. + """ + try: + from openai.types.chat import ChatCompletion, ChatCompletionChunk + + stream = kwargs.get("stream", False) # type: ignore + assert isinstance(stream, bool) + if stream: + return ( + ChatCompletionChunk(**chunk) + for chunk in self.create_chat_completion(*args, **kwargs) + ) # type: ignore + else: + return ChatCompletion(**self.create_chat_completion(*args, **kwargs)) # type: ignore + except ImportError: + raise ImportError( + "To use create_chat_completion_openai_v1, you must install the openai package." + "You can install it with `pip install openai`." + ) def __getstate__(self): return dict( - verbose=self.verbose, model_path=self.model_path, - n_ctx=self.params.n_ctx, - n_gpu_layers=self.params.n_gpu_layers, - seed=self.params.seed, - f16_kv=self.params.f16_kv, - logits_all=self.params.logits_all, - vocab_only=self.params.vocab_only, - use_mmap=self.params.use_mmap, - use_mlock=self.params.use_mlock, - embedding=self.params.embedding, - low_vram=self.params.low_vram, - last_n_tokens_size=self.last_n_tokens_size, + # Model Params + n_gpu_layers=self.model_params.n_gpu_layers, + split_mode=self.model_params.split_mode, + main_gpu=self.model_params.main_gpu, + tensor_split=self.tensor_split, + vocab_only=self.model_params.vocab_only, + use_mmap=self.model_params.use_mmap, + use_mlock=self.model_params.use_mlock, + kv_overrides=self.kv_overrides, + # Context Params + seed=self._seed, + n_ctx=self.context_params.n_ctx, n_batch=self.n_batch, - n_threads=self.n_threads, + n_ubatch=self.context_params.n_ubatch, + n_threads=self.context_params.n_threads, + n_threads_batch=self.context_params.n_threads_batch, + rope_scaling_type=self.context_params.rope_scaling_type, + pooling_type=self.context_params.pooling_type, + attention_type=self.context_params.attention_type, + rope_freq_base=self.context_params.rope_freq_base, + rope_freq_scale=self.context_params.rope_freq_scale, + yarn_ext_factor=self.context_params.yarn_ext_factor, + yarn_attn_factor=self.context_params.yarn_attn_factor, + yarn_beta_fast=self.context_params.yarn_beta_fast, + yarn_beta_slow=self.context_params.yarn_beta_slow, + yarn_orig_ctx=self.context_params.yarn_orig_ctx, + logits_all=self._logits_all, + embedding=self.context_params.embeddings, + offload_kqv=self.context_params.offload_kqv, + flash_attn=( + self.context_params.flash_attn_type + == llama_cpp.LLAMA_FLASH_ATTN_TYPE_ENABLED + ), + op_offload=self.context_params.op_offload, + swa_full=self.context_params.swa_full, + # Sampling Params + no_perf=self.context_params.no_perf, + last_n_tokens_size=self.last_n_tokens_size, + # LoRA Params lora_base=self.lora_base, + lora_scale=self.lora_scale, lora_path=self.lora_path, - ### DEPRECATED ### - n_parts=self.n_parts, - ### DEPRECATED ### + # Backend Params + numa=self.numa, + # Chat Format Params + chat_format=self.chat_format, + chat_handler=self.chat_handler, + # Speculative Decidng + draft_model=self.draft_model, + # KV cache quantization + type_k=self.context_params.type_k, + type_v=self.context_params.type_v, + # Misc + spm_infill=self.spm_infill, + verbose=self.verbose, ) def __setstate__(self, state): - self.__init__( - model_path=state["model_path"], - n_ctx=state["n_ctx"], - n_parts=state["n_parts"], - n_gpu_layers=state["n_gpu_layers"], - seed=state["seed"], - f16_kv=state["f16_kv"], - logits_all=state["logits_all"], - vocab_only=state["vocab_only"], - use_mmap=state["use_mmap"], - use_mlock=state["use_mlock"], - embedding=state["embedding"], - low_vram=state["low_vram"], - n_threads=state["n_threads"], - n_batch=state["n_batch"], - last_n_tokens_size=state["last_n_tokens_size"], - lora_base=state["lora_base"], - lora_path=state["lora_path"], - verbose=state["verbose"], - ) + self.__init__(**state) def save_state(self) -> LlamaState: - assert self.ctx is not None if self.verbose: print("Llama.save_state: saving llama state", file=sys.stderr) - state_size = llama_cpp.llama_get_state_size(self.ctx) + state_size = llama_cpp.llama_state_get_size(self._ctx.ctx) if self.verbose: print(f"Llama.save_state: got state size: {state_size}", file=sys.stderr) - llama_state = (llama_cpp.c_uint8 * int(state_size))() + llama_state = (ctypes.c_uint8 * int(state_size))() if self.verbose: print("Llama.save_state: allocated state", file=sys.stderr) - n_bytes = llama_cpp.llama_copy_state_data(self.ctx, llama_state) + n_bytes = llama_cpp.llama_state_get_data(self._ctx.ctx, llama_state, state_size) if self.verbose: print(f"Llama.save_state: copied llama state: {n_bytes}", file=sys.stderr) if int(n_bytes) > int(state_size): raise RuntimeError("Failed to copy llama state data") - llama_state_compact = (llama_cpp.c_uint8 * int(n_bytes))() + llama_state_compact = (ctypes.c_uint8 * int(n_bytes))() llama_cpp.ctypes.memmove(llama_state_compact, llama_state, int(n_bytes)) if self.verbose: print( @@ -1553,65 +2202,88 @@ def save_state(self) -> LlamaState: file=sys.stderr, ) return LlamaState( - scores=self.scores.copy(), + scores=self._scores.copy(), input_ids=self.input_ids.copy(), n_tokens=self.n_tokens, llama_state=bytes(llama_state_compact), llama_state_size=n_bytes, + seed=self._seed, ) def load_state(self, state: LlamaState) -> None: - assert self.ctx is not None - self.scores = state.scores.copy() + # Only filling in up to `n_tokens` and then zero-ing out the rest + self.scores[: state.n_tokens, :] = state.scores.copy() + rest = self.scores[state.n_tokens :, :] + rest[rest > 0] = 0.0 self.input_ids = state.input_ids.copy() self.n_tokens = state.n_tokens + self._seed = state.seed state_size = state.llama_state_size - LLamaStateArrayType = llama_cpp.c_uint8 * state_size + LLamaStateArrayType = ctypes.c_uint8 * state_size llama_state = LLamaStateArrayType.from_buffer_copy(state.llama_state) - if llama_cpp.llama_set_state_data(self.ctx, llama_state) != state_size: + if ( + llama_cpp.llama_state_set_data(self._ctx.ctx, llama_state, state_size) + != state_size + ): raise RuntimeError("Failed to set llama state data") def n_ctx(self) -> int: """Return the context window size.""" - assert self.ctx is not None - return llama_cpp.llama_n_ctx(self.ctx) + return self._ctx.n_ctx() def n_embd(self) -> int: """Return the embedding size.""" - assert self.ctx is not None - return llama_cpp.llama_n_embd(self.ctx) + return self._model.n_embd() def n_vocab(self) -> int: """Return the vocabulary size.""" - assert self.ctx is not None - return llama_cpp.llama_n_vocab(self.ctx) + return self._model.n_vocab() - def tokenizer(self) -> "LlamaTokenizer": - """Return the tokenizer for this model.""" - assert self.ctx is not None + def tokenizer(self) -> LlamaTokenizer: + """Return the llama tokenizer for this model.""" return LlamaTokenizer(self) - @staticmethod - def token_eos() -> int: + def token_eos(self) -> int: """Return the end-of-sequence token.""" - return llama_cpp.llama_token_eos() + return self._model.token_eos() - @staticmethod - def token_bos() -> int: + def token_bos(self) -> int: """Return the beginning-of-sequence token.""" - return llama_cpp.llama_token_bos() + return self._model.token_bos() - @staticmethod - def token_nl() -> int: + def token_nl(self) -> int: """Return the newline token.""" - return llama_cpp.llama_token_nl() + return self._model.token_nl() + + def pooling_type(self) -> str: + """Return the pooling type.""" + return self._ctx.pooling_type() + + def close(self) -> None: + """Explicitly free the model from memory.""" + self._stack.close() + + def __del__(self) -> None: + self.close() @staticmethod - def logits_to_logprobs(logits: List[float]) -> List[float]: - exps = [math.exp(float(x)) for x in logits] - sum_exps = sum(exps) - return [math.log(x / sum_exps) for x in exps] + def logits_to_logprobs( + logits: Union[npt.NDArray[np.single], List], axis: int = -1 + ) -> npt.NDArray[np.single]: + # https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.log_softmax.html + logits_maxs: np.ndarray = np.amax(logits, axis=axis, keepdims=True) + if logits_maxs.ndim > 0: + logits_maxs[~np.isfinite(logits_maxs)] = 0 + elif not np.isfinite(logits_maxs): + logits_maxs = 0 + subtract_maxs = np.subtract(logits, logits_maxs, dtype=np.single) + exp = np.exp(subtract_maxs) + # Suppress warnings about log of zero + with np.errstate(divide="ignore"): + summed = np.sum(exp, axis=axis, keepdims=True) + out = np.log(summed) + return subtract_maxs - out @staticmethod def longest_token_prefix(a: Sequence[int], b: Sequence[int]): @@ -1623,19 +2295,194 @@ def longest_token_prefix(a: Sequence[int], b: Sequence[int]): break return longest_prefix + @classmethod + def from_pretrained( + cls, + repo_id: str, + filename: Optional[str], + additional_files: Optional[List] = None, + local_dir: Optional[Union[str, os.PathLike[str]]] = None, + local_dir_use_symlinks: Union[bool, Literal["auto"]] = "auto", + cache_dir: Optional[Union[str, os.PathLike[str]]] = None, + **kwargs: Any, + ) -> "Llama": + """Create a Llama model from a pretrained model name or path. + This method requires the huggingface-hub package. + You can install it with `pip install huggingface-hub`. + + Args: + repo_id: The model repo id. + filename: A filename or glob pattern to match the model file in the repo. + additional_files: A list of filenames or glob patterns to match additional model files in the repo. + local_dir: The local directory to save the model to. + local_dir_use_symlinks: Whether to use symlinks when downloading the model. + **kwargs: Additional keyword arguments to pass to the Llama constructor. + + Returns: + A Llama model.""" + try: + from huggingface_hub import hf_hub_download, HfFileSystem + from huggingface_hub.utils import validate_repo_id + except ImportError: + raise ImportError( + "Llama.from_pretrained requires the huggingface-hub package. " + "You can install it with `pip install huggingface-hub`." + ) + + validate_repo_id(repo_id) + + hffs = HfFileSystem() + + files = [ + file["name"] if isinstance(file, dict) else file + for file in hffs.ls(repo_id, recursive=True) + ] + + # split each file into repo_id, subfolder, filename + file_list: List[str] = [] + for file in files: + rel_path = Path(file).relative_to(repo_id) + file_list.append(str(rel_path)) + + # find the only/first shard file: + matching_files = [file for file in file_list if fnmatch.fnmatch(file, filename)] # type: ignore + + if len(matching_files) == 0: + raise ValueError( + f"No file found in {repo_id} that match {filename}\n\n" + f"Available Files:\n{json.dumps(file_list)}" + ) + + if len(matching_files) > 1: + raise ValueError( + f"Multiple files found in {repo_id} matching {filename}\n\n" + f"Available Files:\n{json.dumps(files)}" + ) + + (matching_file,) = matching_files -class LlamaTokenizer: - def __init__(self, llama: Llama): - self.llama = llama + subfolder = str(Path(matching_file).parent) + filename = Path(matching_file).name - def encode(self, text: str, add_bos: bool = True) -> List[int]: - return self.llama.tokenize( - text.encode("utf-8", errors="ignore"), add_bos=add_bos + # download the file + hf_hub_download( + repo_id=repo_id, + filename=filename, + subfolder=subfolder, + local_dir=local_dir, + local_dir_use_symlinks=local_dir_use_symlinks, + cache_dir=cache_dir, ) - def decode(self, tokens: List[int]) -> str: - return self.llama.detokenize(tokens).decode("utf-8", errors="ignore") + if additional_files: + for additonal_file_name in additional_files: + # find the additional shard file: + matching_additional_files = [ + file + for file in file_list + if fnmatch.fnmatch(file, additonal_file_name) + ] + + if len(matching_additional_files) == 0: + raise ValueError( + f"No file found in {repo_id} that match {additonal_file_name}\n\n" + f"Available Files:\n{json.dumps(file_list)}" + ) - @classmethod - def from_ggml_file(cls, path: str) -> "LlamaTokenizer": - return cls(Llama(model_path=path, vocab_only=True)) + if len(matching_additional_files) > 1: + raise ValueError( + f"Multiple files found in {repo_id} matching {additonal_file_name}\n\n" + f"Available Files:\n{json.dumps(files)}" + ) + + (matching_additional_file,) = matching_additional_files + additional_subfolder = str(Path(matching_additional_file).parent) + additional_file_name = Path(matching_additional_file).name + + # Split additional file paths independently to avoid duplicating + # the main model subfolder in Hugging Face download URLs. + hf_hub_download( + repo_id=repo_id, + filename=additional_file_name, + subfolder=additional_subfolder, + local_dir=local_dir, + local_dir_use_symlinks=local_dir_use_symlinks, + cache_dir=cache_dir, + ) + + if local_dir is None: + model_path = hf_hub_download( + repo_id=repo_id, + filename=filename, + subfolder=subfolder, + local_dir=local_dir, + local_dir_use_symlinks=local_dir_use_symlinks, + cache_dir=cache_dir, + local_files_only=True, + ) + else: + model_path = os.path.join(local_dir, filename) + + # loading the first file of a sharded GGUF loads all remaining shard files in the subfolder + return cls( + model_path=model_path, + **kwargs, + ) + + +class LlamaState: + def __init__( + self, + input_ids: npt.NDArray[np.intc], + scores: npt.NDArray[np.single], + n_tokens: int, + llama_state: bytes, + llama_state_size: int, + seed: int, + ): + self.input_ids = input_ids + self.scores = scores + self.n_tokens = n_tokens + self.llama_state = llama_state + self.llama_state_size = llama_state_size + self.seed = seed + + +LogitsProcessor = Callable[ + [npt.NDArray[np.intc], npt.NDArray[np.single]], npt.NDArray[np.single] +] + + +class LogitsProcessorList(List[LogitsProcessor]): + def __call__( + self, input_ids: npt.NDArray[np.intc], scores: npt.NDArray[np.single] + ) -> npt.NDArray[np.single]: + for processor in self: + scores = processor(input_ids, scores) + return scores + + +StoppingCriteria = Callable[[npt.NDArray[np.intc], npt.NDArray[np.single]], bool] + + +class StoppingCriteriaList(List[StoppingCriteria]): + def __call__( + self, input_ids: npt.NDArray[np.intc], logits: npt.NDArray[np.single] + ) -> bool: + return any([stopping_criteria(input_ids, logits) for stopping_criteria in self]) + + +class MinTokensLogitsProcessor(LogitsProcessor): + def __init__(self, min_tokens: int, token_eos: int): + self.min_tokens = min_tokens + self.token_eos = token_eos + self.prompt_tokens = None + + def __call__( + self, input_ids: npt.NDArray[np.intc], scores: npt.NDArray[np.single] + ) -> npt.NDArray[np.single]: + if self.prompt_tokens is None: + self.prompt_tokens = len(input_ids) + if len(input_ids) - self.prompt_tokens < self.min_tokens: + scores[self.token_eos] = -np.inf + return scores diff --git a/llama_cpp/llama_cache.py b/llama_cpp/llama_cache.py new file mode 100644 index 0000000000..5220c79337 --- /dev/null +++ b/llama_cpp/llama_cache.py @@ -0,0 +1,155 @@ +import sys +from abc import ABC, abstractmethod +from typing import ( + Optional, + Sequence, + Tuple, +) +from collections import OrderedDict + +import diskcache + +import llama_cpp.llama + +from .llama_types import * + + +class BaseLlamaCache(ABC): + """Base cache class for a llama.cpp model.""" + + def __init__(self, capacity_bytes: int = (2 << 30)): + self.capacity_bytes = capacity_bytes + + @property + @abstractmethod + def cache_size(self) -> int: + raise NotImplementedError + + def _find_longest_prefix_key( + self, + key: Tuple[int, ...], + ) -> Optional[Tuple[int, ...]]: + pass + + @abstractmethod + def __getitem__(self, key: Sequence[int]) -> "llama_cpp.llama.LlamaState": + raise NotImplementedError + + @abstractmethod + def __contains__(self, key: Sequence[int]) -> bool: + raise NotImplementedError + + @abstractmethod + def __setitem__( + self, key: Sequence[int], value: "llama_cpp.llama.LlamaState" + ) -> None: + raise NotImplementedError + + +class LlamaRAMCache(BaseLlamaCache): + """Cache for a llama.cpp model using RAM.""" + + def __init__(self, capacity_bytes: int = (2 << 30)): + super().__init__(capacity_bytes) + self.capacity_bytes = capacity_bytes + self.cache_state: OrderedDict[Tuple[int, ...], "llama_cpp.llama.LlamaState"] = ( + OrderedDict() + ) + + @property + def cache_size(self): + return sum([state.llama_state_size for state in self.cache_state.values()]) + + def _find_longest_prefix_key( + self, + key: Tuple[int, ...], + ) -> Optional[Tuple[int, ...]]: + min_len = 0 + min_key = None + keys = ( + (k, llama_cpp.llama.Llama.longest_token_prefix(k, key)) + for k in self.cache_state.keys() + ) + for k, prefix_len in keys: + if prefix_len > min_len: + min_len = prefix_len + min_key = k + return min_key + + def __getitem__(self, key: Sequence[int]) -> "llama_cpp.llama.LlamaState": + key = tuple(key) + _key = self._find_longest_prefix_key(key) + if _key is None: + raise KeyError("Key not found") + value = self.cache_state[_key] + self.cache_state.move_to_end(_key) + return value + + def __contains__(self, key: Sequence[int]) -> bool: + return self._find_longest_prefix_key(tuple(key)) is not None + + def __setitem__(self, key: Sequence[int], value: "llama_cpp.llama.LlamaState"): + key = tuple(key) + if key in self.cache_state: + del self.cache_state[key] + self.cache_state[key] = value + while self.cache_size > self.capacity_bytes and len(self.cache_state) > 0: + self.cache_state.popitem(last=False) + + +# Alias for backwards compatibility +LlamaCache = LlamaRAMCache + + +class LlamaDiskCache(BaseLlamaCache): + """Cache for a llama.cpp model using disk.""" + + def __init__( + self, cache_dir: str = ".cache/llama_cache", capacity_bytes: int = (2 << 30) + ): + super().__init__(capacity_bytes) + self.cache = diskcache.Cache(cache_dir) + + @property + def cache_size(self): + return int(self.cache.volume()) # type: ignore + + def _find_longest_prefix_key( + self, + key: Tuple[int, ...], + ) -> Optional[Tuple[int, ...]]: + min_len = 0 + min_key: Optional[Tuple[int, ...]] = None + for k in self.cache.iterkeys(): # type: ignore + prefix_len = llama_cpp.llama.Llama.longest_token_prefix(k, key) + if prefix_len > min_len: + min_len = prefix_len + min_key = k # type: ignore + return min_key + + def __getitem__(self, key: Sequence[int]) -> "llama_cpp.llama.LlamaState": + key = tuple(key) + _key = self._find_longest_prefix_key(key) + if _key is None: + raise KeyError("Key not found") + value: "llama_cpp.llama.LlamaState" = self.cache.pop(_key) # type: ignore + # NOTE: This puts an integer as key in cache, which breaks, + # Llama.longest_token_prefix(k, key) above since k is not a tuple of ints/tokens + # self.cache.push(_key, side="front") # type: ignore + return value + + def __contains__(self, key: Sequence[int]) -> bool: + return self._find_longest_prefix_key(tuple(key)) is not None + + def __setitem__(self, key: Sequence[int], value: "llama_cpp.llama.LlamaState"): + print("LlamaDiskCache.__setitem__: called", file=sys.stderr) + key = tuple(key) + if key in self.cache: + print("LlamaDiskCache.__setitem__: delete", file=sys.stderr) + del self.cache[key] + self.cache[key] = value + print("LlamaDiskCache.__setitem__: set", file=sys.stderr) + while self.cache_size > self.capacity_bytes and len(self.cache) > 0: + key_to_remove = next(iter(self.cache)) + del self.cache[key_to_remove] + print("LlamaDiskCache.__setitem__: trim", file=sys.stderr) diff --git a/llama_cpp/llama_chat_format.py b/llama_cpp/llama_chat_format.py new file mode 100644 index 0000000000..0034bdae98 --- /dev/null +++ b/llama_cpp/llama_chat_format.py @@ -0,0 +1,4574 @@ +from __future__ import annotations + +import os +import sys +import json +import ctypes +import dataclasses +import random +import string + +from datetime import datetime +from contextlib import ExitStack +from typing import ( + Any, + Dict, + Iterator, + List, + Literal, + Optional, + Tuple, + Union, + Protocol, + cast, +) + +import jinja2 +from jinja2.ext import Extension +from jinja2.sandbox import ImmutableSandboxedEnvironment + +import numpy as np +import numpy.typing as npt + +import llama_cpp.llama_cpp as llama_cpp +import llama_cpp.llama as llama +import llama_cpp.llama_types as llama_types +import llama_cpp.llama_grammar as llama_grammar + +from ._logger import logger +from ._utils import suppress_stdout_stderr, Singleton + +### Common Chat Templates and Special Tokens ### + +# Source: https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B/blob/main/tokenizer_config.json +CHATML_CHAT_TEMPLATE = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}" +CHATML_BOS_TOKEN = "<s>" +CHATML_EOS_TOKEN = "<|im_end|>" + +# Source: https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1/blob/main/tokenizer_config.json +MISTRAL_INSTRUCT_CHAT_TEMPLATE = "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ '[INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ message['content'] + eos_token + ' ' }}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}" +MISTRAL_INSTRUCT_BOS_TOKEN = "<s>" +MISTRAL_INSTRUCT_EOS_TOKEN = "</s>" + +# Source: https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1/blob/main/tokenizer_config.json +MIXTRAL_INSTRUCT_CHAT_TEMPLATE = "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ '[INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ message['content'] + eos_token}}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}" + +# Source: https://huggingface.co/meta-llama/Meta-Llama-3-8B-Instruct/blob/main/tokenizer_config.json +LLAMA3_INSTRUCT_CHAT_TEMPLATE = "{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %}" + +### Chat Completion Handler ### + + +class LlamaChatCompletionHandler(Protocol): + """Base Protocol for a llama chat completion handler. + + Very generic protocol that can be used to implement any chat format. + The only hard requirement is that it must return a ChatCompletion when + stream=False and an iterator of ChatCompletionChunks when stream=True.""" + + def __call__( + self, + *, + # llama.cpp instance + llama: llama.Llama, + # openai api parameters + messages: List[llama_types.ChatCompletionRequestMessage], + functions: Optional[List[llama_types.ChatCompletionFunction]] = None, + function_call: Optional[llama_types.ChatCompletionRequestFunctionCall] = None, + tools: Optional[List[llama_types.ChatCompletionTool]] = None, + tool_choice: Optional[llama_types.ChatCompletionToolChoiceOption] = None, + temperature: float = 0.2, + top_p: float = 0.95, + top_k: int = 40, + stream: bool = False, + stop: Optional[Union[str, List[str]]] = [], + seed: Optional[int] = None, + response_format: Optional[ + llama_types.ChatCompletionRequestResponseFormat + ] = None, + max_tokens: Optional[int] = None, + presence_penalty: float = 0.0, + frequency_penalty: float = 0.0, + repeat_penalty: float = 1.1, + model: Optional[str] = None, + logit_bias: Optional[Dict[str, float]] = None, + # llama.cpp parameters + min_p: float = 0.05, + typical_p: float = 1.0, + tfs_z: float = 1.0, + mirostat_mode: int = 0, + mirostat_tau: float = 5.0, + mirostat_eta: float = 0.1, + logits_processor: Optional[llama.LogitsProcessorList] = None, + grammar: Optional[llama.LlamaGrammar] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + **kwargs, # type: ignore + ) -> Union[ + llama_types.CreateChatCompletionResponse, + Iterator[llama_types.CreateChatCompletionStreamResponse], + ]: ... + + +class LlamaChatCompletionHandlerNotFoundException(Exception): + pass + + +class LlamaChatCompletionHandlerRegistry(Singleton): + _chat_handlers: Dict[str, LlamaChatCompletionHandler] = {} + + def register_chat_completion_handler( + self, + name: str, + chat_handler: LlamaChatCompletionHandler, + overwrite: bool = False, + ): + if not overwrite and name in self._chat_handlers: + raise ValueError( + f"Formatter with name '{name}' is already registered. Use `overwrite=True` to overwrite it." + ) + self._chat_handlers[name] = chat_handler + + def unregister_chat_handler(self, name: str): + if name in self._chat_handlers: + del self._chat_handlers[name] + else: + raise ValueError(f"No formatter registered under the name '{name}'.") + + def get_chat_completion_handler_by_name( + self, name: str + ) -> LlamaChatCompletionHandler: + try: + chat_handler = self._chat_handlers[name] + return chat_handler + except KeyError: + raise LlamaChatCompletionHandlerNotFoundException( + f"Invalid chat handler: {name} (valid formats: {list(self._chat_handlers.keys())})" + ) + + +def get_chat_completion_handler(name: str) -> LlamaChatCompletionHandler: + return LlamaChatCompletionHandlerRegistry().get_chat_completion_handler_by_name( + name + ) + + +def register_chat_completion_handler(name: str): + def decorator(f: LlamaChatCompletionHandler): + LlamaChatCompletionHandlerRegistry().register_chat_completion_handler(name, f) + return f + + return decorator + + +### Chat Formatter ### + + +@dataclasses.dataclass +class ChatFormatterResponse: + """Dataclass that stores completion parameters for a given chat format and + create_chat_completion request. + + prompt contains the formatted prompt generated from the chat format and messages. + stop contains the stop token or list of stop tokens to use for the chat format.""" + + prompt: str + stop: Optional[Union[str, List[str]]] = None + stopping_criteria: Optional[llama.StoppingCriteriaList] = None + added_special: bool = False + + +class ChatFormatter(Protocol): + """Base Protocol for a chat formatter. A chat formatter is a function that + takes a list of messages and returns a chat format response which can be used + to generate a completion. The response can also include a stop token or list + of stop tokens to use for the completion.""" + + def __call__( + self, + *, + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, + ) -> ChatFormatterResponse: ... + + +class Jinja2ChatFormatter(ChatFormatter): + class IgnoreGenerationTags(Extension): + """Pass-through for HuggingFace's ``{% generation %}`` chat-template tag.""" + + tags = {"generation"} + + def parse(self, parser: jinja2.parser.Parser): + parser.stream.skip(1) + return parser.parse_statements(("name:endgeneration",), drop_needle=True) + + def __init__( + self, + template: str, + eos_token: str, + bos_token: str, + add_generation_prompt: bool = True, + stop_token_ids: Optional[List[int]] = None, + ): + """A chat formatter that uses jinja2 templates to format the prompt.""" + self.template = template + self.eos_token = eos_token + self.bos_token = bos_token + self.add_generation_prompt = add_generation_prompt + self.stop_token_ids = ( + set(stop_token_ids) if stop_token_ids is not None else None + ) + + environment = ImmutableSandboxedEnvironment( + loader=jinja2.BaseLoader(), + trim_blocks=True, + lstrip_blocks=True, + # Keep this aligned with Transformers' chat-template Jinja extensions. + # https://github.com/huggingface/transformers/blob/39603d0e5cdb6f00e8d473d7fcbb01032d709181/src/transformers/utils/chat_template_utils.py#L489-L490 + extensions=[ + Jinja2ChatFormatter.IgnoreGenerationTags, + jinja2.ext.loopcontrols, + ], + ) + # Match Transformers' chat-template JSON rendering behavior. + # https://github.com/huggingface/transformers/blob/39603d0e5cdb6f00e8d473d7fcbb01032d709181/src/transformers/utils/chat_template_utils.py#L481-L484 + environment.filters["tojson"] = self.tojson + self._environment = environment.from_string(self.template) + + @staticmethod + def strftime_now(f: str) -> str: + return datetime.now().strftime(f) + + @staticmethod + def tojson( + x: Any, + ensure_ascii: bool = False, + indent: Optional[int] = None, + separators: Optional[Tuple[str, str]] = None, + sort_keys: bool = False, + ) -> str: + return json.dumps( + x, + ensure_ascii=ensure_ascii, + indent=indent, + separators=separators, + sort_keys=sort_keys, + ) + + def __call__( + self, + *, + messages: List[llama_types.ChatCompletionRequestMessage], + functions: Optional[List[llama_types.ChatCompletionFunction]] = None, + function_call: Optional[llama_types.ChatCompletionRequestFunctionCall] = None, + tools: Optional[List[llama_types.ChatCompletionTool]] = None, + tool_choice: Optional[llama_types.ChatCompletionToolChoiceOption] = None, + **kwargs: Any, + ) -> ChatFormatterResponse: + def raise_exception(message: str): + raise ValueError(message) + + prompt = self._environment.render( + messages=messages, + eos_token=self.eos_token, + bos_token=self.bos_token, + raise_exception=raise_exception, + add_generation_prompt=self.add_generation_prompt, + functions=functions, + function_call=function_call, + tools=tools, + tool_choice=tool_choice, + strftime_now=self.strftime_now, + **kwargs, + ) + + stopping_criteria = None + if self.stop_token_ids is not None: + + def stop_on_last_token( + tokens: npt.NDArray[np.intc], logits: npt.NDArray[np.single] + ) -> bool: + return tokens[-1] in self.stop_token_ids + + stopping_criteria = llama.StoppingCriteriaList([stop_on_last_token]) + + return ChatFormatterResponse( + prompt=prompt, + stop=[self.eos_token], + stopping_criteria=stopping_criteria, + added_special=True, + ) + + def to_chat_handler(self) -> LlamaChatCompletionHandler: + return chat_formatter_to_chat_completion_handler(self) + + +def _convert_text_completion_logprobs_to_chat( + logprobs: Optional[llama_types.CompletionLogprobs], +) -> llama_types.ChatCompletionLogprobs: + if logprobs is None: + return None + + return { + "content": [ + { + "token": token, + "bytes": None, + "logprob": logprob, + "top_logprobs": [ + { + "token": top_token, + "logprob": top_logprob, + "bytes": None, + } + for top_token, top_logprob in top_logprobs.items() + ], + } + for (token, logprob, top_logprobs) in zip( + logprobs["tokens"], logprobs["token_logprobs"], logprobs["top_logprobs"] + ) + ], + "refusal": None, + } + + +def _convert_text_completion_to_chat( + completion: llama_types.Completion, +) -> llama_types.ChatCompletion: + assert "usage" in completion + return { + "id": "chat" + completion["id"], + "object": "chat.completion", + "created": completion["created"], + "model": completion["model"], + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": completion["choices"][0]["text"], + }, + "logprobs": _convert_text_completion_logprobs_to_chat( + completion["choices"][0]["logprobs"] + ), + "finish_reason": completion["choices"][0]["finish_reason"], + } + ], + "usage": completion["usage"], + } + + +def _convert_text_completion_chunks_to_chat( + chunks: Iterator[llama_types.CreateCompletionStreamResponse], +) -> Iterator[llama_types.ChatCompletionChunk]: + for i, chunk in enumerate(chunks): + if i == 0: + yield { + "id": "chat" + chunk["id"], + "model": chunk["model"], + "created": chunk["created"], + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": { + "role": "assistant", + }, + "logprobs": None, + "finish_reason": None, + } + ], + } + yield { + "id": "chat" + chunk["id"], + "model": chunk["model"], + "created": chunk["created"], + "object": "chat.completion.chunk", + "choices": [ + { + "index": 0, + "delta": ( + { + "content": chunk["choices"][0]["text"], + } + if chunk["choices"][0]["finish_reason"] is None + else {} + ), + "logprobs": _convert_text_completion_logprobs_to_chat( + chunk["choices"][0]["logprobs"] + ), + "finish_reason": chunk["choices"][0]["finish_reason"], + } + ], + } + + +def _convert_completion_to_chat( + completion_or_chunks: Union[ + llama_types.CreateCompletionResponse, + Iterator[llama_types.CreateCompletionStreamResponse], + ], + stream: bool = False, +) -> Union[ + llama_types.CreateChatCompletionResponse, Iterator[llama_types.ChatCompletionChunk] +]: + if stream: + chunks: Iterator[llama_types.CreateCompletionStreamResponse] = ( + completion_or_chunks # type: ignore + ) + return _convert_text_completion_chunks_to_chat(chunks) + else: + completion: llama_types.Completion = completion_or_chunks # type: ignore + return _convert_text_completion_to_chat(completion) + + +def _convert_completion_to_chat_function( + tool_name: str, + completion_or_chunks: Union[ + llama_types.CreateCompletionResponse, + Iterator[llama_types.CreateCompletionStreamResponse], + ], + stream: bool, +): + if not stream: + completion: llama_types.CreateCompletionResponse = completion_or_chunks # type: ignore + assert "usage" in completion + tool_id = "call_" + "_0_" + tool_name + "_" + completion["id"] + # TODO: Fix for legacy function calls + chat_completion: llama_types.CreateChatCompletionResponse = { + "id": "chat" + completion["id"], + "object": "chat.completion", + "created": completion["created"], + "model": completion["model"], + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": None, + "function_call": { + "name": tool_name, + "arguments": completion["choices"][0]["text"], + }, + "tool_calls": [ + { + "id": tool_id, + "type": "function", + "function": { + "name": tool_name, + "arguments": completion["choices"][0]["text"], + }, + } + ], + }, + "logprobs": _convert_text_completion_logprobs_to_chat( + completion["choices"][0]["logprobs"] + ), + "finish_reason": "tool_calls", + } + ], + "usage": completion["usage"], + } + return chat_completion + else: + chunks: Iterator[llama_types.CreateCompletionStreamResponse] = ( + completion_or_chunks # type: ignore + ) + + def _stream_response_to_function_stream( + chunks: Iterator[llama_types.CreateCompletionStreamResponse], + ) -> Iterator[llama_types.CreateChatCompletionStreamResponse]: + # blank first message + first = True + id_ = None + created = None + model = None + tool_id = None + for chunk in chunks: + if first: + id_ = "chat" + chunk["id"] + created = chunk["created"] + model = chunk["model"] + tool_id = "call_" + "_0_" + tool_name + "_" + chunk["id"] + yield { + "id": id_, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": [ + { + "index": 0, + "finish_reason": None, + "logprobs": None, + "delta": { + "role": "assistant", + "content": None, + "function_call": None, + "tool_calls": None, + }, + } + ], + } + yield { + "id": "chat" + chunk["id"], + "object": "chat.completion.chunk", + "created": chunk["created"], + "model": chunk["model"], + "choices": [ + { + "index": 0, + "finish_reason": None, + "logprobs": _convert_text_completion_logprobs_to_chat( + chunk["choices"][0]["logprobs"] + ), + "delta": { + "role": None, + "content": None, + "function_call": { + "name": tool_name, + "arguments": chunk["choices"][0]["text"], + }, + "tool_calls": [ + { + "index": 0, + "id": tool_id, + "type": "function", + "function": { + "name": tool_name, + "arguments": chunk["choices"][0][ + "text" + ], + }, + } + ], + }, + } + ], + } + first = False + continue + assert tool_id is not None + yield { + "id": "chat" + chunk["id"], + "object": "chat.completion.chunk", + "created": chunk["created"], + "model": chunk["model"], + "choices": [ + { + "index": 0, + "finish_reason": None, + "logprobs": _convert_text_completion_logprobs_to_chat( + chunk["choices"][0]["logprobs"] + ), + "delta": { + "role": None, + "content": None, + "function_call": { + "name": tool_name, + "arguments": chunk["choices"][0]["text"], + }, + "tool_calls": [ + { + "index": 0, + "id": tool_id, + "type": "function", + "function": { + "name": tool_name, + "arguments": chunk["choices"][0]["text"], + }, + } + ], + }, + } + ], + } + + if id_ is not None and created is not None and model is not None: + yield { + "id": id_, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "logprobs": None, + "delta": { + "role": None, + "content": None, + "function_call": None, + "tool_calls": None, + }, + } + ], + } + + return _stream_response_to_function_stream(chunks) + + +def chat_formatter_to_chat_completion_handler( + chat_formatter: ChatFormatter, +) -> LlamaChatCompletionHandler: + def chat_completion_handler( + *, + llama: llama.Llama, + messages: List[llama_types.ChatCompletionRequestMessage], + functions: Optional[List[llama_types.ChatCompletionFunction]] = None, + function_call: Optional[llama_types.ChatCompletionRequestFunctionCall] = None, + tools: Optional[List[llama_types.ChatCompletionTool]] = None, + tool_choice: Optional[llama_types.ChatCompletionToolChoiceOption] = None, + temperature: float = 0.2, + top_p: float = 0.95, + top_k: int = 40, + min_p: float = 0.05, + typical_p: float = 1.0, + stream: bool = False, + stop: Optional[Union[str, List[str]]] = [], + seed: Optional[int] = None, + response_format: Optional[ + llama_types.ChatCompletionRequestResponseFormat + ] = None, + max_tokens: Optional[int] = None, + presence_penalty: float = 0.0, + frequency_penalty: float = 0.0, + repeat_penalty: float = 1.1, + tfs_z: float = 1.0, + mirostat_mode: int = 0, + mirostat_tau: float = 5.0, + mirostat_eta: float = 0.1, + model: Optional[str] = None, + logits_processor: Optional[llama.LogitsProcessorList] = None, + grammar: Optional[llama.LlamaGrammar] = None, + logit_bias: Optional[Dict[str, float]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + **kwargs, # type: ignore + ) -> Union[ + llama_types.CreateChatCompletionResponse, + Iterator[llama_types.CreateChatCompletionStreamResponse], + ]: + result = chat_formatter( + messages=messages, + functions=functions, + function_call=function_call, + tools=tools, + tool_choice=tool_choice, + **kwargs, + ) + prompt = llama.tokenize( + result.prompt.encode("utf-8"), + add_bos=not result.added_special, + special=True, + ) + if result.stop is not None: + stop = [] if stop is None else [stop] if isinstance(stop, str) else stop + rstop = result.stop if isinstance(result.stop, list) else [result.stop] + stop = stop + rstop + + stopping_criteria = None + if result.stopping_criteria is not None: + stopping_criteria = result.stopping_criteria + + if response_format is not None and response_format["type"] == "json_object": + grammar = _grammar_for_response_format( + response_format, verbose=llama.verbose + ) + + # Convert legacy functions to tools + if functions is not None: + tools = [ + { + "type": "function", + "function": function, + } + for function in functions + ] + + # Convert legacy function_call to tool_choice + if function_call is not None: + if isinstance(function_call, str) and ( + function_call == "none" or function_call == "auto" + ): + tool_choice = function_call + if isinstance(function_call, dict) and "name" in function_call: + tool_choice = { + "type": "function", + "function": { + "name": function_call["name"], + }, + } + + tool = None + if ( + tool_choice is not None + and isinstance(tool_choice, dict) + and tools is not None + ): + name = tool_choice["function"]["name"] + tool = next((t for t in tools if t["function"]["name"] == name), None) + if tool is None: + raise ValueError(f"Tool choice '{name}' not found in tools.") + schema = tool["function"]["parameters"] + try: + # create grammar from json schema + grammar = llama_grammar.LlamaGrammar.from_json_schema( + json.dumps(schema), verbose=llama.verbose + ) + except Exception as e: + if llama.verbose: + print(str(e), file=sys.stderr) + grammar = llama_grammar.LlamaGrammar.from_string( + llama_grammar.JSON_GBNF, verbose=llama.verbose + ) + + completion_or_chunks = llama.create_completion( + prompt=prompt, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + typical_p=typical_p, + logprobs=top_logprobs if logprobs else None, + stream=stream, + stop=stop, + seed=seed, + max_tokens=max_tokens, + presence_penalty=presence_penalty, + frequency_penalty=frequency_penalty, + repeat_penalty=repeat_penalty, + tfs_z=tfs_z, + mirostat_mode=mirostat_mode, + mirostat_tau=mirostat_tau, + mirostat_eta=mirostat_eta, + model=model, + logits_processor=logits_processor, + stopping_criteria=stopping_criteria, + grammar=grammar, + logit_bias=logit_bias, + ) + if tool is not None: + tool_name = tool["function"]["name"] + return _convert_completion_to_chat_function( + tool_name, completion_or_chunks, stream + ) + return _convert_completion_to_chat(completion_or_chunks, stream=stream) + + return chat_completion_handler + + +def hf_autotokenizer_to_chat_formatter( + pretrained_model_name_or_path: Union[str, os.PathLike[str]], +) -> ChatFormatter: + # https://huggingface.co/docs/transformers/main/chat_templating + # https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1#instruction-format + # https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1/blob/main/tokenizer_config.json + from transformers import AutoTokenizer # type: ignore + + tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path) # type: ignore + + def format_autotokenizer( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, + ) -> ChatFormatterResponse: + tokenizer.use_default_system_prompt = False # type: ignore + prompt: str = tokenizer.apply_chat_template( # type: ignore + messages, tokenize=False, **kwargs + ) + assert isinstance(prompt, str) + # Return formatted prompt and eos token by default + return ChatFormatterResponse( + prompt=prompt, stop=tokenizer.eos_token, added_special=True + ) + + return format_autotokenizer + + +def hf_autotokenizer_to_chat_completion_handler( + pretrained_model_name_or_path: Union[str, os.PathLike[str]], +) -> LlamaChatCompletionHandler: + chat_formatter = hf_autotokenizer_to_chat_formatter(pretrained_model_name_or_path) + return chat_formatter_to_chat_completion_handler(chat_formatter) + + +def hf_tokenizer_config_to_chat_formatter( + tokenizer_config: Dict[str, Any], + add_generation_prompt: bool = True, +) -> ChatFormatter: + assert isinstance(tokenizer_config, dict) + + assert "chat_template" in tokenizer_config + assert isinstance(tokenizer_config["chat_template"], str) + chat_template = tokenizer_config["chat_template"] + + assert "bos_token" in tokenizer_config + assert isinstance(tokenizer_config["bos_token"], str) + bos_token = tokenizer_config["bos_token"] + + assert "eos_token" in tokenizer_config + assert isinstance(tokenizer_config["eos_token"], str) + eos_token = tokenizer_config["eos_token"] + + env = ImmutableSandboxedEnvironment( + trim_blocks=True, + lstrip_blocks=True, + ).from_string(chat_template) + + def format_tokenizer_config( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, + ) -> ChatFormatterResponse: + # TODO: veryify this is correct + # Add a blank assistant message to the end of the messages to prompt the model to generate a response + if add_generation_prompt: + messages = [ + *messages, + llama_types.ChatCompletionRequestAssistantMessage( + role="assistant", content="" + ), + ] + prompt = env.render( + messages=messages, + bos_token=bos_token, + eos_token=eos_token, + **kwargs, + ) + return ChatFormatterResponse( + prompt=prompt, stop=[eos_token, bos_token], added_special=True + ) + + return format_tokenizer_config + + +def hf_tokenizer_config_to_chat_completion_handler( + tokenizer_config: Dict[str, Any], + add_generation_prompt: bool = True, +) -> LlamaChatCompletionHandler: + chat_formatter = hf_tokenizer_config_to_chat_formatter( + tokenizer_config, add_generation_prompt=add_generation_prompt + ) + return chat_formatter_to_chat_completion_handler(chat_formatter) + + +def guess_chat_format_from_gguf_metadata(metadata: Dict[str, str]) -> Optional[str]: + if "tokenizer.chat_template" not in metadata: + return None + + if metadata["tokenizer.chat_template"] == CHATML_CHAT_TEMPLATE: + return "chatml" + + if ( + metadata["tokenizer.chat_template"] == MISTRAL_INSTRUCT_CHAT_TEMPLATE + or metadata["tokenizer.chat_template"] == MIXTRAL_INSTRUCT_CHAT_TEMPLATE + ): + return "mistral-instruct" + + if metadata["tokenizer.chat_template"] == LLAMA3_INSTRUCT_CHAT_TEMPLATE: + return "llama-3" + + return None + + +### Utility functions for formatting chat prompts ### +# TODO: Replace these with jinja2 templates + + +def _get_system_message( + messages: List[llama_types.ChatCompletionRequestMessage], +) -> str: + """Get the first system message.""" + for message in messages: + if message["role"] == "system": + return message["content"] or "" + return "" + + +def _map_roles( + messages: List[llama_types.ChatCompletionRequestMessage], + role_map: Dict[str, str], +) -> List[Tuple[str, Optional[str]]]: + """Map the message roles.""" + output: List[Tuple[str, Optional[str]]] = [] + for message in messages: + role = message["role"] + if role in role_map: + content: str | None = ( + message["content"] if isinstance(message["content"], str) else None + ) + output.append((role_map[role], content)) + return output + + +def _format_llama2( + system_message: str, messages: List[Tuple[str, Optional[str]]], sep: str, sep2: str +) -> str: + """Format the prompt with the llama2 style.""" + seps = [sep, sep2] + ret = system_message + sep + for i, (role, message) in enumerate(messages): + if system_message and i == 0: + m = message or "" + ret += m + seps[i % 2] + elif message: + ret += role + message + " " + seps[i % 2] + else: + ret += role + " " + return ret + + +def _format_add_colon_single( + system_message: str, messages: List[Tuple[str, Optional[str]]], sep: str +) -> str: + """Format the prompt with the add-colon-single style.""" + ret = system_message + sep + for role, message in messages: + if message: + ret += role + ": " + message + sep + else: + ret += role + ":" + return ret + + +def _format_add_colon_two( + system_message: str, messages: List[Tuple[str, Optional[str]]], sep: str, sep2: str +) -> str: + """Format the prompt with the add-colon-two style.""" + seps = [sep, sep2] + ret = system_message + seps[0] + for i, (role, message) in enumerate(messages): + if message: + ret += role + ": " + message + seps[i % 2] + else: + ret += role + ":" + return ret + + +def _format_no_colon_single( + system_message: str, messages: List[Tuple[str, Optional[str]]], sep: str +) -> str: + """Format the prompt with the no-colon-single style.""" + ret = system_message + for role, message in messages: + if message: + ret += role + message + sep + else: + ret += role + return ret + + +def _format_add_colon_space_single( + system_message: str, messages: List[Tuple[str, Optional[str]]], sep: str +) -> str: + """Format the prompt with the add-colon-space-single style.""" + ret = system_message + sep + for role, message in messages: + if message: + ret += role + ": " + message + sep + else: + ret += role + ": " # must be end with a space + return ret + + +def _format_chatml( + system_message: str, messages: List[Tuple[str, Optional[str]]], sep: str +) -> str: + """Format the prompt with the chatml style.""" + ret = "" if system_message == "" else system_message + sep + "\n" + for role, message in messages: + if message: + ret += role + "\n" + message + sep + "\n" + else: + ret += role + "\n" + return ret + + +def _format_chatglm3( + system_message: str, messages: List[Tuple[str, Optional[str]]], sep: str +) -> str: + """Format the prompt with the chatglm3 style.""" + ret = "" + if system_message: + ret += system_message + for role, message in messages: + if message: + ret += role + "\n" + " " + message + else: + ret += role + return ret + + +def _grammar_for_json(verbose: bool = False): + return llama_grammar.LlamaGrammar.from_string( + llama_grammar.JSON_GBNF, verbose=verbose + ) + + +def _grammar_for_json_schema( + schema: str, verbose: bool = False, fallback_to_json: bool = True +): + try: + return llama_grammar.LlamaGrammar.from_json_schema(schema, verbose=verbose) + except Exception as e: + if fallback_to_json: + return _grammar_for_json(verbose=verbose) + else: + raise e + + +def _grammar_for_response_format( + response_format: llama_types.ChatCompletionRequestResponseFormat, + verbose: bool = False, +): + if response_format["type"] != "json_object": + return None + + if "schema" in response_format: + return _grammar_for_json_schema( + json.dumps(response_format["schema"]), verbose=verbose + ) + else: + return _grammar_for_json(verbose=verbose) + + +### Chat Formats ### + + +def register_chat_format(name: str): + def decorator(f: ChatFormatter): + chat_completion_handler = chat_formatter_to_chat_completion_handler(f) + LlamaChatCompletionHandlerRegistry().register_chat_completion_handler( + name, chat_completion_handler + ) + return f + + return decorator + + +# see https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/tokenization_llama.py +# system prompt is "embedded" in the first message +@register_chat_format("llama-2") +def format_llama2( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, +) -> ChatFormatterResponse: + _system_template = "[INST] <<SYS>>\n{system_message}\n<</SYS>>" + _roles = dict(user="<s>[INST]", assistant="[/INST]") + _messages = _map_roles(messages, _roles) + system_message = _get_system_message(messages) + if system_message: + system_message = _system_template.format(system_message=system_message) + _prompt = _format_llama2(system_message, _messages, " ", "</s>") + "[/INST]" + return ChatFormatterResponse(prompt=_prompt) + + +# Chat format for Llama-3 models, see more details at: +# https://github.com/meta-llama/llama3/blob/main/llama/tokenizer.py#L202-L229 +@register_chat_format("llama-3") +def format_llama3( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, +) -> ChatFormatterResponse: + _roles = dict( + system="<|start_header_id|>system<|end_header_id|>\n\n", + user="<|start_header_id|>user<|end_header_id|>\n\n", + assistant="<|start_header_id|>assistant<|end_header_id|>\n\n", + ) + _sep = "<|eot_id|>" + _messages = _map_roles(messages, _roles) + _messages.append((_roles["assistant"], None)) + _prompt = _format_no_colon_single("", _messages, _sep) + return ChatFormatterResponse(prompt=_prompt, stop=_sep) + + +@register_chat_format("alpaca") +def format_alpaca( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, +) -> ChatFormatterResponse: + _roles = dict(user="### Instruction", assistant="### Response") + _sep = "\n\n" + _sep2 = "</s>" + system_message = _get_system_message(messages) + _messages = _map_roles(messages, _roles) + _prompt = _format_add_colon_two(system_message, _messages, _sep, _sep2) + return ChatFormatterResponse(prompt=_prompt) + + +@register_chat_format("qwen") +def format_qwen( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, +) -> ChatFormatterResponse: + _roles = dict(user="<|im_start|>user", assistant="<|im_start|>assistant") + system_message = _get_system_message(messages) or "You are a helpful assistant." + system_template = "<|im_start|>system\n{system_message}" + system_message = system_template.format(system_message=system_message) + _messages = _map_roles(messages, _roles) + _messages.append((_roles["assistant"], None)) + _sep = "<|im_end|>" + _prompt = _format_chatml(system_message, _messages, _sep) + _sep2 = "<|endoftext|>" + return ChatFormatterResponse(prompt=_prompt, stop=_sep2) + + +@register_chat_format("vicuna") +def format( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, +) -> ChatFormatterResponse: + _system_message = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions." + _roles = dict(user="USER", assistant="ASSISTANT") + _sep = " " + _sep2 = "</s>" + system_message = _system_message + _messages = _map_roles(messages, _roles) + _messages.append((_roles["assistant"], None)) + _prompt = _format_add_colon_two(system_message, _messages, _sep, _sep2) + return ChatFormatterResponse(prompt=_prompt) + + +@register_chat_format("oasst_llama") +def format_oasst_llama( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, +) -> ChatFormatterResponse: + _system_template = "[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n" + _roles = dict(user="<|prompter|>", assistant="<|assistant|>") + _sep = "</s>" + system_message = _get_system_message(messages) + system_message = _system_template.format(system_message=system_message) + _messages = _map_roles(messages, _roles) + _messages.append((_roles["assistant"], None)) + _prompt = _format_no_colon_single(system_message, _messages, _sep) + return ChatFormatterResponse(prompt=_prompt) + + +@register_chat_format("baichuan-2") +def format_baichuan2( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, +) -> ChatFormatterResponse: + _system_template = "{system_message}" + _roles = dict(user="<reserved_106>", assistant="<reserved_107>") + _sep = "" + system_message = _get_system_message(messages) + system_message = _system_template.format(system_message=system_message) + _messages = _map_roles(messages, _roles) + _messages.append((_roles["assistant"], None)) + _prompt = _format_no_colon_single(system_message, _messages, _sep) + return ChatFormatterResponse(prompt=_prompt) + + +@register_chat_format("baichuan") +def format_baichuan( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, +) -> ChatFormatterResponse: + _system_template = "{system_message}" + _roles = dict(user="<reserved_102>", assistant="<reserved_103>") + _sep = "" + system_message = _get_system_message(messages) + system_message = _system_template.format(system_message=system_message) + _messages = _map_roles(messages, _roles) + _messages.append((_roles["assistant"], None)) + _prompt = _format_no_colon_single(system_message, _messages, _sep) + return ChatFormatterResponse(prompt=_prompt) + + +@register_chat_format("openbuddy") +def format_openbuddy( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, +) -> ChatFormatterResponse: + _system_message = """You are a helpful, respectful and honest INTP-T AI Assistant named Buddy. You are talking to a human User. +Always answer as helpfully and logically as possible, while being safe. Your answers should not include any harmful, political, religious, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature. +If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information. +You can speak fluently in many languages, for example: English, Chinese. +You cannot access the internet, but you have vast knowledge, cutoff: 2021-09. +You are trained by OpenBuddy team, (https://openbuddy.ai, https://github.com/OpenBuddy/OpenBuddy), you are based on LLaMA and Falcon transformers model, not related to GPT or OpenAI. + +""" + _roles = dict(user="User", assistant="Assistant") + _sep = "\n" + system_message = _system_message + _messages = _map_roles(messages, _roles) + _messages.append((_roles["assistant"], None)) + _prompt = _format_add_colon_single(system_message, _messages, _sep) + return ChatFormatterResponse(prompt=_prompt) + + +@register_chat_format("redpajama-incite") +def format_redpajama_incite( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, +) -> ChatFormatterResponse: + _system_message = _get_system_message(messages) + _roles = dict(user="<human>", assistant="<bot>") + _sep = "\n" + _stop = "<human>" + system_message = _system_message + _messages = _map_roles(messages, _roles) + _messages.append((_roles["assistant"], None)) + _prompt = _format_add_colon_single(system_message, _messages, _sep) + return ChatFormatterResponse(prompt=_prompt, stop=_stop) + + +@register_chat_format("snoozy") +def format_snoozy( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, +) -> ChatFormatterResponse: + system_template = "### Instruction:\n{system_message}" + default_system_message = "The prompt below is a question to answer, a task to complete, or a conversation to respond to; decide which and write an appropriate response." + _system_message = _get_system_message(messages) + _system_message = ( + _system_message if _system_message != "" else default_system_message + ) + system_message = system_template.format(system_message=_system_message) + _roles = dict(user="### Prompt", assistant="### Response") + _sep = "\n" + _stop = "###" + system_message = _system_message + _messages = _map_roles(messages, _roles) + _messages.append((_roles["assistant"], None)) + _prompt = _format_add_colon_single(system_message, _messages, _sep) + return ChatFormatterResponse(prompt=_prompt, stop=_stop) + + +@register_chat_format("phind") +def format_phind( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, +) -> ChatFormatterResponse: + _roles = dict(user="### User Message", assistant="### Assistant") + _sep = "\n\n" + _system_message = "### System Prompt\nYou are an intelligent programming assistant." + _messages = _map_roles(messages, _roles) + _messages.append((_roles["assistant"], None)) + _prompt = _format_add_colon_single(_system_message, _messages, _sep) + return ChatFormatterResponse(prompt=_prompt) + + +@register_chat_format("intel") +def format_intel( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, +) -> ChatFormatterResponse: + _roles = dict(user="### User:", assistant="### Assistant:") + _sep = "\n" + _system_message = "### System:\n{system_message}" + _messages = _map_roles(messages, _roles) + _messages.append((_roles["assistant"], None)) + _prompt = _format_add_colon_single(_system_message, _messages, _sep) + return ChatFormatterResponse(prompt=_prompt) + + +@register_chat_format("open-orca") +def format_open_orca( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, +) -> ChatFormatterResponse: + system_template = "{system_message}" + system_message = ( + "You are a helpful assistant. Please answer truthfully and write out your " + "thinking step by step to be sure you get the right answer. If you make a mistake or encounter " + "an error in your thinking, say so out loud and attempt to correct it. If you don't know or " + "aren't sure about something, say so clearly. You will act as a professional logician, mathematician, " + "and physicist. You will also act as the most appropriate type of expert to answer any particular " + "question or solve the relevant problem; state which expert type your are, if so. Also think of " + "any particular named expert that would be ideal to answer the relevant question or solve the " + "relevant problem; name and act as them, if appropriate." + ) + roles = ("User", "Assistant") + sep = "<|end_of_turn|>\n" + # stop_token_ids=[32000, 32001], # "<|end_of_turn|>" + stop_str = "User" + system_message = system_template.format(system_message=system_message) + _messages = _map_roles(messages, dict(zip(roles, roles))) + _messages.append((roles[1], None)) + _prompt = _format_add_colon_space_single(system_message, _messages, sep) + return ChatFormatterResponse(prompt=_prompt, stop=stop_str) + + +@register_chat_format("mistrallite") +def format_mistrallite( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, +) -> ChatFormatterResponse: + _roles = dict(user="<|prompter|>", assistant="</s>\n<|assistant|>") + _sep = " " + system_template = """<|system|>{system_message}</s>""" + system_message = _get_system_message(messages) + system_message = system_template.format(system_message=system_message) + _messages = _map_roles(messages, _roles) + _messages.append((_roles["assistant"], None)) + _prompt = _format_no_colon_single(system_message, _messages, _sep) + return ChatFormatterResponse(prompt=_prompt) + + +@register_chat_format("zephyr") +def format_zephyr( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, +) -> ChatFormatterResponse: + system_template = """<|system|> +{system_message}""" + system_message = _get_system_message(messages) + system_message = system_template.format(system_message=system_message) + _roles = dict(user="<|user|>\n", assistant="<|assistant|>\n") + _sep = "</s>" + _messages = _map_roles(messages, _roles) + _messages.append((_roles["assistant"], None)) + _prompt = _format_chatml(system_message, _messages, _sep) + return ChatFormatterResponse(prompt=_prompt, stop=_sep) + + +@register_chat_format("pygmalion") +def format_pygmalion( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, +) -> ChatFormatterResponse: + system_template = """<|system|>{system_message}""" + system_message = _get_system_message(messages) + system_message = system_template.format(system_message=system_message) + _roles = dict(user="<|user|>", assistant="<|model|>") + _sep = "\n" + _messages = _map_roles(messages, _roles) + _messages.append((_roles["assistant"], None)) + _prompt = _format_chatml(system_message, _messages, _sep) + return ChatFormatterResponse(prompt=_prompt, stop=_sep) + + +@register_chat_format("chatml") +def format_chatml( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, +) -> ChatFormatterResponse: + system_template = """<|im_start|>system +{system_message}""" + system_message = _get_system_message(messages) + system_message = system_template.format(system_message=system_message) + _roles = dict(user="<|im_start|>user", assistant="<|im_start|>assistant") + _sep = "<|im_end|>" + _messages = _map_roles(messages, _roles) + _messages.append((_roles["assistant"], None)) + _prompt = _format_chatml(system_message, _messages, _sep) + return ChatFormatterResponse(prompt=_prompt, stop=_sep) + + +@register_chat_format("mistral-instruct") +def format_mistral_instruct( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, +) -> ChatFormatterResponse: + eos = "</s>" + stop = eos + prompt = "" + for message in messages: + if ( + message["role"] == "user" + and message["content"] is not None + and isinstance(message["content"], str) + ): + prompt += "[INST] " + message["content"] + elif message["role"] == "assistant" and message["content"] is not None: + prompt += " [/INST]" + message["content"] + eos + prompt += " [/INST]" + return ChatFormatterResponse(prompt=prompt, stop=stop) + + +@register_chat_format("chatglm3") +def format_chatglm3( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, +) -> ChatFormatterResponse: + system_template = """<|system|> +{system_message}""" + system_message = _get_system_message(messages) + system_message = system_template.format(system_message=system_message) + _roles = dict(user="<|user|>", assistant="<|assistant|>") + _sep = "</s>" + _messages = _map_roles(messages, _roles) + _messages.append((_roles["assistant"], None)) + _prompt = _format_chatglm3(system_message, _messages, _sep) + return ChatFormatterResponse(prompt=_prompt, stop=_sep) + + +@register_chat_format("openchat") +def format_openchat( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, +) -> ChatFormatterResponse: + system_template = "{system_message}<|end_of_turn|>" + system_message = _get_system_message(messages) + system_message = system_template.format(system_message=system_message) + _roles = dict( + user="GPT4 Correct User: ", assistant="<|end_of_turn|>GPT4 Correct Assistant: " + ) + _sep = "<|end_of_turn|>" + _messages = _map_roles(messages, _roles) + _messages.append((_roles["assistant"], None)) + _prompt = _format_chatml(system_message, _messages, _sep) + return ChatFormatterResponse(prompt=_prompt, stop=_sep) + + +# Chat format for Saiga models, see more details and available models: +# https://huggingface.co/collections/IlyaGusev/saiga2-saigamistral-6505d4ccc3d1e53166b636cd +@register_chat_format("saiga") +def format_saiga( + messages: list[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, +) -> ChatFormatterResponse: + _message_template = "<s>{role}\n{content}</s>" + _roles = dict(user="user", bot="bot", system="system") + _messages = _map_roles(messages, _roles) + + _prompt = "" + for role, content in _messages: + if content: + _prompt += _message_template.format(role=role, content=content) + else: + _prompt += f"<s>{role}\n" + # Response template + _prompt += "<s>bot" + return ChatFormatterResponse(prompt=_prompt.strip()) + + +# Chat format for Google's Gemma models, see more details and available models: +# https://huggingface.co/collections/google/gemma-release-65d5efbccdbb8c4202ec078b +@register_chat_format("gemma") +def format_gemma( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, +) -> ChatFormatterResponse: + system_message = _get_system_message(messages) + if system_message != "": + logger.debug( + "`role='system'` messages are not allowed on Google's Gemma models." + ) + _roles = dict(user="<start_of_turn>user\n", assistant="<start_of_turn>model\n") + _sep = "<end_of_turn>\n" + _messages = _map_roles(messages, _roles) + _messages.append((_roles["assistant"], None)) + _prompt = _format_no_colon_single(system_message="", messages=_messages, sep=_sep) + return ChatFormatterResponse(prompt=_prompt, stop=_sep) + + +# Tricky chat formats that require custom chat handlers + + +@register_chat_completion_handler("functionary") +def functionary_chat_handler( + llama: llama.Llama, + messages: List[llama_types.ChatCompletionRequestMessage], + functions: Optional[List[llama_types.ChatCompletionFunction]] = None, + function_call: Optional[llama_types.ChatCompletionRequestFunctionCall] = None, + tools: Optional[List[llama_types.ChatCompletionTool]] = None, + tool_choice: Optional[llama_types.ChatCompletionToolChoiceOption] = None, + temperature: float = 0.2, + top_p: float = 0.95, + top_k: int = 40, + min_p: float = 0.05, + typical_p: float = 1.0, + stream: bool = False, + stop: Optional[Union[str, List[str]]] = [], + response_format: Optional[llama_types.ChatCompletionRequestResponseFormat] = None, + max_tokens: Optional[int] = None, + presence_penalty: float = 0.0, + frequency_penalty: float = 0.0, + repeat_penalty: float = 1.1, + tfs_z: float = 1.0, + mirostat_mode: int = 0, + mirostat_tau: float = 5.0, + mirostat_eta: float = 0.1, + model: Optional[str] = None, + logits_processor: Optional[llama.LogitsProcessorList] = None, + grammar: Optional[llama.LlamaGrammar] = None, + **kwargs, # type: ignore +) -> Union[llama_types.ChatCompletion, Iterator[llama_types.ChatCompletionChunk]]: + SYSTEM_MESSAGE = """A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. The assistant calls functions with appropriate input when necessary""" + + def generate_type_definition( + param: Dict[str, llama_types.JsonType], indent_level: int, shared_defs + ) -> str: + indent = " " * indent_level + if "$ref" in param: + # Reference to a shared definition + ref_name = param["$ref"].split("/")[ + -1 + ] # Extract the type name from the reference + return ref_name + elif param.get("type") == "array": + items = param.get("items", {}) + item_type = generate_type_definition(items, indent_level + 1, shared_defs) + return f"Array<{item_type}>" + elif param.get("type") == "object": + properties = param.get("properties", {}) + nested_schema = "{\n" + for nested_param_name, nested_param in properties.items(): + nested_param_type = generate_type_definition( + nested_param, indent_level + 1, shared_defs + ) + nested_schema += ( + f"{indent} {nested_param_name}: {nested_param_type},\n" + ) + nested_schema += indent + "}" + return nested_schema + elif "enum" in param: + # Enum type + return " | ".join([f'"{enum_value}"' for enum_value in param["enum"]]) + else: + # Simple type + return param.get("type", "any") + + def generate_shared_definitions(shared_defs, indent_level: int) -> str: + indent = " " * indent_level + shared_definitions = "" + for def_name, def_properties in shared_defs.items(): + shared_definitions += f"{indent}type {def_name} = " + if def_properties.get("type") == "object": + shared_definitions += generate_type_definition( + def_properties, indent_level, shared_defs + ) + elif "enum" in def_properties: + # Enum type + shared_definitions += " | ".join( + [f'"{enum_value}"' for enum_value in def_properties["enum"]] + ) + shared_definitions += ";\n" + return shared_definitions + + def generate_schema_from_functions(functions, namespace="functions") -> str: + schema = ( + "// Supported function definitions that should be called when necessary.\n" + ) + schema += f"namespace {namespace} {{\n\n" + + # Generate shared definitions + shared_definitions = {} + for function in functions: + parameters = function.get("parameters", {}) + shared_definitions.update(parameters.get("$defs", {})) + + schema += generate_shared_definitions(shared_definitions, 1) + + for function in functions: + function_name = function["name"] + description = function.get("description", "") + parameters = function.get("parameters", {}) + required_params = parameters.get("required", []) + + schema += f" // {description}\n" + schema += f" type {function_name} = (_: {{\n" + + for param_name, param in parameters.get("properties", {}).items(): + param_description = param.get("description", "") + param_type = generate_type_definition(param, 2, shared_definitions) + optional_indicator = "" if param_name in required_params else "?" + schema += f" // {param_description}\n" + schema += f" {param_name}{optional_indicator}: {param_type},\n" + schema += " }) => any;\n\n" + + schema += "}} // namespace {}\n".format(namespace) + return schema + + def prepare_messages_for_inference( + messages: List[llama_types.ChatCompletionRequestMessage], + functions: Optional[List[llama_types.ChatCompletionFunctions]] = None, + tools: Optional[List[llama_types.ChatCompletionTool]] = None, + ): + all_messages: List[llama_types.ChatCompletionRequestMessage] = [] + if functions is not None: + all_messages.append( + llama_types.ChatCompletionRequestSystemMessage( + role="system", content=generate_schema_from_functions(functions) + ) + ) + + if tools is not None: + all_messages.append( + llama_types.ChatCompletionRequestSystemMessage( + role="system", + content=generate_schema_from_functions( + [ + tool["function"] + for tool in tools + if tool["type"] == "function" + ] + ), + ) + ) + + all_messages.append( + llama_types.ChatCompletionRequestSystemMessage( + role="system", content=SYSTEM_MESSAGE + ) + ) + + for message in messages: + # Function call responses + if message["role"] == "function" and "name" in message: + message["name"] = f"functions.{message['name']}" + # Function call requests by assistant + if "function_call" in message: + message["function_call"]["name"] = ( + f"functions.{message['function_call']['name']}" + ) + all_messages.append(message) + + all_messages.append( + llama_types.ChatCompletionRequestAssistantMessage( + role="assistant", content=None + ) + ) + + def message_to_str(msg: llama_types.ChatCompletionRequestMessage): + if msg["role"] == "system": + return f"system:\n{msg['content']}\n" + + elif msg["role"] == "function" and "name" in msg: + return f"function name={msg['name']}:\n{msg['content']}\n" + elif msg["role"] == "function" and "function_call" in msg: + return f"function name={msg['function_call']['name']}:\n{msg['function_call']['arguments']}\n" + elif msg["role"] == "tool": + if msg["content"] is not None: + return f"function name={msg['tool_call_id']}:\n{msg['content']}\n" + else: + return f"function name={msg['tool_call_id']}\n" + elif msg["role"] == "user": + if msg["content"] is None: + return "user:\n</s></s>\n" + else: + return f"user:\n</s>{msg['content']}</s>\n" + elif msg["role"] == "assistant": + if msg["content"] is not None and "function_call" in msg: + return f"assistant:\n{msg['content']}\nassistant to={msg['function_call']['name']}:\n{msg['function_call']['arguments']}</s>\n" + elif "function_call" in msg: + return f"assistant to={msg['function_call']['name']}:\n{msg['function_call']['arguments']}</s>\n" + elif "tool_calls" in msg and len(msg["tool_calls"]) > 0: + for tool_call in msg[ + "tool_calls" + ]: # NOTE: probably doesn't work with the functionary model + return f"assistant to={tool_call['id']}:\n{tool_call['function']['arguments']}</s>\n" + elif msg["content"] is None: + return "assistant" + else: + return f"assistant:\n{msg['content']}\n" + else: + raise ValueError(f"Unsupported role: {msg['role']}") + + return "".join([message_to_str(msg) for msg in all_messages]) + + if tools is not None: + functions = [tool["function"] for tool in tools if tool["type"] == "function"] + + if tool_choice is not None: + function_call = ( + tool_choice if isinstance(tool_choice, str) else tool_choice["function"] + ) + + prompt = prepare_messages_for_inference(messages, functions, tools) + + if function_call is None and (functions is None or len(functions) == 0): + completion_or_completion_chunks = llama.create_completion( + prompt=prompt + ":\n", + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + typical_p=typical_p, + stream=stream, + stop=["user:", "</s>"], + max_tokens=max_tokens, + presence_penalty=presence_penalty, + frequency_penalty=frequency_penalty, + repeat_penalty=repeat_penalty, + tfs_z=tfs_z, + mirostat_mode=mirostat_mode, + mirostat_tau=mirostat_tau, + mirostat_eta=mirostat_eta, + model=model, + logits_processor=logits_processor, + grammar=grammar, + ) + return _convert_completion_to_chat( + completion_or_completion_chunks, stream=stream + ) # type: ignore + + if function_call is None or ( + isinstance(function_call, str) and function_call == "auto" + ): + stop = "\n" + completion: llama_types.Completion = llama.create_completion( + prompt=prompt, stop=stop, stream=False + ) # type: ignore + completion_text = completion["choices"][0]["text"] + # strip " to=functions." and ending ":" + function_call = completion_text.split(".")[-1][:-1] + new_prompt = prompt + completion_text + stop + elif isinstance(function_call, str) and function_call != "none": + new_prompt = prompt + ":\n" + elif isinstance(function_call, dict): + new_prompt = prompt + f" to=functions.{function_call['name']}:\n" + function_call = function_call["name"] + else: + new_prompt = prompt + ":\n" + + function_body = None + for function in functions or []: + if function["name"] == function_call: + function_body = function["parameters"] + break + for tool in tools or []: + if tool["type"] == "function" and tool["function"]["name"] == function_call: + function_body = tool["function"]["parameters"] + break + + if function_body is not None: + try: + with suppress_stdout_stderr(disable=llama.verbose): + grammar_text = llama_grammar.json_schema_to_gbnf( + json.dumps(function_body) + ) + grammar = llama_grammar.LlamaGrammar.from_string( + llama_grammar.json_schema_to_gbnf(json.dumps(function_body)), + verbose=llama.verbose, + ) + print(grammar_text) + except Exception as e: + if llama.verbose: + print( + "Failed to parse function body as JSON schema, falling back to default grammar" + ) + print(e) + with suppress_stdout_stderr(disable=llama.verbose): + grammar = llama_grammar.LlamaGrammar.from_string( + llama_grammar.JSON_GBNF, + verbose=llama.verbose, + ) + else: + with suppress_stdout_stderr(disable=llama.verbose): + grammar = llama_grammar.LlamaGrammar.from_string( + llama_grammar.JSON_GBNF, verbose=llama.verbose + ) + + completion: llama_types.Completion = llama.create_completion( + prompt=new_prompt, + stop=["user:", "</s>"], + stream=False, + grammar=grammar, + max_tokens=max_tokens, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + typical_p=typical_p, + presence_penalty=presence_penalty, + frequency_penalty=frequency_penalty, + repeat_penalty=repeat_penalty, + tfs_z=tfs_z, + mirostat_mode=mirostat_mode, + mirostat_tau=mirostat_tau, + mirostat_eta=mirostat_eta, + model=model, + logits_processor=logits_processor, + ) # type: ignore + + assert "usage" in completion + assert isinstance(function_call, str) + assert stream is False # TODO: support stream mode + + if llama.verbose: + print(new_prompt) + print(completion["choices"][0]["text"]) + + # TODO: support stream mode + return llama_types.CreateChatCompletionResponse( + id="chat" + completion["id"], + object="chat.completion", + created=completion["created"], + model=completion["model"], + choices=[ + { + "index": 0, + "message": { + "role": "assistant", + "content": None, + "function_call": { + "name": function_call, + "arguments": completion["choices"][0]["text"], + }, + "tool_calls": [ + { + "id": function_call, + "type": "function", + "function": { + "name": function_call, + "arguments": completion["choices"][0]["text"], + }, + } + ], + }, + "logprobs": _convert_text_completion_logprobs_to_chat( + completion["choices"][0]["logprobs"] + ), + "finish_reason": "tool_calls", + } + ], + usage=completion["usage"], + ) + + +@register_chat_completion_handler("functionary-v1") +@register_chat_completion_handler("functionary-v2") +def functionary_v1_v2_chat_handler( + llama: llama.Llama, + messages: List[llama_types.ChatCompletionRequestMessage], + functions: Optional[List[llama_types.ChatCompletionFunction]] = None, + function_call: Optional[llama_types.ChatCompletionRequestFunctionCall] = None, + tools: Optional[List[llama_types.ChatCompletionTool]] = None, + tool_choice: Optional[llama_types.ChatCompletionToolChoiceOption] = None, + temperature: float = 0.2, + top_p: float = 0.95, + top_k: int = 40, + min_p: float = 0.05, + typical_p: float = 1.0, + stream: bool = False, + stop: Optional[Union[str, List[str]]] = [], + response_format: Optional[llama_types.ChatCompletionRequestResponseFormat] = None, + max_tokens: Optional[int] = None, + presence_penalty: float = 0.0, + frequency_penalty: float = 0.0, + repeat_penalty: float = 1.1, + tfs_z: float = 1.0, + mirostat_mode: int = 0, + mirostat_tau: float = 5.0, + mirostat_eta: float = 0.1, + model: Optional[str] = None, + logits_processor: Optional[llama.LogitsProcessorList] = None, + grammar: Optional[llama.LlamaGrammar] = None, + **kwargs, # type: ignore +) -> Union[llama_types.ChatCompletion, Iterator[llama_types.ChatCompletionChunk]]: + SYSTEM_MESSAGE = """A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. The assistant calls functions with appropriate input when necessary""" + + tokenizer = llama.tokenizer_ + assert hasattr(tokenizer, "hf_tokenizer"), ( + "Please provide a valid hf_tokenizer_path from https://huggingface.co/meetkai when initializing the Llama class" + ) + from transformers import AutoTokenizer + + if "<|START_OF_FUNCTION_CALL|>" in tokenizer.hf_tokenizer.additional_special_tokens: + version = "v1" + END_SYSTEM_TOKEN = "<|END_OF_SYSTEM|>" + END_USER_TOKEN = "<|END_OF_USER|>" + END_ASSISTANT_TOKEN = "<|END_OF_ASSISTANT|>" + END_FUNCTION_RESULT_TOKEN = "<|END_OF_FUNCTION_RESULT|>" + START_FUNCTION_CALL_TOKEN = "<|START_OF_FUNCTION_CALL|>" + END_FUNCTION_CALL_TOKEN = "<|END_OF_FUNCTION_CALL|>" + else: + version = "v2" + RECIPIENT_TOKEN = "<|recipient|>" + FROM_TOKEN = "<|from|>" + STOP_TOKEN = "<|stop|>" + CONTENT_TOKEN = "<|content|>" + + def generate_type_definition( + param: Dict[str, llama_types.JsonType], indent_level: int, shared_defs + ) -> str: + indent = " " * indent_level + if "$ref" in param: + # Reference to a shared definition + ref_name = param["$ref"].split("/")[ + -1 + ] # Extract the type name from the reference + return ref_name + elif param.get("type") == "array": + items = param.get("items", {}) + item_type = generate_type_definition(items, indent_level + 1, shared_defs) + return f"Array<{item_type}>" + elif param.get("type") == "object": + properties = param.get("properties", {}) + nested_schema = "{\n" + for nested_param_name, nested_param in properties.items(): + nested_param_type = generate_type_definition( + nested_param, indent_level + 1, shared_defs + ) + nested_schema += ( + f"{indent} {nested_param_name}: {nested_param_type},\n" + ) + nested_schema += indent + "}" + return nested_schema + elif "enum" in param: + # Enum type + return " | ".join([f'"{enum_value}"' for enum_value in param["enum"]]) + else: + # Simple type + return param.get("type", "any") + + def generate_shared_definitions(shared_defs, indent_level: int) -> str: + indent = " " * indent_level + shared_definitions = "" + for def_name, def_properties in shared_defs.items(): + shared_definitions += f"{indent}type {def_name} = " + if def_properties.get("type") == "object": + shared_definitions += generate_type_definition( + def_properties, indent_level, shared_defs + ) + elif "enum" in def_properties: + # Enum type + shared_definitions += " | ".join( + [f'"{enum_value}"' for enum_value in def_properties["enum"]] + ) + shared_definitions += ";\n" + return shared_definitions + + def generate_schema_from_functions(functions, namespace="functions") -> str: + schema = ( + "// Supported function definitions that should be called when necessary.\n" + ) + schema += f"namespace {namespace} {{\n\n" + + # Generate shared definitions + shared_definitions = {} + for function in functions: + parameters = function.get("parameters", {}) + shared_definitions.update(parameters.get("$defs", {})) + + schema += generate_shared_definitions(shared_definitions, 1) + + for function in functions: + function_name = function["name"] + description = function.get("description", "") + parameters = function.get("parameters", {}) + required_params = parameters.get("required", []) + + schema += f"// {description}\n" + schema += f"type {function_name} = (_: {{\n" + + for param_name, param in parameters.get("properties", {}).items(): + param_description = param.get("description", "") + param_type = generate_type_definition(param, 2, shared_definitions) + optional_indicator = "" if param_name in required_params else "?" + schema += f"// {param_description}\n" + schema += f"{param_name}{optional_indicator}: {param_type},\n" + schema += "}) => any;\n\n" + + schema += "}} // namespace {}".format(namespace) + return schema + + def prepare_messages_for_inference( + messages: List[llama_types.ChatCompletionRequestMessage], + tokenizer: AutoTokenizer, + version: Literal["v1", "v2"], + functions: Optional[List[llama_types.ChatCompletionFunctions]] = None, + tools: Optional[List[llama_types.ChatCompletionTool]] = None, + tool_choice: Union[Dict, str] = "auto", + ): + all_messages: List[llama_types.ChatCompletionRequestMessage] = [] + if tool_choice == "none": + all_messages.append( + llama_types.ChatCompletionRequestSystemMessage( + role="system", content=generate_schema_from_functions([]) + ) + ) + else: + if functions is not None: + all_messages.append( + llama_types.ChatCompletionRequestSystemMessage( + role="system", content=generate_schema_from_functions(functions) + ) + ) + elif tools is not None and tool_choice != "none": + all_messages.append( + llama_types.ChatCompletionRequestSystemMessage( + role="system", + content=generate_schema_from_functions( + [ + tool["function"] + for tool in tools + if tool["type"] == "function" + ] + ), + ) + ) + + all_messages.append( + llama_types.ChatCompletionRequestSystemMessage( + role="system", content=SYSTEM_MESSAGE + ) + ) + + for message in messages: + # Function call responses + if message["role"] == "function" and "name" in message: + message["name"] = f"functions.{message['name']}" + # Function call requests by assistant + if "function_call" in message: + message["function_call"]["name"] = ( + f"functions.{message['function_call']['name']}" + ) + all_messages.append(message) + + if version == "v1": + suffix = "assistant:\n" + else: + suffix = "<|from|>assistant\n<|recipient|>" + + return ( + tokenizer.hf_tokenizer.apply_chat_template(all_messages, tokenize=False) + + suffix + ) + + if tools is not None: + functions = [tool["function"] for tool in tools if tool["type"] == "function"] + + if tool_choice is not None: + function_call = ( + tool_choice if isinstance(tool_choice, str) else tool_choice["function"] + ) + elif function_call is not None: + pass + else: + function_call = "auto" + + prompt = prepare_messages_for_inference( + messages, tokenizer, version, functions, tools, function_call + ) + + # If no tools/functions are provided + if function_call == "none" or functions is None or len(functions) == 0: + if version == "v1": + stop = END_ASSISTANT_TOKEN + else: + stop = STOP_TOKEN + prompt += "all\n<|content|>" + + completion_or_completion_chunks = llama.create_completion( + prompt=prompt, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + typical_p=typical_p, + stream=stream, + stop=stop, + max_tokens=max_tokens, + presence_penalty=presence_penalty, + frequency_penalty=frequency_penalty, + repeat_penalty=repeat_penalty, + tfs_z=tfs_z, + mirostat_mode=mirostat_mode, + mirostat_tau=mirostat_tau, + mirostat_eta=mirostat_eta, + model=model, + logits_processor=logits_processor, + grammar=grammar, + ) + if stream is False: + completion_or_completion_chunks["choices"][0]["text"] = ( + completion_or_completion_chunks["choices"][0]["text"].lstrip() + ) + return _convert_completion_to_chat( + completion_or_completion_chunks, stream=stream + ) # type: ignore + + def get_grammar(function_call): + function_body = None + for function in functions or []: + if function["name"] == function_call: + function_body = function["parameters"] + break + for tool in tools or []: + if tool["type"] == "function" and tool["function"]["name"] == function_call: + function_body = tool["function"]["parameters"] + break + + try: + with suppress_stdout_stderr(disable=llama.verbose): + grammar_text = llama_grammar.json_schema_to_gbnf( + json.dumps(function_body) + ) + grammar = llama_grammar.LlamaGrammar.from_string( + llama_grammar.json_schema_to_gbnf(json.dumps(function_body)) + ) + print(grammar_text) + except Exception as e: + if llama.verbose: + print( + "Failed to parse function body as JSON schema, falling back to default grammar" + ) + print(e) + with suppress_stdout_stderr(disable=llama.verbose): + grammar = llama_grammar.LlamaGrammar.from_string( + llama_grammar.JSON_GBNF, verbose=llama.verbose + ) + + return grammar + + def create_completion(prompt, stop, grammar): + completion = cast( + llama_types.Completion, + llama.create_completion( + prompt=prompt, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + typical_p=typical_p, + stream=stream, + stop=stop, + max_tokens=max_tokens, + presence_penalty=presence_penalty, + frequency_penalty=frequency_penalty, + repeat_penalty=repeat_penalty, + tfs_z=tfs_z, + mirostat_mode=mirostat_mode, + mirostat_tau=mirostat_tau, + mirostat_eta=mirostat_eta, + model=model, + logits_processor=logits_processor, + grammar=grammar, + ), + ) + + return completion + + content = "" + function_calls, function_bodies = [], [] + completion_tokens = 0 + + def generate_streaming(tools, functions, function_call, prompt): + assert version == "v2", "Streaming for v1 is not supported" + + chunk_id, chunk_created = None, None + + # If tool_choice/function_call is provided + if isinstance(function_call, dict): + prompt += f"{function_call['name']}\n{CONTENT_TOKEN}" + grammar = get_grammar(function_call["name"]) + stops = [STOP_TOKEN, FROM_TOKEN] + tool_id = "".join( + [random.choice(string.ascii_letters + string.digits) for _ in range(24)] + ) + completion = create_completion(prompt=prompt, stop=stops, grammar=grammar) + completion_text = "" + first = True + for chunk in completion: + # Yield the tool/function name first + if first: + if tools is not None: + func_call_dict = { + "tool_calls": [ + { + "index": 0, + "id": "call_" + tool_id, + "type": "function", + "function": { + "name": function_call["name"], + "arguments": "", + }, + } + ] + } + else: + func_call_dict = { + "function_call": { + "name": function_call["name"], + "arguments": "", + } + } + yield llama_types.CreateChatCompletionStreamResponse( + id="chat" + chunk["id"], + object="chat.completion.chunk", + created=chunk["created"], + model=chunk["model"], + choices=[ + { + "index": 0, + "logprobs": None, + "delta": { + "role": None, + "content": None, + **func_call_dict, + }, + } + ], + ) + first = False + if tools is not None: + func_call_dict = { + "tool_calls": [ + { + "index": 0, + "id": "call_" + tool_id, + "type": "function", + "function": { + "name": None, + "arguments": chunk["choices"][0]["text"].rstrip(), + }, + } + ] + } + else: + func_call_dict = { + "function_call": { + "name": None, + "arguments": chunk["choices"][0]["text"].rstrip(), + } + } + if len(chunk["choices"][0]["text"].rstrip()) > 0: + yield llama_types.CreateChatCompletionStreamResponse( + id="chat" + chunk["id"], + object="chat.completion.chunk", + created=chunk["created"], + model=chunk["model"], + choices=[ + { + "index": 0, + "logprobs": _convert_text_completion_logprobs_to_chat( + chunk["choices"][0]["logprobs"] + ), + "delta": { + "role": None, + "content": None, + **func_call_dict, + }, + } + ], + ) + # Yield tool_call/function_call stop message + yield llama_types.CreateChatCompletionStreamResponse( + id="chat" + chunk["id"], + object="chat.completion.chunk", + created=chunk["created"], + model=chunk["model"], + choices=[ + { + "index": 0, + "finish_reason": ( + "tool_calls" if tools is not None else "function_call" + ), + "logprobs": None, + "delta": { + "role": None, + "content": None, + "function_call": None, + "tool_calls": None, + }, + } + ], + ) + # If "auto" or no tool_choice/function_call + elif isinstance(function_call, str) and function_call == "auto": + tool_index = 0 + while True: + # Generate function name first + grammar = None + stops = CONTENT_TOKEN + completion = create_completion( + prompt=prompt, stop=stops, grammar=grammar + ) + completion_text = "" + for chunk in completion: + completion_text += chunk["choices"][0]["text"] + if chunk_id is None: + chunk_id = chunk["id"] + if chunk_created is None: + chunk_created = chunk["created"] + function_name = completion_text.strip() + if function_name == "all": + prompt += "all\n<|content|>" + # Yield the first empty message for content + yield llama_types.CreateChatCompletionStreamResponse( + id="chat" + chunk_id, + model=chunk["model"], + created=chunk_created, + object="chat.completion.chunk", + choices=[ + { + "index": 0, + "delta": {"role": "assistant", "content": ""}, + "logprobs": None, + "finish_reason": None, + } + ], + ) + else: + prompt += f"{function_name}\n<|content|>" + grammar = get_grammar(function_name) + tool_id = "".join( + [ + random.choice(string.ascii_letters + string.digits) + for _ in range(24) + ] + ) + if tools is not None: + func_call_dict = { + "tool_calls": [ + { + "index": tool_index, + "id": "call_" + tool_id, + "type": "function", + "function": { + "name": function_name, + "arguments": "", + }, + } + ] + } + else: + func_call_dict = { + "function_call": {"name": function_name, "arguments": ""} + } + # Stream function name + yield llama_types.CreateChatCompletionStreamResponse( + id="chat" + chunk_id, + object="chat.completion.chunk", + created=chunk_created, + model=chunk["model"], + choices=[ + { + "index": 0, + "logprobs": _convert_text_completion_logprobs_to_chat( + chunk["choices"][0]["logprobs"] + ), + "delta": { + "role": "assistant", + "content": None, + **func_call_dict, + }, + } + ], + ) + # Generate content + stops = [RECIPIENT_TOKEN, STOP_TOKEN] + completion = create_completion( + prompt=prompt, stop=stops, grammar=grammar + ) + if function_name == "all": + completion_text = "" + stop_sequence, buffer, is_end = ( + "\n<|from|>assistant\n<|recipient|>", + [], + False, + ) + for i, chunk in enumerate(completion): + completion_text += chunk["choices"][0]["text"] + if is_end: + buffer.append(chunk["choices"][0]["text"].strip(" ")) + if stop_sequence.startswith("".join(buffer)): + continue + else: + buffer.pop() + while len(buffer) > 0: + yield llama_types.CreateChatCompletionStreamResponse( + id="chat" + chunk_id, + object="chat.completion.chunk", + created=chunk_created, + model=chunk["model"], + choices=[ + { + "index": 0, + "logprobs": _convert_text_completion_logprobs_to_chat( + chunk["choices"][0]["logprobs"] + ), + "delta": { + "role": "assistant", + "content": buffer.pop(0), + }, + } + ], + ) + is_end = False + elif chunk["choices"][0]["text"] == "\n": + is_end = True + buffer.append(chunk["choices"][0]["text"].strip(" ")) + continue + + if len(buffer) == 0 and len(chunk["choices"][0]["text"]) > 0: + yield llama_types.CreateChatCompletionStreamResponse( + id="chat" + chunk_id, + object="chat.completion.chunk", + created=chunk_created, + model=chunk["model"], + choices=[ + { + "index": 0, + "logprobs": _convert_text_completion_logprobs_to_chat( + chunk["choices"][0]["logprobs"] + ), + "delta": { + "role": "assistant", + "content": ( + chunk["choices"][0]["text"] + if i > 0 + else chunk["choices"][0][ + "text" + ].lstrip() + ), + }, + } + ], + ) + # Check whether the model wants to generate another turn + if ( + "<|from|> assistant" in completion_text + or "<|from|>assistant" in completion_text + ): + if completion_text.endswith("\n<|from|>assistant\n"): + cleaned_completion_text = completion_text[ + : -len("\n<|from|>assistant\n") + ].strip() + elif completion_text.endswith("\n<|from|> assistant\n"): + cleaned_completion_text = completion_text[ + : -len("\n<|from|> assistant\n") + ].strip() + else: + cleaned_completion_text = completion_text.strip() + prompt += f"{cleaned_completion_text}\n<|from|>assistant\n<|recipient|>" + else: + # Yield stop message + yield llama_types.CreateChatCompletionStreamResponse( + id="chat" + chunk_id, + model=chunk["model"], + created=chunk_created, + object="chat.completion.chunk", + choices=[ + { + "index": 0, + "delta": {}, + "logprobs": None, + "finish_reason": "stop", + } + ], + ) + break + else: + # Check whether the model wants to generate another turn + completion_text = "" + for chunk in completion: + completion_text += chunk["choices"][0]["text"] + if len(chunk["choices"][0]["text"].rstrip()) > 0: + if tools is not None: + func_call_dict = { + "tool_calls": [ + { + "index": tool_index, + "id": "call_" + tool_id, + "type": "function", + "function": { + "name": None, + "arguments": chunk["choices"][0][ + "text" + ].rstrip(), + }, + } + ] + } + else: + func_call_dict = { + "function_call": { + "name": None, + "arguments": chunk["choices"][0][ + "text" + ].rstrip(), + } + } + yield llama_types.CreateChatCompletionStreamResponse( + id="chat" + chunk_id, + object="chat.completion.chunk", + created=chunk_created, + model=chunk["model"], + choices=[ + { + "index": 0, + "logprobs": _convert_text_completion_logprobs_to_chat( + chunk["choices"][0]["logprobs"] + ), + "delta": { + "role": None, + "content": None, + **func_call_dict, + }, + } + ], + ) + prompt += completion_text.strip() + grammar = None + completion = create_completion( + prompt=prompt, stop=stops, grammar=grammar + ) + completion_text += "".join( + [chunk["choices"][0]["text"] for chunk in completion] + ) + if ( + "<|from|> assistant" in completion_text + or "<|from|>assistant" in completion_text + ) and tools is not None: + prompt += "\n<|from|>assistant\n<|recipient|>" + tool_index += 1 + else: + # Yield tool_call/function_call stop message + yield llama_types.CreateChatCompletionStreamResponse( + id="chat" + chunk_id, + object="chat.completion.chunk", + created=chunk_created, + model=chunk["model"], + choices=[ + { + "index": 0, + "finish_reason": ( + "tool_calls" + if tools is not None + else "function_call" + ), + "logprobs": None, + "delta": { + "role": None, + "content": None, + "function_call": None, + "tool_calls": None, + }, + } + ], + ) + break + + if stream is not False: + return generate_streaming( + tools=tools, functions=functions, function_call=function_call, prompt=prompt + ) + else: + if version == "v1": + # If no or "auto" tool_choice/function_call + if isinstance(function_call, str) and function_call == "auto": + stops = ["\n", END_ASSISTANT_TOKEN] + # If tool_choice/function_call is provided + elif isinstance(function_call, dict): + prompt += f"{START_FUNCTION_CALL_TOKEN}{function_call['name']}:\n" + stops = END_FUNCTION_CALL_TOKEN + function_call = function_call["name"] + function_calls.append(function_call) + grammar = get_grammar(function_call) + else: + prompt = prompt + stops = ["\n", END_ASSISTANT_TOKEN] + + completion = create_completion(prompt=prompt, stop=stops, grammar=grammar) + completion_text = completion["choices"][0]["text"] + completion_tokens += completion["usage"]["completion_tokens"] + + # If the generation does not involve a function call + if ( + START_FUNCTION_CALL_TOKEN not in prompt + and START_FUNCTION_CALL_TOKEN not in completion_text + ): + completion["usage"]["completion_tokens"] = completion_tokens + return _convert_completion_to_chat(completion, stream=stream) # type: ignore + # If the generation involves a function call in completion, generate the parameters + elif ( + START_FUNCTION_CALL_TOKEN not in prompt + and START_FUNCTION_CALL_TOKEN in completion_text + ): + prompt += ( + completion_text.replace( + f"{START_FUNCTION_CALL_TOKEN} ", START_FUNCTION_CALL_TOKEN + ) + + "\n" + ) + function_calls.append( + completion_text.split(START_FUNCTION_CALL_TOKEN)[-1][:-1].strip() + ) + grammar = get_grammar(function_calls[-1]) + completion = create_completion( + prompt=prompt, stop=END_FUNCTION_CALL_TOKEN, grammar=grammar + ) + completion_tokens += completion["usage"]["completion_tokens"] + function_bodies.append(completion["choices"][0]["text"].strip()) + # If the prompt involves a function call, just append generated parameters to function_bodies + else: + function_bodies.append(completion_text.strip()) + else: + # If tool_choice/function_call is provided + if isinstance(function_call, dict): + prompt += f"{function_call['name']}\n{CONTENT_TOKEN}" + function_call = function_call["name"] + function_calls.append(function_call) + grammar = get_grammar(function_call) + stops = [STOP_TOKEN, FROM_TOKEN] + completion = create_completion( + prompt=prompt, stop=stops, grammar=grammar + ) + completion_text = completion["choices"][0]["text"] + completion_tokens += completion["usage"]["completion_tokens"] + function_bodies.append(completion_text.strip()) + # If "auto" or no tool_choice/function_call + elif isinstance(function_call, str) and function_call == "auto": + while True: + # Generate function name first + grammar = None + stops = CONTENT_TOKEN + completion = create_completion( + prompt=prompt, stop=stops, grammar=grammar + ) + completion_text = completion["choices"][0]["text"] + completion_tokens += completion["usage"]["completion_tokens"] + function_name = completion_text.strip() + if function_name == "all": + prompt += "all\n<|content|>" + else: + function_call = completion_text.strip() + prompt += f"{function_call}\n<|content|>" + function_calls.append(function_call) + grammar = get_grammar(function_call) + # Generate content + stops = [RECIPIENT_TOKEN, STOP_TOKEN] + completion = create_completion( + prompt=prompt, stop=stops, grammar=grammar + ) + completion_text = completion["choices"][0]["text"] + completion_tokens += completion["usage"]["completion_tokens"] + if function_name == "all": + if completion_text.endswith("\n<|from|>assistant\n"): + content += completion_text[: -len("\n<|from|>assistant\n")] + if completion_text.endswith("\n<|from|> assistant\n"): + content += completion_text[-len("\n<|from|> assistant\n")] + else: + content += completion_text + content = content.lstrip() + # Check whether the model wants to generate another turn + if ( + "<|from|> assistant" in completion_text + or "<|from|>assistant" in completion_text + ): + if completion_text.endswith("\n<|from|>assistant\n"): + cleaned_completion_text = completion_text[ + : -len("\n<|from|>assistant\n") + ].strip() + elif completion_text.endswith("\n<|from|> assistant\n"): + cleaned_completion_text = completion_text[ + -len("\n<|from|> assistant\n") + ].strip() + else: + cleaned_completion_text = completion_text.strip() + prompt += f"{cleaned_completion_text}\n<|from|>assistant\n<|recipient|>" + else: + break + else: + function_bodies.append(completion_text.strip()) + # Check whether the model wants to generate another turn + prompt += completion_text.strip() + grammar = None + completion = create_completion( + prompt=prompt, stop=stops, grammar=grammar + ) + completion_tokens += completion["usage"]["completion_tokens"] + if ( + "<|from|> assistant" in completion["choices"][0]["text"] + or "<|from|>assistant" in completion["choices"][0]["text"] + ): + prompt += "\n<|from|>assistant\n<|recipient|>" + else: + break + + assert "usage" in completion + assert len(function_calls) == len(function_bodies) + + tool_calls: List[llama_types.ChatCompletionMessageToolCall] = [] + for function_call, function_body in zip(function_calls, function_bodies): + tool_calls.append( + { + "id": "call_" + + "".join( + [ + random.choice(string.ascii_letters + string.digits) + for _ in range(24) + ] + ), + "type": "function", + "function": { + "name": function_call, + "arguments": function_body, + }, + } + ) + + # TODO: support stream mode + function_call_dict: Union[ + Dict[str, str], + Dict[ + Literal["function_call"], + llama_types.ChatCompletionRequestAssistantMessageFunctionCall, + ], + ] = {} + if len(tool_calls) > 0: + if tools is not None: + function_call_dict["tool_calls"] = tool_calls + else: + function_call_dict["function_call"] = { + "name": tool_calls[0]["function"]["name"], + "arguments": tool_calls[0]["function"]["arguments"], + } + completion["usage"]["completion_tokens"] = completion_tokens + return llama_types.CreateChatCompletionResponse( + id="chat" + completion["id"], + object="chat.completion", + created=completion["created"], + model=completion["model"], + choices=[ + { + "index": 0, + "logprobs": _convert_text_completion_logprobs_to_chat( + completion["choices"][0]["logprobs"] + ), + "message": { + "role": "assistant", + "content": None if content == "" else content, + **function_call_dict, + }, + "finish_reason": "tool_calls" if len(tool_calls) > 0 else "stop", + } + ], + usage=completion["usage"], + ) + + +class Llava15ChatHandler: + DEFAULT_SYSTEM_MESSAGE: Optional[str] = ( + "A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions." + ) + + CHAT_FORMAT = ( + "{% for message in messages %}" + "{% if message.role == 'system' %}" + "{{ message.content }}" + "{% endif %}" + "{% if message.role == 'user' %}" + "{% if message.content is string %}" + "\nUSER: {{ message.content }}" + "{% endif %}" + "{% if message.content is iterable %}" + "\nUSER: " + "{% for content in message.content %}" + "{% if content.type == 'image_url' and content.image_url is string %}" + "{{ content.image_url }}" + "{% endif %}" + "{% if content.type == 'image_url' and content.image_url is mapping %}" + "{{ content.image_url.url }}" + "{% endif %}" + "{% endfor %}" + "{% for content in message.content %}" + "{% if content.type == 'text' %}" + "{{ content.text }}" + "{% endif %}" + "{% endfor %}" + "{% endif %}" + "{% endif %}" + "{% if message.role == 'assistant' and message.content is not none %}" + "\nASSISTANT: {{ message.content }}" + "{% endif %}" + "{% endfor %}" + "{% if add_generation_prompt %}" + "\nASSISTANT: " + "{% endif %}" + ) + + def __init__(self, clip_model_path: str, verbose: bool = True): + import llama_cpp.mtmd_cpp as mtmd_cpp + + self.clip_model_path = clip_model_path + self.verbose = verbose + self._mtmd_cpp = mtmd_cpp + self._exit_stack = ExitStack() + self.mtmd_ctx: Optional[mtmd_cpp.mtmd_context_p] = None + + if not os.path.exists(clip_model_path): + raise ValueError(f"Clip model path does not exist: {clip_model_path}") + + def _init_mtmd_context(self, llama_model: llama.Llama): + """Initialize mtmd context with the llama model.""" + if self.mtmd_ctx is not None: + return # Already initialized + + with suppress_stdout_stderr(disable=self.verbose): + # Get default parameters + ctx_params = self._mtmd_cpp.mtmd_context_params_default() + ctx_params.use_gpu = True # TODO: Make this configurable + ctx_params.print_timings = self.verbose + ctx_params.n_threads = llama_model.n_threads + ctx_params.flash_attn_type = ( + llama_cpp.LLAMA_FLASH_ATTN_TYPE_ENABLED + if ( + llama_model.context_params.flash_attn_type + == llama_cpp.LLAMA_FLASH_ATTN_TYPE_ENABLED + ) + else llama_cpp.LLAMA_FLASH_ATTN_TYPE_DISABLED + ) + + # Initialize mtmd context + self.mtmd_ctx = self._mtmd_cpp.mtmd_init_from_file( + self.clip_model_path.encode(), llama_model.model, ctx_params + ) + + if self.mtmd_ctx is None: + raise ValueError( + f"Failed to load mtmd context from: {self.clip_model_path}" + ) + + # Check if vision is supported + if not self._mtmd_cpp.mtmd_support_vision(self.mtmd_ctx): + raise ValueError("Vision is not supported by this model") + + def mtmd_free(): + with suppress_stdout_stderr(disable=self.verbose): + if self.mtmd_ctx is not None: + self._mtmd_cpp.mtmd_free(self.mtmd_ctx) + self.mtmd_ctx = None + + self._exit_stack.callback(mtmd_free) + + def load_image(self, image_url: str) -> bytes: + return self._load_image(image_url) + + def _create_bitmap_from_bytes(self, image_bytes: bytes): + """Create mtmd_bitmap from image bytes.""" + if self.mtmd_ctx is None: + raise ValueError("mtmd context not initialized") + + with suppress_stdout_stderr(disable=self.verbose): + # Create bitmap from buffer using helper function + bitmap = self._mtmd_cpp.mtmd_helper_bitmap_init_from_buf( + self.mtmd_ctx, + (ctypes.c_uint8 * len(image_bytes)).from_buffer(bytearray(image_bytes)), + len(image_bytes), + False, + ) + + if bitmap is None: + raise ValueError("Failed to create bitmap from image bytes") + + return bitmap + + def __call__( + self, + *, + llama: llama.Llama, + messages: List[llama_types.ChatCompletionRequestMessage], + functions: Optional[List[llama_types.ChatCompletionFunction]] = None, + function_call: Optional[llama_types.ChatCompletionRequestFunctionCall] = None, + tools: Optional[List[llama_types.ChatCompletionTool]] = None, + tool_choice: Optional[llama_types.ChatCompletionToolChoiceOption] = None, + temperature: float = 0.2, + top_p: float = 0.95, + top_k: int = 40, + min_p: float = 0.05, + typical_p: float = 1.0, + stream: bool = False, + stop: Optional[Union[str, List[str]]] = [], + seed: Optional[int] = None, + response_format: Optional[ + llama_types.ChatCompletionRequestResponseFormat + ] = None, + max_tokens: Optional[int] = None, + presence_penalty: float = 0.0, + frequency_penalty: float = 0.0, + repeat_penalty: float = 1.1, + tfs_z: float = 1.0, + mirostat_mode: int = 0, + mirostat_tau: float = 5.0, + mirostat_eta: float = 0.1, + model: Optional[str] = None, + logits_processor: Optional[llama.LogitsProcessorList] = None, + grammar: Optional[llama.LlamaGrammar] = None, + logit_bias: Optional[Dict[str, float]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + **kwargs, # type: ignore + ) -> Union[ + llama_types.CreateChatCompletionResponse, + Iterator[llama_types.CreateChatCompletionStreamResponse], + ]: + # Initialize mtmd context + self._init_mtmd_context(llama) + assert self.mtmd_ctx is not None + + system_prompt = _get_system_message(messages) + if system_prompt == "" and self.DEFAULT_SYSTEM_MESSAGE is not None: + messages = [ + llama_types.ChatCompletionRequestSystemMessage( + role="system", content=self.DEFAULT_SYSTEM_MESSAGE + ) + ] + messages + + image_urls = self.get_image_urls(messages) + template = ImmutableSandboxedEnvironment( + trim_blocks=True, + lstrip_blocks=True, + ).from_string(self.CHAT_FORMAT) + + # Get the default media marker + media_marker = self._mtmd_cpp.mtmd_default_marker().decode("utf-8") + + # Replace image URLs with media markers in the template + text = template.render( + messages=messages, + add_generation_prompt=True, + eos_token=llama.detokenize([llama.token_eos()]), + bos_token=llama.detokenize([llama.token_bos()]), + ) + + # Replace image URLs in text with media markers + for image_url in image_urls: + text = text.replace(image_url, media_marker) + + if self.verbose: + print(text, file=sys.stderr) + + # Create bitmaps from images + bitmaps = [] + bitmap_cleanup = [] + try: + for image_url in image_urls: + image_bytes = self.load_image(image_url) + bitmap = self._create_bitmap_from_bytes(image_bytes) + bitmaps.append(bitmap) + bitmap_cleanup.append(bitmap) + + # Create input text structure + input_text = self._mtmd_cpp.mtmd_input_text() + input_text.text = text.encode("utf-8") + input_text.add_special = True + input_text.parse_special = True + + # Create input chunks + chunks = self._mtmd_cpp.mtmd_input_chunks_init() + if chunks is None: + raise ValueError("Failed to create input chunks") + + try: + # Tokenize text and images together + bitmap_array = (self._mtmd_cpp.mtmd_bitmap_p_ctypes * len(bitmaps))( + *bitmaps + ) + result = self._mtmd_cpp.mtmd_tokenize( + self.mtmd_ctx, + chunks, + ctypes.byref(input_text), + bitmap_array, + len(bitmaps), + ) + + if result != 0: + raise ValueError(f"Failed to tokenize input: error code {result}") + + # Reset llama context + llama.reset() + llama._ctx.kv_cache_clear() + + # Process each chunk + n_past = llama_cpp.llama_pos(0) + n_chunks = self._mtmd_cpp.mtmd_input_chunks_size(chunks) + + for i in range(n_chunks): + chunk = self._mtmd_cpp.mtmd_input_chunks_get(chunks, i) + if chunk is None: + continue + + chunk_type = self._mtmd_cpp.mtmd_input_chunk_get_type(chunk) + + if chunk_type == self._mtmd_cpp.MTMD_INPUT_CHUNK_TYPE_TEXT: + # Handle text chunk + n_tokens_out = ctypes.c_size_t() + tokens_ptr = self._mtmd_cpp.mtmd_input_chunk_get_tokens_text( + chunk, ctypes.byref(n_tokens_out) + ) + + if tokens_ptr and n_tokens_out.value > 0: + # Convert ctypes array to Python list + tokens = [tokens_ptr[j] for j in range(n_tokens_out.value)] + + if llama.n_tokens + len(tokens) > llama.n_ctx(): + raise ValueError( + f"Prompt exceeds n_ctx: {llama.n_tokens + len(tokens)} > {llama.n_ctx()}" + ) + llama.eval(tokens) + + elif chunk_type in [ + self._mtmd_cpp.MTMD_INPUT_CHUNK_TYPE_IMAGE, + self._mtmd_cpp.MTMD_INPUT_CHUNK_TYPE_AUDIO, + ]: + # Handle image/audio chunk using helper + chunk_n_tokens = self._mtmd_cpp.mtmd_input_chunk_get_n_tokens( + chunk + ) + + if llama.n_tokens + chunk_n_tokens > llama.n_ctx(): + raise ValueError( + f"Prompt exceeds n_ctx: {llama.n_tokens + chunk_n_tokens} > {llama.n_ctx()}" + ) + + new_n_past = llama_cpp.llama_pos(0) + result = self._mtmd_cpp.mtmd_helper_eval_chunk_single( + self.mtmd_ctx, + llama._ctx.ctx, + chunk, + llama_cpp.llama_pos(llama.n_tokens), + llama_cpp.llama_seq_id(0), + llama.n_batch, + False, # logits_last + ctypes.byref(new_n_past), + ) + + if result != 0: + raise ValueError( + f"Failed to evaluate chunk: error code {result}" + ) + + # Update llama's token count + llama.n_tokens = new_n_past.value + + # Get prompt tokens to avoid a cache miss + prompt = llama.input_ids[: llama.n_tokens].tolist() + + finally: + self._mtmd_cpp.mtmd_input_chunks_free(chunks) + + finally: + # Cleanup bitmaps + for bitmap in bitmap_cleanup: + self._mtmd_cpp.mtmd_bitmap_free(bitmap) + + # Handle response format and tools (same as before) + if response_format is not None and response_format["type"] == "json_object": + grammar = _grammar_for_response_format(response_format) + + # Convert legacy functions to tools + if functions is not None: + tools = [ + { + "type": "function", + "function": function, + } + for function in functions + ] + + # Convert legacy function_call to tool_choice + if function_call is not None: + if isinstance(function_call, str) and ( + function_call == "none" or function_call == "auto" + ): + tool_choice = function_call + if isinstance(function_call, dict) and "name" in function_call: + tool_choice = { + "type": "function", + "function": { + "name": function_call["name"], + }, + } + + tool = None + if ( + tool_choice is not None + and isinstance(tool_choice, dict) + and tools is not None + ): + name = tool_choice["function"]["name"] + tool = next((t for t in tools if t["function"]["name"] == name), None) + if tool is None: + raise ValueError(f"Tool choice '{name}' not found in tools.") + schema = tool["function"]["parameters"] + try: + # create grammar from json schema + grammar = llama_grammar.LlamaGrammar.from_json_schema( + json.dumps(schema), verbose=llama.verbose + ) + except Exception as e: + if llama.verbose: + print(str(e), file=sys.stderr) + grammar = llama_grammar.LlamaGrammar.from_string( + llama_grammar.JSON_GBNF, verbose=llama.verbose + ) + + completion_or_chunks = llama.create_completion( + prompt=prompt, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + typical_p=typical_p, + logprobs=top_logprobs if logprobs else None, + stream=stream, + stop=stop, + seed=seed, + max_tokens=max_tokens, + presence_penalty=presence_penalty, + frequency_penalty=frequency_penalty, + repeat_penalty=repeat_penalty, + tfs_z=tfs_z, + mirostat_mode=mirostat_mode, + mirostat_tau=mirostat_tau, + mirostat_eta=mirostat_eta, + model=model, + logits_processor=logits_processor, + grammar=grammar, + logit_bias=logit_bias, + ) + + if tool is not None: + tool_name = tool["function"]["name"] + return _convert_completion_to_chat_function( + tool_name, completion_or_chunks, stream + ) + return _convert_completion_to_chat(completion_or_chunks, stream=stream) + + @staticmethod + def _load_image(image_url: str) -> bytes: + # TODO: Add Pillow support for other image formats beyond (jpg, png) + if image_url.startswith("data:"): + import base64 + + image_bytes = base64.b64decode(image_url.split(",")[1]) + return image_bytes + else: + import urllib.request + + with urllib.request.urlopen(image_url) as f: + image_bytes = f.read() + return image_bytes + + @staticmethod + def get_image_urls(messages: List[llama_types.ChatCompletionRequestMessage]): + image_urls: List[str] = [] + for message in messages: + if message["role"] == "user": + if message["content"] is None: + continue + for content in message["content"]: + if isinstance(content, dict) and "type" in content: + if content["type"] == "image_url": + if ( + isinstance(content["image_url"], dict) + and "url" in content["image_url"] + ): + image_urls.append(content["image_url"]["url"]) + else: + image_urls.append(content["image_url"]) + return image_urls + + @staticmethod + def split_text_on_image_urls(text: str, image_urls: List[str]): + """This method is no longer used in the new implementation.""" + + def find_first(s: str, substrs: List[str]): + for i, substr in enumerate(substrs): + pos = s.find(substr) + if pos != -1: + return pos, i + return None, None + + split_text: List[Tuple[Literal["text", "image_url"], str]] = [] + remaining = text + while remaining: + # Find first image_url + pos, i = find_first(remaining, image_urls) + if pos is not None and i is not None: + if pos > 0: + split_text.append(("text", remaining[:pos])) + split_text.append(("image_url", image_urls[i])) + remaining = remaining[pos + len(image_urls[i]) :] + else: + split_text.append(("text", remaining)) + remaining = "" + return split_text + + @classmethod + def from_pretrained( + cls, + repo_id: str, + filename: Optional[str], + local_dir: Optional[Union[str, os.PathLike[str]]] = None, + local_dir_use_symlinks: Union[bool, Literal["auto"]] = "auto", + cache_dir: Optional[Union[str, os.PathLike[str]]] = None, + **kwargs: Any, + ) -> "Llava15ChatHandler": + import fnmatch + from pathlib import Path + + try: + from huggingface_hub import hf_hub_download, HfFileSystem # type: ignore + from huggingface_hub.utils import validate_repo_id # type: ignore + except ImportError: + raise ImportError( + "Llama.from_pretrained requires the huggingface-hub package. " + "You can install it with `pip install huggingface-hub`." + ) + + validate_repo_id(repo_id) + + hffs = HfFileSystem() + + files = [ + file["name"] if isinstance(file, dict) else file + for file in hffs.ls(repo_id) # type: ignore + ] + + # split each file into repo_id, subfolder, filename + file_list: List[str] = [] + for file in files: + rel_path = Path(file).relative_to(repo_id) + file_list.append(str(rel_path)) + + matching_files = [file for file in file_list if fnmatch.fnmatch(file, filename)] # type: ignore + + if len(matching_files) == 0: + raise ValueError( + f"No file found in {repo_id} that match {filename}\n\n" + f"Available Files:\n{json.dumps(file_list)}" + ) + + if len(matching_files) > 1: + raise ValueError( + f"Multiple files found in {repo_id} matching {filename}\n\n" + f"Available Files:\n{json.dumps(files)}" + ) + + (matching_file,) = matching_files + + subfolder = str(Path(matching_file).parent) + filename = Path(matching_file).name + + # download the file + hf_hub_download( + repo_id=repo_id, + filename=filename, + subfolder=subfolder, + local_dir=cast(Union[str, Path, None], local_dir), + local_dir_use_symlinks=local_dir_use_symlinks, + cache_dir=cast(Union[str, Path, None], cache_dir), + ) + + if local_dir is None: + model_path = hf_hub_download( + repo_id=repo_id, + filename=filename, + subfolder=subfolder, + local_dir=local_dir, + local_dir_use_symlinks=local_dir_use_symlinks, + cache_dir=cast(Union[str, Path, None], cache_dir), + local_files_only=True, + ) + else: + model_path = os.path.join(local_dir, filename) + + return cls( + clip_model_path=model_path, + **kwargs, + ) + + +class MTMDChatHandler: + def __init__( + self, clip_model_path: str, verbose: bool = True, use_gpu: bool = True + ): + import llama_cpp.mtmd_cpp as mtmd_cpp + + self.clip_model_path = clip_model_path + self.verbose = verbose + self.use_gpu = use_gpu + self._mtmd_cpp = mtmd_cpp + self._exit_stack = ExitStack() + self.mtmd_ctx: Optional[mtmd_cpp.mtmd_context_p] = None + + if not os.path.exists(clip_model_path): + raise ValueError(f"Clip model path does not exist: {clip_model_path}") + + def _init_mtmd_context(self, llama_model: llama.Llama): + self.verbose = llama_model.verbose + if self.mtmd_ctx is not None: + return + + with suppress_stdout_stderr(disable=self.verbose): + ctx_params = self._mtmd_cpp.mtmd_context_params_default() + ctx_params.use_gpu = self.use_gpu + ctx_params.print_timings = self.verbose + ctx_params.n_threads = llama_model.n_threads + ctx_params.flash_attn_type = ( + llama_cpp.LLAMA_FLASH_ATTN_TYPE_ENABLED + if ( + llama_model.context_params.flash_attn_type + == llama_cpp.LLAMA_FLASH_ATTN_TYPE_ENABLED + ) + else llama_cpp.LLAMA_FLASH_ATTN_TYPE_DISABLED + ) + + self.mtmd_ctx = self._mtmd_cpp.mtmd_init_from_file( + self.clip_model_path.encode(), llama_model.model, ctx_params + ) + + if self.mtmd_ctx is None: + raise ValueError( + f"Failed to load mtmd context from: {self.clip_model_path}" + ) + + if not self._mtmd_cpp.mtmd_support_vision(self.mtmd_ctx): + raise ValueError("Vision is not supported by this model") + + def mtmd_free(): + with suppress_stdout_stderr(disable=self.verbose): + if self.mtmd_ctx is not None: + self._mtmd_cpp.mtmd_free(self.mtmd_ctx) + self.mtmd_ctx = None + + self._exit_stack.callback(mtmd_free) + + def load_image(self, image_url: str) -> bytes: + return self._load_image(image_url) + + def _create_bitmap_from_bytes(self, image_bytes: bytes): + if self.mtmd_ctx is None: + raise ValueError("mtmd context not initialized") + + with suppress_stdout_stderr(disable=self.verbose): + bitmap = self._mtmd_cpp.mtmd_helper_bitmap_init_from_buf( + self.mtmd_ctx, + (ctypes.c_uint8 * len(image_bytes)).from_buffer(bytearray(image_bytes)), + len(image_bytes), + False, + ) + + if bitmap is None: + raise ValueError("Failed to create bitmap from image bytes") + + return bitmap + + def _get_chat_template(self, llama_model: llama.Llama) -> str: + chat_template = llama_model.metadata.get("tokenizer.chat_template") + if not isinstance(chat_template, str) or chat_template == "": + raise ValueError( + f"{self.__class__.__name__} requires tokenizer.chat_template metadata" + ) + return chat_template + + def _get_template_messages( + self, + messages: List[llama_types.ChatCompletionRequestMessage], + media_marker: str, + ) -> List[Any]: + return [ + self._convert_message_for_template(message, media_marker) + for message in messages + ] + + @classmethod + def _convert_message_for_template( + cls, + message: llama_types.ChatCompletionRequestMessage, + media_marker: str, + ) -> Dict[str, Any]: + message_dict = dict(message) + content = message_dict.get("content") + if isinstance(content, list): + message_dict["content"] = [ + cls._convert_content_part_for_template(part, media_marker) + for part in content + ] + return message_dict + + @staticmethod + def _convert_content_part_for_template( + part: Any, + media_marker: str, + ) -> Any: + if isinstance(part, dict) and part.get("type") == "image_url": + return {"type": "text", "text": media_marker} + return part + + @staticmethod + def _decode_token_piece(piece: Any) -> str: + if isinstance(piece, bytes): + return piece.decode("utf-8", errors="ignore") + return str(piece) + + def _postprocess_template_text( + self, + text: str, + image_urls: List[str], + media_marker: str, + ) -> str: + for image_url in image_urls: + text = text.replace(image_url, media_marker) + return text + + def __call__( + self, + *, + llama: llama.Llama, + messages: List[llama_types.ChatCompletionRequestMessage], + functions: Optional[List[llama_types.ChatCompletionFunction]] = None, + function_call: Optional[llama_types.ChatCompletionRequestFunctionCall] = None, + tools: Optional[List[llama_types.ChatCompletionTool]] = None, + tool_choice: Optional[llama_types.ChatCompletionToolChoiceOption] = None, + temperature: float = 0.2, + top_p: float = 0.95, + top_k: int = 40, + min_p: float = 0.05, + typical_p: float = 1.0, + stream: bool = False, + stop: Optional[Union[str, List[str]]] = [], + seed: Optional[int] = None, + response_format: Optional[ + llama_types.ChatCompletionRequestResponseFormat + ] = None, + max_tokens: Optional[int] = None, + presence_penalty: float = 0.0, + frequency_penalty: float = 0.0, + repeat_penalty: float = 1.1, + tfs_z: float = 1.0, + mirostat_mode: int = 0, + mirostat_tau: float = 5.0, + mirostat_eta: float = 0.1, + model: Optional[str] = None, + logits_processor: Optional[llama.LogitsProcessorList] = None, + grammar: Optional[llama.LlamaGrammar] = None, + logit_bias: Optional[Dict[str, float]] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + **kwargs, # type: ignore + ) -> Union[ + llama_types.CreateChatCompletionResponse, + Iterator[llama_types.CreateChatCompletionStreamResponse], + ]: + self._init_mtmd_context(llama) + assert self.mtmd_ctx is not None + + image_urls = self.get_image_urls(messages) + media_marker = self._mtmd_cpp.mtmd_default_marker().decode("utf-8") + template_env = ImmutableSandboxedEnvironment( + trim_blocks=True, + lstrip_blocks=True, + extensions=[ + Jinja2ChatFormatter.IgnoreGenerationTags, + jinja2.ext.loopcontrols, + ], + ) + template_env.filters["tojson"] = Jinja2ChatFormatter.tojson + template = template_env.from_string(self._get_chat_template(llama)) + + def raise_exception(message: str): + raise ValueError(message) + + text = template.render( + messages=self._get_template_messages(messages, media_marker), + add_generation_prompt=True, + eos_token=self._decode_token_piece(llama.detokenize([llama.token_eos()])), + bos_token=self._decode_token_piece(llama.detokenize([llama.token_bos()])), + raise_exception=raise_exception, + functions=functions, + function_call=function_call, + tools=tools, + tool_choice=tool_choice, + strftime_now=Jinja2ChatFormatter.strftime_now, + **kwargs, + ) + text = self._postprocess_template_text(text, image_urls, media_marker) + + if self.verbose: + print(text, file=sys.stderr) + + bitmaps = [] + bitmap_cleanup = [] + try: + for image_url in image_urls: + image_bytes = self.load_image(image_url) + bitmap = self._create_bitmap_from_bytes(image_bytes) + bitmaps.append(bitmap) + bitmap_cleanup.append(bitmap) + + input_text = self._mtmd_cpp.mtmd_input_text() + input_text.text = text.encode("utf-8") + input_text.add_special = True + input_text.parse_special = True + + chunks = self._mtmd_cpp.mtmd_input_chunks_init() + if chunks is None: + raise ValueError("Failed to create input chunks") + + try: + bitmap_array = (self._mtmd_cpp.mtmd_bitmap_p_ctypes * len(bitmaps))( + *bitmaps + ) + result = self._mtmd_cpp.mtmd_tokenize( + self.mtmd_ctx, + chunks, + ctypes.byref(input_text), + bitmap_array, + len(bitmaps), + ) + + if result != 0: + raise ValueError(f"Failed to tokenize input: error code {result}") + + llama.reset() + llama._ctx.kv_cache_clear() + + n_chunks = self._mtmd_cpp.mtmd_input_chunks_size(chunks) + + for i in range(n_chunks): + chunk = self._mtmd_cpp.mtmd_input_chunks_get(chunks, i) + if chunk is None: + continue + + chunk_type = self._mtmd_cpp.mtmd_input_chunk_get_type(chunk) + + if chunk_type == self._mtmd_cpp.MTMD_INPUT_CHUNK_TYPE_TEXT: + n_tokens_out = ctypes.c_size_t() + tokens_ptr = self._mtmd_cpp.mtmd_input_chunk_get_tokens_text( + chunk, ctypes.byref(n_tokens_out) + ) + + if tokens_ptr and n_tokens_out.value > 0: + tokens = [tokens_ptr[j] for j in range(n_tokens_out.value)] + + if llama.n_tokens + len(tokens) > llama.n_ctx(): + raise ValueError( + f"Prompt exceeds n_ctx: {llama.n_tokens + len(tokens)} > {llama.n_ctx()}" + ) + llama.eval(tokens) + + elif chunk_type in [ + self._mtmd_cpp.MTMD_INPUT_CHUNK_TYPE_IMAGE, + self._mtmd_cpp.MTMD_INPUT_CHUNK_TYPE_AUDIO, + ]: + chunk_n_tokens = self._mtmd_cpp.mtmd_input_chunk_get_n_tokens( + chunk + ) + + if llama.n_tokens + chunk_n_tokens > llama.n_ctx(): + raise ValueError( + f"Prompt exceeds n_ctx: {llama.n_tokens + chunk_n_tokens} > {llama.n_ctx()}" + ) + + new_n_past = llama_cpp.llama_pos(0) + result = self._mtmd_cpp.mtmd_helper_eval_chunk_single( + self.mtmd_ctx, + llama._ctx.ctx, + chunk, + llama_cpp.llama_pos(llama.n_tokens), + llama_cpp.llama_seq_id(0), + llama.n_batch, + False, # logits_last + ctypes.byref(new_n_past), + ) + + if result != 0: + raise ValueError( + f"Failed to evaluate chunk: error code {result}" + ) + + llama.n_tokens = new_n_past.value + + prompt = llama.input_ids[: llama.n_tokens].tolist() + + finally: + self._mtmd_cpp.mtmd_input_chunks_free(chunks) + + finally: + for bitmap in bitmap_cleanup: + self._mtmd_cpp.mtmd_bitmap_free(bitmap) + + if response_format is not None and response_format["type"] == "json_object": + grammar = _grammar_for_response_format(response_format) + + if functions is not None: + tools = [ + { + "type": "function", + "function": function, + } + for function in functions + ] + + if function_call is not None: + if isinstance(function_call, str) and ( + function_call == "none" or function_call == "auto" + ): + tool_choice = function_call + if isinstance(function_call, dict) and "name" in function_call: + tool_choice = { + "type": "function", + "function": { + "name": function_call["name"], + }, + } + + tool = None + if ( + tool_choice is not None + and isinstance(tool_choice, dict) + and tools is not None + ): + name = tool_choice["function"]["name"] + tool = next((t for t in tools if t["function"]["name"] == name), None) + if tool is None: + raise ValueError(f"Tool choice '{name}' not found in tools.") + schema = tool["function"]["parameters"] + try: + grammar = llama_grammar.LlamaGrammar.from_json_schema( + json.dumps(schema), verbose=llama.verbose + ) + except Exception as e: + if llama.verbose: + print(str(e), file=sys.stderr) + grammar = llama_grammar.LlamaGrammar.from_string( + llama_grammar.JSON_GBNF, verbose=llama.verbose + ) + + completion_or_chunks = llama.create_completion( + prompt=prompt, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + typical_p=typical_p, + logprobs=top_logprobs if logprobs else None, + stream=stream, + stop=stop, + seed=seed, + max_tokens=max_tokens, + presence_penalty=presence_penalty, + frequency_penalty=frequency_penalty, + repeat_penalty=repeat_penalty, + tfs_z=tfs_z, + mirostat_mode=mirostat_mode, + mirostat_tau=mirostat_tau, + mirostat_eta=mirostat_eta, + model=model, + logits_processor=logits_processor, + grammar=grammar, + logit_bias=logit_bias, + ) + + if tool is not None: + tool_name = tool["function"]["name"] + return _convert_completion_to_chat_function( + tool_name, completion_or_chunks, stream + ) + return _convert_completion_to_chat(completion_or_chunks, stream=stream) + + @staticmethod + def _load_image(image_url: str) -> bytes: + if image_url.startswith("data:"): + import base64 + + image_bytes = base64.b64decode(image_url.split(",")[1]) + return image_bytes + else: + import urllib.request + + with urllib.request.urlopen(image_url) as f: + image_bytes = f.read() + return image_bytes + + @staticmethod + def get_image_urls(messages: List[llama_types.ChatCompletionRequestMessage]): + image_urls: List[str] = [] + for message in messages: + if message["role"] == "user": + if message["content"] is None: + continue + for content in message["content"]: + if isinstance(content, dict) and "type" in content: + if content["type"] == "image_url": + if ( + isinstance(content["image_url"], dict) + and "url" in content["image_url"] + ): + image_urls.append(content["image_url"]["url"]) + else: + image_urls.append(content["image_url"]) + return image_urls + + @classmethod + def from_pretrained( + cls, + repo_id: str, + filename: Optional[str], + local_dir: Optional[Union[str, os.PathLike[str]]] = None, + local_dir_use_symlinks: Union[bool, Literal["auto"]] = "auto", + cache_dir: Optional[Union[str, os.PathLike[str]]] = None, + **kwargs: Any, + ) -> "MTMDChatHandler": + import fnmatch + from pathlib import Path + + try: + from huggingface_hub import hf_hub_download, HfFileSystem # type: ignore + from huggingface_hub.utils import validate_repo_id # type: ignore + except ImportError: + raise ImportError( + "Llama.from_pretrained requires the huggingface-hub package. " + "You can install it with `pip install huggingface-hub`." + ) + + validate_repo_id(repo_id) + + hffs = HfFileSystem() + + files = [ + file["name"] if isinstance(file, dict) else file + for file in hffs.ls(repo_id) # type: ignore + ] + + file_list: List[str] = [] + for file in files: + rel_path = Path(file).relative_to(repo_id) + file_list.append(str(rel_path)) + + matching_files = [file for file in file_list if fnmatch.fnmatch(file, filename)] # type: ignore + + if len(matching_files) == 0: + raise ValueError( + f"No file found in {repo_id} that match {filename}\n\n" + f"Available Files:\n{json.dumps(file_list)}" + ) + + if len(matching_files) > 1: + raise ValueError( + f"Multiple files found in {repo_id} matching {filename}\n\n" + f"Available Files:\n{json.dumps(files)}" + ) + + (matching_file,) = matching_files + + subfolder = str(Path(matching_file).parent) + filename = Path(matching_file).name + + hf_hub_download( + repo_id=repo_id, + filename=filename, + subfolder=subfolder, + local_dir=cast(Union[str, Path, None], local_dir), + local_dir_use_symlinks=local_dir_use_symlinks, + cache_dir=cast(Union[str, Path, None], cache_dir), + ) + + if local_dir is None: + model_path = hf_hub_download( + repo_id=repo_id, + filename=filename, + subfolder=subfolder, + local_dir=local_dir, + local_dir_use_symlinks=local_dir_use_symlinks, + cache_dir=cast(Union[str, Path, None], cache_dir), + local_files_only=True, + ) + else: + model_path = os.path.join(local_dir, filename) + + return cls( + clip_model_path=model_path, + **kwargs, + ) + + +class Gemma4ChatHandler(MTMDChatHandler): + pass + + +class ObsidianChatHandler(Llava15ChatHandler): + # Prompt Format + # The model followed ChatML format. However, with ### as the separator + + # <|im_start|>user + # What is this sign about?\n<image> + # ### + # <|im_start|>assistant + # The sign is about bullying, and it is placed on a black background with a red background. + # ### + + CHAT_FORMAT = ( + "{% for message in messages %}" + # System message + "{% if message.role == 'system' %}" + "<|im_start|>system\n" + "{{ message.content }}\n" + "###\n" + "{% endif %}" + # User message + "{% if message.role == 'user' %}" + "<|im_start|>user\n" + "{% if message.content is string %}" + "{{ message.content }}" + "{% endif %}" + "{% if message.content is iterable %}" + "{% for content in message.content %}" + "{% if content.type == 'image_url' and content.image_url is string %}" + "{{ content.image_url }}" + "{% endif %}" + "{% if content.type == 'image_url' and content.image_url is mapping %}" + "{{ content.image_url.url }}" + "{% endif %}" + "{% endfor %}" + "{% for content in message.content %}" + "{% if content.type == 'text' %}" + "{{ content.text }}" + "{% endif %}" + "{% endfor %}" + "{% endif %}" + "###\n" + "{% endif %}" + # Assistant message + "{% if message.role == 'assistant' %}" + "<|im_start|>assistant\n" + "{{ message.content }}" + "###\n" + "{% endif %}" + "{% endfor %}" + # Generation prompt + "{% if add_generation_prompt %}" + "<|im_start|>assistant\n" + "{% endif %}" + ) + + +class MoondreamChatHandler(Llava15ChatHandler): + # Chat Format: + # f"<image>\n\n{chat_history}Question: {question}\n\nAnswer:" + CHAT_FORMAT = ( + "{% for message in messages %}" + "{% if message.role == 'user' %}" + "{% if message.content is iterable %}" + # <image> + "{% for content in message.content %}" + "{% if content.type == 'image_url' %}" + "{% if content.image_url is string %}" + "{{ content.image_url }}\n\n" + "{% endif %}" + "{% if content.image_url is mapping %}" + "{{ content.image_url.url }}\n\n" + "{% endif %}" + "{% endif %}" + "{% endfor %}" + # Question: + "{% for content in message.content %}" + "{% if content.type == 'text' %}" + "Question: {{ content.text }}\n\n" + "{% endif %}" + "{% endfor %}" + "{% endif %}" + # Question: + "{% if message.content is string %}" + "Question: {{ message.content }}\n\n" + "{% endif %}" + "{% endif %}" + # Answer: + "{% if message.role == 'assistant' %}" + "Answer:{{ message.content }}\n\n" + "{% endif %}" + "{% endfor %}" + # Generation prompt + "{% if add_generation_prompt %}" + "Answer:" + "{% endif %}" + ) + + +class Llava16ChatHandler(Llava15ChatHandler): + DEFAULT_SYSTEM_MESSAGE = "A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions. " + + # Example prompt + # "A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions. USER: <image>\nWhat is shown in this image? ASSISTANT:" + + CHAT_FORMAT = ( + "{% for message in messages %}" + "{% if message.role == 'system' %}" + "{{ message.content }}" + "{% endif %}" + "{% if message.role == 'user' %}" + "{% if message.content is iterable %}" + # <image> + "{% for content in message.content %}" + "{% if content.type == 'image_url' %}" + "{% if content.image_url is string %}" + "{{ content.image_url }}\n" + "{% endif %}" + "{% if content.image_url is mapping %}" + "{{ content.image_url.url }}\n" + "{% endif %}" + "{% endif %}" + "{% endfor %}" + # Question: + "{% for content in message.content %}" + "{% if content.type == 'text' %}" + "{{ content.text }}" + "{% endif %}" + "{% endfor %}" + "{% endif %}" + # Question: + "{% if message.content is string %}" + "{{ message.content }}" + "{% endif %}" + "{% endif %}" + # Answer: + "{% if message.role == 'assistant' %}" + "{{ message.content }}" + "{% endif %}" + "{% endfor %}" + # Generation prompt + "{% if add_generation_prompt %}" + "Answer:" + "{% endif %}" + ) + + +class NanoLlavaChatHandler(Llava15ChatHandler): + # Prompt Format + # The model follow the ChatML standard, however, without \n at the end of <|im_end|>: + + # <|im_start|>system + # Answer the question<|im_end|><|im_start|>user + # <image> + # What is the picture about?<|im_end|><|im_start|>assistant + DEFAULT_SYSTEM_MESSAGE = "Answer the question" + + CHAT_FORMAT = ( + "{% for message in messages %}" + # System message + "{% if message.role == 'system' %}" + "<|im_start|>system\n" + "{{ message.content }}" + "<|im_end|>" + "{% endif %}" + # User message + "{% if message.role == 'user' %}" + "<|im_start|>user\n" + "{% if message.content is string %}" + "{{ message.content }}" + "{% endif %}" + "{% if message.content is iterable %}" + "{% for content in message.content %}" + "{% if content.type == 'image_url' and content.image_url is string %}" + "{{ content.image_url }}" + "{% endif %}" + "{% if content.type == 'image_url' and content.image_url is mapping %}" + "{{ content.image_url.url }}" + "{% endif %}" + "{% endfor %}" + "{% for content in message.content %}" + "{% if content.type == 'text' %}" + "{{ content.text }}" + "{% endif %}" + "{% endfor %}" + "{% endif %}" + "<|im_end|>" + "{% endif %}" + # Assistant message + "{% if message.role == 'assistant' %}" + "<|im_start|>assistant\n" + "{{ message.content }}" + "<|im_end|>" + "{% endif %}" + "{% endfor %}" + # Generation prompt + "{% if add_generation_prompt %}" + "<|im_start|>assistant\n" + "{% endif %}" + ) + + +class Llama3VisionAlphaChatHandler(Llava15ChatHandler): + # question = "<image>" + q + + # prompt = f"<|start_header_id|>user<|end_header_id|>\n\n{question}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n" + DEFAULT_SYSTEM_MESSAGE = None + + CHAT_FORMAT = ( + "{% for message in messages %}" + "<|start_header_id|>" + "{% if message.role == 'user' %}" + "user<|end_header_id|>\n\n" + "{% if message.content is iterable %}" + # <image> + "{% for content in message.content %}" + "{% if content.type == 'image_url' %}" + "{% if content.image_url is string %}" + "{{ content.image_url }}" + "{% endif %}" + "{% if content.image_url is mapping %}" + "{{ content.image_url.url }}" + "{% endif %}" + "{% endif %}" + "{% endfor %}" + # Question: + "{% for content in message.content %}" + "{% if content.type == 'text' %}" + "{{ content.text }}" + "{% endif %}" + "{% endfor %}" + "{% endif %}" + # Question: + "{% if message.content is string %}" + "{{ message.content }}" + "{% endif %}" + "{% endif %}" + # Answer: + "{% if message.role == 'assistant' %}" + "assistant<|end_header_id|>\n\n" + "{{ message.content }}" + "{% endif %}" + "<|eot_id|>" + "{% endfor %}" + # Generation prompt + "{% if add_generation_prompt %}" + "<|start_header_id|>assistant<|end_header_id|>\n\n" + "{% endif %}" + ) + + +# alias +Llama3VisionAlpha = Llama3VisionAlphaChatHandler + + +class MiniCPMv26ChatHandler(Llava15ChatHandler): + DEFAULT_SYSTEM_MESSAGE = "You are a helpful assistant." + + CHAT_FORMAT = ( + "{% for message in messages %}" + "{% if loop.first and messages[0]['role'] != 'system' %}" + "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n" + "{% endif %}" + "<|im_start|>{{ message['role'] }}\n" + "{% if message['content'] is iterable %}" + "{% for content in message['content'] %}" + "{% if content.type == 'image_url' %}" + "{% if content.image_url is string %}" + "{{ content.image_url }}" + "{% endif %}" + "{% if content.image_url is mapping %}" + "{{ content.image_url.url }}" + "{% endif %}" + "{% endif %}" + "{% endfor %}" + "{% for content in message['content'] %}" + "{% if content.type == 'text' %}" + "{{ content.text }}" + "{% endif %}" + "{% endfor %}" + "{% endif %}" + "{% if message['content'] is string %}" + "{{ message['content'] }}" + "{% endif %}" + "<|im_end|>\n" + "{% endfor %}" + "{% if add_generation_prompt %}" + "<|im_start|>assistant\n" + "{% endif %}" + ) + + +class Qwen25VLChatHandler(Llava15ChatHandler): + DEFAULT_SYSTEM_MESSAGE = "You are a helpful assistant." + + CHAT_FORMAT = ( + # "{% set image_count = namespace(value=0) %}" + # "{% set video_count = namespace(value=0) %}" + "{% for message in messages %}" + "{% if loop.first and message['role'] != 'system' %}" + "<|im_start|>system\n" + "{{ self.DEFAULT_SYSTEM_MESSAGE }}<|im_end|>\n" + "{% endif %}" + "<|im_start|>{{ message['role'] }}\n" + "{% if message['content'] is string %}" + "{{ message['content'] }}<|im_end|>\n" + "{% else %}" + "{% for content in message['content'] %}" + "{% if content['type'] == 'image_url' %}" + "{% if content.image_url is string %}" + "{{ content.image_url }}" + "{% else %}" + "{{ content.image_url.url }}" + "{% endif %}" + # "{% set image_count.value = image_count.value + 1 %}" + "{% elif content['type'] == 'text' %}" + "{{ content['text'] }}" + "{% endif %}" + "{% endfor %}" + "<|im_end|>\n" + "{% endif %}" + "{% endfor %}" + "<|im_start|>assistant\n" + ) + + def __call__(self, **kwargs): + llama = kwargs["llama"] + + # Clear state for multiple runs + llama.reset() + llama._ctx.kv_cache_clear() + llama.n_tokens = 0 + + if hasattr(llama, "input_ids"): + llama.input_ids.fill(0) + + # Clear any handler state + if hasattr(self, "_last_image_embed"): + self._last_image_embed = None + self._last_image_hash = None + + if self.verbose: + messages = kwargs.get("messages", []) + image_count = len(self.get_image_urls(messages)) + print( + f"Minimal - Cleared state, processing {image_count} images", + file=sys.stderr, + ) + + # Use parent implementation + return super().__call__(**kwargs) + + +@register_chat_completion_handler("chatml-function-calling") +def chatml_function_calling( + llama: llama.Llama, + messages: List[llama_types.ChatCompletionRequestMessage], + functions: Optional[List[llama_types.ChatCompletionFunction]] = None, + function_call: Optional[llama_types.ChatCompletionRequestFunctionCall] = None, + tools: Optional[List[llama_types.ChatCompletionTool]] = None, + tool_choice: Optional[llama_types.ChatCompletionToolChoiceOption] = None, + temperature: float = 0.2, + top_p: float = 0.95, + top_k: int = 40, + min_p: float = 0.05, + typical_p: float = 1.0, + stream: bool = False, + stop: Optional[Union[str, List[str]]] = [], + response_format: Optional[llama_types.ChatCompletionRequestResponseFormat] = None, + max_tokens: Optional[int] = None, + presence_penalty: float = 0.0, + frequency_penalty: float = 0.0, + repeat_penalty: float = 1.1, + tfs_z: float = 1.0, + mirostat_mode: int = 0, + mirostat_tau: float = 5.0, + mirostat_eta: float = 0.1, + model: Optional[str] = None, + logits_processor: Optional[llama.LogitsProcessorList] = None, + grammar: Optional[llama.LlamaGrammar] = None, + logprobs: Optional[bool] = None, + top_logprobs: Optional[int] = None, + **kwargs, # type: ignore +) -> Union[ + llama_types.CreateChatCompletionResponse, + Iterator[llama_types.CreateChatCompletionStreamResponse], +]: + function_calling_template = ( + "{% for message in messages %}" + "<|im_start|>{{ message.role }}\n" + # System message + "{% if message.role == 'system' %}" + "{{ message.content }}" + "{% if tool_calls %}" + "\n\nYou have access to the following functions:\n" + "{% for tool in tools %}" + "\nfunctions.{{ tool.function.name }}:\n" + "{{ tool.function.parameters | tojson }}" + "\n{% endfor %}" + "\n\nYou can respond to users messages with either a single message or one or more function calls." + "\n\nTo respond with a message begin the message with 'message:', use the following format:" + "\n\nmessage:" + "\n<message>" + "\n\nTo respond with one or more function calls begin the message with 'functions.<function_name>:', use the following format:" + "\n\nfunctions.<function_name>:" + '\n{ "arg1": "value1", "arg2": "value2" }' + "\nfunctions.<function_name>:" + '\n{ "arg1": "value1", "arg2": "value2" }' + "{% endif %}" + "<|im_end|>\n" + "{% endif %}" + # User message + "{% if message.role == 'user' %}" + "{{ message.content }}" + "<|im_end|>\n" + "{% endif %}" + # Assistant message + "{% if message.role == 'assistant' %}" + ## Reglar message + "{% if message.content and message.content | length > 0 %}" + "{% if tool_calls %}" + "message:\n" + "{% endif %}" + "{{ message.content }}" + "<|im_end|>\n" + "{% endif %}" + ## Function calls + "{% if 'tool_calls' in message %}" + "{% for tool_call in message.tool_calls %}" + "functions.{{ tool_call.function.name }}:\n" + "{{ tool_call.function.arguments }}" + "{% endfor %}" + "<|im_end|>\n" + "{% endif %}" + "{% endif %}" + "{% endfor %}" + "{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}" + ) + template_renderer = ImmutableSandboxedEnvironment( + autoescape=jinja2.select_autoescape(["html", "xml"]), + undefined=jinja2.StrictUndefined, + ).from_string(function_calling_template) + + # Convert legacy functions to tools + if functions is not None: + tools = [ + { + "type": "function", + "function": function, + } + for function in functions + ] + + # Convert legacy function_call to tool_choice + if function_call is not None: + if isinstance(function_call, str) and ( + function_call == "none" or function_call == "auto" + ): + tool_choice = function_call + if isinstance(function_call, dict) and "name" in function_call: + tool_choice = { + "type": "function", + "function": { + "name": function_call["name"], + }, + } + + stop = ( + [stop, "<|im_end|>"] + if isinstance(stop, str) + else stop + ["<|im_end|>"] + if stop + else ["<|im_end|>"] + ) + + # Case 1: No tool choice by user + if ( + tool_choice is None + or (isinstance(tool_choice, str) and tool_choice == "none") + or tools is None + or len(tools) == 0 + ): + prompt = template_renderer.render( + messages=messages, + tools=[], + tool_calls=None, + add_generation_prompt=True, + ) + + if response_format is not None and response_format["type"] == "json_object": + grammar = _grammar_for_response_format(response_format) + + return _convert_completion_to_chat( + llama.create_completion( + prompt=prompt, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + typical_p=typical_p, + stream=stream, + stop=stop, + max_tokens=max_tokens, + presence_penalty=presence_penalty, + frequency_penalty=frequency_penalty, + repeat_penalty=repeat_penalty, + tfs_z=tfs_z, + mirostat_mode=mirostat_mode, + mirostat_tau=mirostat_tau, + mirostat_eta=mirostat_eta, + model=model, + logits_processor=logits_processor, + grammar=grammar, + logprobs=top_logprobs if logprobs else None, + ), + stream=stream, + ) + + # Case 2: Tool choice by user + if isinstance(tool_choice, dict): + tool_name = tool_choice["function"]["name"] + tool = next( + (tool for tool in tools if tool["function"]["name"] == tool_name), None + ) + if tool is None: + raise ValueError(f"Tool with name '{tool_name}' not found in tools") + prompt = template_renderer.render( + messages=messages, + tools=tools, + tool_calls=True, + add_generation_prompt=True, + ) + prompt += f"functions.{tool_name}:\n" + try: + grammar = llama_grammar.LlamaGrammar.from_json_schema( + json.dumps(tool["function"]["parameters"]), verbose=llama.verbose + ) + except Exception as e: + grammar = llama_grammar.LlamaGrammar.from_string( + llama_grammar.JSON_GBNF, verbose=llama.verbose + ) + if llama.verbose: + print( + "Failed to parse function body as JSON schema, falling back to default grammar" + ) + print(e) + completion_or_chunks = llama.create_completion( + prompt=prompt, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + typical_p=typical_p, + stream=stream, + stop=stop, + max_tokens=max_tokens, + presence_penalty=presence_penalty, + frequency_penalty=frequency_penalty, + repeat_penalty=repeat_penalty, + tfs_z=tfs_z, + mirostat_mode=mirostat_mode, + mirostat_tau=mirostat_tau, + mirostat_eta=mirostat_eta, + model=model, + logits_processor=logits_processor, + grammar=grammar, + ) + return _convert_completion_to_chat_function( + tool_name, completion_or_chunks, stream + ) + + # Case 3: Automatic tool choice + assert isinstance(tool_choice, str) and tool_choice == "auto" + function_names = " | ".join( + [f'''"functions.{tool["function"]["name"]}:"''' for tool in tools] + ) + initial_gbnf_tool_grammar = ( + """root ::= functions | "message:"\n""" + f"""functions ::= {function_names}\n""" + ) + follow_up_gbnf_tool_grammar = ( + """root ::= functions | "<|im_end|>"\n""" + f"""functions ::= {function_names}\n""" + ) + prompt = template_renderer.render( + messages=messages, + tools=tools, + tool_calls=True, + add_generation_prompt=True, + ) + completion_or_chunks = llama.create_completion( + prompt=prompt, + temperature=0, + top_p=top_p, + top_k=top_k, + min_p=min_p, + typical_p=typical_p, + stream=False, + stop=[":"], + max_tokens=None, + presence_penalty=presence_penalty, + frequency_penalty=frequency_penalty, + repeat_penalty=repeat_penalty, + tfs_z=tfs_z, + mirostat_mode=mirostat_mode, + mirostat_tau=mirostat_tau, + mirostat_eta=mirostat_eta, + model=model, + logits_processor=logits_processor, + grammar=llama_grammar.LlamaGrammar.from_string( + initial_gbnf_tool_grammar, verbose=llama.verbose + ), + ) + completion: llama_types.CreateCompletionResponse = completion_or_chunks # type: ignore + text = completion["choices"][0]["text"] + if "message" in text: + return _convert_completion_to_chat( + llama.create_completion( + prompt=prompt + "message:\n", + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + typical_p=typical_p, + stream=stream, + stop=["<|im_end|>"], + logprobs=top_logprobs if logprobs else None, + max_tokens=None, + presence_penalty=presence_penalty, + frequency_penalty=frequency_penalty, + repeat_penalty=repeat_penalty, + tfs_z=tfs_z, + mirostat_mode=mirostat_mode, + mirostat_tau=mirostat_tau, + mirostat_eta=mirostat_eta, + model=model, + logits_processor=logits_processor, + grammar=llama_grammar.LlamaGrammar.from_string( + follow_up_gbnf_tool_grammar, verbose=llama.verbose + ), + ), + stream=stream, + ) + + # One or more function calls + tool_name = text[len("functions.") :] + tool = next((tool for tool in tools if tool["function"]["name"] == tool_name), None) + if not stream: + completions: List[llama_types.CreateCompletionResponse] = [] + completions_tool_name: List[str] = [] + while tool is not None: + prompt += f"functions.{tool_name}:\n" + try: + grammar = llama_grammar.LlamaGrammar.from_json_schema( + json.dumps(tool["function"]["parameters"]), verbose=llama.verbose + ) + except Exception as e: + grammar = llama_grammar.LlamaGrammar.from_string( + llama_grammar.JSON_GBNF, verbose=llama.verbose + ) + if llama.verbose: + print( + "Failed to parse function body as JSON schema, falling back to default grammar" + ) + print(e) + completion_or_chunks = llama.create_completion( + prompt=prompt, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + typical_p=typical_p, + stream=False, + stop=stop, + max_tokens=None, + presence_penalty=presence_penalty, + frequency_penalty=frequency_penalty, + repeat_penalty=repeat_penalty, + tfs_z=tfs_z, + mirostat_mode=mirostat_mode, + mirostat_tau=mirostat_tau, + mirostat_eta=mirostat_eta, + model=model, + logits_processor=logits_processor, + grammar=grammar, + ) + completion_or_chunks = cast( + llama_types.CreateCompletionResponse, completion_or_chunks + ) + completions.append(completion_or_chunks) + completions_tool_name.append(tool_name) + prompt += completion_or_chunks["choices"][0]["text"] + prompt += "\n" + + response = llama.create_completion( + prompt=prompt, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + typical_p=typical_p, + stream=False, + stop=stop, + max_tokens=None, + presence_penalty=presence_penalty, + frequency_penalty=frequency_penalty, + repeat_penalty=repeat_penalty, + tfs_z=tfs_z, + mirostat_mode=mirostat_mode, + mirostat_tau=mirostat_tau, + mirostat_eta=mirostat_eta, + model=model, + logits_processor=logits_processor, + grammar=llama_grammar.LlamaGrammar.from_string( + follow_up_gbnf_tool_grammar, verbose=llama.verbose + ), + ) + response = cast(llama_types.CreateCompletionResponse, response) + + tool_name = response["choices"][0]["text"][len("functions.") :] + tool = next( + (tool for tool in tools if tool["function"]["name"] == tool_name), None + ) + + # Merge completions + function_call_dict: Union[ + Dict[str, str], + Dict[ + Literal["function_call"], + llama_types.ChatCompletionRequestAssistantMessageFunctionCall, + ], + ] = ( + { + "function_call": { + "name": tool_name, + "arguments": completions[0]["choices"][0]["text"], + } + } + if len(completions) == 1 + else {} + ) + return { + "id": "chat" + completion["id"], + "object": "chat.completion", + "created": completion["created"], + "model": completion["model"], + "choices": [ + { + "finish_reason": "tool_calls", + "index": 0, + "logprobs": _convert_text_completion_logprobs_to_chat( + completion["choices"][0]["logprobs"] + ), + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_" + + f"_{i}_" + + tool_name + + "_" + + completion["id"], + "type": "function", + "function": { + "name": tool_name, + "arguments": completion["choices"][0]["text"], + }, + } + for i, (tool_name, completion) in enumerate( + zip(completions_tool_name, completions) + ) + ], + **function_call_dict, + }, + } + ], + "usage": { + "completion_tokens": sum( + ( + completion["usage"]["completion_tokens"] + if "usage" in completion + else 0 + ) + for completion in completions + ), + "prompt_tokens": sum( + completion["usage"]["prompt_tokens"] if "usage" in completion else 0 + for completion in completions + ), + "total_tokens": sum( + completion["usage"]["total_tokens"] if "usage" in completion else 0 + for completion in completions + ), + }, + } + + raise ValueError("Automatic streaming tool choice is not supported") diff --git a/llama_cpp/llama_cpp.py b/llama_cpp/llama_cpp.py index 17c631961d..21f85c81c3 100644 --- a/llama_cpp/llama_cpp.py +++ b/llama_cpp/llama_cpp.py @@ -1,126 +1,527 @@ -import sys +from __future__ import annotations + import os import ctypes -from ctypes import ( - c_double, - c_int, - c_float, - c_char_p, - c_int32, - c_uint32, - c_void_p, - c_bool, - POINTER, - _Pointer, # type: ignore - Structure, - Array, - c_uint8, - c_size_t, -) import pathlib -from typing import List, Union +import warnings + +from typing import ( + Callable, + Union, + NewType, + Optional, + TYPE_CHECKING, +) +from llama_cpp._ctypes_extensions import ( + load_shared_library, + byref, + ctypes_function_for_shared_library, +) -# Load the library -def _load_shared_library(lib_base_name: str): - # Construct the paths to the possible shared library names - _base_path = pathlib.Path(__file__).parent.resolve() - # Searching for the library in the current directory under the name "libllama" (default name - # for llamacpp) and "llama" (default name for this repo) - _lib_paths: List[pathlib.Path] = [] - # Determine the file extension based on the platform - if sys.platform.startswith("linux"): - _lib_paths += [ - _base_path / f"lib{lib_base_name}.so", - ] - elif sys.platform == "darwin": - _lib_paths += [ - _base_path / f"lib{lib_base_name}.so", - _base_path / f"lib{lib_base_name}.dylib", - ] - elif sys.platform == "win32": - _lib_paths += [ - _base_path / f"{lib_base_name}.dll", - ] - else: - raise RuntimeError("Unsupported platform") - - if "LLAMA_CPP_LIB" in os.environ: - lib_base_name = os.environ["LLAMA_CPP_LIB"] - _lib = pathlib.Path(lib_base_name) - _base_path = _lib.parent.resolve() - _lib_paths = [_lib.resolve()] - - cdll_args = dict() # type: ignore - # Add the library directory to the DLL search path on Windows (if needed) - if sys.platform == "win32" and sys.version_info >= (3, 8): - os.add_dll_directory(str(_base_path)) - if "CUDA_PATH" in os.environ: - os.add_dll_directory(os.path.join(os.environ["CUDA_PATH"], "bin")) - os.add_dll_directory(os.path.join(os.environ["CUDA_PATH"], "lib")) - cdll_args["winmode"] = 0 - - # Try to load the shared library, handling potential errors - for _lib_path in _lib_paths: - if _lib_path.exists(): - try: - return ctypes.CDLL(str(_lib_path), **cdll_args) - except Exception as e: - raise RuntimeError(f"Failed to load shared library '{_lib_path}': {e}") - - raise FileNotFoundError( - f"Shared library with base name '{lib_base_name}' not found" +if TYPE_CHECKING: + from llama_cpp._ctypes_extensions import ( + CtypesCData, + CtypesArray, + CtypesPointer, + CtypesVoidPointer, + CtypesRef, + CtypesPointerOrRef, + CtypesFuncPointer, ) # Specify the base name of the shared library to load _lib_base_name = "llama" - +_override_base_path = os.environ.get("LLAMA_CPP_LIB_PATH") +_base_path = ( + pathlib.Path(os.path.abspath(os.path.dirname(__file__))) / "lib" + if _override_base_path is None + else pathlib.Path(_override_base_path) +) # Load the library -_lib = _load_shared_library(_lib_base_name) +_lib = load_shared_library(_lib_base_name, _base_path) + +ctypes_function = ctypes_function_for_shared_library(_lib) -# Misc -c_float_p = POINTER(c_float) -c_uint8_p = POINTER(c_uint8) -c_size_t_p = POINTER(c_size_t) + +def _warn_deprecated(symbol: str, hint: str) -> None: + warnings.warn( + f"{symbol} is deprecated; {hint}", + DeprecationWarning, + stacklevel=2, + ) + + +# from ggml.h +# // NOTE: always add types at the end of the enum to keep backward compatibility +# enum ggml_type { +# GGML_TYPE_F32 = 0, +# GGML_TYPE_F16 = 1, +# GGML_TYPE_Q4_0 = 2, +# GGML_TYPE_Q4_1 = 3, +# // GGML_TYPE_Q4_2 = 4, support has been removed +# // GGML_TYPE_Q4_3 = 5, support has been removed +# GGML_TYPE_Q5_0 = 6, +# GGML_TYPE_Q5_1 = 7, +# GGML_TYPE_Q8_0 = 8, +# GGML_TYPE_Q8_1 = 9, +# GGML_TYPE_Q2_K = 10, +# GGML_TYPE_Q3_K = 11, +# GGML_TYPE_Q4_K = 12, +# GGML_TYPE_Q5_K = 13, +# GGML_TYPE_Q6_K = 14, +# GGML_TYPE_Q8_K = 15, +# GGML_TYPE_IQ2_XXS = 16, +# GGML_TYPE_IQ2_XS = 17, +# GGML_TYPE_IQ3_XXS = 18, +# GGML_TYPE_IQ1_S = 19, +# GGML_TYPE_IQ4_NL = 20, +# GGML_TYPE_IQ3_S = 21, +# GGML_TYPE_IQ2_S = 22, +# GGML_TYPE_IQ4_XS = 23, +# GGML_TYPE_I8 = 24, +# GGML_TYPE_I16 = 25, +# GGML_TYPE_I32 = 26, +# GGML_TYPE_I64 = 27, +# GGML_TYPE_F64 = 28, +# GGML_TYPE_IQ1_M = 29, +# GGML_TYPE_MXFP4 = 39, +# GGML_TYPE_NVFP4 = 40, +# GGML_TYPE_Q1_0 = 41, +# GGML_TYPE_COUNT = 42, +# }; +GGML_TYPE_F32 = 0 +GGML_TYPE_F16 = 1 +GGML_TYPE_Q4_0 = 2 +GGML_TYPE_Q4_1 = 3 +GGML_TYPE_Q5_0 = 6 +GGML_TYPE_Q5_1 = 7 +GGML_TYPE_Q8_0 = 8 +GGML_TYPE_Q8_1 = 9 +GGML_TYPE_Q2_K = 10 +GGML_TYPE_Q3_K = 11 +GGML_TYPE_Q4_K = 12 +GGML_TYPE_Q5_K = 13 +GGML_TYPE_Q6_K = 14 +GGML_TYPE_Q8_K = 15 +GGML_TYPE_IQ2_XXS = 16 +GGML_TYPE_IQ2_XS = 17 +GGML_TYPE_IQ3_XXS = 18 +GGML_TYPE_IQ1_S = 19 +GGML_TYPE_IQ4_NL = 20 +GGML_TYPE_IQ3_S = 21 +GGML_TYPE_IQ2_S = 22 +GGML_TYPE_IQ4_XS = 23 +GGML_TYPE_I8 = 24 +GGML_TYPE_I16 = 25 +GGML_TYPE_I32 = 26 +GGML_TYPE_I64 = 27 +GGML_TYPE_F64 = 28 +GGML_TYPE_IQ1_M = 29 +GGML_TYPE_MXFP4 = 39 +GGML_TYPE_NVFP4 = 40 +GGML_TYPE_Q1_0 = 41 +GGML_TYPE_COUNT = 42 + +# from ggml-backend.h +# typedef bool (*ggml_backend_sched_eval_callback)(struct ggml_tensor * t, bool ask, void * user_data); +ggml_backend_sched_eval_callback = ctypes.CFUNCTYPE( + ctypes.c_bool, ctypes.c_void_p, ctypes.c_bool, ctypes.c_void_p +) + +# // Abort callback +# // If not NULL, called before ggml computation +# // If it returns true, the computation is aborted +# typedef bool (*ggml_abort_callback)(void * data); +ggml_abort_callback = ctypes.CFUNCTYPE(ctypes.c_bool, ctypes.c_void_p) # llama.h bindings -GGML_USE_CUBLAS = hasattr(_lib, "ggml_init_cublas") -GGML_CUDA_MAX_DEVICES = ctypes.c_int(16) -LLAMA_MAX_DEVICES = GGML_CUDA_MAX_DEVICES if GGML_USE_CUBLAS else ctypes.c_int(1) - -# #define LLAMA_FILE_MAGIC_GGJT 0x67676a74u // 'ggjt' -LLAMA_FILE_MAGIC_GGJT = ctypes.c_uint(0x67676A74) -# #define LLAMA_FILE_MAGIC_GGLA 0x67676c61u // 'ggla' -LLAMA_FILE_MAGIC_GGLA = ctypes.c_uint(0x67676C61) -# #define LLAMA_FILE_MAGIC_GGMF 0x67676d66u // 'ggmf' -LLAMA_FILE_MAGIC_GGMF = ctypes.c_uint(0x67676D66) -# #define LLAMA_FILE_MAGIC_GGML 0x67676d6cu // 'ggml' -LLAMA_FILE_MAGIC_GGML = ctypes.c_uint(0x67676D6C) -# #define LLAMA_FILE_MAGIC_GGSN 0x6767736eu // 'ggsn' -LLAMA_FILE_MAGIC_GGSN = ctypes.c_uint(0x6767736E) - -# #define LLAMA_FILE_VERSION 3 -LLAMA_FILE_VERSION = c_int(3) -LLAMA_FILE_MAGIC = LLAMA_FILE_MAGIC_GGJT -LLAMA_FILE_MAGIC_UNVERSIONED = LLAMA_FILE_MAGIC_GGML +_lib.llama_max_devices.argtypes = [] +_lib.llama_max_devices.restype = ctypes.c_size_t + +LLAMA_MAX_DEVICES = _lib.llama_max_devices() + +# define LLAMA_DEFAULT_SEED 0xFFFFFFFF +LLAMA_DEFAULT_SEED = 0xFFFFFFFF + +# define LLAMA_TOKEN_NULL -1 +LLAMA_TOKEN_NULL = -1 + +# define LLAMA_FILE_MAGIC_GGLA 0x67676c61u // 'ggla' +LLAMA_FILE_MAGIC_GGLA = 0x67676C61 + +# define LLAMA_FILE_MAGIC_GGSN 0x6767736eu // 'ggsn' +LLAMA_FILE_MAGIC_GGSN = 0x6767736E + +# define LLAMA_FILE_MAGIC_GGSQ 0x67677371u // 'ggsq' +LLAMA_FILE_MAGIC_GGSQ = 0x67677371 + +# define LLAMA_SESSION_MAGIC LLAMA_FILE_MAGIC_GGSN LLAMA_SESSION_MAGIC = LLAMA_FILE_MAGIC_GGSN -LLAMA_SESSION_VERSION = c_int(1) +# define LLAMA_SESSION_VERSION 9 +LLAMA_SESSION_VERSION = 9 -# #define LLAMA_DEFAULT_SEED 0xFFFFFFFF -LLAMA_DEFAULT_SEED = c_int(0xFFFFFFFF) +# define LLAMA_STATE_SEQ_MAGIC LLAMA_FILE_MAGIC_GGSQ +LLAMA_STATE_SEQ_MAGIC = LLAMA_FILE_MAGIC_GGSQ +# define LLAMA_STATE_SEQ_VERSION 2 +LLAMA_STATE_SEQ_VERSION = 2 + +# struct llama_vocab; +llama_vocab_p = NewType("llama_vocab_p", int) +llama_vocab_p_ctypes = ctypes.c_void_p # struct llama_model; -llama_model_p = c_void_p +llama_model_p = NewType("llama_model_p", int) +llama_model_p_ctypes = ctypes.c_void_p # struct llama_context; -llama_context_p = c_void_p +llama_context_p = NewType("llama_context_p", int) +llama_context_p_ctypes = ctypes.c_void_p + +# typedef struct llama_memory_i * llama_memory_t; +llama_memory_t = NewType("llama_memory_t", int) +llama_memory_t_ctypes = ctypes.c_void_p + +# struct llama_kv_cache; (DEPRECATED) +llama_kv_cache_p = NewType("llama_kv_cache_p", int) +llama_kv_cache_p_ctypes = ctypes.c_void_p + +# struct gguf_context; +gguf_context_p = NewType("gguf_context_p", int) +gguf_context_p_ctypes = ctypes.c_void_p + +# typedef int32_t llama_pos; +llama_pos = ctypes.c_int32 +# typedef int32_t llama_token; +llama_token = ctypes.c_int32 +llama_token_p = ctypes.POINTER(llama_token) +# typedef int32_t llama_seq_id; +llama_seq_id = ctypes.c_int32 +# typedef uint32_t llama_state_seq_flags; +llama_state_seq_flags = ctypes.c_uint32 + + +# enum llama_vocab_type { +# LLAMA_VOCAB_TYPE_NONE = 0, // For models without vocab +# LLAMA_VOCAB_TYPE_SPM = 1, // LLaMA tokenizer based on byte-level BPE with byte fallback +# LLAMA_VOCAB_TYPE_BPE = 2, // GPT-2 tokenizer based on byte-level BPE +# LLAMA_VOCAB_TYPE_WPM = 3, // BERT tokenizer based on WordPiece +# LLAMA_VOCAB_TYPE_UGM = 4, // T5 tokenizer based on Unigram +# LLAMA_VOCAB_TYPE_RWKV = 5, // RWKV tokenizer based on greedy tokenization +# LLAMA_VOCAB_TYPE_PLAMO2 = 6, // PLaMo-2 tokenizer based on Aho-Corasick with dynamic programming +# }; +LLAMA_VOCAB_TYPE_NONE = 0 +"""For models without vocab""" +LLAMA_VOCAB_TYPE_SPM = 1 +"""LLaMA tokenizer based on byte-level BPE with byte fallback""" +LLAMA_VOCAB_TYPE_BPE = 2 +"""GPT-2 tokenizer based on byte-level BPE""" +LLAMA_VOCAB_TYPE_WPM = 3 +"""BERT tokenizer based on WordPiece""" +LLAMA_VOCAB_TYPE_UGM = 4 +"""T5 tokenizer based on Unigram""" +LLAMA_VOCAB_TYPE_RWKV = 5 +"""RWKV tokenizer based on greedy tokenization""" +LLAMA_VOCAB_TYPE_PLAMO2 = 6 +"""PLaMo-2 tokenizer based on Aho-Corasick with dynamic programming""" + + +# NOTE: Deprecated and will be removed in the future. (already gone in llama.cpp) +# // pre-tokenization types +# enum llama_vocab_pre_type { +# LLAMA_VOCAB_PRE_TYPE_DEFAULT = 0, +# LLAMA_VOCAB_PRE_TYPE_LLAMA3 = 1, +# LLAMA_VOCAB_PRE_TYPE_DEEPSEEK_LLM = 2, +# LLAMA_VOCAB_PRE_TYPE_DEEPSEEK_CODER = 3, +# LLAMA_VOCAB_PRE_TYPE_FALCON = 4, +# LLAMA_VOCAB_PRE_TYPE_MPT = 5, +# LLAMA_VOCAB_PRE_TYPE_STARCODER = 6, +# LLAMA_VOCAB_PRE_TYPE_GPT2 = 7, +# LLAMA_VOCAB_PRE_TYPE_REFACT = 8, +# LLAMA_VOCAB_PRE_TYPE_COMMAND_R = 9, +# LLAMA_VOCAB_PRE_TYPE_STABLELM2 = 10, +# LLAMA_VOCAB_PRE_TYPE_QWEN2 = 11, +# LLAMA_VOCAB_PRE_TYPE_OLMO = 12, +# LLAMA_VOCAB_PRE_TYPE_DBRX = 13, +# LLAMA_VOCAB_PRE_TYPE_SMAUG = 14, +# LLAMA_VOCAB_PRE_TYPE_PORO = 15, +# LLAMA_VOCAB_PRE_TYPE_CHATGLM3 = 16, +# LLAMA_VOCAB_PRE_TYPE_CHATGLM4 = 17, +# LLAMA_VOCAB_PRE_TYPE_VIKING = 18, +# LLAMA_VOCAB_PRE_TYPE_JAIS = 19, +# LLAMA_VOCAB_PRE_TYPE_TEKKEN = 20, +# LLAMA_VOCAB_PRE_TYPE_SMOLLM = 21, +# LLAMA_VOCAB_PRE_TYPE_CODESHELL = 22, +# LLAMA_VOCAB_PRE_TYPE_BLOOM = 23, +# LLAMA_VOCAB_PRE_TYPE_GPT3_FINNISH = 24, +# LLAMA_VOCAB_PRE_TYPE_EXAONE = 25, +# LLAMA_VOCAB_PRE_TYPE_CHAMELEON = 26, +# LLAMA_VOCAB_PRE_TYPE_MINERVA = 27, +# LLAMA_VOCAB_PRE_TYPE_DEEPSEEK3_LLM = 28, +# LLAMA_VOCAB_PRE_TYPE_GPT4O = 29, +# LLAMA_VOCAB_PRE_TYPE_SUPERBPE = 30, +# LLAMA_VOCAB_PRE_TYPE_TRILLION = 31, +# LLAMA_VOCAB_PRE_TYPE_BAILINGMOE = 32, +# LLAMA_VOCAB_PRE_TYPE_LLAMA4 = 33, +# LLAMA_VOCAB_PRE_TYPE_PIXTRAL = 34, +# LLAMA_VOCAB_PRE_TYPE_SEED_CODER = 35, +# }; +LLAMA_VOCAB_PRE_TYPE_DEFAULT = 0 +LLAMA_VOCAB_PRE_TYPE_LLAMA3 = 1 +LLAMA_VOCAB_PRE_TYPE_DEEPSEEK_LLM = 2 +LLAMA_VOCAB_PRE_TYPE_DEEPSEEK_CODER = 3 +LLAMA_VOCAB_PRE_TYPE_FALCON = 4 +LLAMA_VOCAB_PRE_TYPE_MPT = 5 +LLAMA_VOCAB_PRE_TYPE_STARCODER = 6 +LLAMA_VOCAB_PRE_TYPE_GPT2 = 7 +LLAMA_VOCAB_PRE_TYPE_REFACT = 8 +LLAMA_VOCAB_PRE_TYPE_COMMAND_R = 9 +LLAMA_VOCAB_PRE_TYPE_STABLELM2 = 10 +LLAMA_VOCAB_PRE_TYPE_QWEN2 = 11 +LLAMA_VOCAB_PRE_TYPE_OLMO = 12 +LLAMA_VOCAB_PRE_TYPE_DBRX = 13 +LLAMA_VOCAB_PRE_TYPE_SMAUG = 14 +LLAMA_VOCAB_PRE_TYPE_PORO = 15 +LLAMA_VOCAB_PRE_TYPE_CHATGLM3 = 16 +LLAMA_VOCAB_PRE_TYPE_CHATGLM4 = 17 +LLAMA_VOCAB_PRE_TYPE_VIKING = 18 +LLAMA_VOCAB_PRE_TYPE_JAIS = 19 +LLAMA_VOCAB_PRE_TYPE_TEKKEN = 20 +LLAMA_VOCAB_PRE_TYPE_SMOLLM = 21 +LLAMA_VOCAB_PRE_TYPE_CODESHELL = 22 +LLAMA_VOCAB_PRE_TYPE_BLOOM = 23 +LLAMA_VOCAB_PRE_TYPE_GPT3_FINNISH = 24 +LLAMA_VOCAB_PRE_TYPE_EXAONE = 25 +LLAMA_VOCAB_PRE_TYPE_CHAMELEON = 26 +LLAMA_VOCAB_PRE_TYPE_MINERVA = 27 +LLAMA_VOCAB_PRE_TYPE_DEEPSEEK3_LLM = 28 +LLAMA_VOCAB_PRE_TYPE_GPT4O = 29 +LLAMA_VOCAB_PRE_TYPE_SUPERBPE = 30 +LLAMA_VOCAB_PRE_TYPE_TRILLION = 31 +LLAMA_VOCAB_PRE_TYPE_BAILINGMOE = 32 +LLAMA_VOCAB_PRE_TYPE_LLAMA4 = 33 +LLAMA_VOCAB_PRE_TYPE_PIXTRAL = 34 +LLAMA_VOCAB_PRE_TYPE_SEED_CODER = 35 + + +# // note: these values should be synchronized with ggml_rope +# // TODO: maybe move this enum to ggml.h (ggml_rope_type) +# enum llama_rope_type { +# LLAMA_ROPE_TYPE_NONE = -1, +# LLAMA_ROPE_TYPE_NORM = 0, +# LLAMA_ROPE_TYPE_NEOX = GGML_ROPE_TYPE_NEOX, +# LLAMA_ROPE_TYPE_MROPE = GGML_ROPE_TYPE_MROPE, +# LLAMA_ROPE_TYPE_IMROPE = GGML_ROPE_TYPE_IMROPE, +# LLAMA_ROPE_TYPE_VISION = GGML_ROPE_TYPE_VISION, +# }; +LLAMA_ROPE_TYPE_NONE = -1 +LLAMA_ROPE_TYPE_NORM = 0 +LLAMA_ROPE_TYPE_NEOX = GGML_ROPE_TYPE_NEOX = 2 +LLAMA_ROPE_TYPE_MROPE = GGML_ROPE_TYPE_MROPE = 8 +LLAMA_ROPE_TYPE_IMROPE = GGML_ROPE_TYPE_IMROPE = 40 +LLAMA_ROPE_TYPE_VISION = GGML_ROPE_TYPE_VISION = 24 + + +# enum llama_token_type { //TODO: remove, required until per token attributes are available from GGUF file +# LLAMA_TOKEN_TYPE_UNDEFINED = 0, +# LLAMA_TOKEN_TYPE_NORMAL = 1, +# LLAMA_TOKEN_TYPE_UNKNOWN = 2, +# LLAMA_TOKEN_TYPE_CONTROL = 3, +# LLAMA_TOKEN_TYPE_USER_DEFINED = 4, +# LLAMA_TOKEN_TYPE_UNUSED = 5, +# LLAMA_TOKEN_TYPE_BYTE = 6, +# }; +LLAMA_TOKEN_TYPE_UNDEFINED = 0 +LLAMA_TOKEN_TYPE_NORMAL = 1 +LLAMA_TOKEN_TYPE_UNKNOWN = 2 +LLAMA_TOKEN_TYPE_CONTROL = 3 +LLAMA_TOKEN_TYPE_USER_DEFINED = 4 +LLAMA_TOKEN_TYPE_UNUSED = 5 +LLAMA_TOKEN_TYPE_BYTE = 6 + + +# enum llama_token_attr { +# LLAMA_TOKEN_ATTR_UNDEFINED = 0, +# LLAMA_TOKEN_ATTR_UNKNOWN = 1 << 0, +# LLAMA_TOKEN_ATTR_UNUSED = 1 << 1, +# LLAMA_TOKEN_ATTR_NORMAL = 1 << 2, +# LLAMA_TOKEN_ATTR_CONTROL = 1 << 3, // SPECIAL? +# LLAMA_TOKEN_ATTR_USER_DEFINED = 1 << 4, +# LLAMA_TOKEN_ATTR_BYTE = 1 << 5, +# LLAMA_TOKEN_ATTR_NORMALIZED = 1 << 6, +# LLAMA_TOKEN_ATTR_LSTRIP = 1 << 7, +# LLAMA_TOKEN_ATTR_RSTRIP = 1 << 8, +# LLAMA_TOKEN_ATTR_SINGLE_WORD = 1 << 9, +# }; +LLAMA_TOKEN_ATTR_UNDEFINED = 0 +LLAMA_TOKEN_ATTR_UNKNOWN = 1 << 0 +LLAMA_TOKEN_ATTR_UNUSED = 1 << 1 +LLAMA_TOKEN_ATTR_NORMAL = 1 << 2 +LLAMA_TOKEN_ATTR_CONTROL = 1 << 3 +LLAMA_TOKEN_ATTR_USER_DEFINED = 1 << 4 +LLAMA_TOKEN_ATTR_BYTE = 1 << 5 +LLAMA_TOKEN_ATTR_NORMALIZED = 1 << 6 +LLAMA_TOKEN_ATTR_LSTRIP = 1 << 7 +LLAMA_TOKEN_ATTR_RSTRIP = 1 << 8 +LLAMA_TOKEN_ATTR_SINGLE_WORD = 1 << 9 + + +# // model file types +# enum llama_ftype { +# LLAMA_FTYPE_ALL_F32 = 0, +# LLAMA_FTYPE_MOSTLY_F16 = 1, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_Q4_0 = 2, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_Q4_1 = 3, // except 1d tensors +# // LLAMA_FTYPE_MOSTLY_Q4_1_SOME_F16 = 4, // tok_embeddings.weight and output.weight are F16 +# // LLAMA_FTYPE_MOSTLY_Q4_2 = 5, // support has been removed +# // LLAMA_FTYPE_MOSTLY_Q4_3 = 6, // support has been removed +# LLAMA_FTYPE_MOSTLY_Q8_0 = 7, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_Q5_0 = 8, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_Q5_1 = 9, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_Q2_K = 10, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_Q3_K_S = 11, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_Q3_K_M = 12, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_Q3_K_L = 13, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_Q4_K_S = 14, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_Q4_K_M = 15, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_Q5_K_S = 16, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_Q5_K_M = 17, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_Q6_K = 18, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_IQ2_XXS = 19, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_IQ2_XS = 20, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_Q2_K_S = 21, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_IQ3_XS = 22, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_IQ3_XXS = 23, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_IQ1_S = 24, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_IQ4_NL = 25, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_IQ3_S = 26, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_IQ3_M = 27, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_IQ2_S = 28, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_IQ2_M = 29, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_IQ4_XS = 30, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_IQ1_M = 31, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_BF16 = 32, // except 1d tensors +# //LLAMA_FTYPE_MOSTLY_Q4_0_4_4 = 33, // removed from gguf files, use Q4_0 and runtime repack +# //LLAMA_FTYPE_MOSTLY_Q4_0_4_8 = 34, // removed from gguf files, use Q4_0 and runtime repack +# //LLAMA_FTYPE_MOSTLY_Q4_0_8_8 = 35, // removed from gguf files, use Q4_0 and runtime repack +# LLAMA_FTYPE_MOSTLY_TQ1_0 = 36, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_TQ2_0 = 37, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_MXFP4_MOE = 38, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_NVFP4 = 39, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_Q1_0 = 40, // except 1d tensors +# +# LLAMA_FTYPE_GUESSED = 1024, // not specified in the model file +# }; +LLAMA_FTYPE_ALL_F32 = 0 +LLAMA_FTYPE_MOSTLY_F16 = 1 +LLAMA_FTYPE_MOSTLY_Q4_0 = 2 +LLAMA_FTYPE_MOSTLY_Q4_1 = 3 +LLAMA_FTYPE_MOSTLY_Q8_0 = 7 +LLAMA_FTYPE_MOSTLY_Q5_0 = 8 +LLAMA_FTYPE_MOSTLY_Q5_1 = 9 +LLAMA_FTYPE_MOSTLY_Q2_K = 10 +LLAMA_FTYPE_MOSTLY_Q3_K_S = 11 +LLAMA_FTYPE_MOSTLY_Q3_K_M = 12 +LLAMA_FTYPE_MOSTLY_Q3_K_L = 13 +LLAMA_FTYPE_MOSTLY_Q4_K_S = 14 +LLAMA_FTYPE_MOSTLY_Q4_K_M = 15 +LLAMA_FTYPE_MOSTLY_Q5_K_S = 16 +LLAMA_FTYPE_MOSTLY_Q5_K_M = 17 +LLAMA_FTYPE_MOSTLY_Q6_K = 18 +LLAMA_FTYPE_MOSTLY_IQ2_XXS = 19 +LLAMA_FTYPE_MOSTLY_IQ2_XS = 20 +LLAMA_FTYPE_MOSTLY_Q2_K_S = 21 +LLAMA_FTYPE_MOSTLY_IQ3_XS = 22 +LLAMA_FTYPE_MOSTLY_IQ3_XXS = 23 +LLAMA_FTYPE_MOSTLY_IQ1_S = 24 +LLAMA_FTYPE_MOSTLY_IQ4_NL = 25 +LLAMA_FTYPE_MOSTLY_IQ3_S = 26 +LLAMA_FTYPE_MOSTLY_IQ3_M = 27 +LLAMA_FTYPE_MOSTLY_IQ2_S = 28 +LLAMA_FTYPE_MOSTLY_IQ2_M = 29 +LLAMA_FTYPE_MOSTLY_IQ4_XS = 30 +LLAMA_FTYPE_MOSTLY_IQ1_M = 31 +LLAMA_FTYPE_MOSTLY_BF16 = 32 +# LLAMA_FTYPE_MOSTLY_Q4_0_4_4 = 33 +# LLAMA_FTYPE_MOSTLY_Q4_0_4_8 = 34 +# LLAMA_FTYPE_MOSTLY_Q4_0_8_8 = 35 +LLAMA_FTYPE_MOSTLY_TQ1_0 = 36 +LLAMA_FTYPE_MOSTLY_TQ2_0 = 37 +LLAMA_FTYPE_MOSTLY_MXFP4_MOE = 38 +LLAMA_FTYPE_MOSTLY_NVFP4 = 39 +LLAMA_FTYPE_MOSTLY_Q1_0 = 40 +LLAMA_FTYPE_GUESSED = 1024 + +# enum llama_rope_scaling_type { +# LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED = -1, +# LLAMA_ROPE_SCALING_TYPE_NONE = 0, +# LLAMA_ROPE_SCALING_TYPE_LINEAR = 1, +# LLAMA_ROPE_SCALING_TYPE_YARN = 2, +# LLAMA_ROPE_SCALING_TYPE_LONGROPE = 3, +# LLAMA_ROPE_SCALING_TYPE_MAX_VALUE = LLAMA_ROPE_SCALING_TYPE_LONGROPE, +# }; +LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED = -1 +LLAMA_ROPE_SCALING_TYPE_NONE = 0 +LLAMA_ROPE_SCALING_TYPE_LINEAR = 1 +LLAMA_ROPE_SCALING_TYPE_YARN = 2 +LLAMA_ROPE_SCALING_TYPE_LONGROPE = 3 +LLAMA_ROPE_SCALING_TYPE_MAX_VALUE = LLAMA_ROPE_SCALING_TYPE_LONGROPE + +# enum llama_pooling_type { +# LLAMA_POOLING_TYPE_UNSPECIFIED = -1, +# LLAMA_POOLING_TYPE_NONE = 0, +# LLAMA_POOLING_TYPE_MEAN = 1, +# LLAMA_POOLING_TYPE_CLS = 2, +# LLAMA_POOLING_TYPE_LAST = 3, +# LLAMA_POOLING_TYPE_RANK = 4, // used by reranking models to attach the classification head to the graph +# }; +LLAMA_POOLING_TYPE_UNSPECIFIED = -1 +LLAMA_POOLING_TYPE_NONE = 0 +LLAMA_POOLING_TYPE_MEAN = 1 +LLAMA_POOLING_TYPE_CLS = 2 +LLAMA_POOLING_TYPE_LAST = 3 +LLAMA_POOLING_TYPE_RANK = 4 + +# enum llama_attention_type { +# LLAMA_ATTENTION_TYPE_UNSPECIFIED = -1, +# LLAMA_ATTENTION_TYPE_CAUSAL = 0, +# LLAMA_ATTENTION_TYPE_NON_CAUSAL = 1, +# }; +LLAMA_ATTENTION_TYPE_UNSPECIFIED = -1 +LLAMA_ATTENTION_TYPE_CAUSAL = 0 +LLAMA_ATTENTION_TYPE_NON_CAUSAL = 1 + + +# enum llama_flash_attn_type { +# LLAMA_FLASH_ATTN_TYPE_AUTO = -1, +# LLAMA_FLASH_ATTN_TYPE_DISABLED = 0, +# LLAMA_FLASH_ATTN_TYPE_ENABLED = 1, +# }; +LLAMA_FLASH_ATTN_TYPE_AUTO = -1 +LLAMA_FLASH_ATTN_TYPE_DISABLED = 0 +LLAMA_FLASH_ATTN_TYPE_ENABLED = 1 + + +# enum llama_split_mode { +# LLAMA_SPLIT_MODE_NONE = 0, // single GPU +# LLAMA_SPLIT_MODE_LAYER = 1, // split layers and KV across GPUs +# LLAMA_SPLIT_MODE_ROW = 2, // split layers and KV across GPUs, use tensor parallelism if supported +# LLAMA_SPLIT_MODE_TENSOR = 3, +# }; +LLAMA_SPLIT_MODE_NONE = 0 +LLAMA_SPLIT_MODE_LAYER = 1 +LLAMA_SPLIT_MODE_ROW = 2 +LLAMA_SPLIT_MODE_TENSOR = 3 -# typedef int llama_token; -llama_token = c_int -llama_token_p = POINTER(llama_token) +# enum llama_context_type { +# LLAMA_CONTEXT_TYPE_DEFAULT = 0, +# LLAMA_CONTEXT_TYPE_MTP = 1, +# }; +LLAMA_CONTEXT_TYPE_DEFAULT = 0 +LLAMA_CONTEXT_TYPE_MTP = 1 # typedef struct llama_token_data { @@ -128,940 +529,4506 @@ def _load_shared_library(lib_base_name: str): # float logit; // log-odds of the token # float p; // probability of the token # } llama_token_data; -class llama_token_data(Structure): +class llama_token_data(ctypes.Structure): + """Used to store token data + + Attributes: + id (llama_token): token id + logit (float): log-odds of the token + p (float): probability of the token""" + + if TYPE_CHECKING: + id: llama_token + logit: float + p: float + _fields_ = [ ("id", llama_token), - ("logit", c_float), - ("p", c_float), + ("logit", ctypes.c_float), + ("p", ctypes.c_float), ] -llama_token_data_p = POINTER(llama_token_data) +llama_token_data_p = ctypes.POINTER(llama_token_data) # typedef struct llama_token_data_array { +# // TODO: consider SoA +# // NOTE: this pointer can be modified by the samplers # llama_token_data * data; # size_t size; +# int64_t selected; // this is the index in the data array (i.e. not the token id) # bool sorted; # } llama_token_data_array; -class llama_token_data_array(Structure): +class llama_token_data_array(ctypes.Structure): + """Used to sample tokens given logits + + Attributes: + data (ctypes.Array[llama_token_data]): token data + size (int): size of the array + selected (int): index in the data array (i.e. not the token id) + sorted (bool): whether the array is sorted""" + + if TYPE_CHECKING: + data: CtypesArray[llama_token_data] + size: int + selected: int + sorted: bool + _fields_ = [ ("data", llama_token_data_p), - ("size", c_size_t), - ("sorted", c_bool), + ("size", ctypes.c_size_t), + ("selected", ctypes.c_int64), + ("sorted", ctypes.c_bool), ] -llama_token_data_array_p = POINTER(llama_token_data_array) +llama_token_data_array_p = ctypes.POINTER(llama_token_data_array) + +# typedef bool (*llama_progress_callback)(float progress, void * user_data); +llama_progress_callback = ctypes.CFUNCTYPE( + ctypes.c_bool, ctypes.c_float, ctypes.c_void_p +) -# typedef void (*llama_progress_callback)(float progress, void *ctx); -llama_progress_callback = ctypes.CFUNCTYPE(None, c_float, c_void_p) +# // Input data for llama_encode/llama_decode +# // A llama_batch object can contain input about one or many sequences +# // The provided arrays (i.e. token, embd, pos, etc.) must have size of n_tokens +# // +# // - token : the token ids of the input (used when embd is NULL) +# // - embd : token embeddings (i.e. float vector of size n_embd) (used when token is NULL) +# // - pos : the positions of the respective token in the sequence +# // (if set to NULL, the token position will be tracked automatically by llama_encode/llama_decode) +# // - seq_id : the sequence to which the respective token belongs +# // (if set to NULL, the sequence ID will be assumed to be 0) +# // - logits : if zero, the logits (and/or the embeddings) for the respective token will not be output +# // (if set to NULL: +# // - if embeddings: all tokens are output +# // - if not: only the last token is output +# // ) +# // +# typedef struct llama_batch { +# int32_t n_tokens; + + +# llama_token * token; +# float * embd; +# llama_pos * pos; +# int32_t * n_seq_id; +# llama_seq_id ** seq_id; +# int8_t * logits; // TODO: rename this to "output" +# } llama_batch; +class llama_batch(ctypes.Structure): + """Input data for llama_encode/llama_decode + + A llama_batch object can contain input about one or many sequences + + The provided arrays (i.e. token, embd, pos, etc.) must have size of n_tokens + + Attributes: + n_tokens (int): number of tokens + token (ctypes.Array[llama_token]): the token ids of the input (used when embd is NULL) + embd (ctypes.Array[ctypes.ctypes.c_float]): token embeddings (i.e. float vector of size n_embd) (used when token is NULL) + pos (ctypes.Array[ctypes.Array[llama_pos]]): the positions of the respective token in the sequence + seq_id (ctypes.Array[ctypes.Array[llama_seq_id]]): the sequence to which the respective token belongs + logits (ctypes.Array[ctypes.ctypes.c_int8]): if zero, the logits for the respective token will not be output + """ + + if TYPE_CHECKING: + n_tokens: int + token: CtypesArray[llama_token] + embd: CtypesArray[ctypes.c_float] + pos: CtypesArray[CtypesArray[llama_pos]] + n_seq_id: CtypesArray[ctypes.c_int] + seq_id: CtypesArray[CtypesArray[llama_seq_id]] + logits: CtypesArray[ctypes.c_int8] -# struct llama_context_params { -# uint32_t seed; // RNG seed, -1 for random -# int32_t n_ctx; // text context -# int32_t n_batch; // prompt processing batch size -# int32_t n_gpu_layers; // number of layers to store in VRAM -# int32_t main_gpu; // the GPU that is used for scratch and small tensors -# float tensor_split[LLAMA_MAX_DEVICES]; // how to split layers across multiple GPUs -# // called with a progress value between 0 and 1, pass NULL to disable -# llama_progress_callback progress_callback; -# // context pointer passed to the progress callback -# void * progress_callback_user_data; + _fields_ = [ + ("n_tokens", ctypes.c_int32), + ("token", ctypes.POINTER(llama_token)), + ("embd", ctypes.POINTER(ctypes.c_float)), + ("pos", ctypes.POINTER(llama_pos)), + ("n_seq_id", ctypes.POINTER(ctypes.c_int32)), + ("seq_id", ctypes.POINTER(ctypes.POINTER(llama_seq_id))), + ("logits", ctypes.POINTER(ctypes.c_int8)), + ] -# // Keep the booleans together to avoid misalignment during copy-by-value. -# bool low_vram; // if true, reduce VRAM usage at the cost of performance -# bool f16_kv; // use fp16 for KV cache -# bool logits_all; // the llama_eval() call computes all logits, not just the last one -# bool vocab_only; // only load the vocabulary, no weights -# bool use_mmap; // use mmap if possible -# bool use_mlock; // force system to keep model in RAM -# bool embedding; // embedding mode only +# enum llama_model_kv_override_type { +# LLAMA_KV_OVERRIDE_TYPE_INT, +# LLAMA_KV_OVERRIDE_TYPE_FLOAT, +# LLAMA_KV_OVERRIDE_TYPE_BOOL, +# LLAMA_KV_OVERRIDE_TYPE_STR, # }; -class llama_context_params(Structure): +LLAMA_KV_OVERRIDE_TYPE_INT = 0 +LLAMA_KV_OVERRIDE_TYPE_FLOAT = 1 +LLAMA_KV_OVERRIDE_TYPE_BOOL = 2 +LLAMA_KV_OVERRIDE_TYPE_STR = 3 + + +# enum llama_model_meta_key { +# LLAMA_MODEL_META_KEY_SAMPLING_SEQUENCE, +# LLAMA_MODEL_META_KEY_SAMPLING_TOP_K, +# LLAMA_MODEL_META_KEY_SAMPLING_TOP_P, +# LLAMA_MODEL_META_KEY_SAMPLING_MIN_P, +# LLAMA_MODEL_META_KEY_SAMPLING_XTC_PROBABILITY, +# LLAMA_MODEL_META_KEY_SAMPLING_XTC_THRESHOLD, +# LLAMA_MODEL_META_KEY_SAMPLING_TEMP, +# LLAMA_MODEL_META_KEY_SAMPLING_PENALTY_LAST_N, +# LLAMA_MODEL_META_KEY_SAMPLING_PENALTY_REPEAT, +# LLAMA_MODEL_META_KEY_SAMPLING_MIROSTAT, +# LLAMA_MODEL_META_KEY_SAMPLING_MIROSTAT_TAU, +# LLAMA_MODEL_META_KEY_SAMPLING_MIROSTAT_ETA, +# }; +LLAMA_MODEL_META_KEY_SAMPLING_SEQUENCE = 0 +LLAMA_MODEL_META_KEY_SAMPLING_TOP_K = 1 +LLAMA_MODEL_META_KEY_SAMPLING_TOP_P = 2 +LLAMA_MODEL_META_KEY_SAMPLING_MIN_P = 3 +LLAMA_MODEL_META_KEY_SAMPLING_XTC_PROBABILITY = 4 +LLAMA_MODEL_META_KEY_SAMPLING_XTC_THRESHOLD = 5 +LLAMA_MODEL_META_KEY_SAMPLING_TEMP = 6 +LLAMA_MODEL_META_KEY_SAMPLING_PENALTY_LAST_N = 7 +LLAMA_MODEL_META_KEY_SAMPLING_PENALTY_REPEAT = 8 +LLAMA_MODEL_META_KEY_SAMPLING_MIROSTAT = 9 +LLAMA_MODEL_META_KEY_SAMPLING_MIROSTAT_TAU = 10 +LLAMA_MODEL_META_KEY_SAMPLING_MIROSTAT_ETA = 11 + + +# struct llama_model_kv_override { +# enum llama_model_kv_override_type tag; + +# char key[128]; + + +# union { +# int64_t val_i64; +# double val_f64; +# bool val_bool; +# char val_str[128]; +# }; +# }; +class llama_model_kv_override_value(ctypes.Union): _fields_ = [ - ("seed", c_uint32), - ("n_ctx", c_int32), - ("n_batch", c_int32), - ("n_gpu_layers", c_int32), - ("main_gpu", c_int32), - ("tensor_split", c_float * LLAMA_MAX_DEVICES.value), - ("progress_callback", llama_progress_callback), - ("progress_callback_user_data", c_void_p), - ("low_vram", c_bool), - ("f16_kv", c_bool), - ("logits_all", c_bool), - ("vocab_only", c_bool), - ("use_mmap", c_bool), - ("use_mlock", c_bool), - ("embedding", c_bool), + ("val_i64", ctypes.c_int64), + ("val_f64", ctypes.c_double), + ("val_bool", ctypes.c_bool), + ("val_str", ctypes.c_char * 128), ] + if TYPE_CHECKING: + val_i64: int + val_f64: float + val_bool: bool + val_str: bytes -llama_context_params_p = POINTER(llama_context_params) -# enum llama_ftype { -# LLAMA_FTYPE_ALL_F32 = 0, -# LLAMA_FTYPE_MOSTLY_F16 = 1, // except 1d tensors -# LLAMA_FTYPE_MOSTLY_Q4_0 = 2, // except 1d tensors -# LLAMA_FTYPE_MOSTLY_Q4_1 = 3, // except 1d tensors -# LLAMA_FTYPE_MOSTLY_Q4_1_SOME_F16 = 4, // tok_embeddings.weight and output.weight are F16 -# // LLAMA_FTYPE_MOSTLY_Q4_2 = 5, // support has been removed -# // LLAMA_FTYPE_MOSTLY_Q4_3 = 6, // support has been removed -# LLAMA_FTYPE_MOSTLY_Q8_0 = 7, // except 1d tensors -# LLAMA_FTYPE_MOSTLY_Q5_0 = 8, // except 1d tensors -# LLAMA_FTYPE_MOSTLY_Q5_1 = 9, // except 1d tensors -# LLAMA_FTYPE_MOSTLY_Q2_K = 10,// except 1d tensors -# LLAMA_FTYPE_MOSTLY_Q3_K_S = 11,// except 1d tensors -# LLAMA_FTYPE_MOSTLY_Q3_K_M = 12,// except 1d tensors -# LLAMA_FTYPE_MOSTLY_Q3_K_L = 13,// except 1d tensors -# LLAMA_FTYPE_MOSTLY_Q4_K_S = 14,// except 1d tensors -# LLAMA_FTYPE_MOSTLY_Q4_K_M = 15,// except 1d tensors -# LLAMA_FTYPE_MOSTLY_Q5_K_S = 16,// except 1d tensors -# LLAMA_FTYPE_MOSTLY_Q5_K_M = 17,// except 1d tensors -# LLAMA_FTYPE_MOSTLY_Q6_K = 18,// except 1d tensors -# }; -LLAMA_FTYPE_ALL_F32 = c_int(0) -LLAMA_FTYPE_MOSTLY_F16 = c_int(1) -LLAMA_FTYPE_MOSTLY_Q4_0 = c_int(2) -LLAMA_FTYPE_MOSTLY_Q4_1 = c_int(3) -LLAMA_FTYPE_MOSTLY_Q4_1_SOME_F16 = c_int(4) -LLAMA_FTYPE_MOSTLY_Q8_0 = c_int(7) -LLAMA_FTYPE_MOSTLY_Q5_0 = c_int(8) -LLAMA_FTYPE_MOSTLY_Q5_1 = c_int(9) -LLAMA_FTYPE_MOSTLY_Q2_K = c_int(10) -LLAMA_FTYPE_MOSTLY_Q3_K_S = c_int(11) -LLAMA_FTYPE_MOSTLY_Q3_K_M = c_int(12) -LLAMA_FTYPE_MOSTLY_Q3_K_L = c_int(13) -LLAMA_FTYPE_MOSTLY_Q4_K_S = c_int(14) -LLAMA_FTYPE_MOSTLY_Q4_K_M = c_int(15) -LLAMA_FTYPE_MOSTLY_Q5_K_S = c_int(16) -LLAMA_FTYPE_MOSTLY_Q5_K_M = c_int(17) -LLAMA_FTYPE_MOSTLY_Q6_K = c_int(18) +class llama_model_kv_override(ctypes.Structure): + _fields_ = [ + ("tag", ctypes.c_int), + ("key", ctypes.c_char * 128), + ("value", llama_model_kv_override_value), + ] + if TYPE_CHECKING: + tag: int + key: bytes + value: Union[int, float, bool, bytes] + + +# struct llama_model_tensor_override { +# const char * pattern; +# enum ggml_type type; +# }; +class llama_model_tensor_override(ctypes.Structure): + """Override the quantization type for tensors matching a pattern.""" -# // model quantization parameters -# typedef struct llama_model_quantize_params { -# int nthread; // number of threads to use for quantizing, if <=0 will use std::thread::hardware_concurrency() -# enum llama_ftype ftype; // quantize to this llama_ftype -# bool allow_requantize; // allow quantizing non-f32/f16 tensors -# bool quantize_output_tensor; // quantize output.weight -# } llama_model_quantize_params; -class llama_model_quantize_params(Structure): _fields_ = [ - ("nthread", c_int), - ("ftype", c_int), - ("allow_requantize", c_bool), - ("quantize_output_tensor", c_bool), + ("pattern", ctypes.c_char_p), + ("type", ctypes.c_int), ] - -# // performance timing information -# struct llama_timings { -# double t_start_ms; -# double t_end_ms; -# double t_load_ms; -# double t_sample_ms; -# double t_p_eval_ms; -# double t_eval_ms; + if TYPE_CHECKING: + pattern: Optional[bytes] + type: int -# int32_t n_sample; -# int32_t n_p_eval; -# int32_t n_eval; +# struct llama_model_imatrix_data { +# const char * name; +# const float * data; +# size_t size; # }; -class llama_timings(Structure): +class llama_model_imatrix_data(ctypes.Structure): + """Importance matrix data for a tensor used during quantization.""" + _fields_ = [ - ("t_start_ms", c_double), - ("t_end_ms", c_double), - ("t_load_ms", c_double), - ("t_sample_ms", c_double), - ("t_p_eval_ms", c_double), - ("t_eval_ms", c_double), - ("n_sample", c_int32), - ("n_p_eval", c_int32), - ("n_eval", c_int32), + ("name", ctypes.c_char_p), + ("data", ctypes.POINTER(ctypes.c_float)), + ("size", ctypes.c_size_t), ] + if TYPE_CHECKING: + name: Optional[bytes] + data: CtypesPointer[ctypes.c_float] + size: int -# LLAMA_API struct llama_context_params llama_context_default_params(); -def llama_context_default_params() -> llama_context_params: - return _lib.llama_context_default_params() +# struct llama_model_tensor_buft_override { +# const char * pattern; +# ggml_backend_buffer_type_t buft; +# }; -_lib.llama_context_default_params.argtypes = [] -_lib.llama_context_default_params.restype = llama_context_params +# struct llama_model_params { +# // NULL-terminated list of devices to use for offloading (if NULL, all available devices are used) +# ggml_backend_dev_t * devices; -# LLAMA_API struct llama_model_quantize_params llama_model_quantize_default_params(); -def llama_model_quantize_default_params() -> llama_model_quantize_params: - return _lib.llama_model_quantize_default_params() +# // NULL-terminated list of buffer types to use for tensors that match a pattern +# const struct llama_model_tensor_buft_override * tensor_buft_overrides; +# int32_t n_gpu_layers; // number of layers to store in VRAM +# enum llama_split_mode split_mode; // how to split the model across multiple GPUs -_lib.llama_model_quantize_default_params.argtypes = [] -_lib.llama_model_quantize_default_params.restype = llama_model_quantize_params +# // the GPU that is used for the entire model when split_mode is LLAMA_SPLIT_MODE_NONE +# int32_t main_gpu; +# // proportion of the model (layers or rows) to offload to each GPU, size: llama_max_devices() +# const float * tensor_split; -# LLAMA_API bool llama_mmap_supported(); -def llama_mmap_supported() -> bool: - return _lib.llama_mmap_supported() +# // Called with a progress value between 0.0 and 1.0. Pass NULL to disable. +# // If the provided progress_callback returns true, model loading continues. +# // If it returns false, model loading is immediately aborted. +# llama_progress_callback progress_callback; +# // context pointer passed to the progress callback +# void * progress_callback_user_data; -_lib.llama_mmap_supported.argtypes = [] -_lib.llama_mmap_supported.restype = c_bool +# // override key-value pairs of the model meta data +# const struct llama_model_kv_override * kv_overrides; -# LLAMA_API bool llama_mlock_supported(); -def llama_mlock_supported() -> bool: - return _lib.llama_mlock_supported() +# // Keep the booleans together to avoid misalignment during copy-by-value. +# bool vocab_only; // only load the vocabulary, no weights +# bool use_mmap; // use mmap if possible +# bool use_direct_io; // use direct io, takes precedence over use_mmap when supported +# bool use_mlock; // force system to keep model in RAM +# bool check_tensors; // validate model tensor data +# bool use_extra_bufts; // use extra buffer types (used for weight repacking) +# bool no_host; // bypass host buffer allowing extra buffers to be used +# bool no_alloc; // only load metadata and simulate memory allocations +# }; +class llama_model_params(ctypes.Structure): + """Parameters for llama_model + + Attributes: + devices (ctypes.Array[ggml_backend_dev_t]): NULL-terminated list of devices to use for offloading (if NULL, all available devices are used) + tensor_buft_overrides (ctypes.Array[llama_model_tensor_buft_override]): NULL-terminated list of buffer types to use for tensors that match a pattern + n_gpu_layers (int): number of layers to store in VRAM + split_mode (int): how to split the model across multiple GPUs + main_gpu (int): the GPU that is used for the entire model when split_mode is LLAMA_SPLIT_MODE_NONE + tensor_split (ctypes.Array[ctypes.ctypes.c_float]): proportion of the model (layers or rows) to offload to each GPU, size: llama_max_devices() + progress_callback (llama_progress_callback): called with a progress value between 0.0 and 1.0. Pass NULL to disable. If the provided progress_callback returns true, model loading continues. If it returns false, model loading is immediately aborted. + progress_callback_user_data (ctypes.ctypes.c_void_p): context pointer passed to the progress callback + kv_overrides (ctypes.Array[llama_model_kv_override]): override key-value pairs of the model meta data + vocab_only (bool): only load the vocabulary, no weights + use_mmap (bool): use mmap if possible + use_direct_io (bool): use direct io, takes precedence over use_mmap when supported + use_mlock (bool): force system to keep model in RAM + check_tensors (bool): validate model tensor data + use_extra_bufts (bool): use extra buffer types (used for weight repacking) + no_host (bool): bypass host buffer allowing extra buffers to be used + no_alloc (bool): only load metadata and simulate memory allocations""" + + if TYPE_CHECKING: + devices: CtypesArray[ctypes.c_void_p] # NOTE: unused + tensor_buft_overrides: CtypesArray[ + llama_model_tensor_buft_override + ] # NOTE: unused + n_gpu_layers: int + split_mode: int + main_gpu: int + tensor_split: CtypesArray[ctypes.c_float] + progress_callback: Callable[[float, ctypes.c_void_p], bool] + progress_callback_user_data: ctypes.c_void_p + kv_overrides: CtypesArray[llama_model_kv_override] + vocab_only: bool + use_mmap: bool + use_direct_io: bool + use_mlock: bool + check_tensors: bool + use_extra_bufts: bool + no_host: bool + no_alloc: bool + _fields_ = [ + ("devices", ctypes.c_void_p), # NOTE: unnused + ("tensor_buft_overrides", ctypes.c_void_p), # NOTE: unused + ("n_gpu_layers", ctypes.c_int32), + ("split_mode", ctypes.c_int), + ("main_gpu", ctypes.c_int32), + ("tensor_split", ctypes.POINTER(ctypes.c_float)), + ("progress_callback", llama_progress_callback), + ("progress_callback_user_data", ctypes.c_void_p), + ("kv_overrides", ctypes.POINTER(llama_model_kv_override)), + ("vocab_only", ctypes.c_bool), + ("use_mmap", ctypes.c_bool), + ("use_direct_io", ctypes.c_bool), + ("use_mlock", ctypes.c_bool), + ("check_tensors", ctypes.c_bool), + ("use_extra_bufts", ctypes.c_bool), + ("no_host", ctypes.c_bool), + ("no_alloc", ctypes.c_bool), + ] -_lib.llama_mlock_supported.argtypes = [] -_lib.llama_mlock_supported.restype = c_bool +# struct llama_sampler_seq_config { +# llama_seq_id seq_id; +# struct llama_sampler * sampler; +# }; +class llama_sampler_seq_config(ctypes.Structure): + if TYPE_CHECKING: + seq_id: int + sampler: ctypes.c_void_p -# // TODO: not great API - very likely to change -# // Initialize the llama + ggml backend -# // If numa is true, use NUMA optimizations -# // Call once at the start of the program -# LLAMA_API void llama_init_backend(bool numa); -def llama_init_backend(numa: c_bool): - return _lib.llama_init_backend(numa) + _fields_ = [ + ("seq_id", llama_seq_id), + ("sampler", ctypes.c_void_p), + ] -_lib.llama_init_backend.argtypes = [c_bool] -_lib.llama_init_backend.restype = None +# // NOTE: changing the default values of parameters marked as [EXPERIMENTAL] may cause crashes or incorrect results in certain configurations +# // https://github.com/ggml-org/llama.cpp/pull/7544 +# struct llama_context_params { +# uint32_t n_ctx; // text context, 0 = from model +# uint32_t n_batch; // logical maximum batch size that can be submitted to llama_decode +# uint32_t n_ubatch; // physical maximum batch size +# uint32_t n_seq_max; // max number of sequences (i.e. distinct states for recurrent models) +# uint32_t n_rs_seq; // number of recurrent-state snapshots per seq for rollback (0 = no rollback) [EXPERIMENTAL] +# uint32_t n_outputs_max; // max outputs in a ubatch (0 = n_batch) +# int32_t n_threads; // number of threads to use for generation +# int32_t n_threads_batch; // number of threads to use for batch processing + +# enum llama_context_type ctx_type; // set the context type (e.g. MTP) +# enum llama_rope_scaling_type rope_scaling_type; // RoPE scaling type, from `enum llama_rope_scaling_type` +# enum llama_pooling_type pooling_type; // whether to pool (sum) embedding results by sequence id +# enum llama_attention_type attention_type; // attention type to use for embeddings +# enum llama_flash_attn_type flash_attn_type; // when to enable Flash Attention + +# // ref: https://github.com/ggml-org/llama.cpp/pull/2054 +# float rope_freq_base; // RoPE base frequency, 0 = from model +# float rope_freq_scale; // RoPE frequency scaling factor, 0 = from model +# float yarn_ext_factor; // YaRN extrapolation mix factor, negative = from model +# float yarn_attn_factor; // YaRN magnitude scaling factor +# float yarn_beta_fast; // YaRN low correction dim +# float yarn_beta_slow; // YaRN high correction dim +# uint32_t yarn_orig_ctx; // YaRN original context size +# float defrag_thold; // defragment the KV cache if holes/size > thold, <= 0 disabled (default) + +# ggml_backend_sched_eval_callback cb_eval; +# void * cb_eval_user_data; + +# enum ggml_type type_k; // data type for K cache [EXPERIMENTAL] +# enum ggml_type type_v; // data type for V cache [EXPERIMENTAL] + +# // Abort callback +# // if it returns true, execution of llama_decode() will be aborted +# // currently works only with CPU execution +# ggml_abort_callback abort_callback; +# void * abort_callback_data; + + +# // Keep the booleans together and at the end of the struct to avoid misalignment during copy-by-value. +# bool embeddings; // if true, extract embeddings (together with logits) +# bool offload_kqv; // offload the KQV ops (including the KV cache) to GPU +# bool no_perf; // measure performance timings +# bool op_offload; // offload host tensor operations to device +# bool swa_full; // use full-size SWA cache (https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055) +# // NOTE: setting to false when n_seq_max > 1 can cause bad performance in some cases +# // ref: https://github.com/ggml-org/llama.cpp/pull/13845#issuecomment-2924800573 +# bool kv_unified; // use a unified buffer across the input sequences when computing the attention +# // try to disable when n_seq_max > 1 for improved performance when the sequences do not share a large prefix +# // ref: https://github.com/ggml-org/llama.cpp/pull/14363 +# struct llama_sampler_seq_config * samplers; +# size_t n_samplers; +# +# // a source/target/parent context +# // can be utilized in various ways, for example by sharing results or llama_memory between 2 contexts +# struct llama_context * ctx_other; +# }; +class llama_context_params(ctypes.Structure): + """Parameters for llama_context + + Attributes: + n_ctx (int): text context, 0 = from model + n_batch (int): logical maximum batch size that can be submitted to llama_decode + n_ubatch (int): physical maximum batch size + n_seq_max (int): max number of sequences (i.e. distinct states for recurrent models) + n_rs_seq (int): number of recurrent-state snapshots per sequence for rollback + n_outputs_max (int): max outputs in a ubatch, 0 = n_batch + n_threads (int): number of threads to use for generation + n_threads_batch (int): number of threads to use for batch processing + ctx_type (int): context type, from `enum llama_context_type` + rope_scaling_type (int): RoPE scaling type, from `enum llama_rope_scaling_type` + pooling_type (int): whether to pool (sum) embedding results by sequence id (ignored if no pooling layer) + attention_type (int): attention type to use for embeddings + flash_attn_type (int): when to enable flash attention + rope_freq_base (float): RoPE base frequency, 0 = from model + rope_freq_scale (float): RoPE frequency scaling factor, 0 = from model + yarn_ext_factor (float): YaRN extrapolation mix factor, negative = from model + yarn_attn_factor (float): YaRN magnitude scaling factor + yarn_beta_fast (float): YaRN low correction dim + yarn_beta_slow (float): YaRN high correction dim + yarn_orig_ctx (int): YaRN original context size + defrag_thold (float): defragment the KV cache if holes/size > thold, <= 0 disabled (default) + cb_eval (ggml_backend_sched_eval_callback): callback for scheduling eval + cb_eval_user_data (ctypes.ctypes.c_void_p): user data for cb_eval + type_k (int): data type for K cache + type_v (int): data type for V cache + abort_callback (ggml_abort_callback): abort callback if it returns true, execution of llama_decode() will be aborted + abort_callback_data (ctypes.ctypes.c_void_p): data for abort_callback + embeddings (bool): if true, extract embeddings (together with logits) + offload_kqv (bool): whether to offload the KQV ops (including the KV cache) to GPU + no_perf (bool): whether to measure performance timings + op_offload (bool): offload host tensor operations to device + swa_full (bool): use full-size SWA cache + kv_unified (bool): use a unified buffer across the input sequences when computing the attention + samplers (ctypes.POINTER(llama_sampler_seq_config)): backend sampler chain configuration + n_samplers (int): number of backend sampler chain configurations + ctx_other (llama_context_p): source, target, or parent context + """ + + if TYPE_CHECKING: + n_ctx: int + n_batch: int + n_ubatch: int + n_seq_max: int + n_rs_seq: int + n_outputs_max: int + n_threads: int + n_threads_batch: int + ctx_type: int + rope_scaling_type: int + pooling_type: int + attention_type: int + flash_attn_type: int + rope_freq_base: float + rope_freq_scale: float + yarn_ext_factor: float + yarn_attn_factor: float + yarn_beta_fast: float + yarn_beta_slow: float + yarn_orig_ctx: int + defrag_thold: float + cb_eval: Callable[[ctypes.c_void_p, bool], bool] + cb_eval_user_data: ctypes.c_void_p + type_k: int + type_v: int + abort_callback: Callable[[ctypes.c_void_p], bool] + abort_callback_data: ctypes.c_void_p + embeddings: bool + offload_kqv: bool + no_perf: bool + op_offload: bool + swa_full: bool + kv_unified: bool + samplers: ctypes.POINTER(llama_sampler_seq_config) + n_samplers: int + ctx_other: llama_context_p + _fields_ = [ + ("n_ctx", ctypes.c_uint32), + ("n_batch", ctypes.c_uint32), + ("n_ubatch", ctypes.c_uint32), + ("n_seq_max", ctypes.c_uint32), + ("n_rs_seq", ctypes.c_uint32), + ("n_outputs_max", ctypes.c_uint32), + ("n_threads", ctypes.c_int32), + ("n_threads_batch", ctypes.c_int32), + ("ctx_type", ctypes.c_int), + ("rope_scaling_type", ctypes.c_int), + ("pooling_type", ctypes.c_int), + ("attention_type", ctypes.c_int), + ("flash_attn_type", ctypes.c_int), + ("rope_freq_base", ctypes.c_float), + ("rope_freq_scale", ctypes.c_float), + ("yarn_ext_factor", ctypes.c_float), + ("yarn_attn_factor", ctypes.c_float), + ("yarn_beta_fast", ctypes.c_float), + ("yarn_beta_slow", ctypes.c_float), + ("yarn_orig_ctx", ctypes.c_uint32), + ("defrag_thold", ctypes.c_float), + ("cb_eval", ggml_backend_sched_eval_callback), + ("cb_eval_user_data", ctypes.c_void_p), + ("type_k", ctypes.c_int), + ("type_v", ctypes.c_int), + ("abort_callback", ggml_abort_callback), + ("abort_callback_data", ctypes.c_void_p), + ("embeddings", ctypes.c_bool), + ("offload_kqv", ctypes.c_bool), + ("no_perf", ctypes.c_bool), + ("op_offload", ctypes.c_bool), + ("swa_full", ctypes.c_bool), + ("kv_unified", ctypes.c_bool), + ("samplers", ctypes.POINTER(llama_sampler_seq_config)), + ("n_samplers", ctypes.c_size_t), + ("ctx_other", llama_context_p_ctypes), + ] -# LLAMA_API struct llama_model * llama_load_model_from_file( -# const char * path_model, -# struct llama_context_params params); -def llama_load_model_from_file( - path_model: bytes, params: llama_context_params -) -> llama_model_p: - return _lib.llama_load_model_from_file(path_model, params) +# // Signature for logging events +# // Note that text includes the new line character at the end for most events. +# // If your logging mechanism cannot handle that, check if the last character is '\n' and strip it +# // if it exists. +# // It might not exist for progress report where '.' is output repeatedly. +# typedef void (*llama_log_callback)(enum llama_log_level level, const char * text, void * user_data); +llama_log_callback = ctypes.CFUNCTYPE( + None, ctypes.c_int, ctypes.c_char_p, ctypes.c_void_p +) +"""Signature for logging events +Note that text includes the new line character at the end for most events. +If your logging mechanism cannot handle that, check if the last character is '\n' and strip it +if it exists. +It might not exist for progress report where '.' is output repeatedly.""" -_lib.llama_load_model_from_file.argtypes = [c_char_p, llama_context_params] -_lib.llama_load_model_from_file.restype = llama_model_p +# // model quantization parameters +# typedef struct llama_model_quantize_params { +# int32_t nthread; // number of threads to use for quantizing, if <=0 will use std::thread::hardware_concurrency() +# enum llama_ftype ftype; // quantize to this llama_ftype +# enum ggml_type output_tensor_type; // output tensor type +# enum ggml_type token_embedding_type; // token embeddings tensor type +# bool allow_requantize; // allow quantizing non-f32/f16 tensors +# bool quantize_output_tensor; // quantize output.weight +# bool only_copy; // only copy tensors - ftype, allow_requantize and quantize_output_tensor are ignored +# bool pure; // quantize all tensors to the default type +# bool keep_split; // quantize to the same number of shards +# bool dry_run; // calculate and show the final quantization size without performing quantization +# const struct llama_model_imatrix_data * imatrix; // pointer to importance matrix data +# const struct llama_model_kv_override * kv_overrides; // pointer to kv overrides +# const struct llama_model_tensor_override * tt_overrides; // pointer to tensor overrides +# const int32_t * prune_layers; // pointer to layer indices to prune +# } llama_model_quantize_params; +class llama_model_quantize_params(ctypes.Structure): + """Parameters for llama_model_quantize + + Attributes: + nthread (int): number of threads to use for quantizing, if <=0 will use std::thread::hardware_concurrency() + ftype (int): quantize to this llama_ftype + output_tensor_type (int): output tensor type + token_embedding_type (int): token embeddings tensor type + allow_requantize (bool): allow quantizing non-f32/f16 tensors + quantize_output_tensor (bool): quantize output.weight + only_copy (bool): only copy tensors - ftype, allow_requantize and quantize_output_tensor are ignored + pure (bool): quantize all tensors to the default type + keep_split (bool): quantize to the same number of shards + dry_run (bool): calculate and show the final quantization size without performing quantization + imatrix (ctypes.Array[llama_model_imatrix_data]): pointer to importance matrix data + kv_overrides (ctypes.Array[llama_model_kv_override]): pointer to kv overrides + tt_overrides (ctypes.Array[llama_model_tensor_override]): pointer to tensor overrides + prune_layers (ctypes.Array[ctypes.c_int32]): pointer to layer indices to prune + """ + + if TYPE_CHECKING: + nthread: int + ftype: int + output_tensor_type: int + token_embedding_type: int + allow_requantize: bool + quantize_output_tensor: bool + only_copy: bool + pure: bool + keep_split: bool + dry_run: bool + imatrix: CtypesPointer[llama_model_imatrix_data] + kv_overrides: CtypesPointer[llama_model_kv_override] + tt_overrides: CtypesPointer[llama_model_tensor_override] + prune_layers: CtypesPointer[ctypes.c_int32] -# LLAMA_API void llama_free_model(struct llama_model * model); -def llama_free_model(model: llama_model_p): - return _lib.llama_free_model(model) + _fields_ = [ + ("nthread", ctypes.c_int32), + ("ftype", ctypes.c_int), + ("output_tensor_type", ctypes.c_int), + ("token_embedding_type", ctypes.c_int), + ("allow_requantize", ctypes.c_bool), + ("quantize_output_tensor", ctypes.c_bool), + ("only_copy", ctypes.c_bool), + ("pure", ctypes.c_bool), + ("keep_split", ctypes.c_bool), + ("dry_run", ctypes.c_bool), + ("imatrix", ctypes.POINTER(llama_model_imatrix_data)), + ("kv_overrides", ctypes.POINTER(llama_model_kv_override)), + ("tt_overrides", ctypes.POINTER(llama_model_tensor_override)), + ("prune_layers", ctypes.POINTER(ctypes.c_int32)), + ] -_lib.llama_free_model.argtypes = [llama_model_p] -_lib.llama_free_model.restype = None +# typedef struct llama_logit_bias { +# llama_token token; +# float bias; +# } llama_logit_bias; +class llama_logit_bias(ctypes.Structure): + """Used to store logit bias + Attributes: + token (llama_token): token id + bias (float): bias""" -# LLAMA_API struct llama_context * llama_new_context_with_model( -# struct llama_model * model, -# struct llama_context_params params); -def llama_new_context_with_model( - model: llama_model_p, params: llama_context_params -) -> llama_context_p: - return _lib.llama_new_context_with_model(model, params) + if TYPE_CHECKING: + token: llama_token + bias: float + _fields_ = [ + ("token", llama_token), + ("bias", ctypes.c_float), + ] -_lib.llama_new_context_with_model.argtypes = [llama_model_p, llama_context_params] -_lib.llama_new_context_with_model.restype = llama_context_p +llama_logit_bias_p = ctypes.POINTER(llama_logit_bias) -# LLAMA_API int64_t llama_time_us(); -def llama_time_us() -> int: - return _lib.llama_time_us() +# typedef struct llama_sampler_chain_params { +# bool no_perf; // whether to measure performance timings +# } llama_sampler_chain_params; +class llama_sampler_chain_params(ctypes.Structure): + """Parameters for llama_sampler_chain -_lib.llama_time_us.argtypes = [] -_lib.llama_time_us.restype = ctypes.c_int64 + Attributes: + no_perf (bool): whether to measure performance timings""" + if TYPE_CHECKING: + no_perf: bool -# // Various functions for loading a ggml llama model. -# // Allocate (almost) all memory needed for the model. -# // Return NULL on failure -# LLAMA_API struct llama_context * llama_init_from_file( -# const char * path_model, -# struct llama_context_params params); -def llama_init_from_file( - path_model: bytes, params: llama_context_params -) -> llama_context_p: - return _lib.llama_init_from_file(path_model, params) + _fields_ = [ + ("no_perf", ctypes.c_bool), + ] -_lib.llama_init_from_file.argtypes = [c_char_p, llama_context_params] -_lib.llama_init_from_file.restype = llama_context_p +# // used in chat template +# typedef struct llama_chat_message { +# const char * role; +# const char * content; +# } llama_chat_message; +class llama_chat_message(ctypes.Structure): + _fields_ = [ + ("role", ctypes.c_char_p), + ("content", ctypes.c_char_p), + ] -# Frees all allocated memory -# LLAMA_API void llama_free(struct llama_context * ctx); -def llama_free(ctx: llama_context_p): - return _lib.llama_free(ctx) +# // lora adapter +# struct llama_adapter_lora; +llama_adapter_lora_p = ctypes.c_void_p +llama_adapter_lora_p_ctypes = ctypes.POINTER(ctypes.c_void_p) -_lib.llama_free.argtypes = [llama_context_p] -_lib.llama_free.restype = None +# // Helpers for getting default parameters +# LLAMA_API struct llama_model_params llama_model_default_params(void); +@ctypes_function( + "llama_model_default_params", + [], + llama_model_params, +) +def llama_model_default_params() -> llama_model_params: + """Get default parameters for llama_model""" + ... -# // Returns 0 on success -# LLAMA_API int llama_model_quantize( -# const char * fname_inp, -# const char * fname_out, -# const llama_model_quantize_params * params); -def llama_model_quantize( - fname_inp: bytes, - fname_out: bytes, - params, # type: POINTER(llama_model_quantize_params) # type: ignore -) -> int: - return _lib.llama_model_quantize(fname_inp, fname_out, params) +# LLAMA_API struct llama_context_params llama_context_default_params(void); +@ctypes_function( + "llama_context_default_params", + [], + llama_context_params, +) +def llama_context_default_params() -> llama_context_params: + """Get default parameters for llama_context""" + ... -_lib.llama_model_quantize.argtypes = [ - c_char_p, - c_char_p, - POINTER(llama_model_quantize_params), -] -_lib.llama_model_quantize.restype = c_int +# LLAMA_API struct llama_sampler_chain_params llama_sampler_chain_default_params(void); +@ctypes_function( + "llama_sampler_chain_default_params", + [], + llama_sampler_chain_params, +) +def llama_sampler_chain_default_params() -> llama_sampler_chain_params: + """Get default parameters for llama_sampler_chain""" + ... -# Apply a LoRA adapter to a loaded model -# path_base_model is the path to a higher quality model to use as a base for -# the layers modified by the adapter. Can be NULL to use the current loaded model. -# The model needs to be reloaded before applying a new adapter, otherwise the adapter -# will be applied on top of the previous one -# Returns 0 on success -# LLAMA_API int llama_apply_lora_from_file( -# struct llama_context * ctx, -# const char * path_lora, -# const char * path_base_model, -# int n_threads); -def llama_apply_lora_from_file( - ctx: llama_context_p, - path_lora: c_char_p, - path_base_model: c_char_p, - n_threads: c_int, -) -> int: - return _lib.llama_apply_lora_from_file(ctx, path_lora, path_base_model, n_threads) +# LLAMA_API struct llama_model_quantize_params llama_model_quantize_default_params(void); +@ctypes_function( + "llama_model_quantize_default_params", + [], + llama_model_quantize_params, +) +def llama_model_quantize_default_params() -> llama_model_quantize_params: + """Get default parameters for llama_model_quantize""" + ... -_lib.llama_apply_lora_from_file.argtypes = [llama_context_p, c_char_p, c_char_p, c_int] -_lib.llama_apply_lora_from_file.restype = c_int +# LLAMA_API const char * llama_flash_attn_type_name(enum llama_flash_attn_type flash_attn_type); +@ctypes_function("llama_flash_attn_type_name", [ctypes.c_int], ctypes.c_char_p) +def llama_flash_attn_type_name(flash_attn_type: int, /) -> Optional[bytes]: + """Get the flash attention type name.""" + ... -# LLAMA_API int llama_model_apply_lora_from_file( -# const struct llama_model * model, -# const char * path_lora, -# const char * path_base_model, -# int n_threads); -def llama_model_apply_lora_from_file( - model: llama_model_p, - path_lora: Union[c_char_p, bytes], - path_base_model: Union[c_char_p, bytes], - n_threads: c_int, -) -> int: - return _lib.llama_model_apply_lora_from_file( - model, path_lora, path_base_model, n_threads - ) +# // Initialize the llama + ggml backend +# // If numa is true, use NUMA optimizations +# // Call once at the start of the program +# LLAMA_API void llama_backend_init(void); +@ctypes_function( + "llama_backend_init", + [], + None, +) +def llama_backend_init(): + """Initialize the llama + ggml backend + Call once at the start of the program""" + ... + + +# // numa strategies +# enum ggml_numa_strategy { +# GGML_NUMA_STRATEGY_DISABLED = 0, +# GGML_NUMA_STRATEGY_DISTRIBUTE = 1, +# GGML_NUMA_STRATEGY_ISOLATE = 2, +# GGML_NUMA_STRATEGY_NUMACTL = 3, +# GGML_NUMA_STRATEGY_MIRROR = 4, +# GGML_NUMA_STRATEGY_COUNT +# }; +GGML_NUMA_STRATEGY_DISABLED = 0 +GGML_NUMA_STRATEGY_DISTRIBUTE = 1 +GGML_NUMA_STRATEGY_ISOLATE = 2 +GGML_NUMA_STRATEGY_NUMACTL = 3 +GGML_NUMA_STRATEGY_MIRROR = 4 +GGML_NUMA_STRATEGY_COUNT = 5 + + +# // Call once at the end of the program - currently only used for MPI +# LLAMA_API void llama_backend_free(void); +@ctypes_function( + "llama_backend_free", + [], + None, +) +def llama_backend_free(): + """Call once at the end of the program - currently only used for MPI""" + ... -_lib.llama_model_apply_lora_from_file.argtypes = [ - llama_model_p, - c_char_p, - c_char_p, - c_int, -] -_lib.llama_model_apply_lora_from_file.restype = c_int +# //optional: +# LLAMA_API void llama_numa_init(enum ggml_numa_strategy numa); +@ctypes_function( + "llama_numa_init", + [ctypes.c_int], + None, +) +def llama_numa_init(numa: int, /): ... -# Returns the number of tokens in the KV cache -# LLAMA_API int llama_get_kv_cache_token_count(const struct llama_context * ctx); -def llama_get_kv_cache_token_count(ctx: llama_context_p) -> int: - return _lib.llama_get_kv_cache_token_count(ctx) +# // Optional: an auto threadpool gets created in ggml if not passed explicitly +# LLAMA_API void llama_attach_threadpool( +# struct llama_context * ctx, +# ggml_threadpool_t threadpool, +# ggml_threadpool_t threadpool_batch); +# TODO: Add llama_attach_threadpool -_lib.llama_get_kv_cache_token_count.argtypes = [llama_context_p] -_lib.llama_get_kv_cache_token_count.restype = c_int +# LLAMA_API void llama_detach_threadpool(struct llama_context * ctx); +# TODO: Add llama_detach_threadpool -# Sets the current rng seed. -# LLAMA_API void llama_set_rng_seed(struct llama_context * ctx, int seed); -def llama_set_rng_seed(ctx: llama_context_p, seed: c_uint32): - return _lib.llama_set_rng_seed(ctx, seed) +# DEPRECATED(LLAMA_API struct llama_model * llama_load_model_from_file( +# const char * path_model, +# struct llama_model_params params), +# "use llama_model_load_from_file instead"); +@ctypes_function( + "llama_load_model_from_file", + [ctypes.c_char_p, llama_model_params], + llama_model_p_ctypes, +) +def llama_load_model_from_file( + path_model: bytes, params: llama_model_params, / +) -> Optional[llama_model_p]: ... -_lib.llama_set_rng_seed.argtypes = [llama_context_p, c_int] -_lib.llama_set_rng_seed.restype = None +_llama_load_model_from_file = llama_load_model_from_file -# Returns the maximum size in bytes of the state (rng, logits, embedding -# and kv_cache) - will often be smaller after compacting tokens -# LLAMA_API size_t llama_get_state_size(const struct llama_context * ctx); -def llama_get_state_size(ctx: llama_context_p) -> int: - return _lib.llama_get_state_size(ctx) +def llama_load_model_from_file( + path_model: bytes, params: llama_model_params, / +) -> Optional[llama_model_p]: + _warn_deprecated( + "llama_load_model_from_file", + "use llama_model_load_from_file instead", + ) + return _llama_load_model_from_file(path_model, params) + + +# // Load the model from a file +# // If the file is split into multiple parts, the file name must follow this pattern: <name>-%05d-of-%05d.gguf +# // If the split file name does not follow this pattern, use llama_model_load_from_splits +# LLAMA_API struct llama_model * llama_model_load_from_file( +# const char * path_model, +# struct llama_model_params params); +@ctypes_function( + "llama_model_load_from_file", + [ctypes.c_char_p, llama_model_params], + llama_model_p_ctypes, +) +def llama_model_load_from_file( + path_model: bytes, params: llama_model_params, / +) -> Optional[llama_model_p]: + """Load the model from a file + + If the file is split into multiple parts, the file name must follow this pattern: <name>-%05d-of-%05d.gguf + + If the split file name does not follow this pattern, use llama_model_load_from_splits""" + ... + + +# // Load the model from multiple splits (support custom naming scheme) +# // The paths must be in the correct order +# LLAMA_API struct llama_model * llama_model_load_from_splits( +# const char ** paths, +# size_t n_paths, +# struct llama_model_params params); +@ctypes_function( + "llama_model_load_from_splits", + [ctypes.POINTER(ctypes.c_char_p), ctypes.c_size_t, llama_model_params], + llama_model_p_ctypes, +) +def llama_model_load_from_splits( + paths: List[bytes], n_paths: int, params: llama_model_params, / +) -> Optional[llama_model_p]: + """Load the model from multiple splits (support custom naming scheme) + + The paths must be in the correct order""" + ... + + +# // Load a model from an open FILE pointer +# LLAMA_API struct llama_model * llama_model_load_from_file_ptr( +# FILE * file, +# struct llama_model_params params); +@ctypes_function( + "llama_model_load_from_file_ptr", + [ctypes.c_void_p, llama_model_params], + llama_model_p_ctypes, +) +def llama_model_load_from_file_ptr( + file: ctypes.c_void_p, params: llama_model_params, / +) -> Optional[llama_model_p]: + """Load a model from an open FILE pointer.""" + ... -_lib.llama_get_state_size.argtypes = [llama_context_p] -_lib.llama_get_state_size.restype = c_size_t +# LLAMA_API void llama_model_save_to_file( +# const struct llama_model * model, +# const char * path_model); +@ctypes_function( + "llama_model_save_to_file", + [llama_model_p_ctypes, ctypes.c_char_p], + None, +) +def llama_model_save_to_file(model: llama_model_p, path_model: bytes, /): + """Save the model to a file""" + ... -# Copies the state to the specified destination address. -# Destination needs to have allocated enough memory. -# Returns the number of bytes copied -# LLAMA_API size_t llama_copy_state_data(struct llama_context * ctx, uint8_t * dst); -def llama_copy_state_data( - ctx: llama_context_p, dst # type: Array[c_uint8] -) -> int: - return _lib.llama_copy_state_data(ctx, dst) +# DEPRECATED(LLAMA_API void llama_free_model(struct llama_model * model), +# "use llama_model_free instead"); +@ctypes_function( + "llama_free_model", + [llama_model_p_ctypes], + None, +) +def llama_free_model(model: llama_model_p, /): ... -_lib.llama_copy_state_data.argtypes = [llama_context_p, c_uint8_p] -_lib.llama_copy_state_data.restype = c_size_t +_llama_free_model = llama_free_model -# Set the state reading from the specified address -# Returns the number of bytes read -# LLAMA_API size_t llama_set_state_data(struct llama_context * ctx, uint8_t * src); -def llama_set_state_data( - ctx: llama_context_p, src # type: Array[c_uint8] -) -> int: - return _lib.llama_set_state_data(ctx, src) +def llama_free_model(model: llama_model_p, /): + _warn_deprecated("llama_free_model", "use llama_model_free instead") + return _llama_free_model(model) -_lib.llama_set_state_data.argtypes = [llama_context_p, c_uint8_p] -_lib.llama_set_state_data.restype = c_size_t +# LLAMA_API void llama_model_free(struct llama_model * model); +@ctypes_function( + "llama_model_free", + [llama_model_p_ctypes], + None, +) +def llama_model_free(model: llama_model_p, /): ... + + +# typedef void (*llama_model_set_tensor_data_t)(struct ggml_tensor * tensor, void * userdata); +llama_model_set_tensor_data_t = ctypes.CFUNCTYPE(None, ctypes.c_void_p, ctypes.c_void_p) + + +# LLAMA_API struct llama_model * llama_model_init_from_user( +# struct gguf_context * metadata, +# llama_model_set_tensor_data_t set_tensor_data, +# void * set_tensor_data_ud, +# struct llama_model_params params); +@ctypes_function( + "llama_model_init_from_user", + [ + gguf_context_p_ctypes, + llama_model_set_tensor_data_t, + ctypes.c_void_p, + llama_model_params, + ], + llama_model_p_ctypes, +) +def llama_model_init_from_user( + metadata: gguf_context_p, + set_tensor_data: llama_model_set_tensor_data_t, + set_tensor_data_ud: ctypes.c_void_p, + params: llama_model_params, + /, +) -> Optional[llama_model_p]: + """Initialize a model from user-provided metadata and tensor data.""" + ... + + +# LLAMA_API struct llama_context * llama_init_from_model( +# struct llama_model * model, +# struct llama_context_params params); +@ctypes_function( + "llama_init_from_model", + [llama_model_p_ctypes, llama_context_params], + llama_context_p_ctypes, +) +def llama_init_from_model( + model: llama_model_p, params: llama_context_params, / +) -> Optional[llama_context_p]: ... + + +# DEPRECATED(LLAMA_API struct llama_context * llama_new_context_with_model( +# struct llama_model * model, +# struct llama_context_params params), +# "use llama_init_from_model instead"); +@ctypes_function( + "llama_new_context_with_model", + [llama_model_p_ctypes, llama_context_params], + llama_context_p_ctypes, +) +def llama_new_context_with_model( + model: llama_model_p, params: llama_context_params, / +) -> Optional[llama_context_p]: ... -# Save/load session file -# LLAMA_API bool llama_load_session_file(struct llama_context * ctx, const char * path_session, llama_token * tokens_out, size_t n_token_capacity, size_t * n_token_count_out); -def llama_load_session_file( - ctx: llama_context_p, - path_session: bytes, - tokens_out, # type: Array[llama_token] - n_token_capacity: c_size_t, - n_token_count_out, # type: _Pointer[c_size_t] -) -> int: - return _lib.llama_load_session_file( - ctx, path_session, tokens_out, n_token_capacity, n_token_count_out - ) +_llama_new_context_with_model = llama_new_context_with_model -_lib.llama_load_session_file.argtypes = [ - llama_context_p, - c_char_p, - llama_token_p, - c_size_t, - c_size_t_p, -] -_lib.llama_load_session_file.restype = c_size_t +def llama_new_context_with_model( + model: llama_model_p, params: llama_context_params, / +) -> Optional[llama_context_p]: + _warn_deprecated( + "llama_new_context_with_model", + "use llama_init_from_model instead", + ) + return _llama_new_context_with_model(model, params) -# LLAMA_API bool llama_save_session_file(struct llama_context * ctx, const char * path_session, const llama_token * tokens, size_t n_token_count); -def llama_save_session_file( - ctx: llama_context_p, - path_session: bytes, - tokens, # type: Array[llama_token] - n_token_count: c_size_t, -) -> int: - return _lib.llama_save_session_file(ctx, path_session, tokens, n_token_count) +# // Frees all allocated memory +# LLAMA_API void llama_free(struct llama_context * ctx); +@ctypes_function( + "llama_free", + [llama_context_p_ctypes], + None, +) +def llama_free(ctx: llama_context_p, /): + """Frees all allocated memory""" + ... -_lib.llama_save_session_file.argtypes = [ - llama_context_p, - c_char_p, - llama_token_p, - c_size_t, -] -_lib.llama_save_session_file.restype = c_size_t +# LLAMA_API int64_t llama_time_us(void); +@ctypes_function( + "llama_time_us", + [], + ctypes.c_int64, +) +def llama_time_us() -> int: ... -# Run the llama inference to obtain the logits and probabilities for the next token. -# tokens + n_tokens is the provided batch of new tokens to process -# n_past is the number of tokens to use from previous eval calls -# Returns 0 on success -# LLAMA_API int llama_eval( -# struct llama_context * ctx, -# const llama_token * tokens, -# int n_tokens, -# int n_past, -# int n_threads); -def llama_eval( - ctx: llama_context_p, - tokens, # type: Array[llama_token] - n_tokens: c_int, - n_past: c_int, - n_threads: c_int, -) -> int: - return _lib.llama_eval(ctx, tokens, n_tokens, n_past, n_threads) +# LLAMA_API size_t llama_max_devices(void); +@ctypes_function("llama_max_devices", [], ctypes.c_size_t) +def llama_max_devices() -> int: ... -_lib.llama_eval.argtypes = [llama_context_p, llama_token_p, c_int, c_int, c_int] -_lib.llama_eval.restype = c_int +# LLAMA_API size_t llama_max_parallel_sequences(void); +@ctypes_function("llama_max_parallel_sequences", [], ctypes.c_size_t) +def llama_max_parallel_sequences() -> int: ... -# // Same as llama_eval, but use float matrix input directly. -# LLAMA_API int llama_eval_embd( -# struct llama_context * ctx, -# const float * embd, -# int n_tokens, -# int n_past, -# int n_threads); -def llama_eval_embd( - ctx: llama_context_p, - embd, # type: Array[c_float] - n_tokens: c_int, - n_past: c_int, - n_threads: c_int, -) -> int: - return _lib.llama_eval_embd(ctx, embd, n_tokens, n_past, n_threads) +# LLAMA_API size_t llama_max_tensor_buft_overrides(void); +@ctypes_function("llama_max_tensor_buft_overrides", [], ctypes.c_size_t) +def llama_max_tensor_buft_overrides() -> int: + """Get the maximum number of tensor buffer type overrides.""" + ... -_lib.llama_eval_embd.argtypes = [llama_context_p, c_float_p, c_int, c_int, c_int] -_lib.llama_eval_embd.restype = c_int +# LLAMA_API bool llama_supports_mmap (void); +@ctypes_function("llama_supports_mmap", [], ctypes.c_bool) +def llama_supports_mmap() -> bool: ... -# Convert the provided text into tokens. -# The tokens pointer must be large enough to hold the resulting tokens. -# Returns the number of tokens on success, no more than n_max_tokens -# Returns a negative number on failure - the number of tokens that would have been returned -# TODO: not sure if correct -# LLAMA_API int llama_tokenize( -# struct llama_context * ctx, -# const char * text, -# llama_token * tokens, -# int n_max_tokens, -# bool add_bos); -def llama_tokenize( - ctx: llama_context_p, - text: bytes, - tokens, # type: Array[llama_token] - n_max_tokens: c_int, - add_bos: c_bool, -) -> int: - return _lib.llama_tokenize(ctx, text, tokens, n_max_tokens, add_bos) +# LLAMA_API bool llama_supports_mlock (void); +@ctypes_function("llama_supports_mlock", [], ctypes.c_bool) +def llama_supports_mlock() -> bool: ... -_lib.llama_tokenize.argtypes = [llama_context_p, c_char_p, llama_token_p, c_int, c_bool] -_lib.llama_tokenize.restype = c_int +# LLAMA_API bool llama_supports_gpu_offload(void); +@ctypes_function("llama_supports_gpu_offload", [], ctypes.c_bool) +def llama_supports_gpu_offload() -> bool: ... -# LLAMA_API int llama_n_vocab(const struct llama_context * ctx); -def llama_n_vocab(ctx: llama_context_p) -> int: - return _lib.llama_n_vocab(ctx) +# LLAMA_API bool llama_supports_rpc (void); +@ctypes_function("llama_supports_rpc", [], ctypes.c_bool) +def llama_supports_rpc() -> bool: ... -_lib.llama_n_vocab.argtypes = [llama_context_p] -_lib.llama_n_vocab.restype = c_int +# LLAMA_API uint32_t llama_n_ctx (const struct llama_context * ctx); +@ctypes_function("llama_n_ctx", [llama_context_p_ctypes], ctypes.c_uint32) +def llama_n_ctx(ctx: llama_context_p, /) -> int: ... -# LLAMA_API int llama_n_ctx (const struct llama_context * ctx); -def llama_n_ctx(ctx: llama_context_p) -> int: - return _lib.llama_n_ctx(ctx) +# LLAMA_API uint32_t llama_n_ctx_seq (const struct llama_context * ctx); +@ctypes_function("llama_n_ctx_seq", [llama_context_p_ctypes], ctypes.c_uint32) +def llama_n_ctx_seq(ctx: llama_context_p, /) -> int: + """Get the context size per sequence.""" + ... -_lib.llama_n_ctx.argtypes = [llama_context_p] -_lib.llama_n_ctx.restype = c_int +# LLAMA_API uint32_t llama_n_batch (const struct llama_context * ctx); +@ctypes_function("llama_n_batch", [llama_context_p_ctypes], ctypes.c_uint32) +def llama_n_batch(ctx: llama_context_p, /) -> int: ... -# LLAMA_API int llama_n_embd (const struct llama_context * ctx); -def llama_n_embd(ctx: llama_context_p) -> int: - return _lib.llama_n_embd(ctx) +# LLAMA_API uint32_t llama_n_ubatch (const struct llama_context * ctx); +@ctypes_function("llama_n_ubatch", [llama_context_p_ctypes], ctypes.c_uint32) +def llama_n_ubatch(ctx: llama_context_p, /) -> int: ... -_lib.llama_n_embd.argtypes = [llama_context_p] -_lib.llama_n_embd.restype = c_int +# LLAMA_API uint32_t llama_n_seq_max (const struct llama_context * ctx); +@ctypes_function("llama_n_seq_max", [llama_context_p_ctypes], ctypes.c_uint32) +def llama_n_seq_max(ctx: llama_context_p, /) -> int: ... -# // Get the vocabulary as output parameters. -# // Returns number of results. -# LLAMA_API int llama_get_vocab( -# const struct llama_context * ctx, -# const char * * strings, -# float * scores, -# int capacity); -def llama_get_vocab( - ctx: llama_context_p, - strings, # type: Array[c_char_p] # type: ignore - scores, # type: Array[c_float] # type: ignore - capacity: c_int, -) -> int: - return _lib.llama_get_vocab(ctx, strings, scores, capacity) +# LLAMA_API uint32_t llama_n_rs_seq (const struct llama_context * ctx); +@ctypes_function("llama_n_rs_seq", [llama_context_p_ctypes], ctypes.c_uint32) +def llama_n_rs_seq(ctx: llama_context_p, /) -> int: ... -_lib.llama_get_vocab.argtypes = [llama_context_p, c_char_p, c_float, c_int] -_lib.llama_get_vocab.restype = c_int +# DEPRECATED(LLAMA_API int32_t llama_n_ctx_train(const struct llama_model * model), "use llama_model_n_ctx_train instead"); +@ctypes_function("llama_n_ctx_train", [llama_model_p_ctypes], ctypes.c_int32) +def llama_n_ctx_train(model: llama_model_p, /) -> int: ... -# Token logits obtained from the last call to llama_eval() -# The logits for the last token are stored in the last row -# Can be mutated in order to change the probabilities of the next token -# Rows: n_tokens -# Cols: n_vocab -# LLAMA_API float * llama_get_logits(struct llama_context * ctx); -def llama_get_logits( - ctx: llama_context_p, -): # type: (...) -> Array[float] # type: ignore - return _lib.llama_get_logits(ctx) +_llama_n_ctx_train = llama_n_ctx_train -_lib.llama_get_logits.argtypes = [llama_context_p] -_lib.llama_get_logits.restype = c_float_p +def llama_n_ctx_train(model: llama_model_p, /) -> int: + _warn_deprecated("llama_n_ctx_train", "use llama_model_n_ctx_train instead") + return _llama_n_ctx_train(model) -# Get the embeddings for the input -# shape: [n_embd] (1-dimensional) -# LLAMA_API float * llama_get_embeddings(struct llama_context * ctx); -def llama_get_embeddings( - ctx: llama_context_p, -): # type: (...) -> Array[float] # type: ignore - return _lib.llama_get_embeddings(ctx) +# DEPRECATED(LLAMA_API int32_t llama_n_embd (const struct llama_model * model), "use llama_model_n_embd instead"); +@ctypes_function("llama_n_embd", [llama_model_p_ctypes], ctypes.c_int32) +def llama_n_embd(model: llama_model_p, /) -> int: ... -_lib.llama_get_embeddings.argtypes = [llama_context_p] -_lib.llama_get_embeddings.restype = c_float_p +_llama_n_embd = llama_n_embd -# Token Id -> String. Uses the vocabulary in the provided context -# LLAMA_API const char * llama_token_to_str(const struct llama_context * ctx, llama_token token); -def llama_token_to_str(ctx: llama_context_p, token: llama_token) -> bytes: - return _lib.llama_token_to_str(ctx, token) +def llama_n_embd(model: llama_model_p, /) -> int: + _warn_deprecated("llama_n_embd", "use llama_model_n_embd instead") + return _llama_n_embd(model) -_lib.llama_token_to_str.argtypes = [llama_context_p, llama_token] -_lib.llama_token_to_str.restype = c_char_p +# DEPRECATED(LLAMA_API int32_t llama_n_layer (const struct llama_model * model), "use llama_model_n_layer instead"); +@ctypes_function("llama_n_layer", [llama_model_p_ctypes], ctypes.c_int32) +def llama_n_layer(model: llama_model_p, /) -> int: ... -# Special tokens +_llama_n_layer = llama_n_layer -# LLAMA_API llama_token llama_token_bos(); // beginning-of-sentence -def llama_token_bos() -> int: - return _lib.llama_token_bos() +def llama_n_layer(model: llama_model_p, /) -> int: + _warn_deprecated("llama_n_layer", "use llama_model_n_layer instead") + return _llama_n_layer(model) -_lib.llama_token_bos.argtypes = [] -_lib.llama_token_bos.restype = llama_token +# DEPRECATED(LLAMA_API int32_t llama_n_head (const struct llama_model * model), "use llama_model_n_head instead"); +@ctypes_function("llama_n_head", [llama_model_p_ctypes], ctypes.c_int32) +def llama_n_head(model: llama_model_p, /) -> int: ... -# LLAMA_API llama_token llama_token_eos(); // end-of-sentence -def llama_token_eos() -> int: - return _lib.llama_token_eos() +_llama_n_head = llama_n_head -_lib.llama_token_eos.argtypes = [] -_lib.llama_token_eos.restype = llama_token +def llama_n_head(model: llama_model_p, /) -> int: + _warn_deprecated("llama_n_head", "use llama_model_n_head instead") + return _llama_n_head(model) -# LLAMA_API llama_token llama_token_nl(); // next-line -def llama_token_nl() -> int: - return _lib.llama_token_nl() +# DEPRECATED(LLAMA_API int32_t llama_n_vocab (const struct llama_vocab * vocab), "use llama_vocab_n_tokens instead"); +@ctypes_function("llama_n_vocab", [llama_vocab_p_ctypes], ctypes.c_int32) +def llama_n_vocab(model: llama_vocab_p, /) -> int: ... -_lib.llama_token_nl.argtypes = [] -_lib.llama_token_nl.restype = llama_token +_llama_n_vocab = llama_n_vocab -# Sampling functions +def llama_n_vocab(model: llama_vocab_p, /) -> int: + _warn_deprecated("llama_n_vocab", "use llama_vocab_n_tokens instead") + return _llama_n_vocab(model) -# @details Repetition penalty described in CTRL academic paper https://arxiv.org/abs/1909.05858, with negative logit fix. -# LLAMA_API void llama_sample_repetition_penalty(struct llama_context * ctx, llama_token_data_array * candidates, const llama_token * last_tokens, size_t last_tokens_size, float penalty); -def llama_sample_repetition_penalty( - ctx: llama_context_p, - candidates, # type: _Pointer[llama_token_data_array] - last_tokens_data, # type: Array[llama_token] - last_tokens_size: c_int, - penalty: c_float, -): - return _lib.llama_sample_repetition_penalty( - ctx, candidates, last_tokens_data, last_tokens_size, penalty - ) +# LLAMA_API const struct llama_model * llama_get_model (const struct llama_context * ctx); +@ctypes_function("llama_get_model", [llama_context_p_ctypes], llama_model_p_ctypes) +def llama_get_model(ctx: llama_context_p, /) -> Optional[llama_model_p]: ... -_lib.llama_sample_repetition_penalty.argtypes = [ - llama_context_p, - llama_token_data_array_p, - llama_token_p, - c_int, - c_float, -] -_lib.llama_sample_repetition_penalty.restype = None +# LLAMA_API llama_memory_t llama_get_memory (const struct llama_context * ctx); +@ctypes_function("llama_get_memory", [llama_context_p_ctypes], llama_memory_t_ctypes) +def llama_get_memory(ctx: llama_context_p, /) -> Optional[llama_memory_t]: + """Get the memory for the context""" + ... -# @details Frequency and presence penalties described in OpenAI API https://platform.openai.com/docs/api-reference/parameter-details. -# LLAMA_API void llama_sample_frequency_and_presence_penalties(struct llama_context * ctx, llama_token_data_array * candidates, const llama_token * last_tokens, size_t last_tokens_size, float alpha_frequency, float alpha_presence); -def llama_sample_frequency_and_presence_penalties( - ctx: llama_context_p, - candidates, # type: _Pointer[llama_token_data_array] - last_tokens_data, # type: Array[llama_token] - last_tokens_size: c_int, - alpha_frequency: c_float, - alpha_presence: c_float, -): - return _lib.llama_sample_frequency_and_presence_penalties( - ctx, - candidates, - last_tokens_data, - last_tokens_size, - alpha_frequency, - alpha_presence, - ) +# LLAMA_API enum llama_pooling_type llama_pooling_type(const struct llama_context * ctx); +@ctypes_function("llama_pooling_type", [llama_context_p_ctypes], ctypes.c_int) +def llama_pooling_type(ctx: llama_context_p, /) -> int: ... -_lib.llama_sample_frequency_and_presence_penalties.argtypes = [ - llama_context_p, - llama_token_data_array_p, - llama_token_p, - c_int, - c_float, - c_float, -] -_lib.llama_sample_frequency_and_presence_penalties.restype = None +# LLAMA_API const struct llama_vocab * llama_model_get_vocab(const struct llama_model * model); +@ctypes_function("llama_model_get_vocab", [llama_model_p_ctypes], llama_vocab_p_ctypes) +def llama_model_get_vocab(model: llama_model_p, /) -> Optional[llama_vocab_p]: ... -# @details Sorts candidate tokens by their logits in descending order and calculate probabilities based on logits. -# LLAMA_API void llama_sample_softmax(struct llama_context * ctx, llama_token_data_array * candidates); -def llama_sample_softmax( - ctx: llama_context_p, candidates # type: _Pointer[llama_token_data] -): - return _lib.llama_sample_softmax(ctx, candidates) +# LLAMA_API enum llama_rope_type llama_model_rope_type(const struct llama_model * model); +@ctypes_function("llama_model_rope_type", [llama_model_p_ctypes], ctypes.c_int) +def llama_model_rope_type(model: llama_model_p, /) -> int: ... -_lib.llama_sample_softmax.argtypes = [ - llama_context_p, - llama_token_data_array_p, -] -_lib.llama_sample_softmax.restype = None +# LLAMA_API int32_t llama_model_n_ctx_train(const struct llama_model * model); +@ctypes_function("llama_model_n_ctx_train", [llama_model_p_ctypes], ctypes.c_int32) +def llama_model_n_ctx_train(model: llama_model_p, /) -> int: ... -# @details Top-K sampling described in academic paper "The Curious Case of Neural Text Degeneration" https://arxiv.org/abs/1904.09751 -# LLAMA_API void llama_sample_top_k(struct llama_context * ctx, llama_token_data_array * candidates, int k, size_t min_keep); -def llama_sample_top_k( - ctx: llama_context_p, - candidates, # type: _Pointer[llama_token_data_array] - k: c_int, - min_keep: c_size_t, -): - return _lib.llama_sample_top_k(ctx, candidates, k, min_keep) +# LLAMA_API int32_t llama_model_n_embd (const struct llama_model * model); +@ctypes_function("llama_model_n_embd", [llama_model_p_ctypes], ctypes.c_int32) +def llama_model_n_embd(model: llama_model_p, /) -> int: ... -_lib.llama_sample_top_k.argtypes = [ - llama_context_p, - llama_token_data_array_p, - c_int, - c_size_t, -] -_lib.llama_sample_top_k.restype = None +# LLAMA_API int32_t llama_model_n_embd_inp (const struct llama_model * model); +@ctypes_function("llama_model_n_embd_inp", [llama_model_p_ctypes], ctypes.c_int32) +def llama_model_n_embd_inp(model: llama_model_p, /) -> int: + """Get the model input embedding size.""" + ... -# @details Nucleus sampling described in academic paper "The Curious Case of Neural Text Degeneration" https://arxiv.org/abs/1904.09751 -# LLAMA_API void llama_sample_top_p(struct llama_context * ctx, llama_token_data_array * candidates, float p, size_t min_keep); -def llama_sample_top_p( - ctx: llama_context_p, - candidates, # type: _Pointer[llama_token_data_array] - p: c_float, - min_keep: c_size_t, -): - return _lib.llama_sample_top_p(ctx, candidates, p, min_keep) +# LLAMA_API int32_t llama_model_n_embd_out (const struct llama_model * model); +@ctypes_function("llama_model_n_embd_out", [llama_model_p_ctypes], ctypes.c_int32) +def llama_model_n_embd_out(model: llama_model_p, /) -> int: + """Get the model output embedding size.""" + ... -_lib.llama_sample_top_p.argtypes = [ - llama_context_p, - llama_token_data_array_p, - c_float, - c_size_t, -] -_lib.llama_sample_top_p.restype = None +# LLAMA_API int32_t llama_model_n_layer (const struct llama_model * model); +@ctypes_function("llama_model_n_layer", [llama_model_p_ctypes], ctypes.c_int32) +def llama_model_n_layer(model: llama_model_p, /) -> int: ... -# @details Tail Free Sampling described in https://www.trentonbricken.com/Tail-Free-Sampling/. -# LLAMA_API void llama_sample_tail_free(struct llama_context * ctx, llama_token_data_array * candidates, float z, size_t min_keep); -def llama_sample_tail_free( - ctx: llama_context_p, - candidates, # type: _Pointer[llama_token_data_array] - z: c_float, - min_keep: c_size_t, -): - return _lib.llama_sample_tail_free(ctx, candidates, z, min_keep) +# LLAMA_API int32_t llama_model_n_head (const struct llama_model * model); +@ctypes_function("llama_model_n_head", [llama_model_p_ctypes], ctypes.c_int32) +def llama_model_n_head(model: llama_model_p, /) -> int: ... -_lib.llama_sample_tail_free.argtypes = [ - llama_context_p, - llama_token_data_array_p, - c_float, - c_size_t, -] -_lib.llama_sample_tail_free.restype = None +# LLAMA_API int32_t llama_model_n_head_kv (const struct llama_model * model); +@ctypes_function("llama_model_n_head_kv", [llama_model_p_ctypes], ctypes.c_int32) +def llama_model_n_head_kv(model: llama_model_p, /) -> int: ... -# @details Locally Typical Sampling implementation described in the paper https://arxiv.org/abs/2202.00666. -# LLAMA_API void llama_sample_typical(struct llama_context * ctx, llama_token_data_array * candidates, float p, size_t min_keep); -def llama_sample_typical( - ctx: llama_context_p, - candidates, # type: _Pointer[llama_token_data_array] - p: c_float, - min_keep: c_size_t, -): - return _lib.llama_sample_typical(ctx, candidates, p, min_keep) +# LLAMA_API int32_t llama_model_n_swa (const struct llama_model * model); +@ctypes_function("llama_model_n_swa", [llama_model_p_ctypes], ctypes.c_int32) +def llama_model_n_swa(model: llama_model_p, /) -> int: ... -_lib.llama_sample_typical.argtypes = [ - llama_context_p, - llama_token_data_array_p, - c_float, - c_size_t, -] -_lib.llama_sample_typical.restype = None +# // Get the model's RoPE frequency scaling factor +# LLAMA_API float llama_model_rope_freq_scale_train(const struct llama_model * model); +@ctypes_function( + "llama_model_rope_freq_scale_train", [llama_model_p_ctypes], ctypes.c_float +) +def llama_model_rope_freq_scale_train(model: llama_model_p, /) -> float: ... -# LLAMA_API void llama_sample_temperature(struct llama_context * ctx, llama_token_data_array * candidates, float temp); -def llama_sample_temperature( - ctx: llama_context_p, - candidates, # type: _Pointer[llama_token_data_array] - temp: c_float, -): - return _lib.llama_sample_temperature(ctx, candidates, temp) +# // Returns the number of classifier outputs (only valid for classifier models) +# // Undefined behavior for non-classifier models +# LLAMA_API uint32_t llama_model_n_cls_out(const struct llama_model * model); +@ctypes_function("llama_model_n_cls_out", [llama_model_p_ctypes], ctypes.c_uint32) +def llama_model_n_cls_out(model: llama_model_p, /) -> int: + """Returns the number of classifier outputs (only valid for classifier models)""" + ... -_lib.llama_sample_temperature.argtypes = [ - llama_context_p, - llama_token_data_array_p, - c_float, -] -_lib.llama_sample_temperature.restype = None +# // Returns label of classifier output by index (<n_cls_out). Returns nullptr if no label provided +# LLAMA_API const char * llama_model_cls_label(const struct llama_model * model, uint32_t i); +@ctypes_function( + "llama_model_cls_label", [llama_model_p_ctypes, ctypes.c_uint32], ctypes.c_char_p +) +def llama_model_cls_label(model: llama_model_p, i: int, /) -> Optional[bytes]: + """Returns label of classifier output by index. Returns None if no label provided""" + ... -# @details Mirostat 1.0 algorithm described in the paper https://arxiv.org/abs/2007.14966. Uses tokens instead of words. -# @param candidates A vector of `llama_token_data` containing the candidate tokens, their probabilities (p), and log-odds (logit) for the current position in the generated text. -# @param tau The target cross-entropy (or surprise) value you want to achieve for the generated text. A higher value corresponds to more surprising or less predictable text, while a lower value corresponds to less surprising or more predictable text. -# @param eta The learning rate used to update `mu` based on the error between the target and observed surprisal of the sampled word. A larger learning rate will cause `mu` to be updated more quickly, while a smaller learning rate will result in slower updates. -# @param m The number of tokens considered in the estimation of `s_hat`. This is an arbitrary value that is used to calculate `s_hat`, which in turn helps to calculate the value of `k`. In the paper, they use `m = 100`, but you can experiment with different values to see how it affects the performance of the algorithm. -# @param mu Maximum cross-entropy. This value is initialized to be twice the target cross-entropy (`2 * tau`) and is updated in the algorithm based on the error between the target and observed surprisal. -# LLAMA_API llama_token llama_sample_token_mirostat(struct llama_context * ctx, llama_token_data_array * candidates, float tau, float eta, int m, float * mu); -def llama_sample_token_mirostat( - ctx: llama_context_p, - candidates, # type: _Pointer[llama_token_data_array] - tau: c_float, - eta: c_float, - m: c_int, - mu, # type: _Pointer[c_float] -) -> int: - return _lib.llama_sample_token_mirostat(ctx, candidates, tau, eta, m, mu) +# LLAMA_API enum llama_vocab_type llama_vocab_type (const struct llama_model * model); +@ctypes_function("llama_vocab_type", [llama_vocab_p_ctypes], ctypes.c_int) +def llama_vocab_type(vocab: llama_vocab_p, /) -> int: ... -_lib.llama_sample_token_mirostat.argtypes = [ - llama_context_p, - llama_token_data_array_p, - c_float, - c_float, - c_int, - c_float_p, -] -_lib.llama_sample_token_mirostat.restype = llama_token +# LLAMA_API int32_t llama_vocab_n_tokens(const struct llama_vocab * vocab); +@ctypes_function("llama_vocab_n_tokens", [llama_vocab_p_ctypes], ctypes.c_int32) +def llama_vocab_n_tokens(vocab: llama_vocab_p, /) -> int: ... -# @details Mirostat 2.0 algorithm described in the paper https://arxiv.org/abs/2007.14966. Uses tokens instead of words. -# @param candidates A vector of `llama_token_data` containing the candidate tokens, their probabilities (p), and log-odds (logit) for the current position in the generated text. -# @param tau The target cross-entropy (or surprise) value you want to achieve for the generated text. A higher value corresponds to more surprising or less predictable text, while a lower value corresponds to less surprising or more predictable text. -# @param eta The learning rate used to update `mu` based on the error between the target and observed surprisal of the sampled word. A larger learning rate will cause `mu` to be updated more quickly, while a smaller learning rate will result in slower updates. -# @param mu Maximum cross-entropy. This value is initialized to be twice the target cross-entropy (`2 * tau`) and is updated in the algorithm based on the error between the target and observed surprisal. -# LLAMA_API llama_token llama_sample_token_mirostat_v2(struct llama_context * ctx, llama_token_data_array * candidates, float tau, float eta, float * mu); -def llama_sample_token_mirostat_v2( - ctx: llama_context_p, - candidates, # type: _Pointer[llama_token_data_array] - tau: c_float, - eta: c_float, - mu, # type: _Pointer[c_float] -) -> int: - return _lib.llama_sample_token_mirostat_v2(ctx, candidates, tau, eta, mu) +# // Functions to access the model's GGUF metadata scalar values +# // - The functions return the length of the string on success, or -1 on failure +# // - The output string is always null-terminated and cleared on failure +# // - When retrieving a string, an extra byte must be allocated to account for the null terminator +# // - GGUF array values are not supported by these functions -_lib.llama_sample_token_mirostat_v2.argtypes = [ - llama_context_p, - llama_token_data_array_p, - c_float, - c_float, - c_float_p, -] -_lib.llama_sample_token_mirostat_v2.restype = llama_token + +# // Get metadata value as a string by key name +# LLAMA_API int32_t llama_model_meta_val_str(const struct llama_model * model, const char * key, char * buf, size_t buf_size); +@ctypes_function( + "llama_model_meta_val_str", + [ + llama_model_p_ctypes, + ctypes.c_char_p, + ctypes.c_char_p, + ctypes.c_size_t, + ], + ctypes.c_int32, +) +def llama_model_meta_val_str( + model: llama_model_p, + key: Union[ctypes.c_char_p, bytes], + buf: bytes, + buf_size: int, + /, +) -> int: + """Get metadata value as a string by key name""" + ... + + +# // Get the number of metadata key/value pairs +# LLAMA_API int32_t llama_model_meta_count(const struct llama_model * model); +@ctypes_function("llama_model_meta_count", [llama_model_p_ctypes], ctypes.c_int32) +def llama_model_meta_count(model: llama_model_p, /) -> int: + """Get the number of metadata key/value pairs""" + ... + + +# // Get sampling metadata key name. Returns nullptr if the key is invalid +# LLAMA_API const char * llama_model_meta_key_str(enum llama_model_meta_key key); +@ctypes_function("llama_model_meta_key_str", [ctypes.c_int], ctypes.c_char_p) +def llama_model_meta_key_str(key: int, /) -> Optional[bytes]: + """Get sampling metadata key name. Returns None if the key is invalid.""" + ... + + +# // Get metadata key name by index +# LLAMA_API int32_t llama_model_meta_key_by_index(const struct llama_model * model, int32_t i, char * buf, size_t buf_size); +@ctypes_function( + "llama_model_meta_key_by_index", + [ + llama_model_p_ctypes, + ctypes.c_int32, + ctypes.c_char_p, + ctypes.c_size_t, + ], + ctypes.c_int32, +) +def llama_model_meta_key_by_index( + model: llama_model_p, + i: Union[ctypes.c_int, int], + buf: Union[bytes, CtypesArray[ctypes.c_char]], + buf_size: int, + /, +) -> int: + """Get metadata key name by index""" + ... + + +# // Get metadata value as a string by index +# LLAMA_API int32_t llama_model_meta_val_str_by_index(const struct llama_model * model, int32_t i, char * buf, size_t buf_size); +@ctypes_function( + "llama_model_meta_val_str_by_index", + [ + llama_model_p_ctypes, + ctypes.c_int32, + ctypes.c_char_p, + ctypes.c_size_t, + ], + ctypes.c_int32, +) +def llama_model_meta_val_str_by_index( + model: llama_model_p, + i: Union[ctypes.c_int, int], + buf: Union[bytes, CtypesArray[ctypes.c_char]], + buf_size: int, + /, +) -> int: + """Get metadata value as a string by index""" + ... -# @details Selects the token with the highest probability. -# LLAMA_API llama_token llama_sample_token_greedy(struct llama_context * ctx, llama_token_data_array * candidates); -def llama_sample_token_greedy( - ctx: llama_context_p, - candidates, # type: _Pointer[llama_token_data_array] +# // Get a string describing the model type +# LLAMA_API int32_t llama_model_desc(const struct llama_model * model, char * buf, size_t buf_size); +@ctypes_function( + "llama_model_desc", + [llama_model_p_ctypes, ctypes.c_char_p, ctypes.c_size_t], + ctypes.c_int32, +) +def llama_model_desc( + model: llama_model_p, + buf: Union[bytes, CtypesArray[ctypes.c_char]], + buf_size: Union[ctypes.c_size_t, int], + /, ) -> int: - return _lib.llama_sample_token_greedy(ctx, candidates) + """Get a string describing the model type""" + ... -_lib.llama_sample_token_greedy.argtypes = [ - llama_context_p, - llama_token_data_array_p, -] -_lib.llama_sample_token_greedy.restype = llama_token +# // Returns the total size of all the tensors in the model in bytes +# LLAMA_API uint64_t llama_model_size(const struct llama_model * model); +@ctypes_function("llama_model_size", [llama_model_p_ctypes], ctypes.c_uint64) +def llama_model_size(model: llama_model_p, /) -> int: + """Returns the total size of all the tensors in the model in bytes""" + ... -# @details Randomly selects a token from the candidates based on their probabilities. -# LLAMA_API llama_token llama_sample_token(struct llama_context * ctx, llama_token_data_array * candidates); -def llama_sample_token( +# // Get the default chat template. Returns nullptr if not available +# // If name is NULL, returns the default chat template +# LLAMA_API const char * llama_model_chat_template(const struct llama_model * model, const char * name); +@ctypes_function( + "llama_model_chat_template", + [llama_model_p_ctypes, ctypes.c_char_p], + ctypes.c_char_p, +) +def llama_model_chat_template( + model: llama_model_p, name: Optional[bytes], / +) -> Optional[bytes]: + """Get the default chat template. Returns None if not available + If name is None, returns the default chat template""" + ... + + +# // Returns the total number of parameters in the model +# LLAMA_API uint64_t llama_model_n_params(const struct llama_model * model); +@ctypes_function("llama_model_n_params", [llama_model_p_ctypes], ctypes.c_uint64) +def llama_model_n_params(model: llama_model_p, /) -> int: + """Returns the total number of parameters in the model""" + ... + + +# // Returns true if the model contains an encoder that requires llama_encode() call +# LLAMA_API bool llama_model_has_encoder(const struct llama_model * model); +@ctypes_function("llama_model_has_encoder", [llama_model_p_ctypes], ctypes.c_bool) +def llama_model_has_encoder(model: llama_model_p, /) -> bool: + """Returns true if the model contains an encoder that requires llama_encode() call""" + ... + + +# // Returns true if the model contains a decoder that requires llama_decode() call +# LLAMA_API bool llama_model_has_decoder(const struct llama_model * model); +@ctypes_function("llama_model_has_decoder", [llama_model_p_ctypes], ctypes.c_bool) +def llama_model_has_decoder(model: llama_model_p, /) -> bool: + """Returns true if the model contains a decoder that requires llama_decode() call""" + ... + + +# // For encoder-decoder models, this function returns id of the token that must be provided +# // to the decoder to start generating output sequence. For other models, it returns -1. +# LLAMA_API llama_token llama_model_decoder_start_token(const struct llama_model * model); +@ctypes_function( + "llama_model_decoder_start_token", [llama_model_p_ctypes], ctypes.c_int32 +) +def llama_model_decoder_start_token(model: llama_model_p, /) -> int: + """For encoder-decoder models, this function returns id of the token that must be provided + to the decoder to start generating output sequence. For other models, it returns -1. + """ + ... + + +# // Returns true if the model is recurrent (like Mamba, RWKV, etc.) +# LLAMA_API bool llama_model_is_recurrent(const struct llama_model * model); +@ctypes_function("llama_model_is_recurrent", [llama_model_p_ctypes], ctypes.c_bool) +def llama_model_is_recurrent(model: llama_model_p, /) -> bool: + """Returns true if the model is recurrent (like Mamba, RWKV, etc.)""" + ... + + +# // Returns true if the model is hybrid (like Jamba, Granite, etc.) +# LLAMA_API bool llama_model_is_hybrid(const struct llama_model * model); +@ctypes_function("llama_model_is_hybrid", [llama_model_p_ctypes], ctypes.c_bool) +def llama_model_is_hybrid(model: llama_model_p, /) -> bool: + """Returns true if the model is hybrid (like Jamba, Granite, etc.)""" + ... + + +# // Returns true if the model is diffusion-based (like LLaDA, Dream, etc.) +# LLAMA_API bool llama_model_is_diffusion(const struct llama_model * model); +@ctypes_function("llama_model_is_diffusion", [llama_model_p_ctypes], ctypes.c_bool) +def llama_model_is_diffusion(model: llama_model_p, /) -> bool: + """Returns true if the model is diffusion-based (like LLaDA, Dream, etc.)""" + ... + + +# // Returns 0 on success +# LLAMA_API uint32_t llama_model_quantize( +# const char * fname_inp, +# const char * fname_out, +# const llama_model_quantize_params * params); +@ctypes_function( + "llama_model_quantize", + [ + ctypes.c_char_p, + ctypes.c_char_p, + ctypes.POINTER(llama_model_quantize_params), + ], + ctypes.c_uint32, +) +def llama_model_quantize( + fname_inp: bytes, + fname_out: bytes, + params: CtypesPointerOrRef[llama_model_quantize_params], + /, +) -> int: + """Returns 0 on success""" + ... + + +# // +# // Adapters +# // + + +# // Load a LoRA adapter from file +# LLAMA_API struct llama_adapter_lora * llama_adapter_lora_init( +# struct llama_model * model, +# const char * path_lora); +@ctypes_function( + "llama_adapter_lora_init", + [llama_model_p_ctypes, ctypes.c_char_p], + llama_adapter_lora_p_ctypes, +) +def llama_adapter_lora_init( + model: llama_model_p, path_lora: bytes, / +) -> Optional[llama_adapter_lora_p]: ... + + +# // Get metadata value as a string by key name +# LLAMA_API int32_t llama_adapter_meta_val_str(const struct llama_adapter_lora * adapter, const char * key, char * buf, size_t buf_size); +@ctypes_function( + "llama_adapter_meta_val_str", + [ + llama_adapter_lora_p_ctypes, + ctypes.c_char_p, + ctypes.c_char_p, + ctypes.c_size_t, + ], + ctypes.c_int32, +) +def llama_adapter_meta_val_str( + adapter: llama_adapter_lora_p, + key: bytes, + buf: Union[bytes, CtypesArray[ctypes.c_char]], + buf_size: int, + /, +) -> int: + """Get adapter metadata value as a string by key name.""" + ... + + +# // Get the number of metadata key/value pairs +# LLAMA_API int32_t llama_adapter_meta_count(const struct llama_adapter_lora * adapter); +@ctypes_function( + "llama_adapter_meta_count", [llama_adapter_lora_p_ctypes], ctypes.c_int32 +) +def llama_adapter_meta_count(adapter: llama_adapter_lora_p, /) -> int: + """Get the number of adapter metadata key/value pairs.""" + ... + + +# // Get metadata key name by index +# LLAMA_API int32_t llama_adapter_meta_key_by_index(const struct llama_adapter_lora * adapter, int32_t i, char * buf, size_t buf_size); +@ctypes_function( + "llama_adapter_meta_key_by_index", + [ + llama_adapter_lora_p_ctypes, + ctypes.c_int32, + ctypes.c_char_p, + ctypes.c_size_t, + ], + ctypes.c_int32, +) +def llama_adapter_meta_key_by_index( + adapter: llama_adapter_lora_p, + i: int, + buf: Union[bytes, CtypesArray[ctypes.c_char]], + buf_size: int, + /, +) -> int: + """Get adapter metadata key name by index.""" + ... + + +# // Get metadata value as a string by index +# LLAMA_API int32_t llama_adapter_meta_val_str_by_index(const struct llama_adapter_lora * adapter, int32_t i, char * buf, size_t buf_size); +@ctypes_function( + "llama_adapter_meta_val_str_by_index", + [ + llama_adapter_lora_p_ctypes, + ctypes.c_int32, + ctypes.c_char_p, + ctypes.c_size_t, + ], + ctypes.c_int32, +) +def llama_adapter_meta_val_str_by_index( + adapter: llama_adapter_lora_p, + i: int, + buf: Union[bytes, CtypesArray[ctypes.c_char]], + buf_size: int, + /, +) -> int: + """Get adapter metadata value as a string by index.""" + ... + + +# // Manually free a LoRA adapter +# // Note: loaded adapters will be free when the associated model is deleted +# LLAMA_API void llama_adapter_lora_free(struct llama_adapter_lora * adapter); +@ctypes_function( + "llama_adapter_lora_free", + [llama_adapter_lora_p_ctypes], + None, +) +def llama_adapter_lora_free(adapter: llama_adapter_lora_p, /): ... + + +# // Get the invocation tokens if the current lora is an alora +# LLAMA_API uint64_t llama_adapter_get_alora_n_invocation_tokens(const struct llama_adapter_lora * adapter); +@ctypes_function( + "llama_adapter_get_alora_n_invocation_tokens", + [llama_adapter_lora_p_ctypes], + ctypes.c_uint64, +) +def llama_adapter_get_alora_n_invocation_tokens( + adapter: llama_adapter_lora_p, / +) -> int: + """Get the invocation token count if the current LoRA is an aLoRA.""" + ... + + +# LLAMA_API const llama_token * llama_adapter_get_alora_invocation_tokens (const struct llama_adapter_lora * adapter); +@ctypes_function( + "llama_adapter_get_alora_invocation_tokens", + [llama_adapter_lora_p_ctypes], + ctypes.POINTER(llama_token), +) +def llama_adapter_get_alora_invocation_tokens( + adapter: llama_adapter_lora_p, / +) -> Optional[CtypesPointer[llama_token]]: + """Get the invocation tokens if the current LoRA is an aLoRA.""" + ... + + +# // The following functions operate on a llama_context, hence the naming: llama_verb_... + + +# // Set LoRa adapters on the context. Will only modify if the adapters currently in context are different. +# LLAMA_API int32_t llama_set_adapters_lora( +# struct llama_context * ctx, +# struct llama_adapter_lora ** adapters, +# size_t n_adapters, +# float * scales); +@ctypes_function( + "llama_set_adapters_lora", + [ + llama_context_p_ctypes, + ctypes.POINTER(llama_adapter_lora_p_ctypes), + ctypes.c_size_t, + ctypes.POINTER(ctypes.c_float), + ], + ctypes.c_int32, +) +def llama_set_adapters_lora( + ctx: llama_context_p, + adapters: Optional[CtypesArray[llama_adapter_lora_p_ctypes]], + n_adapters: int, + scales: Optional[CtypesArray[ctypes.c_float]], + /, +) -> int: + """Set LoRA adapters on the context if they differ from the current adapters.""" + ... + + +# Deprecated compatibility wrapper for the renamed llama_set_adapters_lora(). +def llama_set_adapter_lora( + ctx: llama_context_p, adapter: llama_adapter_lora_p, scale: float, / +) -> int: + warnings.warn( + "llama_set_adapter_lora is deprecated; use llama_set_adapters_lora instead", + DeprecationWarning, + stacklevel=2, + ) + adapters = (llama_adapter_lora_p_ctypes * 1)(adapter) + scales = (ctypes.c_float * 1)(scale) + return llama_set_adapters_lora(ctx, adapters, 1, scales) + + +# // Apply a loaded control vector to a llama_context, or if data is NULL, clear +# // the currently loaded vector. +# // n_embd should be the size of a single layer's control, and data should point +# // to an n_embd x n_layers buffer starting from layer 1. +# // il_start and il_end are the layer range the vector should apply to (both inclusive) +# // See llama_control_vector_load in common to load a control vector. +# LLAMA_API int32_t llama_set_adapter_cvec( +# struct llama_context * ctx, +# const float * data, +# size_t len, +# int32_t n_embd, +# int32_t il_start, +# int32_t il_end); +@ctypes_function( + "llama_set_adapter_cvec", + [ + llama_context_p_ctypes, + ctypes.POINTER(ctypes.c_float), + ctypes.c_size_t, + ctypes.c_int32, + ctypes.c_int32, + ctypes.c_int32, + ], + ctypes.c_int32, +) +def llama_set_adapter_cvec( + ctx: llama_context_p, + data: CtypesPointerOrRef[ctypes.c_float], + len: int, + n_embd: int, + il_start: int, + il_end: int, + /, +) -> int: + """Apply a loaded control vector to a llama_context, or if data is NULL, clear + the currently loaded vector. + n_embd should be the size of a single layer's control, and data should point + to an n_embd x n_layers buffer starting from layer 1. + il_start and il_end are the layer range the vector should apply to (both inclusive) + See llama_control_vector_load in common to load a control vector.""" + ... + + +# Deprecated compatibility wrapper for the renamed llama_set_adapter_cvec(). +def llama_apply_adapter_cvec( + ctx: llama_context_p, + data: CtypesPointerOrRef[ctypes.c_float], + len: int, + n_embd: int, + il_start: int, + il_end: int, + /, +) -> int: + warnings.warn( + "llama_apply_adapter_cvec is deprecated; use llama_set_adapter_cvec instead", + DeprecationWarning, + stacklevel=2, + ) + return llama_set_adapter_cvec(ctx, data, len, n_embd, il_start, il_end) + + +# // +# // Memory +# // + + +# // Clear the memory contents +# // If data == true, the data buffers will also be cleared together with the metadata +# LLAMA_API void llama_memory_clear( +# llama_memory_t mem, +# bool data); +@ctypes_function( + "llama_memory_clear", + [llama_memory_t_ctypes, ctypes.c_bool], + None, +) +def llama_memory_clear(mem: llama_memory_t, data: bool, /): + """Clear the memory contents + If data == true, the data buffers will also be cleared together with the metadata""" + ... + + +# // Removes all tokens that belong to the specified sequence and have positions in [p0, p1) +# // Returns false if a partial sequence cannot be removed. Removing a whole sequence never fails +# // seq_id < 0 : match any sequence +# // p0 < 0 : [0, p1] +# // p1 < 0 : [p0, inf) +# LLAMA_API bool llama_memory_seq_rm( +# llama_memory_t mem, +# llama_seq_id seq_id, +# llama_pos p0, +# llama_pos p1); +@ctypes_function( + "llama_memory_seq_rm", + [ + llama_memory_t_ctypes, + llama_seq_id, + llama_pos, + llama_pos, + ], + ctypes.c_bool, +) +def llama_memory_seq_rm( + mem: llama_memory_t, + seq_id: Union[llama_seq_id, int], + p0: Union[llama_pos, int], + p1: Union[llama_pos, int], + /, +) -> bool: + """Removes all tokens that belong to the specified sequence and have positions in [p0, p1) + + Returns false if a partial sequence cannot be removed. Removing a whole sequence never fails + + seq_id < 0 : match any sequence + p0 < 0 : [0, p1] + p1 < 0 : [p0, inf)""" + ... + + +# // Copy all tokens that belong to the specified sequence to another sequence +# // p0 < 0 : [0, p1] +# // p1 < 0 : [p0, inf) +# LLAMA_API void llama_memory_seq_cp( +# llama_memory_t mem, +# llama_seq_id seq_id_src, +# llama_seq_id seq_id_dst, +# llama_pos p0, +# llama_pos p1); +@ctypes_function( + "llama_memory_seq_cp", + [ + llama_memory_t_ctypes, + llama_seq_id, + llama_seq_id, + llama_pos, + llama_pos, + ], + None, +) +def llama_memory_seq_cp( + mem: llama_memory_t, + seq_id_src: Union[llama_seq_id, int], + seq_id_dst: Union[llama_seq_id, int], + p0: Union[llama_pos, int], + p1: Union[llama_pos, int], + /, +): + """Copy all tokens that belong to the specified sequence to another sequence + p0 < 0 : [0, p1] + p1 < 0 : [p0, inf)""" + ... + + +# // Removes all tokens that do not belong to the specified sequence +# LLAMA_API void llama_memory_seq_keep( +# llama_memory_t mem, +# llama_seq_id seq_id); +@ctypes_function("llama_memory_seq_keep", [llama_memory_t_ctypes, llama_seq_id], None) +def llama_memory_seq_keep(mem: llama_memory_t, seq_id: Union[llama_seq_id, int], /): + """Removes all tokens that do not belong to the specified sequence""" + ... + + +# // Adds relative position "delta" to all tokens that belong to the specified sequence and have positions in [p0, p1) +# // p0 < 0 : [0, p1] +# // p1 < 0 : [p0, inf) +# LLAMA_API void llama_memory_seq_add( +# llama_memory_t mem, +# llama_seq_id seq_id, +# llama_pos p0, +# llama_pos p1, +# llama_pos delta); +@ctypes_function( + "llama_memory_seq_add", + [ + llama_memory_t_ctypes, + llama_seq_id, + llama_pos, + llama_pos, + llama_pos, + ], + None, +) +def llama_memory_seq_add( + mem: llama_memory_t, + seq_id: Union[llama_seq_id, int], + p0: Union[llama_pos, int], + p1: Union[llama_pos, int], + delta: Union[llama_pos, int], + /, +): + """Adds relative position "delta" to all tokens that belong to the specified sequence and have positions in [p0, p1) + p0 < 0 : [0, p1] + p1 < 0 : [p0, inf)""" + ... + + +# // Integer division of the positions by factor of `d > 1` +# // p0 < 0 : [0, p1] +# // p1 < 0 : [p0, inf) +# LLAMA_API void llama_memory_seq_div( +# llama_memory_t mem, +# llama_seq_id seq_id, +# llama_pos p0, +# llama_pos p1, +# int d); +@ctypes_function( + "llama_memory_seq_div", + [ + llama_memory_t_ctypes, + llama_seq_id, + llama_pos, + llama_pos, + ctypes.c_int, + ], + None, +) +def llama_memory_seq_div( + mem: llama_memory_t, + seq_id: Union[llama_seq_id, int], + p0: Union[llama_pos, int], + p1: Union[llama_pos, int], + d: Union[ctypes.c_int, int], + /, +): + """Integer division of the positions by factor of `d > 1` + p0 < 0 : [0, p1] + p1 < 0 : [p0, inf)""" + ... + + +# // Returns the smallest position present in the memory for the specified sequence +# // This is typically non-zero only for SWA caches +# // Note that all positions in the range [pos_min, pos_max] are guaranteed to be present in the memory +# // Return -1 if the sequence is empty +# LLAMA_API llama_pos llama_memory_seq_pos_min( +# llama_memory_t mem, +# llama_seq_id seq_id); +@ctypes_function( + "llama_memory_seq_pos_min", [llama_memory_t_ctypes, llama_seq_id], llama_pos +) +def llama_memory_seq_pos_min( + mem: llama_memory_t, seq_id: Union[llama_seq_id, int], / +) -> int: + """Returns the smallest position present in the memory for the specified sequence + This is typically non-zero only for SWA caches + Return -1 if the sequence is empty""" + ... + + +# // Returns the largest position present in the memory for the specified sequence +# // Note that all positions in the range [pos_min, pos_max] are guaranteed to be present in the memory +# // Return -1 if the sequence is empty +# LLAMA_API llama_pos llama_memory_seq_pos_max( +# llama_memory_t mem, +# llama_seq_id seq_id); +@ctypes_function( + "llama_memory_seq_pos_max", [llama_memory_t_ctypes, llama_seq_id], llama_pos +) +def llama_memory_seq_pos_max( + mem: llama_memory_t, seq_id: Union[llama_seq_id, int], / +) -> int: + """Returns the largest position present in the memory for the specified sequence + Return -1 if the sequence is empty""" + ... + + +# // Check if the memory supports shifting +# LLAMA_API bool llama_memory_can_shift(llama_memory_t mem); +@ctypes_function("llama_memory_can_shift", [llama_memory_t_ctypes], ctypes.c_bool) +def llama_memory_can_shift(mem: llama_memory_t, /) -> bool: + """Check if the memory supports shifting""" + ... + + +# // +# // State / sessions +# // + + +# // Returns the *actual* size in bytes of the state +# // (logits, embedding and memory) +# // Only use when saving the state, not when restoring it, otherwise the size may be too small. +# LLAMA_API size_t llama_state_get_size(struct llama_context * ctx); +@ctypes_function("llama_state_get_size", [llama_context_p_ctypes], ctypes.c_size_t) +def llama_state_get_size(ctx: llama_context_p, /) -> int: + """Returns the *actual* size in bytes of the state (logits, embedding and memory)""" + ... + + +# LLAMA_API DEPRECATED(size_t llama_get_state_size(struct llama_context * ctx), +# "use llama_state_get_size instead"); +@ctypes_function("llama_get_state_size", [llama_context_p_ctypes], ctypes.c_size_t) +def llama_get_state_size(ctx: llama_context_p, /) -> int: + """Returns the size in bytes of the state (DEPRECATED)""" + ... + + +_llama_get_state_size = llama_get_state_size + + +def llama_get_state_size(ctx: llama_context_p, /) -> int: + _warn_deprecated("llama_get_state_size", "use llama_state_get_size instead") + return _llama_get_state_size(ctx) + + +# // Copies the state to the specified destination address. +# // Destination needs to have allocated enough memory. +# // Returns the number of bytes copied +# LLAMA_API size_t llama_state_get_data( +# struct llama_context * ctx, +# uint8_t * dst, +# size_t size); +@ctypes_function( + "llama_state_get_data", + [ + llama_context_p_ctypes, + ctypes.POINTER(ctypes.c_uint8), + ctypes.c_size_t, + ], + ctypes.c_size_t, +) +def llama_state_get_data( + ctx: llama_context_p, + dst: CtypesArray[ctypes.c_uint8], + size: Union[ctypes.c_size_t, int], + /, +) -> int: + """Copies the state to the specified destination address. + Destination needs to have allocated enough memory. + Returns the number of bytes copied""" + ... + + +# LLAMA_API DEPRECATED(size_t llama_copy_state_data( +# struct llama_context * ctx, +# uint8_t * dst), +# "use llama_state_get_data instead"); +@ctypes_function( + "llama_copy_state_data", + [ + llama_context_p_ctypes, + ctypes.POINTER(ctypes.c_uint8), + ], + ctypes.c_size_t, +) +def llama_copy_state_data( + ctx: llama_context_p, dst: CtypesArray[ctypes.c_uint8], / +) -> int: + """Copies the state to the specified destination address (DEPRECATED)""" + ... + + +_llama_copy_state_data = llama_copy_state_data + + +def llama_copy_state_data( + ctx: llama_context_p, dst: CtypesArray[ctypes.c_uint8], / +) -> int: + _warn_deprecated("llama_copy_state_data", "use llama_state_get_data instead") + return _llama_copy_state_data(ctx, dst) + + +# // Set the state reading from the specified address +# // Returns the number of bytes read +# LLAMA_API size_t llama_state_set_data( +# struct llama_context * ctx, +# const uint8_t * src, +# size_t size); +@ctypes_function( + "llama_state_set_data", + [llama_context_p_ctypes, ctypes.POINTER(ctypes.c_uint8), ctypes.c_size_t], + ctypes.c_size_t, +) +def llama_state_set_data( + ctx: llama_context_p, + src: CtypesArray[ctypes.c_uint8], + size: Union[ctypes.c_size_t, int], + /, +) -> int: + """Set the state reading from the specified address + Returns the number of bytes read""" + ... + + +# LLAMA_API DEPRECATED(size_t llama_set_state_data( +# struct llama_context * ctx, +# const uint8_t * src), +# "use llama_state_set_data instead"); +@ctypes_function( + "llama_set_state_data", + [llama_context_p_ctypes, ctypes.POINTER(ctypes.c_uint8)], + ctypes.c_size_t, +) +def llama_set_state_data( + ctx: llama_context_p, src: CtypesArray[ctypes.c_uint8], / +) -> int: + """Set the state reading from the specified address (DEPRECATED)""" + ... + + +_llama_set_state_data = llama_set_state_data + + +def llama_set_state_data( + ctx: llama_context_p, src: CtypesArray[ctypes.c_uint8], / +) -> int: + _warn_deprecated("llama_set_state_data", "use llama_state_set_data instead") + return _llama_set_state_data(ctx, src) + + +# Save/load session file +# LLAMA_API bool llama_state_load_file( +# struct llama_context * ctx, +# const char * path_session, +# llama_token * tokens_out, +# size_t n_token_capacity, +# size_t * n_token_count_out); +@ctypes_function( + "llama_state_load_file", + [ + llama_context_p_ctypes, + ctypes.c_char_p, + llama_token_p, + ctypes.c_size_t, + ctypes.POINTER(ctypes.c_size_t), + ], + ctypes.c_bool, +) +def llama_state_load_file( ctx: llama_context_p, - candidates, # type: _Pointer[llama_token_data_array] + path_session: bytes, + tokens_out: CtypesArray[llama_token], + n_token_capacity: Union[ctypes.c_size_t, int], + n_token_count_out: CtypesPointerOrRef[ctypes.c_size_t], + /, +) -> bool: ... + + +# LLAMA_API DEPRECATED(bool llama_load_session_file( +# struct llama_context * ctx, +# const char * path_session, +# llama_token * tokens_out, +# size_t n_token_capacity, +# size_t * n_token_count_out), +# "use llama_state_load_file instead"); +@ctypes_function( + "llama_load_session_file", + [ + llama_context_p_ctypes, + ctypes.c_char_p, + llama_token_p, + ctypes.c_size_t, + ctypes.POINTER(ctypes.c_size_t), + ], + ctypes.c_bool, +) +def llama_load_session_file( + ctx: llama_context_p, + path_session: bytes, + tokens_out: CtypesArray[llama_token], + n_token_capacity: Union[ctypes.c_size_t, int], + n_token_count_out: CtypesPointerOrRef[ctypes.c_size_t], + /, +) -> bool: ... + + +_llama_load_session_file = llama_load_session_file + + +def llama_load_session_file( + ctx: llama_context_p, + path_session: bytes, + tokens_out: CtypesArray[llama_token], + n_token_capacity: Union[ctypes.c_size_t, int], + n_token_count_out: CtypesPointerOrRef[ctypes.c_size_t], + /, +) -> bool: + _warn_deprecated("llama_load_session_file", "use llama_state_load_file instead") + return _llama_load_session_file( + ctx, path_session, tokens_out, n_token_capacity, n_token_count_out + ) + + +# LLAMA_API bool llama_state_save_file( +# struct llama_context * ctx, +# const char * path_session, +# const llama_token * tokens, +# size_t n_token_count); +@ctypes_function( + "llama_state_save_file", + [ + llama_context_p_ctypes, + ctypes.c_char_p, + llama_token_p, + ctypes.c_size_t, + ], + ctypes.c_bool, +) +def llama_state_save_file( + ctx: llama_context_p, + path_session: bytes, + tokens: CtypesArray[llama_token], + n_token_count: Union[ctypes.c_size_t, int], + /, +) -> bool: ... + + +# LLAMA_API DEPRECATED(bool llama_save_session_file( +# struct llama_context * ctx, +# const char * path_session, +# const llama_token * tokens, +# size_t n_token_count), +# "use llama_state_save_file instead"); +@ctypes_function( + "llama_save_session_file", + [ + llama_context_p_ctypes, + ctypes.c_char_p, + llama_token_p, + ctypes.c_size_t, + ], + ctypes.c_bool, +) +def llama_save_session_file( + ctx: llama_context_p, + path_session: bytes, + tokens: CtypesArray[llama_token], + n_token_count: Union[ctypes.c_size_t, int], + /, +) -> bool: ... + + +_llama_save_session_file = llama_save_session_file + + +def llama_save_session_file( + ctx: llama_context_p, + path_session: bytes, + tokens: CtypesArray[llama_token], + n_token_count: Union[ctypes.c_size_t, int], + /, +) -> bool: + _warn_deprecated("llama_save_session_file", "use llama_state_save_file instead") + return _llama_save_session_file(ctx, path_session, tokens, n_token_count) + + +# // Get the exact size needed to copy the state of a single sequence +# LLAMA_API size_t llama_state_seq_get_size( +# struct llama_context * ctx, +# llama_seq_id seq_id); +@ctypes_function( + "llama_state_seq_get_size", + [llama_context_p_ctypes, llama_seq_id], + ctypes.c_size_t, +) +def llama_state_seq_get_size(ctx: llama_context_p, seq_id: llama_seq_id, /) -> int: + """Get the exact size needed to copy the state of a single sequence""" + ... + + +# // Copy the state of a single sequence into the specified buffer +# LLAMA_API size_t llama_state_seq_get_data( +# struct llama_context * ctx, +# uint8_t * dst, +# size_t size, +# llama_seq_id seq_id); +@ctypes_function( + "llama_state_seq_get_data", + [ + llama_context_p_ctypes, + ctypes.POINTER(ctypes.c_uint8), + ctypes.c_size_t, + llama_seq_id, + ], + ctypes.c_size_t, +) +def llama_state_seq_get_data( + ctx: llama_context_p, + dst: CtypesArray[ctypes.c_uint8], + size: Union[ctypes.c_size_t, int], + seq_id: llama_seq_id, + /, +) -> int: + """Copy the state of a single sequence into the specified buffer""" + ... + + +# // Copy the sequence data (originally copied with `llama_state_seq_get_data`) into the specified sequence +# // Returns: +# // - Positive: Ok +# // - Zero: Failed to load +# LLAMA_API size_t llama_state_seq_set_data( +# struct llama_context * ctx, +# const uint8_t * src, +# size_t size, +# llama_seq_id dest_seq_id); +@ctypes_function( + "llama_state_seq_set_data", + [ + llama_context_p_ctypes, + ctypes.POINTER(ctypes.c_uint8), + ctypes.c_size_t, + llama_seq_id, + ], + ctypes.c_size_t, +) +def llama_state_seq_set_data( + ctx: llama_context_p, + src: CtypesArray[ctypes.c_uint8], + size: Union[ctypes.c_size_t, int], + dest_seq_id: llama_seq_id, + /, +) -> int: + """Copy the sequence data into the specified sequence""" + ... + + +# LLAMA_API size_t llama_state_seq_save_file( +# struct llama_context * ctx, +# const char * filepath, +# llama_seq_id seq_id, +# const llama_token * tokens, +# size_t n_token_count); +@ctypes_function( + "llama_state_seq_save_file", + [ + llama_context_p_ctypes, + ctypes.c_char_p, + llama_seq_id, + llama_token_p, + ctypes.c_size_t, + ], + ctypes.c_size_t, +) +def llama_state_seq_save_file( + ctx: llama_context_p, + filepath: bytes, + seq_id: llama_seq_id, + tokens: CtypesArray[llama_token], + n_token_count: Union[ctypes.c_size_t, int], + /, +) -> int: ... + + +# LLAMA_API size_t llama_state_seq_load_file( +# struct llama_context * ctx, +# const char * filepath, +# llama_seq_id dest_seq_id, +# llama_token * tokens_out, +# size_t n_token_capacity, +# size_t * n_token_count_out); +@ctypes_function( + "llama_state_seq_load_file", + [ + llama_context_p_ctypes, + ctypes.c_char_p, + llama_seq_id, + llama_token_p, + ctypes.c_size_t, + ctypes.POINTER(ctypes.c_size_t), + ], + ctypes.c_size_t, +) +def llama_state_seq_load_file( + ctx: llama_context_p, + filepath: bytes, + dest_seq_id: llama_seq_id, + tokens_out: CtypesArray[llama_token], + n_token_capacity: Union[ctypes.c_size_t, int], + n_token_count_out: CtypesPointerOrRef[ctypes.c_size_t], + /, +) -> int: ... + + +# define LLAMA_STATE_SEQ_FLAGS_NONE 0 +LLAMA_STATE_SEQ_FLAGS_NONE = 0 + +# for backwards-compat +# define LLAMA_STATE_SEQ_FLAGS_SWA_ONLY 1 +LLAMA_STATE_SEQ_FLAGS_SWA_ONLY = 1 + +# work only with partial states, such as SWA KV cache or recurrent cache +# (e.g. Mamba) +# define LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY 1 +LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY = 1 + +# keeps the tensor data on device buffers +# (i.e. not accessible in host memory, but faster save/load) +# define LLAMA_STATE_SEQ_FLAGS_ON_DEVICE 2 +LLAMA_STATE_SEQ_FLAGS_ON_DEVICE = 2 + + +# LLAMA_API size_t llama_state_seq_get_size_ext( +# struct llama_context * ctx, +# llama_seq_id seq_id, +# llama_state_seq_flags flags); +@ctypes_function( + "llama_state_seq_get_size_ext", + [llama_context_p_ctypes, llama_seq_id, llama_state_seq_flags], + ctypes.c_size_t, +) +def llama_state_seq_get_size_ext( + ctx: llama_context_p, + seq_id: llama_seq_id, + flags: llama_state_seq_flags, + /, +) -> int: ... + + +# LLAMA_API size_t llama_state_seq_get_data_ext( +# struct llama_context * ctx, +# uint8_t * dst, +# size_t size, +# llama_seq_id seq_id, +# llama_state_seq_flags flags); +@ctypes_function( + "llama_state_seq_get_data_ext", + [ + llama_context_p_ctypes, + ctypes.POINTER(ctypes.c_uint8), + ctypes.c_size_t, + llama_seq_id, + llama_state_seq_flags, + ], + ctypes.c_size_t, +) +def llama_state_seq_get_data_ext( + ctx: llama_context_p, + dst: CtypesArray[ctypes.c_uint8], + size: Union[ctypes.c_size_t, int], + seq_id: llama_seq_id, + flags: llama_state_seq_flags, + /, +) -> int: ... + + +# LLAMA_API size_t llama_state_seq_set_data_ext( +# struct llama_context * ctx, +# const uint8_t * src, +# size_t size, +# llama_seq_id dest_seq_id, +# llama_state_seq_flags flags); +@ctypes_function( + "llama_state_seq_set_data_ext", + [ + llama_context_p_ctypes, + ctypes.POINTER(ctypes.c_uint8), + ctypes.c_size_t, + llama_seq_id, + llama_state_seq_flags, + ], + ctypes.c_size_t, +) +def llama_state_seq_set_data_ext( + ctx: llama_context_p, + src: CtypesArray[ctypes.c_uint8], + size: Union[ctypes.c_size_t, int], + dest_seq_id: llama_seq_id, + flags: llama_state_seq_flags, + /, +) -> int: ... + + +# // +# // Decoding +# // + + +# // Return batch for single sequence of tokens +# // The sequence ID will be fixed to 0 +# // The position of the tokens will be tracked automatically by llama_decode +# // +# // NOTE: this is a helper function to facilitate transition to the new batch API - avoid using it +# // +# LLAMA_API struct llama_batch llama_batch_get_one( +# llama_token * tokens, +# int32_t n_tokens); +@ctypes_function( + "llama_batch_get_one", + [ + llama_token_p, + ctypes.c_int32, + ], + llama_batch, +) +def llama_batch_get_one( + tokens: CtypesArray[llama_token], + n_tokens: Union[ctypes.c_int, int], + /, +) -> llama_batch: + """Return batch for single sequence of tokens + + NOTE: this is a helper function to facilitate transition to the new batch API - avoid using it + """ + ... + + +# // Allocates a batch of tokens on the heap that can hold a maximum of n_tokens +# // Each token can be assigned up to n_seq_max sequence ids +# // The batch has to be freed with llama_batch_free() +# // If embd != 0, llama_batch.embd will be allocated with size of n_tokens * embd * sizeof(float) +# // Otherwise, llama_batch.token will be allocated to store n_tokens llama_token +# // The rest of the llama_batch members are allocated with size n_tokens +# // All members are left uninitialized +# LLAMA_API struct llama_batch llama_batch_init( +# int32_t n_tokens, +# int32_t embd, +# int32_t n_seq_max); +@ctypes_function( + "llama_batch_init", [ctypes.c_int32, ctypes.c_int32, ctypes.c_int32], llama_batch +) +def llama_batch_init( + n_tokens: Union[ctypes.c_int32, int], + embd: Union[ctypes.c_int32, int], + n_seq_max: Union[ctypes.c_int32, int], + /, +) -> llama_batch: + """Allocates a batch of tokens on the heap that can hold a maximum of n_tokens + Each token can be assigned up to n_seq_max sequence ids + The batch has to be freed with llama_batch_free() + If embd != 0, llama_batch.embd will be allocated with size of n_tokens * embd * sizeof(float) + Otherwise, llama_batch.token will be allocated to store n_tokens llama_token + The rest of the llama_batch members are allocated with size n_tokens + All members are left uninitialized""" + ... + + +# // Frees a batch of tokens allocated with llama_batch_init() +# LLAMA_API void llama_batch_free(struct llama_batch batch); +@ctypes_function("llama_batch_free", [llama_batch], None) +def llama_batch_free(batch: llama_batch, /): + """Frees a batch of tokens allocated with llama_batch_init()""" + ... + + +# // Process a batch of tokens. +# // In contrast to llama_decode() - this call does not use KV cache. +# // For encode-decoder contexts, processes the batch using the encoder. +# // Can store the encoder output internally for later use by the decoder's cross-attention layers. +# // 0 - success +# // < 0 - error. the memory state is restored to the state before this call +# LLAMA_API int32_t llama_encode( +# struct llama_context * ctx, +# struct llama_batch batch); +@ctypes_function("llama_encode", [llama_context_p_ctypes, llama_batch], ctypes.c_int32) +def llama_encode(ctx: llama_context_p, batch: llama_batch, /) -> int: + """Process a batch of tokens using the encoder. + 0 - success + < 0 - error""" + ... + + +# // Process a batch of tokens. +# // Requires the context to have a memory. +# // For encode-decoder contexts, processes the batch using the decoder. +# // Positive return values does not mean a fatal error, but rather a warning. +# // Upon fatal-error or abort, the ubatches that managed to be been processed will remain in the memory state of the context +# // To handle this correctly, query the memory state using llama_memory_seq_pos_min() and llama_memory_seq_pos_max() +# // Upon other return values, the memory state is restored to the state before this call +# // 0 - success +# // 1 - could not find a KV slot for the batch (try reducing the size of the batch or increase the context) +# // 2 - aborted (processed ubatches will remain in the context's memory) +# // -1 - invalid input batch +# // < -1 - fatal error (processed ubatches will remain in the context's memory) +# LLAMA_API int32_t llama_decode( +# struct llama_context * ctx, +# struct llama_batch batch); +@ctypes_function("llama_decode", [llama_context_p_ctypes, llama_batch], ctypes.c_int32) +def llama_decode(ctx: llama_context_p, batch: llama_batch, /) -> int: + """Process a batch of tokens. + 0 - success + 1 - could not find a KV slot for the batch (try reducing the size of the batch or increase the context) + 2 - aborted (processed ubatches will remain in the context's memory) + -1 - invalid input batch + < -1 - fatal error (processed ubatches will remain in the context's memory)""" + ... + + +# // Set the number of threads used for decoding +# // n_threads is the number of threads used for generation (single token) +# // n_threads_batch is the number of threads used for prompt and batch processing (multiple tokens) +# LLAMA_API void llama_set_n_threads(struct llama_context * ctx, int32_t n_threads, int32_t n_threads_batch); +@ctypes_function( + "llama_set_n_threads", + [ + llama_context_p_ctypes, + ctypes.c_int32, + ctypes.c_int32, + ], + None, +) +def llama_set_n_threads( + ctx: llama_context_p, + n_threads: Union[ctypes.c_int32, int], + n_threads_batch: Union[ctypes.c_int32, int], + /, +): + """Set the number of threads used for decoding + n_threads is the number of threads used for generation (single token) + n_threads_batch is the number of threads used for prompt and batch processing (multiple tokens) + """ + ... + + +# // Get the number of threads used for generation of a single token. +# LLAMA_API int32_t llama_n_threads(struct llama_context * ctx); +@ctypes_function("llama_n_threads", [llama_context_p_ctypes], ctypes.c_int32) +def llama_n_threads(ctx: llama_context_p, /) -> int: + """Get the number of threads used for generation of a single token""" + ... + + +# // Get the number of threads used for prompt and batch processing (multiple token). +# LLAMA_API int32_t llama_n_threads_batch(struct llama_context * ctx); +@ctypes_function("llama_n_threads_batch", [llama_context_p_ctypes], ctypes.c_int32) +def llama_n_threads_batch(ctx: llama_context_p, /) -> int: + """Get the number of threads used for prompt and batch processing (multiple token)""" + ... + + +# // Set whether the context outputs embeddings or not +# // TODO: rename to avoid confusion with llama_get_embeddings() +# LLAMA_API void llama_set_embeddings(struct llama_context * ctx, bool embeddings); +@ctypes_function("llama_set_embeddings", [llama_context_p_ctypes, ctypes.c_bool], None) +def llama_set_embeddings(ctx: llama_context_p, embeddings: bool, /): + """Set whether the context outputs embeddings or not""" + ... + + +# // Set whether to use causal attention or not +# // If set to true, the model will only attend to the past tokens +# LLAMA_API void llama_set_causal_attn(struct llama_context * ctx, bool causal_attn); +@ctypes_function("llama_set_causal_attn", [llama_context_p_ctypes, ctypes.c_bool], None) +def llama_set_causal_attn(ctx: llama_context_p, causal_attn: bool, /): + """Set whether to use causal attention or not + If set to true, the model will only attend to the past tokens""" + ... + + +# // Set whether the model is in warmup mode or not +# // If true, all model tensors are activated during llama_decode() to load and cache their weights. +# LLAMA_API void llama_set_warmup(struct llama_context * ctx, bool warmup); +@ctypes_function("llama_set_warmup", [llama_context_p_ctypes, ctypes.c_bool], None) +def llama_set_warmup(ctx: llama_context_p, warmup: bool, /): + """Set whether the model is in warmup mode or not + If true, all model tensors are activated during llama_decode() to load and cache their weights.""" + ... + + +# // Set abort callback +# LLAMA_API void llama_set_abort_callback(struct llama_context * ctx, ggml_abort_callback abort_callback, void * abort_callback_data); +@ctypes_function( + "llama_set_abort_callback", + [llama_context_p_ctypes, ggml_abort_callback, ctypes.c_void_p], + None, +) +def llama_set_abort_callback( + ctx: llama_context_p, + abort_callback: Callable[[ctypes.c_void_p], None], + abort_callback_data: ctypes.c_void_p, + /, +): + """Set abort callback""" + ... + + +# // Wait until all computations are finished +# // This is automatically done when using one of the functions below to obtain the computation results +# // and is not necessary to call it explicitly in most cases +# LLAMA_API void llama_synchronize(struct llama_context * ctx); +@ctypes_function("llama_synchronize", [llama_context_p_ctypes], None) +def llama_synchronize(ctx: llama_context_p, /): + """Wait until all computations are finished + This is automatically done when using one of the functions below to obtain the computation results + and is not necessary to call it explicitly in most cases""" + ... + + +# // Token logits obtained from the last call to llama_decode() +# // The logits for which llama_batch.logits[i] != 0 are stored contiguously +# // in the order they have appeared in the batch. +# // Rows: number of tokens for which llama_batch.logits[i] != 0 +# // Cols: n_vocab +# // TODO: deprecate in favor of llama_get_logits_ith() (ref: https://github.com/ggml-org/llama.cpp/pull/14853#issuecomment-3113143522) +# LLAMA_API float * llama_get_logits(struct llama_context * ctx); +@ctypes_function( + "llama_get_logits", [llama_context_p_ctypes], ctypes.POINTER(ctypes.c_float) +) +def llama_get_logits(ctx: llama_context_p, /) -> CtypesArray[ctypes.c_float]: + """Token logits obtained from the last call to llama_decode() + The logits for which llama_batch.logits[i] != 0 are stored contiguously + in the order they have appeared in the batch. + Rows: number of tokens for which llama_batch.logits[i] != 0 + Cols: n_vocab + + Returns: + Pointer to the logits buffer of shape (n_tokens, n_vocab)""" + ... + + +# // Logits for the ith token. For positive indices, Equivalent to: +# // llama_get_logits(ctx) + ctx->output_ids[i]*n_vocab +# // Negative indicies can be used to access logits in reverse order, -1 is the last logit. +# // returns NULL for invalid ids. +# LLAMA_API float * llama_get_logits_ith(struct llama_context * ctx, int32_t i); +@ctypes_function( + "llama_get_logits_ith", + [llama_context_p_ctypes, ctypes.c_int32], + ctypes.POINTER(ctypes.c_float), +) +def llama_get_logits_ith( + ctx: llama_context_p, i: Union[ctypes.c_int32, int], / +) -> CtypesArray[ctypes.c_float]: + """Logits for the ith token. Equivalent to: + llama_get_logits(ctx) + i*n_vocab""" + ... + + +# // Get all output token embeddings. +# // when pooling_type == LLAMA_POOLING_TYPE_NONE or when using a generative model, +# // the embeddings for which llama_batch.logits[i] != 0 are stored contiguously +# // in the order they have appeared in the batch. +# // shape: [n_outputs*n_embd] +# // Otherwise, returns NULL. +# // TODO: deprecate in favor of llama_get_embeddings_ith() (ref: https://github.com/ggml-org/llama.cpp/pull/14853#issuecomment-3113143522) +# LLAMA_API float * llama_get_embeddings(struct llama_context * ctx); +@ctypes_function( + "llama_get_embeddings", [llama_context_p_ctypes], ctypes.POINTER(ctypes.c_float) +) +def llama_get_embeddings(ctx: llama_context_p, /) -> CtypesArray[ctypes.c_float]: + """Get the embeddings for the input + shape: [n_embd] (1-dimensional)""" + ... + + +# // Get the embeddings for the ith token. For positive indices, Equivalent to: +# // llama_get_embeddings(ctx) + ctx->output_ids[i]*n_embd +# // Negative indicies can be used to access embeddings in reverse order, -1 is the last embedding. +# // shape: [n_embd] (1-dimensional) +# // returns NULL for invalid ids. +# LLAMA_API float * llama_get_embeddings_ith(struct llama_context * ctx, int32_t i); +@ctypes_function( + "llama_get_embeddings_ith", + [llama_context_p_ctypes, ctypes.c_int32], + ctypes.POINTER(ctypes.c_float), +) +def llama_get_embeddings_ith( + ctx: llama_context_p, i: Union[ctypes.c_int32, int], / +) -> CtypesArray[ctypes.c_float]: + """Get the embeddings for the ith sequence + llama_get_embeddings(ctx) + i*n_embd""" + ... + + +# // Get the embeddings for a sequence id +# // Returns NULL if pooling_type is LLAMA_POOLING_TYPE_NONE +# // when pooling_type == LLAMA_POOLING_TYPE_RANK, returns float[n_cls_out] with the rank(s) of the sequence +# // otherwise: float[n_embd] (1-dimensional) +# LLAMA_API float * llama_get_embeddings_seq(struct llama_context * ctx, llama_seq_id seq_id); +@ctypes_function( + "llama_get_embeddings_seq", + [llama_context_p_ctypes, llama_seq_id], + ctypes.POINTER(ctypes.c_float), +) +def llama_get_embeddings_seq( + ctx: llama_context_p, seq_id: Union[llama_seq_id, int], / +) -> CtypesArray[ctypes.c_float]: + """Get the embeddings for a sequence id + Returns NULL if pooling_type is LLAMA_POOLING_TYPE_NONE + shape: [n_embd] (1-dimensional)""" + ... + + +# // Get the backend sampled token for the ith token. +# // Returns LLAMA_TOKEN_NULL if no token was sampled. +# LLAMA_API llama_token llama_get_sampled_token_ith(struct llama_context * ctx, int32_t i); +@ctypes_function( + "llama_get_sampled_token_ith", [llama_context_p_ctypes, ctypes.c_int32], llama_token +) +def llama_get_sampled_token_ith( + ctx: llama_context_p, i: Union[ctypes.c_int32, int], / +) -> int: + """Get the backend sampled token for the ith token.""" + ... + + +# // Get the backend sampled probabilities for the ith token +# LLAMA_API float * llama_get_sampled_probs_ith (struct llama_context * ctx, int32_t i); +@ctypes_function( + "llama_get_sampled_probs_ith", + [llama_context_p_ctypes, ctypes.c_int32], + ctypes.POINTER(ctypes.c_float), +) +def llama_get_sampled_probs_ith( + ctx: llama_context_p, i: Union[ctypes.c_int32, int], / +) -> Optional[CtypesPointer[ctypes.c_float]]: + """Get the backend sampled probabilities for the ith token.""" + ... + + +# LLAMA_API uint32_t llama_get_sampled_probs_count_ith(struct llama_context * ctx, int32_t i); +@ctypes_function( + "llama_get_sampled_probs_count_ith", + [llama_context_p_ctypes, ctypes.c_int32], + ctypes.c_uint32, +) +def llama_get_sampled_probs_count_ith( + ctx: llama_context_p, i: Union[ctypes.c_int32, int], / +) -> int: + """Get the backend sampled probability count for the ith token.""" + ... + + +# // Get the backend sampled logits for the ith token +# LLAMA_API float * llama_get_sampled_logits_ith (struct llama_context * ctx, int32_t i); +@ctypes_function( + "llama_get_sampled_logits_ith", + [llama_context_p_ctypes, ctypes.c_int32], + ctypes.POINTER(ctypes.c_float), +) +def llama_get_sampled_logits_ith( + ctx: llama_context_p, i: Union[ctypes.c_int32, int], / +) -> Optional[CtypesPointer[ctypes.c_float]]: + """Get the backend sampled logits for the ith token.""" + ... + + +# LLAMA_API uint32_t llama_get_sampled_logits_count_ith(struct llama_context * ctx, int32_t i); +@ctypes_function( + "llama_get_sampled_logits_count_ith", + [llama_context_p_ctypes, ctypes.c_int32], + ctypes.c_uint32, +) +def llama_get_sampled_logits_count_ith( + ctx: llama_context_p, i: Union[ctypes.c_int32, int], / +) -> int: + """Get the backend sampled logit count for the ith token.""" + ... + + +# // Get the backend sampled candidates for the ith token +# LLAMA_API llama_token * llama_get_sampled_candidates_ith (struct llama_context * ctx, int32_t i); +@ctypes_function( + "llama_get_sampled_candidates_ith", + [llama_context_p_ctypes, ctypes.c_int32], + ctypes.POINTER(llama_token), +) +def llama_get_sampled_candidates_ith( + ctx: llama_context_p, i: Union[ctypes.c_int32, int], / +) -> Optional[CtypesPointer[llama_token]]: + """Get the backend sampled candidates for the ith token.""" + ... + + +# LLAMA_API uint32_t llama_get_sampled_candidates_count_ith(struct llama_context * ctx, int32_t i); +@ctypes_function( + "llama_get_sampled_candidates_count_ith", + [llama_context_p_ctypes, ctypes.c_int32], + ctypes.c_uint32, +) +def llama_get_sampled_candidates_count_ith( + ctx: llama_context_p, i: Union[ctypes.c_int32, int], / +) -> int: + """Get the backend sampled candidate count for the ith token.""" + ... + + +# // +# // Vocab +# // + + +# LLAMA_API const char * llama_vocab_get_text(const struct llama_vocab * vocab, llama_token token); +@ctypes_function( + "llama_vocab_get_text", [llama_vocab_p_ctypes, llama_token], ctypes.c_char_p +) +def llama_vocab_get_text( + vocab: llama_vocab_p, token: Union[llama_token, int], / +) -> bytes: ... + + +# LLAMA_API float llama_vocab_get_score(const struct llama_vocab * vocab, llama_token token); +@ctypes_function( + "llama_vocab_get_score", [llama_vocab_p_ctypes, llama_token], ctypes.c_float +) +def llama_vocab_get_score( + vocab: llama_vocab_p, token: Union[llama_token, int], / +) -> float: ... + + +# LLAMA_API enum llama_token_attr llama_vocab_get_attr(const struct llama_vocab * vocab, llama_token token); +@ctypes_function( + "llama_vocab_get_attr", [llama_vocab_p_ctypes, llama_token], ctypes.c_int +) +def llama_vocab_get_attr( + vocab: llama_vocab_p, token: Union[llama_token, int], / +) -> int: ... + + +# // Check if the token is supposed to end generation (end-of-generation, eg. EOS, EOT, etc.) +# LLAMA_API bool llama_vocab_is_eog(const struct llama_vocab * vocab, llama_token token); +@ctypes_function( + "llama_vocab_is_eog", [llama_vocab_p_ctypes, llama_token], ctypes.c_bool +) +def llama_vocab_is_eog(vocab: llama_vocab_p, token: Union[llama_token, int], /) -> bool: + """Check if the token is supposed to end generation (end-of-generation, eg. EOS, EOT, etc.)""" + ... + + +# // Identify if Token Id is a control token or a render-able token +# LLAMA_API bool llama_vocab_is_control(const struct llama_vocab * vocab, llama_token token); +@ctypes_function( + "llama_vocab_is_control", [llama_vocab_p_ctypes, llama_token], ctypes.c_bool +) +def llama_vocab_is_control( + vocab: llama_vocab_p, token: Union[llama_token, int], / +) -> bool: + """Identify if Token Id is a control token or a render-able token""" + ... + + +# // Special tokens +# LLAMA_API llama_token llama_vocab_bos(const struct llama_vocab * vocab); // beginning-of-sentence +@ctypes_function("llama_vocab_bos", [llama_vocab_p_ctypes], llama_token) +def llama_vocab_bos(vocab: llama_vocab_p, /) -> llama_token: + """beginning-of-sentence""" + ... + + +# LLAMA_API llama_token llama_vocab_eos(const struct llama_vocab * vocab); // end-of-sentence +@ctypes_function("llama_vocab_eos", [llama_vocab_p_ctypes], llama_token) +def llama_vocab_eos(vocab: llama_vocab_p, /) -> llama_token: + """end-of-sentence""" + ... + + +# LLAMA_API llama_token llama_vocab_eot(const struct llama_vocab * vocab); // end-of-turn +@ctypes_function("llama_vocab_eot", [llama_vocab_p_ctypes], llama_token) +def llama_vocab_eot(vocab: llama_vocab_p, /) -> llama_token: + """end-of-turn""" + ... + + +# LLAMA_API llama_token llama_vocab_sep(const struct llama_vocab * vocab); // sentence separator +@ctypes_function("llama_vocab_sep", [llama_vocab_p_ctypes], llama_token) +def llama_vocab_sep(vocab: llama_vocab_p, /) -> llama_token: + """sentence separator""" + ... + + +# LLAMA_API llama_token llama_vocab_nl (const struct llama_vocab * vocab); // next-line +@ctypes_function("llama_vocab_nl", [llama_vocab_p_ctypes], llama_token) +def llama_vocab_nl(vocab: llama_vocab_p, /) -> llama_token: + """next-line""" + ... + + +# LLAMA_API llama_token llama_vocab_pad(const struct llama_vocab * vocab); // padding +@ctypes_function("llama_vocab_pad", [llama_vocab_p_ctypes], llama_token) +def llama_vocab_pad(vocab: llama_vocab_p, /) -> llama_token: + """padding""" + ... + + +# LLAMA_API llama_token llama_vocab_mask(const struct llama_vocab * vocab); // mask +@ctypes_function("llama_vocab_mask", [llama_vocab_p_ctypes], llama_token) +def llama_vocab_mask(vocab: llama_vocab_p, /) -> llama_token: + """mask""" + ... + + +# LLAMA_API bool llama_vocab_get_add_bos(const struct llama_vocab * vocab); +@ctypes_function( + "llama_vocab_get_add_bos", + [llama_vocab_p_ctypes], + ctypes.c_bool, +) +def llama_vocab_get_add_bos(vocab: llama_vocab_p, /) -> bool: ... + + +# LLAMA_API bool llama_vocab_get_add_eos(const struct llama_vocab * vocab); +@ctypes_function( + "llama_vocab_get_add_eos", + [llama_vocab_p_ctypes], + ctypes.c_bool, +) +def llama_vocab_get_add_eos(vocab: llama_vocab_p, /) -> bool: ... + + +# LLAMA_API bool llama_vocab_get_add_sep(const struct llama_vocab * vocab); +@ctypes_function( + "llama_vocab_get_add_sep", + [llama_vocab_p_ctypes], + ctypes.c_bool, +) +def llama_vocab_get_add_sep(vocab: llama_vocab_p, /) -> bool: ... + + +# LLAMA_API llama_token llama_vocab_fim_pre(const struct llama_vocab * vocab); +@ctypes_function( + "llama_vocab_fim_pre", + [llama_vocab_p_ctypes], + llama_token, +) +def llama_vocab_fim_pre(vocab: llama_vocab_p, /) -> llama_token: ... + + +# LLAMA_API llama_token llama_vocab_fim_suf(const struct llama_vocab * vocab); +@ctypes_function( + "llama_vocab_fim_suf", + [llama_vocab_p_ctypes], + llama_token, +) +def llama_vocab_fim_suf(vocab: llama_vocab_p, /) -> llama_token: ... + + +# LLAMA_API llama_token llama_vocab_fim_mid(const struct llama_vocab * vocab); +@ctypes_function( + "llama_vocab_fim_mid", + [llama_vocab_p_ctypes], + llama_token, +) +def llama_vocab_fim_mid(vocab: llama_vocab_p, /) -> llama_token: ... + + +# LLAMA_API llama_token llama_vocab_fim_pad(const struct llama_vocab * vocab); +@ctypes_function( + "llama_vocab_fim_pad", + [llama_vocab_p_ctypes], + llama_token, +) +def llama_vocab_fim_pad(vocab: llama_vocab_p, /) -> llama_token: ... + + +# LLAMA_API llama_token llama_vocab_fim_rep(const struct llama_vocab * vocab); +@ctypes_function( + "llama_vocab_fim_rep", + [llama_vocab_p_ctypes], + llama_token, +) +def llama_vocab_fim_rep(vocab: llama_vocab_p, /) -> llama_token: ... + + +# LLAMA_API llama_token llama_vocab_fim_sep(const struct llama_vocab * vocab); +@ctypes_function( + "llama_vocab_fim_sep", + [llama_vocab_p_ctypes], + llama_token, +) +def llama_vocab_fim_sep(vocab: llama_vocab_p, /) -> llama_token: ... + + +# DEPRECATED functions +# DEPRECATED(LLAMA_API const char * llama_token_get_text(const struct llama_vocab * vocab, llama_token token), "use llama_vocab_get_text instead"); +@ctypes_function( + "llama_token_get_text", + [llama_vocab_p_ctypes, llama_token], + ctypes.c_char_p, +) +def llama_token_get_text( + vocab: llama_vocab_p, token: Union[llama_token, int], / +) -> bytes: ... + + +_llama_token_get_text = llama_token_get_text + + +def llama_token_get_text( + vocab: llama_vocab_p, token: Union[llama_token, int], / +) -> bytes: + _warn_deprecated("llama_token_get_text", "use llama_vocab_get_text instead") + return _llama_token_get_text(vocab, token) + + +# DEPRECATED(LLAMA_API float llama_token_get_score(const struct llama_vocab * vocab, llama_token token), "use llama_vocab_get_score instead"); +@ctypes_function( + "llama_token_get_score", + [llama_vocab_p_ctypes, llama_token], + ctypes.c_float, +) +def llama_token_get_score( + vocab: llama_vocab_p, token: Union[llama_token, int], / +) -> float: ... + + +_llama_token_get_score = llama_token_get_score + + +def llama_token_get_score( + vocab: llama_vocab_p, token: Union[llama_token, int], / +) -> float: + _warn_deprecated("llama_token_get_score", "use llama_vocab_get_score instead") + return _llama_token_get_score(vocab, token) + + +# DEPRECATED(LLAMA_API enum llama_token_attr llama_token_get_attr(const struct llama_vocab * vocab, llama_token token), "use llama_vocab_get_attr instead"); +@ctypes_function( + "llama_token_get_attr", + [llama_vocab_p_ctypes, llama_token], + ctypes.c_int, +) +def llama_token_get_attr( + vocab: llama_vocab_p, token: Union[llama_token, int], / +) -> int: ... + + +_llama_token_get_attr = llama_token_get_attr + + +def llama_token_get_attr( + vocab: llama_vocab_p, token: Union[llama_token, int], / +) -> int: + _warn_deprecated("llama_token_get_attr", "use llama_vocab_get_attr instead") + return _llama_token_get_attr(vocab, token) + + +# DEPRECATED(LLAMA_API bool llama_token_is_eog(const struct llama_vocab * vocab, llama_token token), "use llama_vocab_is_eog instead"); +@ctypes_function( + "llama_token_is_eog", + [llama_vocab_p_ctypes, llama_token], + ctypes.c_bool, +) +def llama_token_is_eog( + vocab: llama_vocab_p, token: Union[llama_token, int], / +) -> bool: ... + + +_llama_token_is_eog = llama_token_is_eog + + +def llama_token_is_eog(vocab: llama_vocab_p, token: Union[llama_token, int], /) -> bool: + _warn_deprecated("llama_token_is_eog", "use llama_vocab_is_eog instead") + return _llama_token_is_eog(vocab, token) + + +# DEPRECATED(LLAMA_API bool llama_token_is_control(const struct llama_vocab * vocab, llama_token token), "use llama_vocab_is_control instead"); +@ctypes_function( + "llama_token_is_control", + [llama_vocab_p_ctypes, llama_token], + ctypes.c_bool, +) +def llama_token_is_control( + vocab: llama_vocab_p, token: Union[llama_token, int], / +) -> bool: ... + + +_llama_token_is_control = llama_token_is_control + + +def llama_token_is_control( + vocab: llama_vocab_p, token: Union[llama_token, int], / +) -> bool: + _warn_deprecated( + "llama_token_is_control", + "use llama_vocab_is_control instead", + ) + return _llama_token_is_control(vocab, token) + + +# DEPRECATED(LLAMA_API llama_token llama_token_bos(const struct llama_vocab * vocab), "use llama_vocab_bos instead"); +@ctypes_function( + "llama_token_bos", + [llama_vocab_p_ctypes], + llama_token, +) +def llama_token_bos(vocab: llama_vocab_p, /) -> int: ... + + +_llama_token_bos = llama_token_bos + + +def llama_token_bos(vocab: llama_vocab_p, /) -> int: + _warn_deprecated("llama_token_bos", "use llama_vocab_bos instead") + return _llama_token_bos(vocab) + + +# DEPRECATED(LLAMA_API llama_token llama_token_eos(const struct llama_vocab * vocab), "use llama_vocab_eos instead"); +@ctypes_function( + "llama_token_eos", + [llama_vocab_p_ctypes], + llama_token, +) +def llama_token_eos(vocab: llama_vocab_p, /) -> int: ... + + +_llama_token_eos = llama_token_eos + + +def llama_token_eos(vocab: llama_vocab_p, /) -> int: + _warn_deprecated("llama_token_eos", "use llama_vocab_eos instead") + return _llama_token_eos(vocab) + + +# DEPRECATED(LLAMA_API llama_token llama_token_eot(const struct llama_vocab * vocab), "use llama_vocab_eot instead"); +@ctypes_function( + "llama_token_eot", + [llama_vocab_p_ctypes], + llama_token, +) +def llama_token_eot(vocab: llama_vocab_p, /) -> int: ... + + +_llama_token_eot = llama_token_eot + + +def llama_token_eot(vocab: llama_vocab_p, /) -> int: + _warn_deprecated("llama_token_eot", "use llama_vocab_eot instead") + return _llama_token_eot(vocab) + + +# DEPRECATED(LLAMA_API llama_token llama_token_cls(const struct llama_vocab * vocab), "use llama_vocab_cls instead"); +@ctypes_function( + "llama_token_cls", + [llama_vocab_p_ctypes], + llama_token, +) +def llama_token_cls(vocab: llama_vocab_p, /) -> int: ... + + +_llama_token_cls = llama_token_cls + + +def llama_token_cls(vocab: llama_vocab_p, /) -> int: + _warn_deprecated("llama_token_cls", "use llama_vocab_cls instead") + return _llama_token_cls(vocab) + + +# DEPRECATED(LLAMA_API llama_token llama_token_sep(const struct llama_vocab * vocab), "use llama_vocab_sep instead"); +@ctypes_function( + "llama_token_sep", + [llama_vocab_p_ctypes], + llama_token, +) +def llama_token_sep(vocab: llama_vocab_p, /) -> int: ... + + +_llama_token_sep = llama_token_sep + + +def llama_token_sep(vocab: llama_vocab_p, /) -> int: + _warn_deprecated("llama_token_sep", "use llama_vocab_sep instead") + return _llama_token_sep(vocab) + + +# DEPRECATED(LLAMA_API llama_token llama_token_nl (const struct llama_vocab * vocab), "use llama_vocab_nl instead"); +@ctypes_function( + "llama_token_nl", + [llama_vocab_p_ctypes], + llama_token, +) +def llama_token_nl(vocab: llama_vocab_p, /) -> int: ... + + +_llama_token_nl = llama_token_nl + + +def llama_token_nl(vocab: llama_vocab_p, /) -> int: + _warn_deprecated("llama_token_nl", "use llama_vocab_nl instead") + return _llama_token_nl(vocab) + + +# DEPRECATED(LLAMA_API llama_token llama_token_pad(const struct llama_vocab * vocab), "use llama_vocab_pad instead"); +@ctypes_function( + "llama_token_pad", + [llama_vocab_p_ctypes], + llama_token, +) +def llama_token_pad(vocab: llama_vocab_p, /) -> int: ... + + +_llama_token_pad = llama_token_pad + + +def llama_token_pad(vocab: llama_vocab_p, /) -> int: + _warn_deprecated("llama_token_pad", "use llama_vocab_pad instead") + return _llama_token_pad(vocab) + + +# DEPRECATED(LLAMA_API bool llama_add_bos_token(const struct llama_vocab * vocab), "use llama_vocab_get_add_bos instead"); +@ctypes_function( + "llama_add_bos_token", + [llama_vocab_p_ctypes], + ctypes.c_bool, +) +def llama_add_bos_token(vocab: llama_vocab_p, /) -> bool: ... + + +_llama_add_bos_token = llama_add_bos_token + + +def llama_add_bos_token(vocab: llama_vocab_p, /) -> bool: + _warn_deprecated("llama_add_bos_token", "use llama_vocab_get_add_bos instead") + return _llama_add_bos_token(vocab) + + +# DEPRECATED(LLAMA_API bool llama_add_eos_token(const struct llama_vocab * vocab), "use llama_vocab_get_add_eos instead"); +@ctypes_function( + "llama_add_eos_token", + [llama_vocab_p_ctypes], + ctypes.c_bool, +) +def llama_add_eos_token(vocab: llama_vocab_p, /) -> bool: ... + + +_llama_add_eos_token = llama_add_eos_token + + +def llama_add_eos_token(vocab: llama_vocab_p, /) -> bool: + _warn_deprecated("llama_add_eos_token", "use llama_vocab_get_add_eos instead") + return _llama_add_eos_token(vocab) + + +# DEPRECATED(LLAMA_API llama_token llama_token_fim_pre(const struct llama_vocab * vocab), "use llama_vocab_fim_pre instead"); +@ctypes_function( + "llama_token_fim_pre", + [llama_vocab_p_ctypes], + llama_token, +) +def llama_token_fim_pre(vocab: llama_vocab_p, /) -> llama_token: ... + + +_llama_token_fim_pre = llama_token_fim_pre + + +def llama_token_fim_pre(vocab: llama_vocab_p, /) -> llama_token: + _warn_deprecated("llama_token_fim_pre", "use llama_vocab_fim_pre instead") + return _llama_token_fim_pre(vocab) + + +# DEPRECATED(LLAMA_API llama_token llama_token_fim_suf(const struct llama_vocab * vocab), "use llama_vocab_fim_suf instead"); +@ctypes_function( + "llama_token_fim_suf", + [llama_vocab_p_ctypes], + llama_token, +) +def llama_token_fim_suf(vocab: llama_vocab_p, /) -> llama_token: ... + + +_llama_token_fim_suf = llama_token_fim_suf + + +def llama_token_fim_suf(vocab: llama_vocab_p, /) -> llama_token: + _warn_deprecated("llama_token_fim_suf", "use llama_vocab_fim_suf instead") + return _llama_token_fim_suf(vocab) + + +# DEPRECATED(LLAMA_API llama_token llama_token_fim_mid(const struct llama_vocab * vocab), "use llama_vocab_fim_mid instead"); +@ctypes_function( + "llama_token_fim_mid", + [llama_vocab_p_ctypes], + llama_token, +) +def llama_token_fim_mid(vocab: llama_vocab_p, /) -> llama_token: ... + + +_llama_token_fim_mid = llama_token_fim_mid + + +def llama_token_fim_mid(vocab: llama_vocab_p, /) -> llama_token: + _warn_deprecated("llama_token_fim_mid", "use llama_vocab_fim_mid instead") + return _llama_token_fim_mid(vocab) + + +# DEPRECATED(LLAMA_API llama_token llama_token_fim_pad(const struct llama_vocab * vocab), "use llama_vocab_fim_pad instead"); +@ctypes_function( + "llama_token_fim_pad", + [llama_vocab_p_ctypes], + llama_token, +) +def llama_token_fim_pad(vocab: llama_vocab_p, /) -> llama_token: ... + + +_llama_token_fim_pad = llama_token_fim_pad + + +def llama_token_fim_pad(vocab: llama_vocab_p, /) -> llama_token: + _warn_deprecated("llama_token_fim_pad", "use llama_vocab_fim_pad instead") + return _llama_token_fim_pad(vocab) + + +# DEPRECATED(LLAMA_API llama_token llama_token_fim_rep(const struct llama_vocab * vocab), "use llama_vocab_fim_rep instead"); +@ctypes_function( + "llama_token_fim_rep", + [llama_vocab_p_ctypes], + llama_token, +) +def llama_token_fim_rep(vocab: llama_vocab_p, /) -> llama_token: ... + + +_llama_token_fim_rep = llama_token_fim_rep + + +def llama_token_fim_rep(vocab: llama_vocab_p, /) -> llama_token: + _warn_deprecated("llama_token_fim_rep", "use llama_vocab_fim_rep instead") + return _llama_token_fim_rep(vocab) + + +# DEPRECATED(LLAMA_API llama_token llama_token_fim_sep(const struct llama_vocab * vocab), "use llama_vocab_fim_sep instead"); +@ctypes_function( + "llama_token_fim_sep", + [llama_vocab_p_ctypes], + llama_token, +) +def llama_token_fim_sep(vocab: llama_vocab_p, /) -> llama_token: ... + + +_llama_token_fim_sep = llama_token_fim_sep + + +def llama_token_fim_sep(vocab: llama_vocab_p, /) -> llama_token: + _warn_deprecated("llama_token_fim_sep", "use llama_vocab_fim_sep instead") + return _llama_token_fim_sep(vocab) + + +# // CLS is equivalent to BOS +# DEPRECATED(LLAMA_API llama_token llama_vocab_cls(const struct llama_vocab * vocab), // classification +# "use llama_vocab_bos instead"); +@ctypes_function( + "llama_vocab_cls", + [llama_vocab_p_ctypes], + llama_token, +) +def llama_vocab_cls(vocab: llama_vocab_p, /) -> llama_token: ... + + +_llama_vocab_cls = llama_vocab_cls + + +def llama_vocab_cls(vocab: llama_vocab_p, /) -> llama_token: + _warn_deprecated("llama_vocab_cls", "use llama_vocab_bos instead") + return _llama_vocab_cls(vocab) + + +# // +# // Tokenization +# // +# // The API is thread-safe. +# // + + +# /// @details Convert the provided text into tokens. +# /// @param tokens The tokens pointer must be large enough to hold the resulting tokens. +# /// @return Returns the number of tokens on success, no more than n_tokens_max +# /// @return Returns a negative number on failure - the number of tokens that would have been returned +# /// @return Returns INT32_MIN on overflow (e.g., tokenization result size exceeds int32_t limit) +# /// @param add_special Allow to add BOS and EOS tokens if model is configured to do so. +# /// @param parse_special Allow tokenizing special and/or control tokens which otherwise are not exposed and treated +# /// as plaintext. Does not insert a leading space. +# LLAMA_API int32_t llama_tokenize( +# const struct llama_vocab * vocab, +# const char * text, +# int32_t text_len, +# llama_token * tokens, +# int32_t n_tokens_max, +# bool add_special, +# bool parse_special); +@ctypes_function( + "llama_tokenize", + [ + llama_vocab_p_ctypes, + ctypes.c_char_p, + ctypes.c_int32, + llama_token_p, + ctypes.c_int32, + ctypes.c_bool, + ctypes.c_bool, + ], + ctypes.c_int32, +) +def llama_tokenize( + vocab: llama_vocab_p, + text: bytes, + text_len: Union[ctypes.c_int, int], + tokens: CtypesArray[llama_token], + n_tokens_max: Union[ctypes.c_int, int], + add_special: Union[ctypes.c_bool, bool], + parse_special: Union[ctypes.c_bool, bool], + /, +) -> int: + """Convert the provided text into tokens. + + Args: + vocab: The vocabulary to use for tokenization. + text: The text to tokenize. + text_len: The length of the text. + tokens: The tokens pointer must be large enough to hold the resulting tokens. + n_max_tokens: The maximum number of tokens to return. + add_special: Allow adding special tokens if the model is configured to do so. + parse_special: Allow parsing special tokens. + + Returns: + Returns the number of tokens on success, no more than n_tokens_max + Returns a negative number on failure - the number of tokens that would have been returned + """ + ... + + +# // Token Id -> Piece. +# // Uses the vocabulary in the provided context. +# // Does not write null terminator to the buffer. +# // User can skip up to 'lstrip' leading spaces before copying (useful when encoding/decoding multiple tokens with 'add_space_prefix') +# // @param special If true, special tokens are rendered in the output. +# LLAMA_API int32_t llama_token_to_piece( +# const struct llama_vocab * vocab, +# llama_token token, +# char * buf, +# int32_t length, +# int32_t lstrip, +# bool special); +@ctypes_function( + "llama_token_to_piece", + [ + llama_vocab_p_ctypes, + llama_token, + ctypes.c_char_p, + ctypes.c_int32, + ctypes.c_int32, + ctypes.c_bool, + ], + ctypes.c_int32, +) +def llama_token_to_piece( + vocab: llama_vocab_p, + token: Union[llama_token, int], + buf: Union[ctypes.c_char_p, bytes, CtypesArray[ctypes.c_char]], + length: Union[ctypes.c_int, int], + lstrip: Union[ctypes.c_int, int], + special: Union[ctypes.c_bool, bool], + /, +) -> int: + """Token Id -> Piece. + Uses the vocabulary in the provided context. + Does not write null terminator to the buffer. + User code is responsible to remove the leading whitespace of the first non-BOS token when decoding multiple tokens. + + Args: + vocab: The vocabulary to use for tokenization. + token: The token to convert. + buf: The buffer to write the token to. + length: The length of the buffer. + lstrip: The number of leading spaces to skip. + special: If true, special tokens are rendered in the output.""" + ... + + +# /// @details Convert the provided tokens into text (inverse of llama_tokenize()). +# /// @param text The char pointer must be large enough to hold the resulting text. +# /// @return Returns the number of chars/bytes on success, no more than text_len_max. +# /// @return Returns a negative number on failure - the number of chars/bytes that would have been returned. +# /// @param remove_special Allow to remove BOS and EOS tokens if model is configured to do so. +# /// @param unparse_special If true, special tokens are rendered in the output. +# LLAMA_API int32_t llama_detokenize( +# const struct llama_vocab * vocab, +# const llama_token * tokens, +# int32_t n_tokens, +# char * text, +# int32_t text_len_max, +# bool remove_special, +# bool unparse_special); +@ctypes_function( + "llama_detokenize", + [ + llama_vocab_p_ctypes, + ctypes.POINTER(llama_token), + ctypes.c_int32, + ctypes.c_char_p, + ctypes.c_int32, + ctypes.c_bool, + ctypes.c_bool, + ], + ctypes.c_int32, +) +def llama_detokenize( + vocab: llama_vocab_p, + tokens: CtypesArray[llama_token], + n_tokens: Union[ctypes.c_int, int], + text: bytes, + text_len_max: Union[ctypes.c_int, int], + remove_special: Union[ctypes.c_bool, bool], + unparse_special: Union[ctypes.c_bool, bool], + /, ) -> int: - return _lib.llama_sample_token(ctx, candidates) + """Convert the provided tokens into text (inverse of llama_tokenize()). + + Args: + vocab: The vocabulary to use for tokenization. + tokens: The tokens to convert. + n_tokens: The number of tokens. + text: The buffer to write the text to. + text_len_max: The length of the buffer. + remove_special: Allow to remove BOS and EOS tokens if model is configured to do so. + unparse_special: If true, special tokens are rendered in the output.""" + ... + + +# // +# // Chat templates +# // + + +# /// Apply chat template. Inspired by hf apply_chat_template() on python. +# /// Both "model" and "custom_template" are optional, but at least one is required. "custom_template" has higher precedence than "model" +# /// NOTE: This function does not use a jinja parser. It only support a pre-defined list of template. See more: https://github.com/ggml-org/llama.cpp/wiki/Templates-supported-by-llama_chat_apply_template +# /// @param tmpl A Jinja template to use for this chat. If this is nullptr, the model's default chat template will be used instead. +# /// @param chat Pointer to a list of multiple llama_chat_message +# /// @param n_msg Number of llama_chat_message in this chat +# /// @param add_ass Whether to end the prompt with the token(s) that indicate the start of an assistant message. +# /// @param buf A buffer to hold the output formatted prompt. The recommended alloc size is 2 * (total number of characters of all messages) +# /// @param length The size of the allocated buffer +# /// @return The total number of bytes of the formatted prompt. If is it larger than the size of buffer, you may need to re-alloc it and then re-apply the template. +# LLAMA_API int32_t llama_chat_apply_template( +# const char * tmpl, +# const struct llama_chat_message * chat, +# size_t n_msg, +# bool add_ass, +# char * buf, +# int32_t length); +@ctypes_function( + "llama_chat_apply_template", + [ + ctypes.c_char_p, # tmpl + ctypes.POINTER(llama_chat_message), # chat + ctypes.c_size_t, # n_msg + ctypes.c_bool, # add_ass (added) + ctypes.c_char_p, # buf + ctypes.c_int32, # length + ], + ctypes.c_int32, +) +def llama_chat_apply_template( + tmpl: bytes, + chat: CtypesArray[llama_chat_message], + n_msg: int, + add_ass: bool, # Added parameter + buf: bytes, + length: int, + /, +) -> int: + """Apply chat template. + + Args: + tmpl: Template to use. If None, uses model's default + chat: Array of chat messages + n_msg: Number of messages + add_ass: Whether to end prompt with assistant token + buf: Output buffer + length: Buffer length + + Returns: + Number of bytes written, or needed if buffer too small + """ + ... + + +# // Get list of built-in chat templates +# LLAMA_API int32_t llama_chat_builtin_templates(const char ** output, size_t len); +@ctypes_function( + "llama_chat_builtin_templates", + [ + ctypes.POINTER(ctypes.c_char_p), + ctypes.c_size_t, + ], + ctypes.c_int32, +) +def llama_chat_builtin_templates( + output: CtypesArray[bytes], + len: Union[ctypes.c_size_t, int], + /, +) -> int: + """Get list of built-in chat templates. + + Args: + output: Output buffer to store template names. + len: Length of the output buffer. + + Returns: + Number of templates available. + Returns a negative number on error. + """ + ... + + +# // +# // Sampling API +# // + +# typedef void * llama_sampler_context_t; +llama_sampler_context_t = ctypes.c_void_p + + +# struct llama_sampler_data { +# struct ggml_tensor * logits; +# struct ggml_tensor * probs; +# struct ggml_tensor * sampled; +# struct ggml_tensor * candidates; +# }; +class llama_sampler_data(ctypes.Structure): + _fields_ = [ + ("logits", ctypes.c_void_p), + ("probs", ctypes.c_void_p), + ("sampled", ctypes.c_void_p), + ("candidates", ctypes.c_void_p), + ] + + +# // user code can implement the interface below in order to create custom llama_sampler +# struct llama_sampler_i { +# const char * (*name) (const struct llama_sampler * smpl); // can be NULL +# void (*accept)( struct llama_sampler * smpl, llama_token token); // can be NULL +# void (*apply) ( struct llama_sampler * smpl, llama_token_data_array * cur_p); // required +# void (*reset) ( struct llama_sampler * smpl); // can be NULL +# struct llama_sampler * (*clone) (const struct llama_sampler * smpl); // can be NULL if ctx is NULL +# void (*free) ( struct llama_sampler * smpl); // can be NULL if ctx is NULL -_lib.llama_sample_token.argtypes = [ - llama_context_p, - llama_token_data_array_p, +# // TODO: API for internal libllama usage for appending the sampling to an existing ggml_cgraph +# //void (*apply_ggml) (struct llama_sampler * smpl, ...); +# }; +class llama_sampler_i(ctypes.Structure): ... + + +# struct llama_sampler { +# const struct llama_sampler_i * iface; +# llama_sampler_context_t ctx; +# }; +class llama_sampler(ctypes.Structure): + _fields_ = [ + ("iface", ctypes.POINTER(llama_sampler_i)), + ("ctx", llama_sampler_context_t), + ] + + +if TYPE_CHECKING: + llama_sampler_p = CtypesPointer[llama_sampler] + +llama_sampler_p_ctypes = ctypes.POINTER(llama_sampler) + +llama_sampler_i_name = ctypes.CFUNCTYPE(ctypes.c_char_p, llama_sampler_p_ctypes) +llama_sampler_i_accept = ctypes.CFUNCTYPE(None, llama_sampler_p_ctypes, llama_token) +llama_sampler_i_apply = ctypes.CFUNCTYPE( + None, llama_sampler_p_ctypes, llama_token_data_array_p +) +llama_sampler_i_reset = ctypes.CFUNCTYPE(None, llama_sampler_p_ctypes) +llama_sampler_i_clone = ctypes.CFUNCTYPE(llama_sampler_p_ctypes, llama_sampler_p_ctypes) +llama_sampler_i_free = ctypes.CFUNCTYPE(None, llama_sampler_p_ctypes) +llama_sampler_i_backend_init = ctypes.CFUNCTYPE( + ctypes.c_bool, llama_sampler_p_ctypes, ctypes.c_void_p +) +llama_sampler_i_backend_accept = ctypes.CFUNCTYPE( + None, + llama_sampler_p_ctypes, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, +) +llama_sampler_i_backend_apply = ctypes.CFUNCTYPE( + None, + llama_sampler_p_ctypes, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.POINTER(llama_sampler_data), +) +llama_sampler_i_backend_set_input = ctypes.CFUNCTYPE(None, llama_sampler_p_ctypes) + +llama_sampler_i._fields_ = [ + ("name", llama_sampler_i_name), + ("accept", llama_sampler_i_accept), + ("apply", llama_sampler_i_apply), + ("reset", llama_sampler_i_reset), + ("clone", llama_sampler_i_clone), + ("free", llama_sampler_i_free), + ("backend_init", llama_sampler_i_backend_init), + ("backend_accept", llama_sampler_i_backend_accept), + ("backend_apply", llama_sampler_i_backend_apply), + ("backend_set_input", llama_sampler_i_backend_set_input), ] -_lib.llama_sample_token.restype = llama_token -# Performance information +# // attach a sampler to the context +# LLAMA_API bool llama_set_sampler(struct llama_context * ctx, llama_seq_id seq_id, struct llama_sampler * smpl); +@ctypes_function( + "llama_set_sampler", + [llama_context_p_ctypes, llama_seq_id, llama_sampler_p_ctypes], + ctypes.c_bool, +) +def llama_set_sampler( + ctx: llama_context_p, seq_id: Union[llama_seq_id, int], smpl: llama_sampler_p, / +) -> bool: + """Attach a sampler to the context.""" + ... + + +# // mirror of llama_sampler_i: +# LLAMA_API struct llama_sampler * llama_sampler_init (const struct llama_sampler_i * iface, llama_sampler_context_t ctx); +@ctypes_function( + "llama_sampler_init", + [ctypes.POINTER(llama_sampler_i), llama_sampler_context_t], + llama_sampler_p_ctypes, +) +def llama_sampler_init( + iface: ctypes.POINTER(llama_sampler_i), ctx: llama_sampler_context_t, / +) -> llama_sampler_p: ... + + +# LLAMA_API const char * llama_sampler_name (const struct llama_sampler * smpl); +@ctypes_function( + "llama_sampler_name", + [llama_sampler_p_ctypes], + ctypes.c_char_p, +) +def llama_sampler_name(smpl: llama_sampler_p, /) -> bytes: ... + + +# LLAMA_API void llama_sampler_accept( struct llama_sampler * smpl, llama_token token); +@ctypes_function( + "llama_sampler_accept", + [llama_sampler_p_ctypes, llama_token], + None, +) +def llama_sampler_accept(smpl: llama_sampler_p, token: Union[llama_token, int], /): ... + + +# LLAMA_API void llama_sampler_apply ( struct llama_sampler * smpl, llama_token_data_array * cur_p); +@ctypes_function( + "llama_sampler_apply", + [llama_sampler_p_ctypes, llama_token_data_array_p], + None, +) +def llama_sampler_apply( + smpl: llama_sampler_p, cur_p: CtypesArray[llama_token_data_array], / +): ... + + +# LLAMA_API void llama_sampler_reset ( struct llama_sampler * smpl); +@ctypes_function( + "llama_sampler_reset", + [llama_sampler_p_ctypes], + None, +) +def llama_sampler_reset(smpl: llama_sampler_p, /): ... + + +# LLAMA_API struct llama_sampler * llama_sampler_clone (const struct llama_sampler * smpl); +@ctypes_function( + "llama_sampler_clone", + [llama_sampler_p_ctypes], + llama_sampler_p_ctypes, +) +def llama_sampler_clone(smpl: llama_sampler_p, /) -> llama_sampler_p: ... + + +# // important: do not free if the sampler has been added to a llama_sampler_chain (via llama_sampler_chain_add) +# LLAMA_API void llama_sampler_free ( struct llama_sampler * smpl); +@ctypes_function( + "llama_sampler_free", + [llama_sampler_p_ctypes], + None, +) +def llama_sampler_free(smpl: llama_sampler_p, /): ... + + +# // llama_sampler_chain +# // a type of llama_sampler that can chain multiple samplers one after another + + +# LLAMA_API struct llama_sampler * llama_sampler_chain_init(struct llama_sampler_chain_params params); +@ctypes_function( + "llama_sampler_chain_init", + [llama_sampler_chain_params], + llama_sampler_p_ctypes, +) +def llama_sampler_chain_init( + params: llama_sampler_chain_params, / +) -> llama_sampler_p: ... + + +# // important: takes ownership of the sampler object and will free it when llama_sampler_free is called +# LLAMA_API void llama_sampler_chain_add( struct llama_sampler * chain, struct llama_sampler * smpl); +@ctypes_function( + "llama_sampler_chain_add", + [llama_sampler_p_ctypes, llama_sampler_p_ctypes], + None, +) +def llama_sampler_chain_add(chain: llama_sampler_p, smpl: llama_sampler_p, /): ... + + +# LLAMA_API struct llama_sampler * llama_sampler_chain_get(const struct llama_sampler * chain, int32_t i); +@ctypes_function( + "llama_sampler_chain_get", + [llama_sampler_p_ctypes, ctypes.c_int32], + llama_sampler_p_ctypes, +) +def llama_sampler_chain_get( + chain: llama_sampler_p, i: Union[ctypes.c_int32, int], / +) -> llama_sampler_p: ... + + +# LLAMA_API int llama_sampler_chain_n (const struct llama_sampler * chain); +@ctypes_function( + "llama_sampler_chain_n", + [llama_sampler_p_ctypes], + ctypes.c_int, +) +def llama_sampler_chain_n(chain: llama_sampler_p, /) -> int: ... + + +# // after removing a sampler, the chain will no longer own it, and it will not be freed when the chain is freed +# LLAMA_API struct llama_sampler * llama_sampler_chain_remove( struct llama_sampler * chain, int32_t i); +@ctypes_function( + "llama_sampler_chain_remove", + [llama_sampler_p_ctypes, ctypes.c_int32], + llama_sampler_p_ctypes, +) +def llama_sampler_chain_remove( + chain: llama_sampler_p, i: Union[ctypes.c_int32, int], / +) -> llama_sampler_p: ... + + +# // available samplers: + + +# LLAMA_API struct llama_sampler * llama_sampler_init_greedy(void); +@ctypes_function("llama_sampler_init_greedy", [], llama_sampler_p_ctypes) +def llama_sampler_init_greedy() -> llama_sampler_p: ... + + +# LLAMA_API struct llama_sampler * llama_sampler_init_dist (uint32_t seed); +@ctypes_function("llama_sampler_init_dist", [ctypes.c_uint32], llama_sampler_p_ctypes) +def llama_sampler_init_dist(seed: int) -> llama_sampler_p: ... + + +# /// @details Top-K sampling described in academic paper "The Curious Case of Neural Text Degeneration" https://arxiv.org/abs/1904.09751 +# /// Setting k <= 0 makes this a noop +# LLAMA_API struct llama_sampler * llama_sampler_init_top_k (int32_t k); +@ctypes_function("llama_sampler_init_top_k", [ctypes.c_int32], llama_sampler_p_ctypes) +def llama_sampler_init_top_k(k: int) -> llama_sampler_p: ... + + +# /// @details Nucleus sampling described in academic paper "The Curious Case of Neural Text Degeneration" https://arxiv.org/abs/1904.09751 +# LLAMA_API struct llama_sampler * llama_sampler_init_top_p (float p, size_t min_keep); +@ctypes_function( + "llama_sampler_init_top_p", + [ctypes.c_float, ctypes.c_size_t], + llama_sampler_p_ctypes, +) +def llama_sampler_init_top_p(p: float, min_keep: int) -> llama_sampler_p: ... + + +# /// @details Minimum P sampling as described in https://github.com/ggml-org/llama.cpp/pull/3841 +# LLAMA_API struct llama_sampler * llama_sampler_init_min_p (float p, size_t min_keep); +@ctypes_function( + "llama_sampler_init_min_p", + [ctypes.c_float, ctypes.c_size_t], + llama_sampler_p_ctypes, +) +def llama_sampler_init_min_p(p: float, min_keep: int) -> llama_sampler_p: ... + + +# /// @details Locally Typical Sampling implementation described in the paper https://arxiv.org/abs/2202.00666. +# LLAMA_API struct llama_sampler * llama_sampler_init_typical (float p, size_t min_keep); +@ctypes_function( + "llama_sampler_init_typical", + [ctypes.c_float, ctypes.c_size_t], + llama_sampler_p_ctypes, +) +def llama_sampler_init_typical(p: float, min_keep: int) -> llama_sampler_p: ... + + +# /// #details Updates the logits l_i` = l_i/t. When t <= 0.0f, the maximum logit is kept at it's original value, the rest are set to -inf +# LLAMA_API struct llama_sampler * llama_sampler_init_temp (float t); +@ctypes_function("llama_sampler_init_temp", [ctypes.c_float], llama_sampler_p_ctypes) +def llama_sampler_init_temp(t: float) -> llama_sampler_p: ... + + +# /// @details Dynamic temperature implementation (a.k.a. entropy) described in the paper https://arxiv.org/abs/2309.02772. +# LLAMA_API struct llama_sampler * llama_sampler_init_temp_ext (float t, float delta, float exponent); +@ctypes_function( + "llama_sampler_init_temp_ext", + [ctypes.c_float, ctypes.c_float, ctypes.c_float], + llama_sampler_p_ctypes, +) +def llama_sampler_init_temp_ext( + t: float, delta: float, exponent: float +) -> llama_sampler_p: ... + + +# /// @details XTC sampler as described in https://github.com/oobabooga/text-generation-webui/pull/6335 +# LLAMA_API struct llama_sampler * llama_sampler_init_xtc (float p, float t, size_t min_keep, uint32_t seed); +@ctypes_function( + "llama_sampler_init_xtc", + [ctypes.c_float, ctypes.c_float, ctypes.c_size_t, ctypes.c_uint32], + llama_sampler_p_ctypes, +) +def llama_sampler_init_xtc( + p: float, t: float, min_keep: int, seed: int, / +) -> llama_sampler_p: ... + + +# /// @details Top n sigma sampling as described in academic paper "Top-nσ: Not All Logits Are You Need" https://arxiv.org/pdf/2411.07641 +# LLAMA_API struct llama_sampler * llama_sampler_init_top_n_sigma(float n); +@ctypes_function( + "llama_sampler_init_top_n_sigma", + [ctypes.c_float], + llama_sampler_p_ctypes, +) +def llama_sampler_init_top_n_sigma(n: float, /) -> llama_sampler_p: ... + + +# /// @details Mirostat 1.0 algorithm described in the paper https://arxiv.org/abs/2007.14966. Uses tokens instead of words. +# LLAMA_API struct llama_sampler * llama_sampler_init_mirostat( +# int32_t n_vocab, +# uint32_t seed, +# float tau, +# float eta, +# int32_t m); +@ctypes_function( + "llama_sampler_init_mirostat", + [ctypes.c_int32, ctypes.c_uint32, ctypes.c_float, ctypes.c_float, ctypes.c_int32], + llama_sampler_p_ctypes, +) +def llama_sampler_init_mirostat( + n_vocab: int, seed: int, tau: float, eta: float, m: int, / +) -> llama_sampler_p: ... + + +# /// @details Mirostat 2.0 algorithm described in the paper https://arxiv.org/abs/2007.14966. Uses tokens instead of words. +# LLAMA_API struct llama_sampler * llama_sampler_init_mirostat_v2( +# uint32_t seed, +# float tau, +# float eta); +@ctypes_function( + "llama_sampler_init_mirostat_v2", + [ctypes.c_uint32, ctypes.c_float, ctypes.c_float], + llama_sampler_p_ctypes, +) +def llama_sampler_init_mirostat_v2( + seed: int, tau: float, eta: float, / +) -> llama_sampler_p: ... + + +# /// @details Intializes a GBNF grammar, see grammars/README.md for details. +# LLAMA_API struct llama_sampler * llama_sampler_init_grammar( +# const struct llama_vocab * vocab, +# const char * grammar_str, +# const char * grammar_root); +@ctypes_function( + "llama_sampler_init_grammar", + [llama_vocab_p_ctypes, ctypes.c_char_p, ctypes.c_char_p], + llama_sampler_p_ctypes, +) +def llama_sampler_init_grammar( + vocab: llama_vocab_p, grammar_str: bytes, grammar_root: bytes, / +) -> llama_sampler_p: ... + + +# DEPRECATED(LLAMA_API struct llama_sampler * llama_sampler_init_grammar_lazy( +# const struct llama_vocab * vocab, +# const char * grammar_str, +# const char * grammar_root, +# const char ** trigger_words, +# size_t num_trigger_words, +# const llama_token * trigger_tokens, +# size_t num_trigger_tokens), +# "use llama_sampler_init_grammar_lazy_patterns instead"); +@ctypes_function( + "llama_sampler_init_grammar_lazy", + [ + llama_vocab_p_ctypes, + ctypes.c_char_p, + ctypes.c_char_p, + ctypes.POINTER(ctypes.c_char_p), + ctypes.c_size_t, + ctypes.POINTER(llama_token), + ctypes.c_size_t, + ], + llama_sampler_p_ctypes, +) +def llama_sampler_init_grammar_lazy( + vocab: llama_vocab_p, + grammar_str: bytes, + grammar_root: bytes, + trigger_words: CtypesArray[bytes], + num_trigger_words: int, + trigger_tokens: CtypesArray[llama_token], + num_trigger_tokens: int, + /, +) -> llama_sampler_p: ... + + +_llama_sampler_init_grammar_lazy = llama_sampler_init_grammar_lazy + + +def llama_sampler_init_grammar_lazy( + vocab: llama_vocab_p, + grammar_str: bytes, + grammar_root: bytes, + trigger_words: CtypesArray[bytes], + num_trigger_words: int, + trigger_tokens: CtypesArray[llama_token], + num_trigger_tokens: int, + /, +) -> llama_sampler_p: + _warn_deprecated( + "llama_sampler_init_grammar_lazy", + "use llama_sampler_init_grammar_lazy_patterns instead", + ) + return _llama_sampler_init_grammar_lazy( + vocab, + grammar_str, + grammar_root, + trigger_words, + num_trigger_words, + trigger_tokens, + num_trigger_tokens, + ) + + +# /// @details Lazy grammar sampler, introduced in https://github.com/ggml-org/llama.cpp/pull/9639 +# LLAMA_API struct llama_sampler * llama_sampler_init_grammar_lazy_patterns( +# const struct llama_vocab * vocab, +# const char * grammar_str, +# const char * grammar_root, +# const char ** trigger_patterns, +# size_t num_trigger_patterns, +# const llama_token * trigger_tokens, +# size_t num_trigger_tokens); +@ctypes_function( + "llama_sampler_init_grammar_lazy_patterns", + [ + llama_vocab_p_ctypes, + ctypes.c_char_p, + ctypes.c_char_p, + ctypes.POINTER(ctypes.c_char_p), + ctypes.c_size_t, + ctypes.POINTER(llama_token), + ctypes.c_size_t, + ], + llama_sampler_p_ctypes, +) +def llama_sampler_init_grammar_lazy_patterns( + vocab: llama_vocab_p, + grammar_str: bytes, + grammar_root: bytes, + trigger_patterns: CtypesArray[bytes], + num_trigger_patterns: int, + trigger_tokens: CtypesArray[llama_token], + num_trigger_tokens: int, + /, +) -> llama_sampler_p: ... + + +# /// NOTE: Avoid using on the full vocabulary as searching for repeated tokens can become slow. For example, apply top-k or top-p sampling first. +# LLAMA_API struct llama_sampler * llama_sampler_init_penalties( +# int32_t penalty_last_n, // last n tokens to penalize (0 = disable penalty, -1 = context size) +# float penalty_repeat, // 1.0 = disabled +# float penalty_freq, // 0.0 = disabled +# float penalty_present); // 0.0 = disabled +@ctypes_function( + "llama_sampler_init_penalties", + [ctypes.c_int32, ctypes.c_float, ctypes.c_float, ctypes.c_float], + llama_sampler_p_ctypes, +) +def llama_sampler_init_penalties( + penalty_last_n: int, + penalty_repeat: float, + penalty_freq: float, + penalty_present: float, + /, +) -> llama_sampler_p: ... + + +# /// @details DRY sampler, designed by p-e-w, as described in: https://github.com/oobabooga/text-generation-webui/pull/5677, porting Koboldcpp implementation authored by pi6am: https://github.com/LostRuins/koboldcpp/pull/982 +# LLAMA_API struct llama_sampler * llama_sampler_init_dry( +# const struct llama_vocab * vocab, +# int32_t n_ctx_train, +# float dry_multiplier, +# float dry_base, +# int32_t dry_allowed_length, +# int32_t dry_penalty_last_n, +# const char ** seq_breakers, +# size_t num_breakers); +@ctypes_function( + "llama_sampler_init_dry", + [ + llama_vocab_p_ctypes, + ctypes.c_int32, + ctypes.c_float, + ctypes.c_float, + ctypes.c_int32, + ctypes.c_int32, + ctypes.POINTER(ctypes.c_char_p), + ctypes.c_size_t, + ], + llama_sampler_p_ctypes, +) +def llama_sampler_init_dry( + vocab: llama_vocab_p, + n_ctx_train: int, + dry_multiplier: float, + dry_base: float, + dry_allowed_length: int, + dry_penalty_last_n: int, + seq_breakers, + num_breakers: int, + /, +) -> llama_sampler_p: ... + + +# LLAMA_API struct llama_sampler * llama_sampler_init_adaptive_p( +# float target, +# float decay, +# uint32_t seed); +@ctypes_function( + "llama_sampler_init_adaptive_p", + [ctypes.c_float, ctypes.c_float, ctypes.c_uint32], + llama_sampler_p_ctypes, +) +def llama_sampler_init_adaptive_p( + target: float, decay: float, seed: int, / +) -> llama_sampler_p: + """Initialize an adaptive-p sampler.""" + ... + + +# LLAMA_API struct llama_sampler * llama_sampler_init_logit_bias( +# int32_t n_vocab, +# int32_t n_logit_bias, +# const llama_logit_bias * logit_bias); +@ctypes_function( + "llama_sampler_init_logit_bias", + [ctypes.c_int32, ctypes.c_int32, llama_logit_bias_p], + llama_sampler_p_ctypes, +) +def llama_sampler_init_logit_bias( + n_vocab: int, n_logit_bias: int, logit_bias: CtypesArray[llama_logit_bias], / +) -> llama_sampler_p: ... -# LLAMA_API struct llama_timings llama_get_timings(struct llama_context * ctx); -def llama_get_timings(ctx: llama_context_p) -> llama_timings: - return _lib.llama_get_timings(ctx) +# // this sampler is meant to be used for fill-in-the-middle infilling +# LLAMA_API struct llama_sampler * llama_sampler_init_infill(const struct llama_vocab * vocab); +@ctypes_function( + "llama_sampler_init_infill", + [llama_vocab_p_ctypes], + llama_sampler_p_ctypes, +) +def llama_sampler_init_infill(vocab: llama_vocab_p, /) -> llama_sampler_p: ... -_lib.llama_get_timings.argtypes = [llama_context_p] -_lib.llama_get_timings.restype = llama_timings +# // Returns the seed used by the sampler if applicable, LLAMA_DEFAULT_SEED otherwise +# LLAMA_API uint32_t llama_sampler_get_seed(const struct llama_sampler * smpl); +@ctypes_function( + "llama_sampler_get_seed", + [llama_sampler_p_ctypes], + ctypes.c_uint32, +) +def llama_sampler_get_seed(smpl: llama_sampler_p, /) -> int: ... -# LLAMA_API void llama_print_timings(struct llama_context * ctx); -def llama_print_timings(ctx: llama_context_p): - _lib.llama_print_timings(ctx) +# /// @details Sample and accept a token from the idx-th output of the last evaluation +# LLAMA_API llama_token llama_sampler_sample(struct llama_sampler * smpl, struct llama_context * ctx, int32_t idx); +@ctypes_function( + "llama_sampler_sample", + [llama_sampler_p_ctypes, llama_context_p_ctypes, ctypes.c_int32], + llama_token, +) +def llama_sampler_sample( + smpl: llama_sampler_p, ctx: llama_context_p, idx: int, / +) -> int: ... -_lib.llama_print_timings.argtypes = [llama_context_p] -_lib.llama_print_timings.restype = None +# // +# // Model split +# // -# LLAMA_API void llama_reset_timings(struct llama_context * ctx); -def llama_reset_timings(ctx: llama_context_p): - _lib.llama_reset_timings(ctx) +# /// @details Build a split GGUF final path for this chunk. +# LLAMA_API int llama_split_path(char * split_path, size_t maxlen, const char * path_prefix, int split_no, int split_count); +@ctypes_function( + "llama_split_path", + [ctypes.c_char_p, ctypes.c_size_t, ctypes.c_char_p, ctypes.c_int, ctypes.c_int], + ctypes.c_int, +) +def llama_split_path( + split_path: bytes, + maxlen: Union[ctypes.c_size_t, int], + path_prefix: bytes, + split_no: Union[ctypes.c_int, int], + split_count: Union[ctypes.c_int, int], + /, +) -> int: + """Build a split GGUF final path for this chunk.""" + ... -_lib.llama_reset_timings.argtypes = [llama_context_p] -_lib.llama_reset_timings.restype = None +# /// @details Extract the path prefix from the split_path if and only if the split_no and split_count match. +# LLAMA_API int llama_split_prefix(char * split_prefix, size_t maxlen, const char * split_path, int split_no, int split_count); +@ctypes_function( + "llama_split_prefix", + [ctypes.c_char_p, ctypes.c_size_t, ctypes.c_char_p, ctypes.c_int, ctypes.c_int], + ctypes.c_int, +) +def llama_split_prefix( + split_prefix: bytes, + maxlen: Union[ctypes.c_size_t, int], + split_path: bytes, + split_no: Union[ctypes.c_int, int], + split_count: Union[ctypes.c_int, int], + /, +) -> int: + """Extract the path prefix from the split_path if and only if the split_no and split_count match.""" + ... -# Print system information +# // Print system information # LLAMA_API const char * llama_print_system_info(void); -def llama_print_system_info() -> bytes: - return _lib.llama_print_system_info() +@ctypes_function("llama_print_system_info", [], ctypes.c_char_p) +def llama_print_system_info() -> bytes: ... + + +# // Set callback for all future logging events. +# // If this is not called, or NULL is supplied, everything is output on stderr. +# // The logger state is global so these functions are NOT thread safe. +# LLAMA_API void llama_log_get(ggml_log_callback * log_callback, void ** user_data); +@ctypes_function( + "llama_log_get", + [ctypes.POINTER(llama_log_callback), ctypes.POINTER(ctypes.c_void_p)], + None, +) +def llama_log_get( + log_callback: CtypesPointerOrRef[llama_log_callback], + user_data: CtypesPointerOrRef[ctypes.c_void_p], + /, +): + """Get the current logging callback and user data.""" + ... + + +# LLAMA_API void llama_log_set(ggml_log_callback log_callback, void * user_data); +@ctypes_function( + "llama_log_set", + [llama_log_callback, ctypes.c_void_p], + None, +) +def llama_log_set( + log_callback: Optional[CtypesFuncPointer], + user_data: ctypes.c_void_p, + /, +): + """Set callback for all future logging events. + + If this is not called, or NULL is supplied, everything is output on stderr.""" + ... + + +# // +# // Performance utils +# // + +# struct llama_perf_context_data { +# double t_start_ms; +# double t_load_ms; +# double t_p_eval_ms; +# double t_eval_ms; + + +# int32_t n_p_eval; +# int32_t n_eval; +# int32_t n_reused; // number of times a ggml compute graph had been reused +# }; +class llama_perf_context_data(ctypes.Structure): + _fields_ = [ + ("t_start_ms", ctypes.c_double), + ("t_load_ms", ctypes.c_double), + ("t_p_eval_ms", ctypes.c_double), + ("t_eval_ms", ctypes.c_double), + ("n_p_eval", ctypes.c_int32), + ("n_eval", ctypes.c_int32), + ("n_reused", ctypes.c_int32), + ] + + +# struct llama_perf_sampler_data { +# double t_sample_ms; + + +# int32_t n_sample; +# }; +class llama_perf_sampler_data(ctypes.Structure): + _fields_ = [ + ("t_sample_ms", ctypes.c_double), + ("n_sample", ctypes.c_int32), + ] + + +# LLAMA_API struct llama_perf_context_data llama_perf_context (const struct llama_context * ctx); +@ctypes_function( + "llama_perf_context", + [llama_context_p_ctypes], + llama_perf_context_data, +) +def llama_perf_context(ctx: llama_context_p, /) -> llama_perf_context_data: ... + + +# LLAMA_API void llama_perf_context_print(const struct llama_context * ctx); +@ctypes_function( + "llama_perf_context_print", + [llama_context_p_ctypes], + None, +) +def llama_perf_context_print(ctx: llama_context_p, /): ... + + +# LLAMA_API void llama_perf_context_reset( struct llama_context * ctx); +@ctypes_function( + "llama_perf_context_reset", + [llama_context_p_ctypes], + None, +) +def llama_perf_context_reset(ctx: llama_context_p, /): ... + + +# // NOTE: the following work only with samplers constructed via llama_sampler_chain_init +# LLAMA_API struct llama_perf_sampler_data llama_perf_sampler (const struct llama_sampler * chain); +@ctypes_function( + "llama_perf_sampler", + [llama_sampler_p_ctypes], + llama_perf_sampler_data, +) +def llama_perf_sampler(chain: llama_sampler_p, /) -> llama_perf_sampler_data: ... + + +# LLAMA_API void llama_perf_sampler_print(const struct llama_sampler * chain); +@ctypes_function( + "llama_perf_sampler_print", + [llama_sampler_p_ctypes], + None, +) +def llama_perf_sampler_print(chain: llama_sampler_p, /): ... -_lib.llama_print_system_info.argtypes = [] -_lib.llama_print_system_info.restype = c_char_p +# LLAMA_API void llama_perf_sampler_reset( struct llama_sampler * chain); +@ctypes_function( + "llama_perf_sampler_reset", + [llama_sampler_p_ctypes], + None, +) +def llama_perf_sampler_reset(chain: llama_sampler_p, /): ... + + +# // +# // training +# // + +# // function that returns whether or not a given tensor contains trainable parameters +# typedef bool (*llama_opt_param_filter)(const struct ggml_tensor * tensor, void * userdata); +llama_opt_param_filter = ctypes.CFUNCTYPE( + ctypes.c_bool, ctypes.c_void_p, ctypes.c_void_p +) + + +# // always returns true +# LLAMA_API bool llama_opt_param_filter_all(const struct ggml_tensor * tensor, void * userdata); +@ctypes_function( + "llama_opt_param_filter_all", + [ctypes.c_void_p, ctypes.c_void_p], + ctypes.c_bool, +) +def llama_opt_param_filter_all( + tensor: ctypes.c_void_p, userdata: ctypes.c_void_p, / +) -> bool: ... -################################################################################################### +# struct llama_opt_params { +# uint32_t n_ctx_train; // assumed context size post training, use context size specified in llama_context if 0 -_llama_initialized = False +# llama_opt_param_filter param_filter; // callback for determining which tensors contain trainable parameters +# void * param_filter_ud; // userdata for determining which tensors contain trainable parameters -if not _llama_initialized: - llama_init_backend(c_bool(False)) - _llama_initialized = True + +# ggml_opt_get_optimizer_params get_opt_pars; // callback for calculating optimizer parameters +# void * get_opt_pars_ud; // userdata for calculating optimizer parameters +# }; +class llama_opt_params(ctypes.Structure): + _fields_ = [ + ("n_ctx_train", ctypes.c_uint32), + ("param_filter", llama_opt_param_filter), + ("param_filter_ud", ctypes.c_void_p), + ( + "get_opt_pars", + ctypes.c_void_p, + ), # ggml_opt_get_optimizer_params - not implemented here + ("get_opt_pars_ud", ctypes.c_void_p), + ] + + +# LLAMA_API void llama_opt_init(struct llama_context * lctx, struct llama_model * model, struct llama_opt_params lopt_params); +@ctypes_function( + "llama_opt_init", + [llama_context_p_ctypes, llama_model_p_ctypes, llama_opt_params], + None, +) +def llama_opt_init( + lctx: llama_context_p, model: llama_model_p, lopt_params: llama_opt_params, / +): ... + + +# LLAMA_API void llama_opt_epoch( +# struct llama_context * lctx, +# ggml_opt_dataset_t dataset, +# ggml_opt_result_t result_train, +# ggml_opt_result_t result_eval, +# int64_t idata_split, +# ggml_opt_epoch_callback callback_train, +# ggml_opt_epoch_callback callback_eval); +@ctypes_function( + "llama_opt_epoch", + [ + llama_context_p_ctypes, + ctypes.c_void_p, # ggml_opt_dataset_t + ctypes.c_void_p, # ggml_opt_result_t + ctypes.c_void_p, # ggml_opt_result_t + ctypes.c_int64, + ctypes.c_void_p, # ggml_opt_epoch_callback + ctypes.c_void_p, # ggml_opt_epoch_callback + ], + None, +) +def llama_opt_epoch( + lctx: llama_context_p, + dataset: ctypes.c_void_p, + result_train: ctypes.c_void_p, + result_eval: ctypes.c_void_p, + idata_split: int, + callback_train: ctypes.c_void_p, + callback_eval: ctypes.c_void_p, + /, +): ... diff --git a/llama_cpp/llama_cpp_ext.py b/llama_cpp/llama_cpp_ext.py new file mode 100644 index 0000000000..284811086a --- /dev/null +++ b/llama_cpp/llama_cpp_ext.py @@ -0,0 +1,117 @@ +"""Experimental bindings for non-public llama.cpp APIs from `llama-ext.h`. + +This module is not part of the stable llama-cpp-python public API. +Downstream code should not import or depend on it directly. +""" + +from __future__ import annotations + +import ctypes +import functools + +from typing import Any, Iterable, Union + +from . import llama_cpp + +_lib = llama_cpp._lib + + +def _ctypes_function_from_names( + names: Iterable[str], + argtypes: list[Any], + restype: Any, +): + """Decorator for extension functions whose exported symbol name can vary by ABI.""" + + def decorator(f): + missing: list[str] = [] + for name in names: + try: + func = getattr(_lib, name) + except AttributeError: + missing.append(name) + continue + func.argtypes = argtypes + func.restype = restype + functools.wraps(f)(func) + return func + raise AttributeError( + f"None of the shared library symbols were found: {', '.join(missing)}" + ) + + return decorator + + +# LLAMA_API void llama_set_embeddings_nextn(struct llama_context * ctx, bool value, bool masked); +@_ctypes_function_from_names( + ( + "llama_set_embeddings_nextn", + "_Z26llama_set_embeddings_nextnP13llama_contextbb", + "?llama_set_embeddings_nextn@@YAXPEAUllama_context@@_N1@Z", + ), + [llama_cpp.llama_context_p_ctypes, ctypes.c_bool, ctypes.c_bool], + None, +) +def llama_set_embeddings_nextn( + ctx: llama_cpp.llama_context_p, + value: bool, + masked: bool, + /, +): + """Set whether the context outputs nextn embeddings or not.""" + ... + + +# LLAMA_API float * llama_get_embeddings_nextn(struct llama_context * ctx); +@_ctypes_function_from_names( + ( + "llama_get_embeddings_nextn", + "_Z26llama_get_embeddings_nextnP13llama_context", + "?llama_get_embeddings_nextn@@YAPEAMPEAUllama_context@@@Z", + ), + [llama_cpp.llama_context_p_ctypes], + ctypes.POINTER(ctypes.c_float), +) +def llama_get_embeddings_nextn( + ctx: llama_cpp.llama_context_p, + /, +): + """Get the nextn embeddings from the last evaluation.""" + ... + + +# LLAMA_API float * llama_get_embeddings_nextn_ith(struct llama_context * ctx, int32_t i); +@_ctypes_function_from_names( + ( + "llama_get_embeddings_nextn_ith", + "_Z30llama_get_embeddings_nextn_ithP13llama_contexti", + "?llama_get_embeddings_nextn_ith@@YAPEAMPEAUllama_context@@H@Z", + ), + [llama_cpp.llama_context_p_ctypes, ctypes.c_int32], + ctypes.POINTER(ctypes.c_float), +) +def llama_get_embeddings_nextn_ith( + ctx: llama_cpp.llama_context_p, + i: Union[ctypes.c_int32, int], + /, +): + """Get the nextn embeddings for the ith output row from the last evaluation.""" + ... + + +# LLAMA_API llama_context * llama_get_ctx_other(struct llama_context * ctx); +@_ctypes_function_from_names( + ( + "llama_get_ctx_other", + "_Z19llama_get_ctx_otherP13llama_context", + "?llama_get_ctx_other@@YAPEAUllama_context@@PEAU1@@Z", + ), + [llama_cpp.llama_context_p_ctypes], + llama_cpp.llama_context_p_ctypes, +) +def llama_get_ctx_other( + ctx: llama_cpp.llama_context_p, + /, +): + """Get the context linked through llama_context_params.ctx_other.""" + ... diff --git a/llama_cpp/llama_grammar.py b/llama_cpp/llama_grammar.py new file mode 100644 index 0000000000..ba34dda831 --- /dev/null +++ b/llama_cpp/llama_grammar.py @@ -0,0 +1,953 @@ +"""Python implementation of llama grammar parser directly translated from C++ source file in vendor/llama.cpp/common/grammar-parser.cpp.""" + +# flake8: noqa +from pathlib import Path + +from itertools import groupby +from typing import ( + Any, + Set, + List, + Optional, + Tuple, + Union, +) + +LLAMA_GRAMMAR_DEFAULT_ROOT = "root" + + +class LlamaGrammar: + def __init__(self, *args, _grammar: str, **kwargs): + self._grammar = _grammar + self._root = LLAMA_GRAMMAR_DEFAULT_ROOT + + @classmethod + def from_string(cls, grammar: str, verbose: bool = True) -> "LlamaGrammar": + return cls(_grammar=grammar) + + @classmethod + def from_file(cls, file: Union[str, Path], verbose: bool = True) -> "LlamaGrammar": + try: + with open(file) as f: + grammar = f.read() + except Exception as err: + raise Exception( + f"{cls.from_file.__name__}: error reading grammar file: {err}" + ) + + if grammar: + return cls.from_string(grammar, verbose=verbose) + + raise ValueError( + f"{cls.from_file.__name__}: error parsing grammar file: params_grammer is empty" + ) + + @classmethod + def from_json_schema(cls, json_schema: str, verbose: bool = True) -> "LlamaGrammar": + return cls.from_string(json_schema_to_gbnf(json_schema), verbose=verbose) + + +"""llama.cpp gbnf rules from vendor/llama.cpp/grammars""" + +ARITHMETIC_GBNF = r""" +root ::= (expr "=" ws term "\n")+ +expr ::= term ([-+*/] term)* +term ::= ident | num | "(" ws expr ")" ws +ident ::= [a-z] [a-z0-9_]* ws +num ::= [0-9]+ ws +ws ::= [ \t\n]* +""" + +C_GBNF = r""" +root ::= (declaration)* + +declaration ::= dataType identifier "(" parameter? ")" "{" statement* "}" + +dataType ::= "int" ws | "float" ws | "char" ws +identifier ::= [a-zA-Z_] [a-zA-Z_0-9]* + +parameter ::= dataType identifier + +statement ::= + ( dataType identifier ws "=" ws expression ";" ) | + ( identifier ws "=" ws expression ";" ) | + ( identifier ws "(" argList? ")" ";" ) | + ( "return" ws expression ";" ) | + ( "while" "(" condition ")" "{" statement* "}" ) | + ( "for" "(" forInit ";" ws condition ";" ws forUpdate ")" "{" statement* "}" ) | + ( "if" "(" condition ")" "{" statement* "}" ("else" "{" statement* "}")? ) | + ( singleLineComment ) | + ( multiLineComment ) + +forInit ::= dataType identifier ws "=" ws expression | identifier ws "=" ws expression +forUpdate ::= identifier ws "=" ws expression + +condition ::= expression relationOperator expression +relationOperator ::= ("<=" | "<" | "==" | "!=" | ">=" | ">") + +expression ::= term (("+" | "-") term)* +term ::= factor(("*" | "/") factor)* + +factor ::= identifier | number | unaryTerm | funcCall | parenExpression +unaryTerm ::= "-" factor +funcCall ::= identifier "(" argList? ")" +parenExpression ::= "(" ws expression ws ")" + +argList ::= expression ("," ws expression)* + +number ::= [0-9]+ + +singleLineComment ::= "//" [^\n]* "\n" +multiLineComment ::= "/*" ( [^*] | ("*" [^/]) )* "*/" + +ws ::= ([ \t\n]+) +""" + +CHESS_GBNF = r""" +root ::= object +value ::= object | array | string | number | ("true" | "false" | "null") ws + +object ::= + "{" ws ( + string ":" ws value + ("," ws string ":" ws value)* + )? "}" ws + +array ::= + "[" ws ( + value + ("," ws value)* + )? "]" ws + +string ::= + "\"" ( + [^"\\] | + "\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F]) # escapes + )* "\"" ws + +number ::= ("-"? ([0-9] | [1-9] [0-9]*)) ("." [0-9]+)? ([eE] [-+]? [0-9]+)? ws + +# Optional space: by convention, applied in this grammar after literal chars when allowed +ws ::= ([ \t\n] ws)? +""" + +JAPANESE_GBNF = r""" +root ::= object +value ::= object | array | string | number | ("true" | "false" | "null") ws + +object ::= + "{" ws ( + string ":" ws value + ("," ws string ":" ws value)* + )? "}" ws + +array ::= + "[" ws ( + value + ("," ws value)* + )? "]" ws + +string ::= + "\"" ( + [^"\\] | + "\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F]) # escapes + )* "\"" ws + +number ::= ("-"? ([0-9] | [1-9] [0-9]*)) ("." [0-9]+)? ([eE] [-+]? [0-9]+)? ws + +# Optional space: by convention, applied in this grammar after literal chars when allowed +ws ::= ([ \t\n] ws)? +""" + +JSON_ARR_GBNF = r""" +# This is the same as json.gbnf but we restrict whitespaces at the end of the root array +# Useful for generating JSON arrays + +root ::= arr +value ::= object | array | string | number | ("true" | "false" | "null") ws + +arr ::= + "[\n" ws ( + value + (",\n" ws value)* + )? "]" + +object ::= + "{" ws ( + string ":" ws value + ("," ws string ":" ws value)* + )? "}" ws + +array ::= + "[" ws ( + value + ("," ws value)* + )? "]" ws + +string ::= + "\"" ( + [^"\\\x7F\x00-\x1F] | + "\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F]) # escapes + )* "\"" ws + +number ::= ("-"? ([0-9] | [1-9] [0-9]*)) ("." [0-9]+)? ([eE] [-+]? [0-9]+)? ws + +# Optional space: by convention, applied in this grammar after literal chars when allowed +ws ::= ([ \t\n] ws)? +""" + + +JSON_GBNF = r""" +root ::= object +value ::= object | array | string | number | ("true" | "false" | "null") ws + +object ::= + "{" ws ( + string ":" ws value + ("," ws string ":" ws value)* + )? "}" ws + +array ::= + "[" ws ( + value + ("," ws value)* + )? "]" ws + +string ::= + "\"" ( + [^"\\\x7F\x00-\x1F] | + "\\" (["\\bfnrt] | "u" [0-9a-fA-F]{4}) # escapes + )* "\"" ws + +number ::= ("-"? ([0-9] | [1-9] [0-9]{0,15})) ("." [0-9]+)? ([eE] [-+]? [0-9] [1-9]{0,15})? ws + +# Optional space: by convention, applied in this grammar after literal chars when allowed +ws ::= | " " | "\n" [ \t]{0,20} +""" + +LIST_GBNF = r""" +root ::= item+ + +# Excludes various line break characters +item ::= "- " [^\r\n\x0b\x0c\x85\u2028\u2029]+ "\n" +""" + +"""llama.cpp json-schema to grammar converter from vendor/llama.cpp/examples/json-schema-to-grammar.py""" +import json +import re +from typing import List, Optional + +# whitespace is constrained to a single space char to prevent model "running away" in +# whitespace. Also maybe improves generation quality? +SPACE_RULE = '" "?' + + +INVALID_RULE_CHARS_RE = re.compile(r"[^a-zA-Z0-9-]+") +GRAMMAR_LITERAL_ESCAPE_RE = re.compile(r'[\r\n"]') +GRAMMAR_LITERAL_ESCAPES = {"\r": "\\r", "\n": "\\n", '"': '\\"'} + +# whitespace is constrained to a single space char to prevent model "running away" in +# whitespace. Also maybe improves generation quality? +SPACE_RULE = '" "?' + + +def _build_repetition( + item_rule, min_items, max_items, separator_rule=None, item_rule_is_literal=False +): + if not separator_rule: + if min_items == 0 and max_items == 1: + return f"{item_rule}?" + elif min_items == 1 and max_items is None: + return f"{item_rule}+" + + result = "" + + if min_items > 0: + if item_rule_is_literal and separator_rule is None: + result = '"' + (item_rule[1:-1] * min_items) + '"' + else: + result = (f" {separator_rule} " if separator_rule else " ").join( + [item_rule] * min_items + ) + + def opt_repetitions(up_to_n, prefix_with_sep=False): + """ + - n=4, no sep: '(a (a (a (a)?)?)?)?' + - n=4, sep=',', prefix: '("," a ("," a ("," a ("," a)?)?)?)?' + - n=4, sep=',', no prefix: '(a ("," a ("," a ("," a)?)?)?)?' + """ + + content = ( + f"{separator_rule} {item_rule}" + if prefix_with_sep and separator_rule + else item_rule + ) + if up_to_n == 0: + return "" + elif up_to_n == 1: + return f"({content})?" + elif separator_rule and not prefix_with_sep: + return f"({content} {opt_repetitions(up_to_n - 1, prefix_with_sep=True)})?" + else: + return (f"({content} " * up_to_n).rstrip() + (")?" * up_to_n) + + if min_items > 0 and max_items != min_items: + result += " " + + if max_items is not None: + result += opt_repetitions(max_items - min_items, prefix_with_sep=min_items > 0) + else: + item_operator = f"({separator_rule + ' ' if separator_rule else ''}{item_rule})" + + if min_items == 0 and separator_rule: + result = f"({item_rule} {item_operator}*)?" + else: + result += f"{item_operator}*" + + return result + + +class BuiltinRule: + def __init__(self, content: str, deps: list = None): + self.content = content + self.deps = deps or [] + + +_up_to_15_digits = _build_repetition("[0-9]", 0, 15) + +PRIMITIVE_RULES = { + "boolean": BuiltinRule('("true" | "false") space', []), + "decimal-part": BuiltinRule("[0-9] " + _up_to_15_digits, []), + "integral-part": BuiltinRule("[0-9] | [1-9] " + _up_to_15_digits, []), + "number": BuiltinRule( + '("-"? integral-part) ("." decimal-part)? ([eE] [-+]? integral-part)? space', + ["integral-part", "decimal-part"], + ), + "integer": BuiltinRule('("-"? integral-part) space', ["integral-part"]), + "value": BuiltinRule( + "object | array | string | number | boolean | null", + ["object", "array", "string", "number", "boolean", "null"], + ), + "object": BuiltinRule( + '"{" space ( string ":" space value ("," space string ":" space value)* )? "}" space', + ["string", "value"], + ), + "array": BuiltinRule( + '"[" space ( value ("," space value)* )? "]" space', ["value"] + ), + "uuid": BuiltinRule( + r'"\"" ' + + ' "-" '.join("[0-9a-fA-F]" * n for n in [8, 4, 4, 4, 12]) + + r' "\"" space', + [], + ), + "char": BuiltinRule( + r'[^"\\] | "\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F])', + [], + ), + "string": BuiltinRule(r'"\"" char* "\"" space', ["char"]), + "null": BuiltinRule('"null" space', []), +} + +# TODO: support "uri", "email" string formats +STRING_FORMAT_RULES = { + "date": BuiltinRule( + '[0-9] [0-9] [0-9] [0-9] "-" ( "0" [1-9] | "1" [0-2] ) "-" ( "0" [1-9] | [1-2] [0-9] | "3" [0-1] )', + [], + ), + "time": BuiltinRule( + '([01] [0-9] | "2" [0-3]) ":" [0-5] [0-9] ":" [0-5] [0-9] ( "." [0-9] [0-9] [0-9] )? ( "Z" | ( "+" | "-" ) ( [01] [0-9] | "2" [0-3] ) ":" [0-5] [0-9] )', + [], + ), + "date-time": BuiltinRule('date "T" time', ["date", "time"]), + "date-string": BuiltinRule('"\\"" date "\\"" space', ["date"]), + "time-string": BuiltinRule('"\\"" time "\\"" space', ["time"]), + "date-time-string": BuiltinRule('"\\"" date-time "\\"" space', ["date-time"]), +} + +DOTALL = "[\\U00000000-\\U0010FFFF]" +DOT = "[^\\x0A\\x0D]" + +RESERVED_NAMES = set( + ["root", "dot", *PRIMITIVE_RULES.keys(), *STRING_FORMAT_RULES.keys()] +) + + +NON_LITERAL_SET = set("|.()[]{}*+?") +ESCAPED_IN_REGEXPS_BUT_NOT_IN_LITERALS = set("[]()|{}*+?") + + +class SchemaConverter: + def __init__(self, *, prop_order, allow_fetch, dotall, raw_pattern): + self._prop_order = prop_order + self._allow_fetch = allow_fetch + self._dotall = dotall + self._raw_pattern = raw_pattern + self._rules = { + "space": SPACE_RULE, + } + self._refs = {} + self._refs_being_resolved = set() + + def _format_literal(self, literal): + escaped = GRAMMAR_LITERAL_ESCAPE_RE.sub( + lambda m: GRAMMAR_LITERAL_ESCAPES.get(m.group(0)), literal + ) + return f'"{escaped}"' + + def not_literal( + self, literal: str, dotall: bool = True, maybe_escaped_underscores=False + ) -> str: + """ + not_literal('a') -> '[^a]' + not_literal('abc') -> '([^a] | "a" ([^b] | "b" ([^c])?)?)?' + """ + assert len(literal) > 0, "Empty literal not supported" + + def recurse(i: int): + c = literal[i] + if maybe_escaped_underscores and c == "_": + yield f"[^{c}\\\\]" + yield " | " + yield f'"\\\\"? "{c}"' + else: + yield f"[^{c}]" + if i < len(literal) - 1: + yield " | " + yield self._format_literal(c) + yield " (" + yield from recurse(i + 1) + yield ")?" + + return "".join(("(", *recurse(0), ")")) + + def _add_rule(self, name, rule): + esc_name = INVALID_RULE_CHARS_RE.sub("-", name) + if esc_name not in self._rules or self._rules[esc_name] == rule: + key = esc_name + else: + i = 0 + while ( + f"{esc_name}{i}" in self._rules + and self._rules[f"{esc_name}{i}"] != rule + ): + i += 1 + key = f"{esc_name}{i}" + self._rules[key] = rule + return key + + def resolve_refs(self, schema: dict, url: str): + """ + Resolves all $ref fields in the given schema, fetching any remote schemas, + replacing $ref with absolute reference URL and populating self._refs with the + respective referenced (sub)schema dictionaries. + """ + + def visit(n: dict): + if isinstance(n, list): + return [visit(x) for x in n] + elif isinstance(n, dict): + ref = n.get("$ref") + if ref is not None and ref not in self._refs: + if ref.startswith("https://"): + assert self._allow_fetch, ( + "Fetching remote schemas is not allowed (use --allow-fetch for force)" + ) + import requests + + frag_split = ref.split("#") + base_url = frag_split[0] + + target = self._refs.get(base_url) + if target is None: + target = self.resolve_refs( + requests.get(ref).json(), base_url + ) + self._refs[base_url] = target + + if len(frag_split) == 1 or frag_split[-1] == "": + return target + elif ref.startswith("#/"): + target = schema + ref = f"{url}{ref}" + n["$ref"] = ref + else: + raise ValueError(f"Unsupported ref {ref}") + + for sel in ref.split("#")[-1].split("/")[1:]: + assert target is not None and sel in target, ( + f"Error resolving ref {ref}: {sel} not in {target}" + ) + target = target[sel] + + self._refs[ref] = target + else: + for v in n.values(): + visit(v) + + return n + + return visit(schema) + + def _generate_union_rule(self, name, alt_schemas): + return " | ".join( + ( + self.visit(alt_schema, f"{name}{'-' if name else 'alternative-'}{i}") + for i, alt_schema in enumerate(alt_schemas) + ) + ) + + def _visit_pattern(self, pattern, name): + """ + Transforms a regular expression pattern into a GBNF rule. + + Input: https://json-schema.org/understanding-json-schema/reference/regular_expressions + Output: https://github.com/ggerganov/llama.cpp/blob/master/grammars/README.md + + Unsupported features: negative/positive lookaheads, greedy/non-greedy modifiers. + + Mostly a 1:1 translation, except for {x} / {x,} / {x,y} quantifiers for which + we define sub-rules to keep the output lean. + """ + + assert pattern.startswith("^") and pattern.endswith("$"), ( + 'Pattern must start with "^" and end with "$"' + ) + pattern = pattern[1:-1] + sub_rule_ids = {} + + i = 0 + length = len(pattern) + + def to_rule(s: Tuple[str, bool]) -> str: + (txt, is_literal) = s + return '"' + txt + '"' if is_literal else txt + + def transform() -> Tuple[str, bool]: + """ + Parse a unit at index i (advancing it), and return its string representation + whether it's a literal. + """ + nonlocal i + nonlocal pattern + nonlocal sub_rule_ids + + start = i + # For each component of this sequence, store its string representation and whether it's a literal. + # We only need a flat structure here to apply repetition operators to the last item, and + # to merge literals at the and (we're parsing grouped ( sequences ) recursively and don't treat '|' specially + # (GBNF's syntax is luckily very close to regular expressions!) + seq: list[Tuple[str, bool]] = [] + + def get_dot(): + if self._dotall: + rule = DOTALL + else: + # Accept any character... except \n and \r line break chars (\x0A and \xOD) + rule = DOT + return self._add_rule(f"dot", rule) + + def join_seq(): + nonlocal seq + ret = [] + for is_literal, g in groupby(seq, lambda x: x[1]): + if is_literal: + ret.append(("".join(x[0] for x in g), True)) + else: + ret.extend(g) + if len(ret) == 1: + return ret[0] + return (" ".join(to_rule(x) for x in seq), False) + + while i < length: + c = pattern[i] + if c == ".": + seq.append((get_dot(), False)) + i += 1 + elif c == "(": + i += 1 + if i < length: + assert pattern[i] != "?", ( + f'Unsupported pattern syntax "{pattern[i]}" at index {i} of /{pattern}/' + ) + seq.append((f"({to_rule(transform())})", False)) + elif c == ")": + i += 1 + assert start > 0 and pattern[start - 1] == "(", ( + f"Unbalanced parentheses; start = {start}, i = {i}, pattern = {pattern}" + ) + return join_seq() + elif c == "[": + square_brackets = c + i += 1 + while i < length and pattern[i] != "]": + if pattern[i] == "\\": + square_brackets += pattern[i : i + 2] + i += 2 + else: + square_brackets += pattern[i] + i += 1 + assert i < length, ( + f"Unbalanced square brackets; start = {start}, i = {i}, pattern = {pattern}" + ) + square_brackets += "]" + i += 1 + seq.append((square_brackets, False)) + elif c == "|": + seq.append(("|", False)) + i += 1 + elif c in ("*", "+", "?"): + seq[-1] = (to_rule(seq[-1]) + c, False) + i += 1 + elif c == "{": + curly_brackets = c + i += 1 + while i < length and pattern[i] != "}": + curly_brackets += pattern[i] + i += 1 + assert i < length, ( + f"Unbalanced curly brackets; start = {start}, i = {i}, pattern = {pattern}" + ) + curly_brackets += "}" + i += 1 + nums = [s.strip() for s in curly_brackets[1:-1].split(",")] + min_times = 0 + max_times = None + try: + if len(nums) == 1: + min_times = int(nums[0]) + max_times = min_times + else: + assert len(nums) == 2 + min_times = int(nums[0]) if nums[0] else 0 + max_times = int(nums[1]) if nums[1] else None + except ValueError: + raise ValueError( + f"Invalid quantifier {curly_brackets} in /{pattern}/" + ) + + (sub, sub_is_literal) = seq[-1] + + if not sub_is_literal: + id = sub_rule_ids.get(sub) + if id is None: + id = self._add_rule(f"{name}-{len(sub_rule_ids) + 1}", sub) + sub_rule_ids[sub] = id + sub = id + + seq[-1] = ( + _build_repetition( + f'"{sub}"' if sub_is_literal else sub, + min_times, + max_times, + item_rule_is_literal=sub_is_literal, + ), + False, + ) + else: + literal = "" + while i < length: + if pattern[i] == "\\" and i < length - 1: + next = pattern[i + 1] + if next in ESCAPED_IN_REGEXPS_BUT_NOT_IN_LITERALS: + i += 1 + literal += pattern[i] + i += 1 + else: + literal += pattern[i : i + 2] + i += 2 + elif pattern[i] == '"' and not self._raw_pattern: + literal += '\\"' + i += 1 + elif pattern[i] not in NON_LITERAL_SET and ( + i == length - 1 + or literal == "" + or pattern[i + 1] == "." + or pattern[i + 1] not in NON_LITERAL_SET + ): + literal += pattern[i] + i += 1 + else: + break + if literal: + seq.append((literal, True)) + + return join_seq() + + return self._add_rule( + name, + ( + to_rule(transform()) + if self._raw_pattern + else '"\\"" ' + to_rule(transform()) + ' "\\"" space' + ), + ) + + def _resolve_ref(self, ref): + ref_name = ref.split("/")[-1] + if ref_name not in self._rules and ref not in self._refs_being_resolved: + self._refs_being_resolved.add(ref) + resolved = self._refs[ref] + ref_name = self.visit(resolved, ref_name) + self._refs_being_resolved.remove(ref) + return ref_name + + def _generate_constant_rule(self, value): + return self._format_literal(json.dumps(value)) + + def visit(self, schema, name): + schema_type = schema.get("type") + schema_format = schema.get("format") + rule_name = name + "-" if name in RESERVED_NAMES else name or "root" + + if (ref := schema.get("$ref")) is not None: + return self._add_rule(rule_name, self._resolve_ref(ref)) + + elif "oneOf" in schema or "anyOf" in schema: + return self._add_rule( + rule_name, + self._generate_union_rule(name, schema.get("oneOf") or schema["anyOf"]), + ) + + elif isinstance(schema_type, list): + return self._add_rule( + rule_name, + self._generate_union_rule(name, [{"type": t} for t in schema_type]), + ) + + elif "const" in schema: + return self._add_rule( + rule_name, self._generate_constant_rule(schema["const"]) + ) + + elif "enum" in schema: + rule = " | ".join((self._generate_constant_rule(v) for v in schema["enum"])) + return self._add_rule(rule_name, rule) + + elif schema_type in (None, "object") and ( + "properties" in schema + or ( + "additionalProperties" in schema + and schema["additionalProperties"] is not True + ) + ): + required = set(schema.get("required", [])) + properties = list(schema.get("properties", {}).items()) + return self._add_rule( + rule_name, + self._build_object_rule( + properties, required, name, schema.get("additionalProperties") + ), + ) + + elif schema_type in (None, "object") and "allOf" in schema: + required = set() + properties = [] + hybrid_name = name + + def add_component(comp_schema, is_required): + if (ref := comp_schema.get("$ref")) is not None: + comp_schema = self._refs[ref] + + if "properties" in comp_schema: + for prop_name, prop_schema in comp_schema["properties"].items(): + properties.append((prop_name, prop_schema)) + if is_required: + required.add(prop_name) + + for t in schema["allOf"]: + if "anyOf" in t: + for tt in t["anyOf"]: + add_component(tt, is_required=False) + else: + add_component(t, is_required=True) + + return self._add_rule( + rule_name, + self._build_object_rule( + properties, required, hybrid_name, additional_properties=[] + ), + ) + + elif schema_type in (None, "array") and ( + "items" in schema or "prefixItems" in schema + ): + items = schema.get("items") or schema["prefixItems"] + if isinstance(items, list): + return self._add_rule( + rule_name, + '"[" space ' + + ' "," space '.join( + self.visit(item, f"{name}{'-' if name else ''}tuple-{i}") + for i, item in enumerate(items) + ) + + ' "]" space', + ) + else: + item_rule_name = self.visit(items, f"{name}{'-' if name else ''}item") + min_items = schema.get("minItems", 0) + max_items = schema.get("maxItems") + return self._add_rule( + rule_name, + '"[" space ' + + _build_repetition( + item_rule_name, min_items, max_items, separator_rule='"," space' + ) + + ' "]" space', + ) + + elif schema_type in (None, "string") and "pattern" in schema: + return self._visit_pattern(schema["pattern"], rule_name) + + elif schema_type in (None, "string") and re.match( + r"^uuid[1-5]?$", schema_format or "" + ): + return self._add_primitive( + "root" if rule_name == "root" else schema_format, + PRIMITIVE_RULES["uuid"], + ) + + elif ( + schema_type in (None, "string") + and f"{schema_format}-string" in STRING_FORMAT_RULES + ): + prim_name = f"{schema_format}-string" + return self._add_rule( + rule_name, + self._add_primitive(prim_name, STRING_FORMAT_RULES[prim_name]), + ) + + elif schema_type == "string" and ( + "minLength" in schema or "maxLength" in schema + ): + char_rule = self._add_primitive("char", PRIMITIVE_RULES["char"]) + min_len = schema.get("minLength", 0) + max_len = schema.get("maxLength") + + return self._add_rule( + rule_name, + r'"\"" ' + + _build_repetition(char_rule, min_len, max_len) + + r' "\"" space', + ) + + elif (schema_type == "object") or (len(schema) == 0): + return self._add_rule( + rule_name, self._add_primitive("object", PRIMITIVE_RULES["object"]) + ) + + else: + assert schema_type in PRIMITIVE_RULES, f"Unrecognized schema: {schema}" + # TODO: support minimum, maximum, exclusiveMinimum, exclusiveMaximum at least for zero + return self._add_primitive( + "root" if rule_name == "root" else schema_type, + PRIMITIVE_RULES[schema_type], + ) + + def _add_primitive(self, name: str, rule: BuiltinRule): + n = self._add_rule(name, rule.content) + + for dep in rule.deps: + dep_rule = PRIMITIVE_RULES.get(dep) or STRING_FORMAT_RULES.get(dep) + assert dep_rule, f"Rule {dep} not known" + if dep not in self._rules: + self._add_primitive(dep, dep_rule) + return n + + def _build_object_rule( + self, + properties: List[Tuple[str, Any]], + required: Set[str], + name: str, + additional_properties: Union[bool, Any], + ): + prop_order = self._prop_order + # sort by position in prop_order (if specified) then by original order + sorted_props = [ + kv[0] + for _, kv in sorted( + enumerate(properties), + key=lambda ikv: (prop_order.get(ikv[1][0], len(prop_order)), ikv[0]), + ) + ] + + prop_kv_rule_names = {} + for prop_name, prop_schema in properties: + prop_rule_name = self.visit( + prop_schema, f"{name}{'-' if name else ''}{prop_name}" + ) + prop_kv_rule_names[prop_name] = self._add_rule( + f"{name}{'-' if name else ''}{prop_name}-kv", + rf'{self._format_literal(json.dumps(prop_name))} space ":" space {prop_rule_name}', + ) + required_props = [k for k in sorted_props if k in required] + optional_props = [k for k in sorted_props if k not in required] + + if additional_properties == True or isinstance(additional_properties, dict): + sub_name = f"{name}{'-' if name else ''}additional" + value_rule = self.visit( + {} if additional_properties == True else additional_properties, + f"{sub_name}-value", + ) + prop_kv_rule_names["*"] = self._add_rule( + f"{sub_name}-kv", + self._add_primitive("string", PRIMITIVE_RULES["string"]) + + f' ":" space {value_rule}', + ) + optional_props.append("*") + + rule = '"{" space ' + rule += ' "," space '.join(prop_kv_rule_names[k] for k in required_props) + + if optional_props: + rule += " (" + if required_props: + rule += ' "," space ( ' + + def get_recursive_refs(ks, first_is_optional): + [k, *rest] = ks + kv_rule_name = prop_kv_rule_names[k] + if k == "*": + res = self._add_rule( + f"{name}{'-' if name else ''}additional-kvs", + f'{kv_rule_name} ( "," space ' + kv_rule_name + " )*", + ) + elif first_is_optional: + res = f'( "," space {kv_rule_name} )?' + else: + res = kv_rule_name + if len(rest) > 0: + res += " " + self._add_rule( + f"{name}{'-' if name else ''}{k}-rest", + get_recursive_refs(rest, first_is_optional=True), + ) + return res + + rule += " | ".join( + get_recursive_refs(optional_props[i:], first_is_optional=False) + for i in range(len(optional_props)) + ) + if required_props: + rule += " )" + rule += " )?" + + rule += ' "}" space' + + return rule + + def format_grammar(self): + return "\n".join( + f"{name} ::= {rule}" + for name, rule in sorted(self._rules.items(), key=lambda kv: kv[0]) + ) + + +def json_schema_to_gbnf(schema: str, prop_order: Optional[List[str]] = None): + prop_order = prop_order or [] + schema = json.loads(schema) + prop_order = {name: idx for idx, name in enumerate(prop_order)} + converter = SchemaConverter( + prop_order=prop_order, allow_fetch=False, dotall=False, raw_pattern=False + ) + schema = converter.resolve_refs(schema, "stdin") + converter.visit(schema, "") + return converter.format_grammar() diff --git a/llama_cpp/llama_speculative.py b/llama_cpp/llama_speculative.py new file mode 100644 index 0000000000..39dfb903ba --- /dev/null +++ b/llama_cpp/llama_speculative.py @@ -0,0 +1,64 @@ +import abc + +from typing import Any + +import numpy as np +import numpy.typing as npt + + +class LlamaDraftModel(abc.ABC): + @abc.abstractmethod + def __call__( + self, input_ids: npt.NDArray[np.intc], /, **kwargs: Any + ) -> npt.NDArray[np.intc]: + raise NotImplementedError() + + +class LlamaPromptLookupDecoding(LlamaDraftModel): + """Based on https://github.com/apoorvumang/prompt-lookup-decoding""" + + def __init__(self, max_ngram_size: int = 2, num_pred_tokens: int = 10): + self.max_ngram_size = max_ngram_size + self.num_pred_tokens = num_pred_tokens + + @staticmethod + def find_candidate_pred_tokens( + input_ids: npt.NDArray[np.intc], + max_ngram_size: int, + num_pred_tokens: int, + ): + input_length = input_ids.shape[0] + + for ngram_size in range(min(max_ngram_size, input_length - 1), 0, -1): + # Create sliding windows of size ngram_size + windows = np.lib.stride_tricks.sliding_window_view(input_ids, (ngram_size,)) + + # Convert ngram to an array for comparison + ngram_array = input_ids[-ngram_size:] + + # Find where the windows match the ngram + matches = np.all(windows == ngram_array, axis=1) + + # Get the indices of matches + match_indices = np.nonzero(matches)[0] + + # Iterate through match indices to find a valid continuation + for idx in match_indices: + start_idx = idx + ngram_size + end_idx = start_idx + num_pred_tokens + end_idx = min(end_idx, input_length) + + if start_idx < end_idx: + return input_ids[start_idx:end_idx] + + # If no match is found, return an empty array + return np.array([], dtype=np.intc) + + def __call__( + self, input_ids: npt.NDArray[np.intc], /, **kwargs: Any + ) -> npt.NDArray[np.intc]: + return self.find_candidate_pred_tokens( + input_ids=input_ids, + max_ngram_size=self.max_ngram_size, + num_pred_tokens=self.num_pred_tokens, + ) diff --git a/llama_cpp/llama_tokenizer.py b/llama_cpp/llama_tokenizer.py new file mode 100644 index 0000000000..1375e1392d --- /dev/null +++ b/llama_cpp/llama_tokenizer.py @@ -0,0 +1,120 @@ +from __future__ import annotations + +import abc +from typing import ( + List, + Optional, + Any, +) + +import llama_cpp +from llama_cpp.llama_types import List + + +class BaseLlamaTokenizer(abc.ABC): + @abc.abstractmethod + def tokenize( + self, text: bytes, add_bos: bool = True, special: bool = True + ) -> List[int]: + """Tokenize the text into tokens. + + Args: + text: The utf-8 encoded string to tokenize. + add_bos: Whether to add a beginning of sequence token. + special: Whether to tokenize special tokens. + """ + raise NotImplementedError + + @abc.abstractmethod + def detokenize( + self, + tokens: List[int], + prev_tokens: Optional[List[int]] = None, + special: bool = False, + ) -> bytes: + """Detokenize the tokens into text. + + Args: + tokens: The list of tokens to detokenize. + prev_tokens: The list of previous tokens. Offset mapping will be performed if provided. + special: Whether to detokenize special tokens. + """ + raise NotImplementedError + + +class LlamaTokenizer(BaseLlamaTokenizer): + def __init__(self, llama: llama_cpp.Llama): + self._model = llama._model # type: ignore + + def tokenize( + self, text: bytes, add_bos: bool = True, special: bool = True + ) -> List[int]: + return self._model.tokenize(text, add_bos=add_bos, special=special) + + def detokenize( + self, + tokens: List[int], + prev_tokens: Optional[List[int]] = None, + special: bool = False, + ) -> bytes: + return self._model.detokenize(tokens, special=special) + + def encode( + self, text: str, add_bos: bool = True, special: bool = True + ) -> List[int]: + return self.tokenize( + text.encode("utf-8", errors="ignore"), add_bos=add_bos, special=special + ) + + def decode(self, tokens: List[int]) -> str: + return self.detokenize(tokens).decode("utf-8", errors="ignore") + + @classmethod + def from_ggml_file(cls, path: str) -> "LlamaTokenizer": + return cls(llama_cpp.Llama(model_path=path, vocab_only=True)) + + +class LlamaHFTokenizer(BaseLlamaTokenizer): + def __init__(self, hf_tokenizer: Any): + self.hf_tokenizer = hf_tokenizer + + def tokenize( + self, text: bytes, add_bos: bool = True, special: bool = True + ) -> List[int]: + return self.hf_tokenizer.encode( + text.decode("utf-8", errors="ignore"), add_special_tokens=special + ) + + def detokenize( + self, + tokens: List[int], + prev_tokens: Optional[List[int]] = None, + special: bool = False, + ) -> bytes: + skip_special_tokens = not special + if prev_tokens is not None: + text = self.hf_tokenizer.decode( + prev_tokens + tokens, skip_special_tokens=skip_special_tokens + ).encode("utf-8", errors="ignore") + prev_text = self.hf_tokenizer.decode( + prev_tokens, skip_special_tokens=skip_special_tokens + ).encode("utf-8", errors="ignore") + return text[len(prev_text) :] + else: + return self.hf_tokenizer.decode( + tokens, skip_special_tokens=skip_special_tokens + ).encode("utf-8", errors="ignore") + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path: str) -> "LlamaHFTokenizer": + try: + from transformers import AutoTokenizer + except ImportError: + raise ImportError( + "The `transformers` library is required to use the `HFTokenizer`." + "You can install it with `pip install transformers`." + ) + hf_tokenizer = AutoTokenizer.from_pretrained( + pretrained_model_name_or_path=pretrained_model_name_or_path + ) + return cls(hf_tokenizer) diff --git a/llama_cpp/llama_types.py b/llama_cpp/llama_types.py index 6ba8023bd8..f647822ff5 100644 --- a/llama_cpp/llama_types.py +++ b/llama_cpp/llama_types.py @@ -1,22 +1,37 @@ +"""Types and request signatures for OpenAI compatibility + +NOTE: These types may change to match the OpenAI OpenAPI specification. + +Based on the OpenAI OpenAPI specification: +https://github.com/openai/openai-openapi/blob/master/openapi.yaml + +""" + from typing import Any, List, Optional, Dict, Union from typing_extensions import TypedDict, NotRequired, Literal +# NOTE: Defining this correctly using annotations seems to break pydantic validation. +# This is a workaround until we can figure out how to do this correctly +# JsonType = Union[None, int, str, bool, List["JsonType"], Dict[str, "JsonType"]] +JsonType = Union[None, int, str, bool, List[Any], Dict[str, Any]] + + class EmbeddingUsage(TypedDict): prompt_tokens: int total_tokens: int -class EmbeddingData(TypedDict): +class Embedding(TypedDict): index: int object: str - embedding: List[float] + embedding: Union[List[float], List[List[float]]] -class Embedding(TypedDict): +class CreateEmbeddingResponse(TypedDict): object: Literal["list"] model: str - data: List[EmbeddingData] + data: List[Embedding] usage: EmbeddingUsage @@ -31,7 +46,7 @@ class CompletionChoice(TypedDict): text: str index: int logprobs: Optional[CompletionLogprobs] - finish_reason: Optional[str] + finish_reason: Optional[Literal["stop", "length"]] class CompletionUsage(TypedDict): @@ -40,60 +55,262 @@ class CompletionUsage(TypedDict): total_tokens: int -class CompletionChunk(TypedDict): +class CreateCompletionResponse(TypedDict): id: str object: Literal["text_completion"] created: int model: str choices: List[CompletionChoice] + usage: NotRequired[CompletionUsage] -class Completion(TypedDict): - id: str - object: Literal["text_completion"] - created: int - model: str - choices: List[CompletionChoice] - usage: CompletionUsage +class ChatCompletionResponseFunctionCall(TypedDict): + name: str + arguments: str + + +class ChatCompletionResponseMessage(TypedDict): + content: Optional[str] + tool_calls: NotRequired["ChatCompletionMessageToolCalls"] + role: Literal["assistant", "function"] # NOTE: "function" may be incorrect here + function_call: NotRequired[ChatCompletionResponseFunctionCall] # DEPRECATED + + +class ChatCompletionFunction(TypedDict): + name: str + description: NotRequired[str] + parameters: Dict[str, JsonType] # TODO: make this more specific + + +class ChatCompletionTopLogprobToken(TypedDict): + token: str + logprob: float + bytes: Optional[List[int]] + + +class ChatCompletionLogprobToken(ChatCompletionTopLogprobToken): + token: str + logprob: float + bytes: Optional[List[int]] + top_logprobs: List[ChatCompletionTopLogprobToken] -class ChatCompletionMessage(TypedDict): - role: Literal["assistant", "user", "system"] - content: str - user: NotRequired[str] +class ChatCompletionLogprobs(TypedDict): + content: Optional[List[ChatCompletionLogprobToken]] + refusal: Optional[List[ChatCompletionLogprobToken]] -class ChatCompletionChoice(TypedDict): +class ChatCompletionResponseChoice(TypedDict): index: int - message: ChatCompletionMessage + message: "ChatCompletionResponseMessage" + logprobs: Optional[ChatCompletionLogprobs] finish_reason: Optional[str] -class ChatCompletion(TypedDict): +class CreateChatCompletionResponse(TypedDict): id: str object: Literal["chat.completion"] created: int model: str - choices: List[ChatCompletionChoice] + choices: List["ChatCompletionResponseChoice"] usage: CompletionUsage -class ChatCompletionChunkDeltaEmpty(TypedDict): + +class ChatCompletionMessageToolCallChunkFunction(TypedDict): + name: Optional[str] + arguments: str + + +class ChatCompletionMessageToolCallChunk(TypedDict): + index: int + id: NotRequired[str] + type: Literal["function"] + function: ChatCompletionMessageToolCallChunkFunction + + +class ChatCompletionStreamResponseDeltaEmpty(TypedDict): pass -class ChatCompletionChunkDelta(TypedDict): - role: NotRequired[Literal["assistant"]] - content: NotRequired[str] +class ChatCompletionStreamResponseDeltaFunctionCall(TypedDict): + name: str + arguments: str + + +class ChatCompletionStreamResponseDelta(TypedDict): + content: NotRequired[Optional[str]] + function_call: NotRequired[ + Optional[ChatCompletionStreamResponseDeltaFunctionCall] + ] # DEPRECATED + tool_calls: NotRequired[Optional[List[ChatCompletionMessageToolCallChunk]]] + role: NotRequired[Optional[Literal["system", "user", "assistant", "tool"]]] -class ChatCompletionChunkChoice(TypedDict): + +class ChatCompletionStreamResponseChoice(TypedDict): index: int - delta: Union[ChatCompletionChunkDelta, ChatCompletionChunkDeltaEmpty] - finish_reason: Optional[str] + delta: Union[ + ChatCompletionStreamResponseDelta, ChatCompletionStreamResponseDeltaEmpty + ] + finish_reason: Optional[Literal["stop", "length", "tool_calls", "function_call"]] + logprobs: NotRequired[Optional[ChatCompletionLogprobs]] -class ChatCompletionChunk(TypedDict): +class CreateChatCompletionStreamResponse(TypedDict): id: str model: str object: Literal["chat.completion.chunk"] created: int - choices: List[ChatCompletionChunkChoice] + choices: List[ChatCompletionStreamResponseChoice] + + +class ChatCompletionFunctions(TypedDict): + name: str + description: NotRequired[str] + parameters: Dict[str, JsonType] # TODO: make this more specific + + +class ChatCompletionFunctionCallOption(TypedDict): + name: str + + +class ChatCompletionRequestResponseFormat(TypedDict): + type: Literal["text", "json_object"] + schema: NotRequired[ + JsonType + ] # https://docs.endpoints.anyscale.com/guides/json_mode/ + + +class ChatCompletionRequestMessageContentPartText(TypedDict): + type: Literal["text"] + text: str + + +class ChatCompletionRequestMessageContentPartImageImageUrl(TypedDict): + url: str + detail: NotRequired[Literal["auto", "low", "high"]] + + +class ChatCompletionRequestMessageContentPartImage(TypedDict): + type: Literal["image_url"] + image_url: Union[str, ChatCompletionRequestMessageContentPartImageImageUrl] + + +ChatCompletionRequestMessageContentPart = Union[ + ChatCompletionRequestMessageContentPartText, + ChatCompletionRequestMessageContentPartImage, +] + + +class ChatCompletionRequestSystemMessage(TypedDict): + role: Literal["system"] + content: Optional[str] + + +class ChatCompletionRequestUserMessage(TypedDict): + role: Literal["user"] + content: Optional[Union[str, List[ChatCompletionRequestMessageContentPart]]] + + +class ChatCompletionMessageToolCallFunction(TypedDict): + name: str + arguments: str + + +class ChatCompletionMessageToolCall(TypedDict): + id: str + type: Literal["function"] + function: ChatCompletionMessageToolCallFunction + + +ChatCompletionMessageToolCalls = List[ChatCompletionMessageToolCall] + + +class ChatCompletionRequestAssistantMessageFunctionCall(TypedDict): + name: str + arguments: str + + +class ChatCompletionRequestAssistantMessage(TypedDict): + role: Literal["assistant"] + content: NotRequired[str] + tool_calls: NotRequired[ChatCompletionMessageToolCalls] + function_call: NotRequired[ + ChatCompletionRequestAssistantMessageFunctionCall + ] # DEPRECATED + + +class ChatCompletionRequestToolMessage(TypedDict): + role: Literal["tool"] + content: Optional[str] + tool_call_id: str + + +class ChatCompletionRequestFunctionMessage(TypedDict): + role: Literal["function"] + content: Optional[str] + name: str + + +ChatCompletionRequestMessage = Union[ + ChatCompletionRequestSystemMessage, + ChatCompletionRequestUserMessage, + ChatCompletionRequestAssistantMessage, + ChatCompletionRequestUserMessage, + ChatCompletionRequestToolMessage, + ChatCompletionRequestFunctionMessage, +] + + +class ChatCompletionRequestFunctionCallOption(TypedDict): + name: str + + +ChatCompletionRequestFunctionCall = Union[ + Literal["none", "auto"], ChatCompletionRequestFunctionCallOption +] + +ChatCompletionFunctionParameters = Dict[str, JsonType] # TODO: make this more specific + + +class ChatCompletionToolFunction(TypedDict): + name: str + description: NotRequired[str] + parameters: ChatCompletionFunctionParameters + + +class ChatCompletionTool(TypedDict): + type: Literal["function"] + function: ChatCompletionToolFunction + + +class ChatCompletionNamedToolChoiceFunction(TypedDict): + name: str + + +class ChatCompletionNamedToolChoice(TypedDict): + type: Literal["function"] + function: ChatCompletionNamedToolChoiceFunction + + +ChatCompletionToolChoiceOption = Union[ + Literal["none", "auto", "required"], ChatCompletionNamedToolChoice +] + + +# NOTE: The following type names are not part of the OpenAI OpenAPI specification +# and will be removed in a future major release. + +EmbeddingData = Embedding +CompletionChunk = CreateCompletionResponse +Completion = CreateCompletionResponse +CreateCompletionStreamResponse = CreateCompletionResponse +ChatCompletionMessage = ChatCompletionResponseMessage +ChatCompletionChoice = ChatCompletionResponseChoice +ChatCompletion = CreateChatCompletionResponse +ChatCompletionChunkDeltaEmpty = ChatCompletionStreamResponseDeltaEmpty +ChatCompletionChunkChoice = ChatCompletionStreamResponseChoice +ChatCompletionChunkDelta = ChatCompletionStreamResponseDelta +ChatCompletionChunk = CreateChatCompletionStreamResponse +ChatCompletionStreamResponse = CreateChatCompletionStreamResponse +ChatCompletionResponseFunction = ChatCompletionFunction +ChatCompletionFunctionCall = ChatCompletionResponseFunctionCall diff --git a/llama_cpp/llava_cpp.py b/llama_cpp/llava_cpp.py new file mode 100644 index 0000000000..ea9b216f7d --- /dev/null +++ b/llama_cpp/llava_cpp.py @@ -0,0 +1,154 @@ +from __future__ import annotations + +import os +from ctypes import ( + c_bool, + c_char_p, + c_int, + c_uint8, + c_float, + c_void_p, + POINTER, + _Pointer, # type: ignore + Structure, +) +import pathlib +from typing import ( + Union, + NewType, + Optional, + TYPE_CHECKING, +) + +import llama_cpp.llama_cpp as llama_cpp + +from llama_cpp._ctypes_extensions import ( + load_shared_library, + ctypes_function_for_shared_library, +) + +if TYPE_CHECKING: + from llama_cpp._ctypes_extensions import ( + CtypesArray, + ) + + +# Specify the base name of the shared library to load +_libllava_base_name = "llava" +_libllava_override_path = os.environ.get("LLAVA_CPP_LIB") +_libllava_base_path = ( + pathlib.Path(os.path.abspath(os.path.dirname(__file__))) / "lib" + if _libllava_override_path is None + else pathlib.Path(_libllava_override_path) +) + +# Load the library +_libllava = load_shared_library(_libllava_base_name, _libllava_base_path) + +ctypes_function = ctypes_function_for_shared_library(_libllava) + + +################################################ +# llava.h +################################################ + +# struct clip_ctx; +clip_ctx_p = NewType("clip_ctx_p", int) +clip_ctx_p_ctypes = c_void_p + + +# struct llava_image_embed { +# float * embed; +# int n_image_pos; +# }; +class llava_image_embed(Structure): + _fields_ = [ + ("embed", POINTER(c_float)), + ("n_image_pos", c_int), + ] + + +# /** sanity check for clip <-> llava embed size match */ +# LLAVA_API bool llava_validate_embed_size(const llama_context * ctx_llama, const clip_ctx * ctx_clip); +@ctypes_function( + "llava_validate_embed_size", + [llama_cpp.llama_context_p_ctypes, clip_ctx_p_ctypes], + c_bool, +) +def llava_validate_embed_size( + ctx_llama: llama_cpp.llama_context_p, ctx_clip: clip_ctx_p, / +) -> bool: ... + + +# /** build an image embed from image file bytes */ +# LLAVA_API struct llava_image_embed * llava_image_embed_make_with_bytes(struct clip_ctx * ctx_clip, int n_threads, const unsigned char * image_bytes, int image_bytes_length); +@ctypes_function( + "llava_image_embed_make_with_bytes", + [clip_ctx_p_ctypes, c_int, POINTER(c_uint8), c_int], + POINTER(llava_image_embed), +) +def llava_image_embed_make_with_bytes( + ctx_clip: clip_ctx_p, + n_threads: Union[c_int, int], + image_bytes: CtypesArray[c_uint8], + image_bytes_length: Union[c_int, int], + /, +) -> "_Pointer[llava_image_embed]": ... + + +# /** build an image embed from a path to an image filename */ +# LLAVA_API struct llava_image_embed * llava_image_embed_make_with_filename(struct clip_ctx * ctx_clip, int n_threads, const char * image_path); +@ctypes_function( + "llava_image_embed_make_with_filename", + [clip_ctx_p_ctypes, c_int, c_char_p], + POINTER(llava_image_embed), +) +def llava_image_embed_make_with_filename( + ctx_clip: clip_ctx_p, n_threads: Union[c_int, int], image_path: bytes, / +) -> "_Pointer[llava_image_embed]": ... + + +# LLAVA_API void llava_image_embed_free(struct llava_image_embed * embed); +# /** free an embedding made with llava_image_embed_make_* */ +@ctypes_function("llava_image_embed_free", [POINTER(llava_image_embed)], None) +def llava_image_embed_free(embed: "_Pointer[llava_image_embed]", /): ... + + +# /** write the image represented by embed into the llama context with batch size n_batch, starting at context pos n_past. on completion, n_past points to the next position in the context after the image embed. */ +# LLAVA_API bool llava_eval_image_embed(struct llama_context * ctx_llama, const struct llava_image_embed * embed, int n_batch, int * n_past); +@ctypes_function( + "llava_eval_image_embed", + [ + llama_cpp.llama_context_p_ctypes, + POINTER(llava_image_embed), + c_int, + POINTER(c_int), + ], + c_bool, +) +def llava_eval_image_embed( + ctx_llama: llama_cpp.llama_context_p, + embed: "_Pointer[llava_image_embed]", + n_batch: Union[c_int, int], + n_past: "_Pointer[c_int]", + /, +) -> bool: ... + + +################################################ +# clip.h +################################################ + + +# /** load mmproj model */ +# CLIP_API struct clip_ctx * clip_model_load (const char * fname, int verbosity); +@ctypes_function("clip_model_load", [c_char_p, c_int], clip_ctx_p_ctypes) +def clip_model_load( + fname: bytes, verbosity: Union[c_int, int], / +) -> Optional[clip_ctx_p]: ... + + +# /** free mmproj model */ +# CLIP_API void clip_free(struct clip_ctx * ctx); +@ctypes_function("clip_free", [clip_ctx_p_ctypes], None) +def clip_free(ctx: clip_ctx_p, /): ... diff --git a/llama_cpp/mtmd_cpp.py b/llama_cpp/mtmd_cpp.py new file mode 100644 index 0000000000..919cefb357 --- /dev/null +++ b/llama_cpp/mtmd_cpp.py @@ -0,0 +1,946 @@ +from __future__ import annotations + +import os +import warnings +from ctypes import ( + CFUNCTYPE, + c_bool, + c_char_p, + c_int, + c_int64, + c_uint8, + c_uint32, + c_size_t, + c_float, + c_void_p, + POINTER, + _Pointer, # type: ignore + Structure, + byref, +) +import pathlib +from typing import ( + Union, + NewType, + Optional, + TYPE_CHECKING, +) + +import llama_cpp.llama_cpp as llama_cpp + +from llama_cpp._ctypes_extensions import ( + load_shared_library, + ctypes_function_for_shared_library, +) + +if TYPE_CHECKING: + from llama_cpp._ctypes_extensions import ( + CtypesArray, + ) + + +# Specify the base name of the shared library to load +_libmtmd_base_name = "mtmd" +_libmtmd_override_path = os.environ.get("MTMD_CPP_LIB") +_libmtmd_base_path = ( + pathlib.Path(os.path.abspath(os.path.dirname(__file__))) / "lib" + if _libmtmd_override_path is None + else pathlib.Path(_libmtmd_override_path) +) + +# Load the library +_libmtmd = load_shared_library(_libmtmd_base_name, _libmtmd_base_path) + +ctypes_function = ctypes_function_for_shared_library(_libmtmd) + +################################################ +# mtmd.h types +################################################ + +# Opaque types +mtmd_context_p = NewType("mtmd_context_p", int) +mtmd_context_p_ctypes = c_void_p + +mtmd_bitmap_p = NewType("mtmd_bitmap_p", int) +mtmd_bitmap_p_ctypes = c_void_p + +mtmd_helper_video_p = NewType("mtmd_helper_video_p", int) +mtmd_helper_video_p_ctypes = c_void_p + +mtmd_image_tokens_p = NewType("mtmd_image_tokens_p", int) +mtmd_image_tokens_p_ctypes = c_void_p + +mtmd_input_chunk_p = NewType("mtmd_input_chunk_p", int) +mtmd_input_chunk_p_ctypes = c_void_p + +mtmd_input_chunks_p = NewType("mtmd_input_chunks_p", int) +mtmd_input_chunks_p_ctypes = c_void_p + +# Enums +MTMD_INPUT_CHUNK_TYPE_TEXT = 0 +MTMD_INPUT_CHUNK_TYPE_IMAGE = 1 +MTMD_INPUT_CHUNK_TYPE_AUDIO = 2 + + +# Structures +class mtmd_context_params(Structure): + """Context parameters for MTMD initialization. + + `image_marker` is deprecated upstream and kept for compatibility; use + `media_marker` for multimodal prompt placeholders. + """ + + if TYPE_CHECKING: + use_gpu: bool + print_timings: bool + n_threads: int + image_marker: Optional[bytes] + media_marker: Optional[bytes] + flash_attn_type: int + warmup: bool + image_min_tokens: int + image_max_tokens: int + cb_eval: llama_cpp.ggml_backend_sched_eval_callback + cb_eval_user_data: c_void_p + + _fields_ = [ + ("use_gpu", c_bool), + ("print_timings", c_bool), + ("n_threads", c_int), + ("image_marker", c_char_p), + ("media_marker", c_char_p), + ("flash_attn_type", c_int), + ("warmup", c_bool), + ("image_min_tokens", c_int), + ("image_max_tokens", c_int), + ("cb_eval", llama_cpp.ggml_backend_sched_eval_callback), + ("cb_eval_user_data", c_void_p), + ] + + +class mtmd_input_text(Structure): + """Text input passed to `mtmd_tokenize`.""" + + _fields_ = [ + ("text", c_char_p), + ("add_special", c_bool), + ("parse_special", c_bool), + ] + + +class mtmd_decoder_pos(Structure): + """Decoder attention position for M-RoPE models.""" + + _fields_ = [ + ("t", c_uint32), + ("x", c_uint32), + ("y", c_uint32), + ("z", c_uint32), + ] + + +# struct mtmd_caps { +# bool inp_vision; +# bool inp_audio; +# }; +class mtmd_caps(Structure): + """Capabilities exposed by an mmproj file.""" + + if TYPE_CHECKING: + inp_vision: bool + inp_audio: bool + + _fields_ = [ + ("inp_vision", c_bool), + ("inp_audio", c_bool), + ] + + +mtmd_bitmap_lazy_callback = CFUNCTYPE( + c_int, + c_size_t, + c_void_p, + POINTER(mtmd_bitmap_p_ctypes), + POINTER(c_char_p), +) + + +class mtmd_helper_bitmap_wrapper(Structure): + """Bitmap wrapper returned by MTMD helper media loaders.""" + + if TYPE_CHECKING: + bitmap: Optional[mtmd_bitmap_p] + video_ctx: Optional[mtmd_helper_video_p] + + _fields_ = [ + ("bitmap", mtmd_bitmap_p_ctypes), + ("video_ctx", mtmd_helper_video_p_ctypes), + ] + + +class mtmd_helper_video_info(Structure): + """Metadata for a decoded video stream.""" + + if TYPE_CHECKING: + width: int + height: int + fps: float + n_frames: int + + _fields_ = [ + ("width", c_uint32), + ("height", c_uint32), + ("fps", c_float), + ("n_frames", c_int), + ] + + +class mtmd_helper_video_init_params(Structure): + """Parameters for initializing an MTMD helper video stream.""" + + if TYPE_CHECKING: + fps_target: float + ffmpeg_bin_dir: Optional[bytes] + timestamp_interval_ms: int + + _fields_ = [ + ("fps_target", c_float), + ("ffmpeg_bin_dir", c_char_p), + ("timestamp_interval_ms", c_int64), + ] + + +################################################ +# mtmd.h functions +################################################ + + +# MTMD_API const char * mtmd_default_marker(void); +@ctypes_function("mtmd_default_marker", [], c_char_p) +def mtmd_default_marker() -> bytes: + """Return the default media marker.""" + ... + + +# MTMD_API struct mtmd_context_params mtmd_context_params_default(void); +@ctypes_function("mtmd_context_params_default", [], mtmd_context_params) +def mtmd_context_params_default() -> mtmd_context_params: + """Return the default MTMD context parameters.""" + ... + + +# MTMD_API mtmd_context * mtmd_init_from_file(const char * mmproj_fname, +# const struct llama_model * text_model, +# const struct mtmd_context_params ctx_params); +@ctypes_function( + "mtmd_init_from_file", + [c_char_p, llama_cpp.llama_model_p_ctypes, mtmd_context_params], + mtmd_context_p_ctypes, +) +def mtmd_init_from_file( + mmproj_fname: bytes, + text_model: llama_cpp.llama_model_p, + ctx_params: mtmd_context_params, + /, +) -> Optional[mtmd_context_p]: + """Initialize the MTMD context from a projector file. Returns None on failure.""" + ... + + +# MTMD_API void mtmd_free(mtmd_context * ctx); +@ctypes_function("mtmd_free", [mtmd_context_p_ctypes], None) +def mtmd_free(ctx: mtmd_context_p, /): ... + + +# MTMD_API bool mtmd_decode_use_non_causal(const mtmd_context * ctx, const mtmd_input_chunk * chunk); +@ctypes_function( + "mtmd_decode_use_non_causal", + [mtmd_context_p_ctypes, mtmd_input_chunk_p_ctypes], + c_bool, +) +def mtmd_decode_use_non_causal( + ctx: mtmd_context_p, chunk: Optional[mtmd_input_chunk_p], / +) -> bool: + """Check whether MTMD decoding uses non-causal attention.""" + ... + + +# MTMD_API bool mtmd_decode_use_mrope(const mtmd_context * ctx); +@ctypes_function("mtmd_decode_use_mrope", [mtmd_context_p_ctypes], c_bool) +def mtmd_decode_use_mrope(ctx: mtmd_context_p, /) -> bool: + """Check whether MTMD decoding uses mRoPE.""" + ... + + +# MTMD_API bool mtmd_support_vision(const mtmd_context * ctx); +@ctypes_function("mtmd_support_vision", [mtmd_context_p_ctypes], c_bool) +def mtmd_support_vision(ctx: mtmd_context_p, /) -> bool: + """Check whether the current model supports vision input.""" + ... + + +# MTMD_API bool mtmd_support_audio(const mtmd_context * ctx); +@ctypes_function("mtmd_support_audio", [mtmd_context_p_ctypes], c_bool) +def mtmd_support_audio(ctx: mtmd_context_p, /) -> bool: + """Check whether MTMD supports audio.""" + ... + + +# MTMD_API int mtmd_get_audio_sample_rate(const mtmd_context * ctx); +@ctypes_function("mtmd_get_audio_sample_rate", [mtmd_context_p_ctypes], c_int) +def mtmd_get_audio_sample_rate(ctx: mtmd_context_p, /) -> int: + """Get the audio sample rate in Hz. Returns -1 if audio is not supported.""" + ... + + +# MTMD_API const char * mtmd_get_marker(const mtmd_context * ctx); +@ctypes_function("mtmd_get_marker", [mtmd_context_p_ctypes], c_char_p) +def mtmd_get_marker(ctx: mtmd_context_p, /) -> Optional[bytes]: + """Get the current media marker string.""" + ... + + +# Deprecated compatibility wrapper for the renamed mtmd_get_audio_sample_rate(). +def mtmd_get_audio_bitrate(ctx: mtmd_context_p, /) -> int: + warnings.warn( + "mtmd_get_audio_bitrate is deprecated; use mtmd_get_audio_sample_rate instead", + DeprecationWarning, + stacklevel=2, + ) + return mtmd_get_audio_sample_rate(ctx) + + +# MTMD_API mtmd_bitmap * mtmd_bitmap_init(uint32_t nx, uint32_t ny, const unsigned char * data); +@ctypes_function( + "mtmd_bitmap_init", [c_uint32, c_uint32, POINTER(c_uint8)], mtmd_bitmap_p_ctypes +) +def mtmd_bitmap_init( + nx: Union[c_uint32, int], + ny: Union[c_uint32, int], + data: CtypesArray[c_uint8], + /, +) -> Optional[mtmd_bitmap_p]: ... + + +# MTMD_API mtmd_bitmap * mtmd_bitmap_init_from_audio(size_t n_samples, const float * data); +@ctypes_function( + "mtmd_bitmap_init_from_audio", + [c_size_t, POINTER(c_float)], + mtmd_bitmap_p_ctypes, +) +def mtmd_bitmap_init_from_audio( + n_samples: Union[c_size_t, int], + data: CtypesArray[c_float], + /, +) -> Optional[mtmd_bitmap_p]: + """Initialize an MTMD bitmap from audio samples.""" + ... + + +# MTMD_API void mtmd_bitmap_free(mtmd_bitmap * bitmap); +@ctypes_function("mtmd_bitmap_free", [mtmd_bitmap_p_ctypes], None) +def mtmd_bitmap_free(bitmap: mtmd_bitmap_p, /): ... + + +# MTMD_API uint32_t mtmd_bitmap_get_nx(const mtmd_bitmap * bitmap); +@ctypes_function("mtmd_bitmap_get_nx", [mtmd_bitmap_p_ctypes], c_uint32) +def mtmd_bitmap_get_nx(bitmap: mtmd_bitmap_p, /) -> int: + """Get the bitmap width in pixels.""" + ... + + +# MTMD_API uint32_t mtmd_bitmap_get_ny(const mtmd_bitmap * bitmap); +@ctypes_function("mtmd_bitmap_get_ny", [mtmd_bitmap_p_ctypes], c_uint32) +def mtmd_bitmap_get_ny(bitmap: mtmd_bitmap_p, /) -> int: + """Get the bitmap height in pixels.""" + ... + + +# MTMD_API const unsigned char * mtmd_bitmap_get_data(const mtmd_bitmap * bitmap); +@ctypes_function("mtmd_bitmap_get_data", [mtmd_bitmap_p_ctypes], POINTER(c_uint8)) +def mtmd_bitmap_get_data(bitmap: mtmd_bitmap_p, /) -> Optional[CtypesArray[c_uint8]]: + """Get the raw bitmap data buffer.""" + ... + + +# MTMD_API size_t mtmd_bitmap_get_n_bytes(const mtmd_bitmap * bitmap); +@ctypes_function("mtmd_bitmap_get_n_bytes", [mtmd_bitmap_p_ctypes], c_size_t) +def mtmd_bitmap_get_n_bytes(bitmap: mtmd_bitmap_p, /) -> int: + """Get the bitmap data size in bytes.""" + ... + + +# MTMD_API bool mtmd_bitmap_is_audio(const mtmd_bitmap * bitmap); +@ctypes_function("mtmd_bitmap_is_audio", [mtmd_bitmap_p_ctypes], c_bool) +def mtmd_bitmap_is_audio(bitmap: mtmd_bitmap_p, /) -> bool: + """Check whether the bitmap contains audio data.""" + ... + + +# MTMD_API const char * mtmd_bitmap_get_id(const mtmd_bitmap * bitmap); +@ctypes_function("mtmd_bitmap_get_id", [mtmd_bitmap_p_ctypes], c_char_p) +def mtmd_bitmap_get_id(bitmap: mtmd_bitmap_p, /) -> Optional[bytes]: + """Get the optional bitmap identifier.""" + ... + + +# MTMD_API void mtmd_bitmap_set_id(mtmd_bitmap * bitmap, const char * id); +@ctypes_function("mtmd_bitmap_set_id", [mtmd_bitmap_p_ctypes, c_char_p], None) +def mtmd_bitmap_set_id(bitmap: mtmd_bitmap_p, id: Optional[bytes], /): + """Set the optional bitmap identifier.""" + ... + + +# MTMD_API mtmd_bitmap * mtmd_bitmap_init_lazy(mtmd_context * ctx, +# const char * id, +# void * user_data, +# mtmd_bitmap_lazy_callback callback); +@ctypes_function( + "mtmd_bitmap_init_lazy", + [mtmd_context_p_ctypes, c_char_p, c_void_p, mtmd_bitmap_lazy_callback], + mtmd_bitmap_p_ctypes, +) +def mtmd_bitmap_init_lazy( + ctx: mtmd_context_p, + id: Optional[bytes], + user_data: c_void_p, + callback: mtmd_bitmap_lazy_callback, + /, +) -> Optional[mtmd_bitmap_p]: + """Initialize a lazy MTMD bitmap.""" + ... + + +# MTMD_API mtmd_input_chunks * mtmd_input_chunks_init(void); +@ctypes_function("mtmd_input_chunks_init", [], mtmd_input_chunks_p_ctypes) +def mtmd_input_chunks_init() -> Optional[mtmd_input_chunks_p]: ... + + +# MTMD_API void mtmd_input_chunks_free(mtmd_input_chunks * chunks); +@ctypes_function("mtmd_input_chunks_free", [mtmd_input_chunks_p_ctypes], None) +def mtmd_input_chunks_free(chunks: mtmd_input_chunks_p, /): ... + + +# MTMD_API size_t mtmd_input_chunks_size(const mtmd_input_chunks * chunks); +@ctypes_function("mtmd_input_chunks_size", [mtmd_input_chunks_p_ctypes], c_size_t) +def mtmd_input_chunks_size(chunks: mtmd_input_chunks_p, /) -> int: ... + + +# MTMD_API const mtmd_input_chunk * mtmd_input_chunks_get(const mtmd_input_chunks * chunks, size_t idx); +@ctypes_function( + "mtmd_input_chunks_get", + [mtmd_input_chunks_p_ctypes, c_size_t], + mtmd_input_chunk_p_ctypes, +) +def mtmd_input_chunks_get( + chunks: mtmd_input_chunks_p, idx: Union[c_size_t, int], / +) -> Optional[mtmd_input_chunk_p]: ... + + +# MTMD_API int32_t mtmd_tokenize(mtmd_context * ctx, +# mtmd_input_chunks * output, +# const mtmd_input_text * text, +# const mtmd_bitmap ** bitmaps, +# size_t n_bitmaps); +@ctypes_function( + "mtmd_tokenize", + [ + mtmd_context_p_ctypes, + mtmd_input_chunks_p_ctypes, + POINTER(mtmd_input_text), + POINTER(mtmd_bitmap_p_ctypes), + c_size_t, + ], + c_int, +) +def mtmd_tokenize( + ctx: mtmd_context_p, + output: mtmd_input_chunks_p, + text: "_Pointer[mtmd_input_text]", + bitmaps: CtypesArray[mtmd_bitmap_p_ctypes], + n_bitmaps: Union[c_size_t, int], + /, +) -> int: ... + + +# MTMD_API size_t mtmd_input_chunk_get_n_tokens(const mtmd_input_chunk * chunk); +@ctypes_function("mtmd_input_chunk_get_n_tokens", [mtmd_input_chunk_p_ctypes], c_size_t) +def mtmd_input_chunk_get_n_tokens(chunk: mtmd_input_chunk_p, /) -> int: ... + + +# MTMD_API enum mtmd_input_chunk_type mtmd_input_chunk_get_type(const mtmd_input_chunk * chunk); +@ctypes_function("mtmd_input_chunk_get_type", [mtmd_input_chunk_p_ctypes], c_int) +def mtmd_input_chunk_get_type(chunk: mtmd_input_chunk_p, /) -> int: ... + + +# MTMD_API const llama_token * mtmd_input_chunk_get_tokens_text(const mtmd_input_chunk * chunk, size_t * n_tokens_output); +@ctypes_function( + "mtmd_input_chunk_get_tokens_text", + [mtmd_input_chunk_p_ctypes, POINTER(c_size_t)], + POINTER(llama_cpp.llama_token), +) +def mtmd_input_chunk_get_tokens_text( + chunk: mtmd_input_chunk_p, n_tokens_output: "_Pointer[c_size_t]", / +) -> Optional["_Pointer[llama_cpp.llama_token]"]: ... + + +# MTMD_API const mtmd_image_tokens * mtmd_input_chunk_get_tokens_image(const mtmd_input_chunk * chunk); +@ctypes_function( + "mtmd_input_chunk_get_tokens_image", + [mtmd_input_chunk_p_ctypes], + mtmd_image_tokens_p_ctypes, +) +def mtmd_input_chunk_get_tokens_image( + chunk: mtmd_input_chunk_p, / +) -> Optional[mtmd_image_tokens_p]: ... + + +# MTMD_API const char * mtmd_input_chunk_get_id(const mtmd_input_chunk * chunk); +@ctypes_function("mtmd_input_chunk_get_id", [mtmd_input_chunk_p_ctypes], c_char_p) +def mtmd_input_chunk_get_id(chunk: mtmd_input_chunk_p, /) -> Optional[bytes]: + """Get the optional chunk identifier.""" + ... + + +# MTMD_API llama_pos mtmd_input_chunk_get_n_pos(const mtmd_input_chunk * chunk); +@ctypes_function( + "mtmd_input_chunk_get_n_pos", + [mtmd_input_chunk_p_ctypes], + llama_cpp.llama_pos, +) +def mtmd_input_chunk_get_n_pos(chunk: mtmd_input_chunk_p, /) -> int: + """Get the number of positions consumed by the chunk.""" + ... + + +# MTMD_API mtmd_input_chunk * mtmd_input_chunk_copy(const mtmd_input_chunk * chunk); +@ctypes_function( + "mtmd_input_chunk_copy", [mtmd_input_chunk_p_ctypes], mtmd_input_chunk_p_ctypes +) +def mtmd_input_chunk_copy(chunk: mtmd_input_chunk_p, /) -> Optional[mtmd_input_chunk_p]: + """Copy an input chunk and transfer ownership to the caller.""" + ... + + +# MTMD_API void mtmd_input_chunk_free(mtmd_input_chunk * chunk); +@ctypes_function("mtmd_input_chunk_free", [mtmd_input_chunk_p_ctypes], None) +def mtmd_input_chunk_free(chunk: mtmd_input_chunk_p, /): + """Free an owned input chunk.""" + ... + + +# MTMD_API size_t mtmd_image_tokens_get_n_tokens(const mtmd_image_tokens * image_tokens); +@ctypes_function( + "mtmd_image_tokens_get_n_tokens", [mtmd_image_tokens_p_ctypes], c_size_t +) +def mtmd_image_tokens_get_n_tokens(image_tokens: mtmd_image_tokens_p, /) -> int: + """Get the number of image tokens.""" + ... + + +# DEPRECATED(MTMD_API size_t mtmd_image_tokens_get_nx(const mtmd_image_tokens * image_tokens), +# "use mtmd_image_tokens_get_decoder_pos() instead"); +@ctypes_function("mtmd_image_tokens_get_nx", [mtmd_image_tokens_p_ctypes], c_size_t) +def mtmd_image_tokens_get_nx(image_tokens: mtmd_image_tokens_p, /) -> int: + """Get the image token grid width.""" + ... + + +# DEPRECATED(MTMD_API size_t mtmd_image_tokens_get_ny(const mtmd_image_tokens * image_tokens), +# "use mtmd_image_tokens_get_decoder_pos() instead"); +@ctypes_function("mtmd_image_tokens_get_ny", [mtmd_image_tokens_p_ctypes], c_size_t) +def mtmd_image_tokens_get_ny(image_tokens: mtmd_image_tokens_p, /) -> int: + """Get the image token grid height.""" + ... + + +# MTMD_API const char * mtmd_image_tokens_get_id(const mtmd_image_tokens * image_tokens); +@ctypes_function("mtmd_image_tokens_get_id", [mtmd_image_tokens_p_ctypes], c_char_p) +def mtmd_image_tokens_get_id(image_tokens: mtmd_image_tokens_p, /) -> Optional[bytes]: + """Get the optional image token identifier.""" + ... + + +# MTMD_API llama_pos mtmd_image_tokens_get_n_pos(const mtmd_image_tokens * image_tokens); +@ctypes_function( + "mtmd_image_tokens_get_n_pos", + [mtmd_image_tokens_p_ctypes], + llama_cpp.llama_pos, +) +def mtmd_image_tokens_get_n_pos(image_tokens: mtmd_image_tokens_p, /) -> int: + """Get the number of positions consumed by the image tokens.""" + ... + + +# MTMD_API struct mtmd_decoder_pos mtmd_image_tokens_get_decoder_pos( +# const mtmd_image_tokens * image_tokens, llama_pos pos_0, size_t i); +@ctypes_function( + "mtmd_image_tokens_get_decoder_pos", + [mtmd_image_tokens_p_ctypes, llama_cpp.llama_pos, c_size_t], + mtmd_decoder_pos, +) +def mtmd_image_tokens_get_decoder_pos( + image_tokens: mtmd_image_tokens_p, + pos_0: llama_cpp.llama_pos, + i: Union[c_size_t, int], + /, +) -> mtmd_decoder_pos: + """Get decoder attention position for an image embedding token.""" + ... + + +# MTMD_API int32_t mtmd_encode(mtmd_context * ctx, const mtmd_image_tokens * image_tokens); +@ctypes_function( + "mtmd_encode", + [mtmd_context_p_ctypes, mtmd_image_tokens_p_ctypes], + c_int, +) +def mtmd_encode(ctx: mtmd_context_p, image_tokens: mtmd_image_tokens_p, /) -> int: + """Run an MTMD encode pass for image tokens.""" + ... + + +# MTMD_API int32_t mtmd_encode_chunk(mtmd_context * ctx, const mtmd_input_chunk * chunk); +@ctypes_function( + "mtmd_encode_chunk", + [mtmd_context_p_ctypes, mtmd_input_chunk_p_ctypes], + c_int, +) +def mtmd_encode_chunk(ctx: mtmd_context_p, chunk: mtmd_input_chunk_p, /) -> int: + """Run an MTMD encode pass for a single chunk.""" + ... + + +# MTMD_API float * mtmd_get_output_embd(mtmd_context * ctx); +@ctypes_function("mtmd_get_output_embd", [mtmd_context_p_ctypes], POINTER(c_float)) +def mtmd_get_output_embd(ctx: mtmd_context_p, /) -> Optional[CtypesArray[c_float]]: + """Get output embeddings from the last encode pass.""" + ... + + +# MTMD_API struct mtmd_caps mtmd_get_cap_from_file(const char * mmproj_fname); +@ctypes_function("mtmd_get_cap_from_file", [c_char_p], mtmd_caps) +def mtmd_get_cap_from_file(mmproj_fname: bytes, /) -> mtmd_caps: + """Get mmproj capabilities without initializing a full MTMD context.""" + ... + + +# MTMD_API mtmd_input_chunks * mtmd_test_create_input_chunks(void); +@ctypes_function("mtmd_test_create_input_chunks", [], mtmd_input_chunks_p_ctypes) +def mtmd_test_create_input_chunks() -> Optional[mtmd_input_chunks_p]: + """Create MTMD test chunks for the C API tests.""" + ... + + +################################################ +# mtmd-helper.h functions +################################################ + + +# MTMD_API bool mtmd_helper_support_video(mtmd_context * ctx); +@ctypes_function( + "mtmd_helper_support_video", + [mtmd_context_p_ctypes], + c_bool, +) +def mtmd_helper_support_video(ctx: mtmd_context_p, /) -> bool: + """Check whether MTMD helper video support is available.""" + ... + + +# MTMD_API struct mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_file(mtmd_context * ctx, const char * fname, bool placeholder); +@ctypes_function( + "mtmd_helper_bitmap_init_from_file", + [mtmd_context_p_ctypes, c_char_p, c_bool], + mtmd_helper_bitmap_wrapper, +) +def mtmd_helper_bitmap_init_from_file_wrapper( + ctx: mtmd_context_p, fname: bytes, placeholder: Union[c_bool, bool], / +) -> mtmd_helper_bitmap_wrapper: + """Initialize an MTMD bitmap wrapper from a file.""" + ... + + +def mtmd_helper_bitmap_init_from_file( + ctx: mtmd_context_p, fname: bytes, placeholder: Union[c_bool, bool], / +) -> Optional[mtmd_bitmap_p]: + """Initialize an MTMD bitmap from a file.""" + return mtmd_helper_bitmap_init_from_file_wrapper(ctx, fname, placeholder).bitmap + + +# MTMD_API struct mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, const unsigned char * buf, size_t len, bool placeholder); +@ctypes_function( + "mtmd_helper_bitmap_init_from_buf", + [mtmd_context_p_ctypes, POINTER(c_uint8), c_size_t, c_bool], + mtmd_helper_bitmap_wrapper, +) +def mtmd_helper_bitmap_init_from_buf_wrapper( + ctx: mtmd_context_p, + buf: CtypesArray[c_uint8], + length: Union[c_size_t, int], + placeholder: Union[c_bool, bool], + /, +) -> mtmd_helper_bitmap_wrapper: ... + + +def mtmd_helper_bitmap_init_from_buf( + ctx: mtmd_context_p, + buf: CtypesArray[c_uint8], + length: Union[c_size_t, int], + placeholder: Union[c_bool, bool], + /, +) -> Optional[mtmd_bitmap_p]: + """Initialize an MTMD bitmap from a buffer.""" + return mtmd_helper_bitmap_init_from_buf_wrapper( + ctx, buf, length, placeholder + ).bitmap + + +# MTMD_API size_t mtmd_helper_get_n_tokens(const mtmd_input_chunks * chunks); +@ctypes_function("mtmd_helper_get_n_tokens", [mtmd_input_chunks_p_ctypes], c_size_t) +def mtmd_helper_get_n_tokens(chunks: mtmd_input_chunks_p, /) -> int: ... + + +# MTMD_API llama_pos mtmd_helper_get_n_pos(const mtmd_input_chunks * chunks); +@ctypes_function( + "mtmd_helper_get_n_pos", + [mtmd_input_chunks_p_ctypes], + llama_cpp.llama_pos, +) +def mtmd_helper_get_n_pos(chunks: mtmd_input_chunks_p, /) -> int: + """Count the total positions consumed by the chunks.""" + ... + + +# MTMD_API void mtmd_helper_image_get_decoder_pos( +# const mtmd_image_tokens * image, llama_pos pos_0, struct mtmd_decoder_pos * out_pos); +@ctypes_function( + "mtmd_helper_image_get_decoder_pos", + [mtmd_image_tokens_p_ctypes, llama_cpp.llama_pos, POINTER(mtmd_decoder_pos)], + None, +) +def mtmd_helper_image_get_decoder_pos( + image: mtmd_image_tokens_p, + pos_0: llama_cpp.llama_pos, + out_pos: "_Pointer[mtmd_decoder_pos]", + /, +): + """Fill decoder attention positions for all image embedding tokens.""" + ... + + +# MTMD_API int32_t mtmd_helper_eval_chunks(mtmd_context * ctx, +# struct llama_context * lctx, +# const mtmd_input_chunks * chunks, +# llama_pos n_past, +# llama_seq_id seq_id, +# int32_t n_batch, +# bool logits_last, +# llama_pos * new_n_past); +@ctypes_function( + "mtmd_helper_eval_chunks", + [ + mtmd_context_p_ctypes, + llama_cpp.llama_context_p_ctypes, + mtmd_input_chunks_p_ctypes, + llama_cpp.llama_pos, + llama_cpp.llama_seq_id, + c_int, + c_bool, + POINTER(llama_cpp.llama_pos), + ], + c_int, +) +def mtmd_helper_eval_chunks( + ctx: mtmd_context_p, + lctx: llama_cpp.llama_context_p, + chunks: mtmd_input_chunks_p, + n_past: llama_cpp.llama_pos, + seq_id: llama_cpp.llama_seq_id, + n_batch: Union[c_int, int], + logits_last: Union[c_bool, bool], + new_n_past: "_Pointer[llama_cpp.llama_pos]", + /, +) -> int: ... + + +# MTMD_API int32_t mtmd_helper_eval_chunk_single(mtmd_context * ctx, +# struct llama_context * lctx, +# const mtmd_input_chunk * chunk, +# llama_pos n_past, +# llama_seq_id seq_id, +# int32_t n_batch, +# bool logits_last, +# llama_pos * new_n_past); +@ctypes_function( + "mtmd_helper_eval_chunk_single", + [ + mtmd_context_p_ctypes, + llama_cpp.llama_context_p_ctypes, + mtmd_input_chunk_p_ctypes, + llama_cpp.llama_pos, + llama_cpp.llama_seq_id, + c_int, + c_bool, + POINTER(llama_cpp.llama_pos), + ], + c_int, +) +def mtmd_helper_eval_chunk_single( + ctx: mtmd_context_p, + lctx: llama_cpp.llama_context_p, + chunk: mtmd_input_chunk_p, + n_past: llama_cpp.llama_pos, + seq_id: llama_cpp.llama_seq_id, + n_batch: Union[c_int, int], + logits_last: Union[c_bool, bool], + new_n_past: "_Pointer[llama_cpp.llama_pos]", + /, +) -> int: ... + + +# MTMD_API int32_t mtmd_helper_decode_image_chunk(mtmd_context * ctx, +# struct llama_context * lctx, +# const mtmd_input_chunk * chunk, +# float * encoded_embd, +# llama_pos n_past, +# llama_seq_id seq_id, +# int32_t n_batch, +# llama_pos * new_n_past); +@ctypes_function( + "mtmd_helper_decode_image_chunk", + [ + mtmd_context_p_ctypes, + llama_cpp.llama_context_p_ctypes, + mtmd_input_chunk_p_ctypes, + POINTER(c_float), + llama_cpp.llama_pos, + llama_cpp.llama_seq_id, + c_int, + POINTER(llama_cpp.llama_pos), + ], + c_int, +) +def mtmd_helper_decode_image_chunk( + ctx: mtmd_context_p, + lctx: llama_cpp.llama_context_p, + chunk: mtmd_input_chunk_p, + encoded_embd: CtypesArray[c_float], + n_past: llama_cpp.llama_pos, + seq_id: llama_cpp.llama_seq_id, + n_batch: Union[c_int, int], + new_n_past: "_Pointer[llama_cpp.llama_pos]", + /, +) -> int: + """Decode a pre-encoded image chunk.""" + ... + + +# MTMD_API struct mtmd_helper_video_init_params mtmd_helper_video_init_params_default(void); +@ctypes_function( + "mtmd_helper_video_init_params_default", [], mtmd_helper_video_init_params +) +def mtmd_helper_video_init_params_default() -> mtmd_helper_video_init_params: + """Return the default MTMD helper video initialization parameters.""" + ... + + +# MTMD_API mtmd_helper_video * mtmd_helper_video_init( +# struct mtmd_context * mctx, +# const char * path, +# struct mtmd_helper_video_init_params params); +@ctypes_function( + "mtmd_helper_video_init", + [mtmd_context_p_ctypes, c_char_p, mtmd_helper_video_init_params], + mtmd_helper_video_p_ctypes, +) +def mtmd_helper_video_init( + ctx: mtmd_context_p, + path: bytes, + params: mtmd_helper_video_init_params, + /, +) -> Optional[mtmd_helper_video_p]: + """Initialize an MTMD helper video stream from a file path.""" + ... + + +# MTMD_API mtmd_helper_video * mtmd_helper_video_init_from_buf( +# struct mtmd_context * mctx, +# const unsigned char * buf, size_t len, +# struct mtmd_helper_video_init_params params); +@ctypes_function( + "mtmd_helper_video_init_from_buf", + [mtmd_context_p_ctypes, POINTER(c_uint8), c_size_t, mtmd_helper_video_init_params], + mtmd_helper_video_p_ctypes, +) +def mtmd_helper_video_init_from_buf( + ctx: mtmd_context_p, + buf: CtypesArray[c_uint8], + length: Union[c_size_t, int], + params: mtmd_helper_video_init_params, + /, +) -> Optional[mtmd_helper_video_p]: + """Initialize an MTMD helper video stream from a buffer.""" + ... + + +# MTMD_API void mtmd_helper_video_free(mtmd_helper_video * ctx); +@ctypes_function("mtmd_helper_video_free", [mtmd_helper_video_p_ctypes], None) +def mtmd_helper_video_free(ctx: mtmd_helper_video_p, /): + """Free an MTMD helper video stream.""" + ... + + +# MTMD_API struct mtmd_helper_video_info mtmd_helper_video_get_info(const mtmd_helper_video * ctx); +@ctypes_function( + "mtmd_helper_video_get_info", + [mtmd_helper_video_p_ctypes], + mtmd_helper_video_info, +) +def mtmd_helper_video_get_info(ctx: mtmd_helper_video_p, /) -> mtmd_helper_video_info: + """Get metadata for an MTMD helper video stream.""" + ... + + +# MTMD_API int32_t mtmd_helper_video_read_next(mtmd_helper_video * ctx, +# mtmd_bitmap ** out_bitmap, +# char ** out_text); +@ctypes_function( + "mtmd_helper_video_read_next", + [ + mtmd_helper_video_p_ctypes, + POINTER(mtmd_bitmap_p_ctypes), + POINTER(c_char_p), + ], + c_int, +) +def mtmd_helper_video_read_next( + ctx: mtmd_helper_video_p, + out_bitmap: "_Pointer[mtmd_bitmap_p_ctypes]", + out_text: "_Pointer[c_char_p]", + /, +) -> int: + """Read the next bitmap or text chunk from an MTMD helper video stream.""" + ... + + +# MTMD_API void mtmd_log_set(ggml_log_callback log_callback, void * user_data); +@ctypes_function( + "mtmd_log_set", + [llama_cpp.llama_log_callback, c_void_p], + None, +) +def mtmd_log_set(log_callback, user_data: c_void_p, /): + """Set the MTMD logging callback.""" + ... + + +# MTMD_API void mtmd_helper_log_set(ggml_log_callback log_callback, void * user_data); +@ctypes_function( + "mtmd_helper_log_set", + [llama_cpp.llama_log_callback, c_void_p], + None, +) +def mtmd_helper_log_set(log_callback, user_data: c_void_p, /): + """Set the MTMD helper logging callback.""" + ... diff --git a/llama_cpp/py.typed b/llama_cpp/py.typed new file mode 100644 index 0000000000..e69de29bb2 diff --git a/llama_cpp/server/__main__.py b/llama_cpp/server/__main__.py index 995dd44490..bbac4957e9 100644 --- a/llama_cpp/server/__main__.py +++ b/llama_cpp/server/__main__.py @@ -9,7 +9,7 @@ Then run: ``` -uvicorn llama_cpp.server.app:app --reload +uvicorn llama_cpp.server.app:create_app --reload ``` or @@ -21,30 +21,80 @@ Then visit http://localhost:8000/docs to see the interactive API docs. """ + +from __future__ import annotations + import os +import sys import argparse import uvicorn -from llama_cpp.server.app import create_app, Settings +from llama_cpp.server.app import create_app +from llama_cpp.server.settings import ( + Settings, + ServerSettings, + ModelSettings, + ConfigFileSettings, +) +from llama_cpp.server.cli import add_args_from_model, parse_model_from_args -if __name__ == "__main__": - parser = argparse.ArgumentParser() - for name, field in Settings.model_fields.items(): - description = field.description - if field.default is not None and description is not None: - description += f" (default: {field.default})" - parser.add_argument( - f"--{name}", - dest=name, - type=field.annotation if field.annotation is not None else str, - help=description, - ) +def main(): + description = "🦙 Llama.cpp python server. Host your own LLMs!🚀" + parser = argparse.ArgumentParser(description=description) + + add_args_from_model(parser, Settings) + parser.add_argument( + "--config_file", + type=str, + help="Path to a config file to load.", + ) + server_settings: ServerSettings | None = None + model_settings: list[ModelSettings] = [] args = parser.parse_args() - settings = Settings(**{k: v for k, v in vars(args).items() if v is not None}) - app = create_app(settings=settings) + try: + # Load server settings from config_file if provided + config_file = os.environ.get("CONFIG_FILE", args.config_file) + if config_file: + if not os.path.exists(config_file): + raise ValueError(f"Config file {config_file} not found!") + with open(config_file, "rb") as f: + # Check if yaml file + if config_file.endswith(".yaml") or config_file.endswith(".yml"): + import yaml + import json + config_file_settings = ConfigFileSettings.model_validate_json( + json.dumps(yaml.safe_load(f)) + ) + else: + config_file_settings = ConfigFileSettings.model_validate_json( + f.read() + ) + server_settings = ServerSettings.model_validate(config_file_settings) + model_settings = config_file_settings.models + else: + server_settings = parse_model_from_args(ServerSettings, args) + model_settings = [parse_model_from_args(ModelSettings, args)] + except Exception as e: + print(e, file=sys.stderr) + parser.print_help() + sys.exit(1) + assert server_settings is not None + assert model_settings is not None + app = create_app( + server_settings=server_settings, + model_settings=model_settings, + ) uvicorn.run( - app, host=os.getenv("HOST", settings.host), port=int(os.getenv("PORT", settings.port)) + app, + host=os.getenv("HOST", server_settings.host), + port=int(os.getenv("PORT", server_settings.port)), + ssl_keyfile=server_settings.ssl_keyfile, + ssl_certfile=server_settings.ssl_certfile, ) + + +if __name__ == "__main__": + main() diff --git a/llama_cpp/server/app.py b/llama_cpp/server/app.py index ffd07fa6b7..f776fe159c 100644 --- a/llama_cpp/server/app.py +++ b/llama_cpp/server/app.py @@ -1,172 +1,87 @@ +from __future__ import annotations + +import os import json -import multiprocessing -from threading import Lock +import typing +import contextlib + +from anyio import Lock from functools import partial -from typing import Iterator, List, Optional, Union, Dict -from typing_extensions import TypedDict, Literal +from typing import List, Optional, Union, Dict import llama_cpp import anyio from anyio.streams.memory import MemoryObjectSendStream from starlette.concurrency import run_in_threadpool, iterate_in_threadpool -from fastapi import Depends, FastAPI, APIRouter, Request +from fastapi import Depends, FastAPI, APIRouter, Request, HTTPException, status, Body +from fastapi.middleware import Middleware from fastapi.middleware.cors import CORSMiddleware -from pydantic import BaseModel, Field -from pydantic_settings import BaseSettings +from fastapi.security import HTTPBearer from sse_starlette.sse import EventSourceResponse +from starlette_context.plugins import RequestIdPlugin # type: ignore +from starlette_context.middleware import RawContextMiddleware +from llama_cpp.server.model import ( + LlamaProxy, +) +from llama_cpp.server.settings import ( + ConfigFileSettings, + Settings, + ModelSettings, + ServerSettings, +) +from llama_cpp.server.types import ( + CreateCompletionRequest, + CreateEmbeddingRequest, + CreateChatCompletionRequest, + ModelList, + TokenizeInputRequest, + TokenizeInputResponse, + TokenizeInputCountResponse, + DetokenizeInputRequest, + DetokenizeInputResponse, +) +from llama_cpp.server.errors import RouteErrorHandler -class Settings(BaseSettings): - model: str = Field( - description="The path to the model to use for generating completions." - ) - model_alias: Optional[str] = Field( - default=None, - description="The alias of the model to use for generating completions.", - ) - n_ctx: int = Field(default=2048, ge=1, description="The context size.") - n_gpu_layers: int = Field( - default=0, - ge=0, - description="The number of layers to put on the GPU. The rest will be on the CPU.", - ) - seed: int = Field( - default=1337, description="Random seed. -1 for random." - ) - n_batch: int = Field( - default=512, ge=1, description="The batch size to use per eval." - ) - n_threads: int = Field( - default=max(multiprocessing.cpu_count() // 2, 1), - ge=1, - description="The number of threads to use.", - ) - f16_kv: bool = Field(default=True, description="Whether to use f16 key/value.") - use_mlock: bool = Field( - default=llama_cpp.llama_mlock_supported(), - description="Use mlock.", - ) - use_mmap: bool = Field( - default=llama_cpp.llama_mmap_supported(), - description="Use mmap.", - ) - embedding: bool = Field(default=True, description="Whether to use embeddings.") - low_vram: bool = Field( - default=False, - description="Whether to use less VRAM. This will reduce performance.", - ) - last_n_tokens_size: int = Field( - default=64, - ge=0, - description="Last n tokens to keep for repeat penalty calculation.", - ) - logits_all: bool = Field(default=True, description="Whether to return logits.") - cache: bool = Field( - default=False, - description="Use a cache to reduce processing times for evaluated prompts.", - ) - cache_type: Literal["ram", "disk"] = Field( - default="ram", - description="The type of cache to use. Only used if cache is True.", - ) - cache_size: int = Field( - default=2 << 30, - description="The size of the cache in bytes. Only used if cache is True.", - ) - vocab_only: bool = Field( - default=False, description="Whether to only return the vocabulary." - ) - verbose: bool = Field( - default=True, description="Whether to print debug information." - ) - host: str = Field( - default="localhost", description="Listen address" - ) - port: int = Field( - default=8000, description="Listen port" - ) - interrupt_requests: bool = Field( - default=True, - description="Whether to interrupt requests when a new request is received.", - ) +router = APIRouter(route_class=RouteErrorHandler) -router = APIRouter() +_server_settings: Optional[ServerSettings] = None -settings: Optional[Settings] = None -llama: Optional[llama_cpp.Llama] = None +def set_server_settings(server_settings: ServerSettings): + global _server_settings + _server_settings = server_settings -def create_app(settings: Optional[Settings] = None): - if settings is None: - settings = Settings() - app = FastAPI( - title="🦙 llama.cpp Python API", - version="0.0.1", - ) - app.add_middleware( - CORSMiddleware, - allow_origins=["*"], - allow_credentials=True, - allow_methods=["*"], - allow_headers=["*"], - ) - app.include_router(router) - global llama - llama = llama_cpp.Llama( - model_path=settings.model, - n_gpu_layers=settings.n_gpu_layers, - seed=settings.seed, - f16_kv=settings.f16_kv, - use_mlock=settings.use_mlock, - use_mmap=settings.use_mmap, - embedding=settings.embedding, - logits_all=settings.logits_all, - n_threads=settings.n_threads, - n_batch=settings.n_batch, - n_ctx=settings.n_ctx, - last_n_tokens_size=settings.last_n_tokens_size, - vocab_only=settings.vocab_only, - verbose=settings.verbose, - ) - if settings.cache: - if settings.cache_type == "disk": - if settings.verbose: - print(f"Using disk cache with size {settings.cache_size}") - cache = llama_cpp.LlamaDiskCache(capacity_bytes=settings.cache_size) - else: - if settings.verbose: - print(f"Using ram cache with size {settings.cache_size}") - cache = llama_cpp.LlamaRAMCache(capacity_bytes=settings.cache_size) - cache = llama_cpp.LlamaCache(capacity_bytes=settings.cache_size) - llama.set_cache(cache) +def get_server_settings(): + yield _server_settings - def set_settings(_settings: Settings): - global settings - settings = _settings - - set_settings(settings) - return app +_llama_proxy: Optional[LlamaProxy] = None llama_outer_lock = Lock() llama_inner_lock = Lock() -def get_llama(): +def set_llama_proxy(model_settings: List[ModelSettings]): + global _llama_proxy + _llama_proxy = LlamaProxy(models=model_settings) + + +async def get_llama_proxy(): # NOTE: This double lock allows the currently streaming llama model to # check if any other requests are pending in the same thread and cancel # the stream if so. - llama_outer_lock.acquire() + await llama_outer_lock.acquire() release_outer_lock = True try: - llama_inner_lock.acquire() + await llama_inner_lock.acquire() try: llama_outer_lock.release() release_outer_lock = False - yield llama + yield _llama_proxy finally: llama_inner_lock.release() finally: @@ -174,395 +89,509 @@ def get_llama(): llama_outer_lock.release() -def get_settings(): - yield settings +_ping_message_factory: typing.Optional[typing.Callable[[], bytes]] = None -model_field = Field(description="The model to use for generating completions.") +def set_ping_message_factory(factory: typing.Callable[[], bytes]): + global _ping_message_factory + _ping_message_factory = factory -max_tokens_field = Field( - default=16, ge=1, le=2048, description="The maximum number of tokens to generate." -) - -temperature_field = Field( - default=0.8, - ge=0.0, - le=2.0, - description="Adjust the randomness of the generated text.\n\n" - + "Temperature is a hyperparameter that controls the randomness of the generated text. It affects the probability distribution of the model's output tokens. A higher temperature (e.g., 1.5) makes the output more random and creative, while a lower temperature (e.g., 0.5) makes the output more focused, deterministic, and conservative. The default value is 0.8, which provides a balance between randomness and determinism. At the extreme, a temperature of 0 will always pick the most likely next token, leading to identical outputs in each run.", -) - -top_p_field = Field( - default=0.95, - ge=0.0, - le=1.0, - description="Limit the next token selection to a subset of tokens with a cumulative probability above a threshold P.\n\n" - + "Top-p sampling, also known as nucleus sampling, is another text generation method that selects the next token from a subset of tokens that together have a cumulative probability of at least p. This method provides a balance between diversity and quality by considering both the probabilities of tokens and the number of tokens to sample from. A higher value for top_p (e.g., 0.95) will lead to more diverse text, while a lower value (e.g., 0.5) will generate more focused and conservative text.", -) -stop_field = Field( - default=None, - description="A list of tokens at which to stop generation. If None, no stop tokens are used.", -) +def create_app( + settings: Settings | None = None, + server_settings: ServerSettings | None = None, + model_settings: List[ModelSettings] | None = None, +): + config_file = os.environ.get("CONFIG_FILE", None) + if config_file is not None: + if not os.path.exists(config_file): + raise ValueError(f"Config file {config_file} not found!") + with open(config_file, "rb") as f: + # Check if yaml file + if config_file.endswith(".yaml") or config_file.endswith(".yml"): + import yaml + + config_file_settings = ConfigFileSettings.model_validate_json( + json.dumps(yaml.safe_load(f)) + ) + else: + config_file_settings = ConfigFileSettings.model_validate_json(f.read()) + server_settings = ServerSettings.model_validate(config_file_settings) + model_settings = config_file_settings.models + + if server_settings is None and model_settings is None: + if settings is None: + settings = Settings() + server_settings = ServerSettings.model_validate(settings) + model_settings = [ModelSettings.model_validate(settings)] + + assert server_settings is not None and model_settings is not None, ( + "server_settings and model_settings must be provided together" + ) -stream_field = Field( - default=False, - description="Whether to stream the results as they are generated. Useful for chatbots.", -) + set_server_settings(server_settings) + middleware = [Middleware(RawContextMiddleware, plugins=(RequestIdPlugin(),))] + app = FastAPI( + middleware=middleware, + title="🦙 llama.cpp Python API", + version=llama_cpp.__version__, + root_path=server_settings.root_path, + ) + app.add_middleware( + CORSMiddleware, + allow_origins=["*"], + allow_credentials=True, + allow_methods=["*"], + allow_headers=["*"], + ) + app.include_router(router) -top_k_field = Field( - default=40, - ge=0, - description="Limit the next token selection to the K most probable tokens.\n\n" - + "Top-k sampling is a text generation method that selects the next token only from the top k most likely tokens predicted by the model. It helps reduce the risk of generating low-probability or nonsensical tokens, but it may also limit the diversity of the output. A higher value for top_k (e.g., 100) will consider more tokens and lead to more diverse text, while a lower value (e.g., 10) will focus on the most probable tokens and generate more conservative text.", -) + assert model_settings is not None + set_llama_proxy(model_settings=model_settings) -repeat_penalty_field = Field( - default=1.1, - ge=0.0, - description="A penalty applied to each token that is already generated. This helps prevent the model from repeating itself.\n\n" - + "Repeat penalty is a hyperparameter used to penalize the repetition of token sequences during text generation. It helps prevent the model from generating repetitive or monotonous text. A higher value (e.g., 1.5) will penalize repetitions more strongly, while a lower value (e.g., 0.9) will be more lenient.", -) + if server_settings.disable_ping_events: + set_ping_message_factory(lambda: bytes()) -presence_penalty_field = Field( - default=0.0, - ge=-2.0, - le=2.0, - description="Positive values penalize new tokens based on whether they appear in the text so far, increasing the model's likelihood to talk about new topics.", -) + return app -frequency_penalty_field = Field( - default=0.0, - ge=-2.0, - le=2.0, - description="Positive values penalize new tokens based on their existing frequency in the text so far, decreasing the model's likelihood to repeat the same line verbatim.", -) -mirostat_mode_field = Field( - default=0, - ge=0, - le=2, - description="Enable Mirostat constant-perplexity algorithm of the specified version (1 or 2; 0 = disabled)" -) +def prepare_request_resources( + body: CreateCompletionRequest | CreateChatCompletionRequest, + llama_proxy: LlamaProxy, + body_model: str | None, + kwargs, +) -> llama_cpp.Llama: + if llama_proxy is None: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail="Service is not available", + ) + llama = llama_proxy(body_model) + if body.logit_bias is not None: + kwargs["logit_bias"] = ( + _logit_bias_tokens_to_input_ids(llama, body.logit_bias) + if body.logit_bias_type == "tokens" + else body.logit_bias + ) -mirostat_tau_field = Field( - default=5.0, - ge=0.0, - le=10.0, - description="Mirostat target entropy, i.e. the target perplexity - lower values produce focused and coherent text, larger values produce more diverse and less coherent text" -) + if body.grammar is not None: + kwargs["grammar"] = llama_cpp.LlamaGrammar.from_string(body.grammar) -mirostat_eta_field = Field( - default=0.1, - ge=0.001, - le=1.0, - description="Mirostat learning rate" -) + if body.min_tokens > 0: + _min_tokens_logits_processor = llama_cpp.LogitsProcessorList( + [llama_cpp.MinTokensLogitsProcessor(body.min_tokens, llama.token_eos())] + ) + if "logits_processor" not in kwargs: + kwargs["logits_processor"] = _min_tokens_logits_processor + else: + kwargs["logits_processor"].extend(_min_tokens_logits_processor) + return llama -class CreateCompletionRequest(BaseModel): - prompt: Union[str, List[str]] = Field( - default="", description="The prompt to generate completions for." - ) - suffix: Optional[str] = Field( - default=None, - description="A suffix to append to the generated text. If None, no suffix is appended. Useful for chatbots.", - ) - max_tokens: int = max_tokens_field - temperature: float = temperature_field - top_p: float = top_p_field - mirostat_mode: int = mirostat_mode_field - mirostat_tau: float = mirostat_tau_field - mirostat_eta: float = mirostat_eta_field - echo: bool = Field( - default=False, - description="Whether to echo the prompt in the generated text. Useful for chatbots.", - ) - stop: Optional[Union[str, List[str]]] = stop_field - stream: bool = stream_field - logprobs: Optional[int] = Field( - default=None, - ge=0, - description="The number of logprobs to generate. If None, no logprobs are generated.", +async def get_event_publisher( + request: Request, + inner_send_chan: MemoryObjectSendStream[typing.Any], + body: CreateCompletionRequest | CreateChatCompletionRequest, + body_model: str | None, + llama_call, + kwargs, +): + server_settings = next(get_server_settings()) + interrupt_requests = ( + server_settings.interrupt_requests if server_settings else False ) - presence_penalty: Optional[float] = presence_penalty_field - frequency_penalty: Optional[float] = frequency_penalty_field - logit_bias: Optional[Dict[str, float]] = Field(None) - logprobs: Optional[int] = Field(None) - - # ignored or currently unsupported - model: Optional[str] = model_field - n: Optional[int] = 1 - best_of: Optional[int] = 1 - user: Optional[str] = Field(None) - - # llama.cpp specific parameters - top_k: int = top_k_field - repeat_penalty: float = repeat_penalty_field - logit_bias_type: Optional[Literal["input_ids", "tokens"]] = Field(None) - - class Config: - schema_extra = { - "example": { - "prompt": "\n\n### Instructions:\nWhat is the capital of France?\n\n### Response:\n", - "stop": ["\n", "###"], - } - } - - - - -def make_logit_bias_processor( + async with contextlib.asynccontextmanager(get_llama_proxy)() as llama_proxy: + llama = prepare_request_resources(body, llama_proxy, body_model, kwargs) + async with inner_send_chan: + try: + iterator = await run_in_threadpool(llama_call, llama, **kwargs) + async for chunk in iterate_in_threadpool(iterator): + await inner_send_chan.send(dict(data=json.dumps(chunk))) + if await request.is_disconnected(): + raise anyio.get_cancelled_exc_class()() + if interrupt_requests and llama_outer_lock.locked(): + await inner_send_chan.send(dict(data="[DONE]")) + raise anyio.get_cancelled_exc_class()() + await inner_send_chan.send(dict(data="[DONE]")) + except anyio.get_cancelled_exc_class() as e: + print("disconnected") + with anyio.move_on_after(1, shield=True): + print( + f"Disconnected from client (via refresh/close) {request.client}" + ) + raise e + + +def _logit_bias_tokens_to_input_ids( llama: llama_cpp.Llama, logit_bias: Dict[str, float], - logit_bias_type: Optional[Literal["input_ids", "tokens"]], -): - if logit_bias_type is None: - logit_bias_type = "input_ids" +) -> Dict[str, float]: + to_bias: Dict[str, float] = {} + for token, score in logit_bias.items(): + token = token.encode("utf-8") + for input_id in llama.tokenize(token, add_bos=False, special=True): + to_bias[str(input_id)] = score + return to_bias - to_bias: Dict[int, float] = {} - if logit_bias_type == "input_ids": - for input_id, score in logit_bias.items(): - input_id = int(input_id) - to_bias[input_id] = score - elif logit_bias_type == "tokens": - for token, score in logit_bias.items(): - token = token.encode('utf-8') - for input_id in llama.tokenize(token, add_bos=False): - to_bias[input_id] = score +# Setup Bearer authentication scheme +bearer_scheme = HTTPBearer(auto_error=False) - def logit_bias_processor( - input_ids: List[int], - scores: List[float], - ) -> List[float]: - new_scores = [None] * len(scores) - for input_id, score in enumerate(scores): - new_scores[input_id] = score + to_bias.get(input_id, 0.0) - return new_scores +async def authenticate( + settings: Settings = Depends(get_server_settings), + authorization: Optional[str] = Depends(bearer_scheme), +): + # Skip API key check if it's not set in settings + if settings.api_key is None: + return True + + # check bearer credentials against the api_key + if authorization and authorization.credentials == settings.api_key: + # api key is valid + return authorization.credentials + + # raise http error 401 + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Invalid API key", + ) - return logit_bias_processor + +openai_v1_tag = "OpenAI V1" @router.post( "/v1/completions", + summary="Completion", + dependencies=[Depends(authenticate)], + response_model=Union[ + llama_cpp.CreateCompletionResponse, + str, + ], + responses={ + "200": { + "description": "Successful Response", + "content": { + "application/json": { + "schema": { + "anyOf": [ + {"$ref": "#/components/schemas/CreateCompletionResponse"} + ], + "title": "Completion response, when stream=False", + } + }, + "text/event-stream": { + "schema": { + "type": "string", + "title": "Server Side Streaming response, when stream=True. " + + "See SSE format: https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events#Event_stream_format", # noqa: E501 + "example": """data: {... see CreateCompletionResponse ...} \\n\\n data: ... \\n\\n ... data: [DONE]""", + } + }, + }, + } + }, + tags=[openai_v1_tag], +) +@router.post( + "/v1/engines/copilot-codex/completions", + include_in_schema=False, + dependencies=[Depends(authenticate)], + tags=[openai_v1_tag], ) async def create_completion( request: Request, body: CreateCompletionRequest, - llama: llama_cpp.Llama = Depends(get_llama), -): +) -> llama_cpp.Completion: if isinstance(body.prompt, list): assert len(body.prompt) <= 1 body.prompt = body.prompt[0] if len(body.prompt) > 0 else "" + body_model = ( + body.model + if request.url.path != "/v1/engines/copilot-codex/completions" + else "copilot-codex" + ) + exclude = { "n", "best_of", - "logit_bias", "logit_bias_type", "user", + "min_tokens", } - kwargs = body.dict(exclude=exclude) + kwargs = body.model_dump(exclude=exclude) - if body.logit_bias is not None: - kwargs['logits_processor'] = llama_cpp.LogitsProcessorList([ - make_logit_bias_processor(llama, body.logit_bias, body.logit_bias_type), - ]) - - if body.stream: + # handle streaming request + if kwargs.get("stream", False): send_chan, recv_chan = anyio.create_memory_object_stream(10) - - async def event_publisher(inner_send_chan: MemoryObjectSendStream): - async with inner_send_chan: - try: - iterator: Iterator[llama_cpp.CompletionChunk] = await run_in_threadpool(llama, **kwargs) # type: ignore - async for chunk in iterate_in_threadpool(iterator): - await inner_send_chan.send(dict(data=json.dumps(chunk))) - if await request.is_disconnected(): - raise anyio.get_cancelled_exc_class()() - if settings.interrupt_requests and llama_outer_lock.locked(): - await inner_send_chan.send(dict(data="[DONE]")) - raise anyio.get_cancelled_exc_class()() - await inner_send_chan.send(dict(data="[DONE]")) - except anyio.get_cancelled_exc_class() as e: - print("disconnected") - with anyio.move_on_after(1, shield=True): - print( - f"Disconnected from client (via refresh/close) {request.client}" - ) - raise e - return EventSourceResponse( - recv_chan, data_sender_callable=partial(event_publisher, send_chan) + recv_chan, + data_sender_callable=partial( # type: ignore + get_event_publisher, + request=request, + inner_send_chan=send_chan, + body=body, + body_model=body_model, + llama_call=llama_cpp.Llama.__call__, + kwargs=kwargs, + ), + sep="\n", + ping_message_factory=_ping_message_factory, ) - else: - completion: llama_cpp.Completion = await run_in_threadpool(llama, **kwargs) # type: ignore - return completion - -class CreateEmbeddingRequest(BaseModel): - model: Optional[str] = model_field - input: Union[str, List[str]] = Field(description="The input to embed.") - user: Optional[str] - - class Config: - schema_extra = { - "example": { - "input": "The food was delicious and the waiter...", - } - } + # handle regular request + async with contextlib.asynccontextmanager(get_llama_proxy)() as llama_proxy: + llama = prepare_request_resources(body, llama_proxy, body_model, kwargs) + if await request.is_disconnected(): + print( + f"Disconnected from client (via refresh/close) before llm invoked {request.client}" + ) + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Client closed request", + ) + return await run_in_threadpool(llama, **kwargs) @router.post( "/v1/embeddings", + summary="Embedding", + dependencies=[Depends(authenticate)], + tags=[openai_v1_tag], ) async def create_embedding( - request: CreateEmbeddingRequest, llama: llama_cpp.Llama = Depends(get_llama) + request: CreateEmbeddingRequest, + llama_proxy: LlamaProxy = Depends(get_llama_proxy), ): return await run_in_threadpool( - llama.create_embedding, **request.dict(exclude={"user"}) - ) - - -class ChatCompletionRequestMessage(BaseModel): - role: Literal["system", "user", "assistant"] = Field( - default="user", description="The role of the message." + llama_proxy(request.model).create_embedding, + **request.model_dump(exclude={"user"}), ) - content: str = Field(default="", description="The content of the message.") - - -class CreateChatCompletionRequest(BaseModel): - messages: List[ChatCompletionRequestMessage] = Field( - default=[], description="A list of messages to generate completions for." - ) - max_tokens: int = max_tokens_field - temperature: float = temperature_field - top_p: float = top_p_field - mirostat_mode: int = mirostat_mode_field - mirostat_tau: float = mirostat_tau_field - mirostat_eta: float = mirostat_eta_field - stop: Optional[List[str]] = stop_field - stream: bool = stream_field - presence_penalty: Optional[float] = presence_penalty_field - frequency_penalty: Optional[float] = frequency_penalty_field - logit_bias: Optional[Dict[str, float]] = Field(None) - - # ignored or currently unsupported - model: Optional[str] = model_field - n: Optional[int] = 1 - user: Optional[str] = Field(None) - - # llama.cpp specific parameters - top_k: int = top_k_field - repeat_penalty: float = repeat_penalty_field - logit_bias_type: Optional[Literal["input_ids", "tokens"]] = Field(None) - - class Config: - schema_extra = { - "example": { - "messages": [ - ChatCompletionRequestMessage( - role="system", content="You are a helpful assistant." - ), - ChatCompletionRequestMessage( - role="user", content="What is the capital of France?" - ), - ] - } - } - - @router.post( "/v1/chat/completions", + summary="Chat", + dependencies=[Depends(authenticate)], + response_model=Union[llama_cpp.ChatCompletion, str], + responses={ + "200": { + "description": "Successful Response", + "content": { + "application/json": { + "schema": { + "anyOf": [ + { + "$ref": "#/components/schemas/CreateChatCompletionResponse" + } + ], + "title": "Completion response, when stream=False", + } + }, + "text/event-stream": { + "schema": { + "type": "string", + "title": "Server Side Streaming response, when stream=True" + + "See SSE format: https://developer.mozilla.org/en-US/docs/Web/API/Server-sent_events/Using_server-sent_events#Event_stream_format", # noqa: E501 + "example": """data: {... see CreateChatCompletionResponse ...} \\n\\n data: ... \\n\\n ... data: [DONE]""", + } + }, + }, + } + }, + tags=[openai_v1_tag], ) async def create_chat_completion( request: Request, - body: CreateChatCompletionRequest, - llama: llama_cpp.Llama = Depends(get_llama), - settings: Settings = Depends(get_settings), -) -> Union[llama_cpp.ChatCompletion]: # type: ignore + body: CreateChatCompletionRequest = Body( + openapi_examples={ + "normal": { + "summary": "Chat Completion", + "value": { + "model": "gpt-3.5-turbo", + "messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "What is the capital of France?"}, + ], + }, + }, + "json_mode": { + "summary": "JSON Mode", + "value": { + "model": "gpt-3.5-turbo", + "messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Who won the world series in 2020"}, + ], + "response_format": {"type": "json_object"}, + }, + }, + "tool_calling": { + "summary": "Tool Calling", + "value": { + "model": "gpt-3.5-turbo", + "messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Extract Jason is 30 years old."}, + ], + "tools": [ + { + "type": "function", + "function": { + "name": "User", + "description": "User record", + "parameters": { + "type": "object", + "properties": { + "name": {"type": "string"}, + "age": {"type": "number"}, + }, + "required": ["name", "age"], + }, + }, + } + ], + "tool_choice": { + "type": "function", + "function": { + "name": "User", + }, + }, + }, + }, + "logprobs": { + "summary": "Logprobs", + "value": { + "model": "gpt-3.5-turbo", + "messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "What is the capital of France?"}, + ], + "logprobs": True, + "top_logprobs": 10, + }, + }, + } + ), +) -> llama_cpp.ChatCompletion: + # This is a workaround for an issue in FastAPI dependencies + # where the dependency is cleaned up before a StreamingResponse + # is complete. + # https://github.com/tiangolo/fastapi/issues/11143 + + body_model = body.model exclude = { "n", - "logit_bias", "logit_bias_type", "user", + "min_tokens", } - kwargs = body.dict(exclude=exclude) + kwargs = body.model_dump(exclude=exclude) - if body.logit_bias is not None: - kwargs['logits_processor'] = llama_cpp.LogitsProcessorList([ - make_logit_bias_processor(llama, body.logit_bias, body.logit_bias_type), - ]) - - if body.stream: + # handle streaming request + if kwargs.get("stream", False): send_chan, recv_chan = anyio.create_memory_object_stream(10) - - async def event_publisher(inner_send_chan: MemoryObjectSendStream): - async with inner_send_chan: - try: - iterator: Iterator[llama_cpp.ChatCompletionChunk] = await run_in_threadpool(llama.create_chat_completion, **kwargs) # type: ignore - async for chat_chunk in iterate_in_threadpool(iterator): - await inner_send_chan.send(dict(data=json.dumps(chat_chunk))) - if await request.is_disconnected(): - raise anyio.get_cancelled_exc_class()() - if settings.interrupt_requests and llama_outer_lock.locked(): - await inner_send_chan.send(dict(data="[DONE]")) - raise anyio.get_cancelled_exc_class()() - await inner_send_chan.send(dict(data="[DONE]")) - except anyio.get_cancelled_exc_class() as e: - print("disconnected") - with anyio.move_on_after(1, shield=True): - print( - f"Disconnected from client (via refresh/close) {request.client}" - ) - raise e - return EventSourceResponse( recv_chan, - data_sender_callable=partial(event_publisher, send_chan), - ) - else: - completion: llama_cpp.ChatCompletion = await run_in_threadpool( - llama.create_chat_completion, **kwargs # type: ignore + data_sender_callable=partial( # type: ignore + get_event_publisher, + request=request, + inner_send_chan=send_chan, + body=body, + body_model=body_model, + llama_call=llama_cpp.Llama.create_chat_completion, + kwargs=kwargs, + ), + sep="\n", + ping_message_factory=_ping_message_factory, ) - return completion + # handle regular request + async with contextlib.asynccontextmanager(get_llama_proxy)() as llama_proxy: + llama = prepare_request_resources(body, llama_proxy, body_model, kwargs) -class ModelData(TypedDict): - id: str - object: Literal["model"] - owned_by: str - permissions: List[str] + if await request.is_disconnected(): + print( + f"Disconnected from client (via refresh/close) before llm invoked {request.client}" + ) + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Client closed request", + ) + return await run_in_threadpool(llama.create_chat_completion, **kwargs) -class ModelList(TypedDict): - object: Literal["list"] - data: List[ModelData] - - - -@router.get("/v1/models") +@router.get( + "/v1/models", + summary="Models", + dependencies=[Depends(authenticate)], + tags=[openai_v1_tag], +) async def get_models( - settings: Settings = Depends(get_settings), + llama_proxy: LlamaProxy = Depends(get_llama_proxy), ) -> ModelList: - assert llama is not None return { "object": "list", "data": [ { - "id": settings.model_alias - if settings.model_alias is not None - else llama.model_path, + "id": model_alias, "object": "model", "owned_by": "me", "permissions": [], } + for model_alias in llama_proxy ], } + + +extras_tag = "Extras" + + +@router.post( + "/extras/tokenize", + summary="Tokenize", + dependencies=[Depends(authenticate)], + tags=[extras_tag], +) +async def tokenize( + body: TokenizeInputRequest, + llama_proxy: LlamaProxy = Depends(get_llama_proxy), +) -> TokenizeInputResponse: + tokens = llama_proxy(body.model).tokenize(body.input.encode("utf-8"), special=True) + + return TokenizeInputResponse(tokens=tokens) + + +@router.post( + "/extras/tokenize/count", + summary="Tokenize Count", + dependencies=[Depends(authenticate)], + tags=[extras_tag], +) +async def count_query_tokens( + body: TokenizeInputRequest, + llama_proxy: LlamaProxy = Depends(get_llama_proxy), +) -> TokenizeInputCountResponse: + tokens = llama_proxy(body.model).tokenize(body.input.encode("utf-8"), special=True) + + return TokenizeInputCountResponse(count=len(tokens)) + + +@router.post( + "/extras/detokenize", + summary="Detokenize", + dependencies=[Depends(authenticate)], + tags=[extras_tag], +) +async def detokenize( + body: DetokenizeInputRequest, + llama_proxy: LlamaProxy = Depends(get_llama_proxy), +) -> DetokenizeInputResponse: + text = llama_proxy(body.model).detokenize(body.tokens).decode("utf-8") + + return DetokenizeInputResponse(text=text) diff --git a/llama_cpp/server/cli.py b/llama_cpp/server/cli.py new file mode 100644 index 0000000000..171b8db30c --- /dev/null +++ b/llama_cpp/server/cli.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +import argparse +import json + +from typing import List, Literal, Union, Any, Type, TypeVar, Dict + +from pydantic import BaseModel + + +def _get_base_type(annotation: Type[Any]) -> Type[Any]: + if getattr(annotation, "__origin__", None) is Literal: + assert hasattr(annotation, "__args__") and len(annotation.__args__) >= 1 # type: ignore + return type(annotation.__args__[0]) # type: ignore + elif getattr(annotation, "__origin__", None) is Union: + assert hasattr(annotation, "__args__") and len(annotation.__args__) >= 1 # type: ignore + non_optional_args: List[Type[Any]] = [ + arg + for arg in annotation.__args__ + if arg is not type(None) # type: ignore + ] + if non_optional_args: + return _get_base_type(non_optional_args[0]) + elif ( + getattr(annotation, "__origin__", None) is list + or getattr(annotation, "__origin__", None) is List + ): + assert hasattr(annotation, "__args__") and len(annotation.__args__) >= 1 # type: ignore + return _get_base_type(annotation.__args__[0]) # type: ignore + return annotation + + +def _contains_list_type(annotation: Type[Any] | None) -> bool: + origin = getattr(annotation, "__origin__", None) + + if origin is list or origin is List: + return True + elif origin in (Literal, Union): + return any(_contains_list_type(arg) for arg in annotation.__args__) # type: ignore + else: + return False + + +def _contains_dict_type(annotation: Type[Any] | None) -> bool: + origin = getattr(annotation, "__origin__", None) + + if origin is dict or origin is Dict: + return True + elif origin in (Literal, Union): + return any(_contains_dict_type(arg) for arg in annotation.__args__) # type: ignore + else: + return False + + +def _parse_bool_arg(arg: str | bytes | bool) -> bool: + if isinstance(arg, bytes): + arg = arg.decode("utf-8") + + true_values = {"1", "on", "t", "true", "y", "yes"} + false_values = {"0", "off", "f", "false", "n", "no"} + + arg_str = str(arg).lower().strip() + + if arg_str in true_values: + return True + elif arg_str in false_values: + return False + else: + raise ValueError(f"Invalid boolean argument: {arg}") + + +def _parse_json_object_arg(arg: str | bytes) -> dict[str, Any]: + if isinstance(arg, bytes): + arg = arg.decode("utf-8") + + value = json.loads(arg) + if not isinstance(value, dict): + raise ValueError(f"Invalid JSON object argument: {arg}") + return value + + +def add_args_from_model(parser: argparse.ArgumentParser, model: Type[BaseModel]): + """Add arguments from a pydantic model to an argparse parser.""" + + for name, field in model.model_fields.items(): + description = field.description + if field.default and description and not field.is_required(): + description += f" (default: {field.default})" + base_type = ( + _get_base_type(field.annotation) if field.annotation is not None else str + ) + list_type = _contains_list_type(field.annotation) + dict_type = _contains_dict_type(field.annotation) + if dict_type: + parser.add_argument( + f"--{name}", + dest=name, + type=_parse_json_object_arg, + help=description, + ) + elif base_type is not bool: + parser.add_argument( + f"--{name}", + dest=name, + nargs="*" if list_type else None, + type=base_type, + help=description, + ) + if base_type is bool: + parser.add_argument( + f"--{name}", + dest=name, + type=_parse_bool_arg, + help=f"{description}", + ) + + +T = TypeVar("T", bound=Type[BaseModel]) + + +def parse_model_from_args(model: T, args: argparse.Namespace) -> T: + """Parse a pydantic model from an argparse namespace.""" + return model( + **{ + k: v + for k, v in vars(args).items() + if v is not None and k in model.model_fields + } + ) diff --git a/llama_cpp/server/errors.py b/llama_cpp/server/errors.py new file mode 100644 index 0000000000..d0eda5664b --- /dev/null +++ b/llama_cpp/server/errors.py @@ -0,0 +1,212 @@ +from __future__ import annotations + +import sys +import traceback +import time +from re import compile, Match, Pattern +from typing import Callable, Coroutine, Optional, Tuple, Union, Dict +from typing_extensions import TypedDict + + +from fastapi import ( + Request, + Response, + HTTPException, +) +from fastapi.responses import JSONResponse +from fastapi.routing import APIRoute + +from llama_cpp.server.types import ( + CreateCompletionRequest, + CreateEmbeddingRequest, + CreateChatCompletionRequest, +) + + +class ErrorResponse(TypedDict): + """OpenAI style error response""" + + message: str + type: str + param: Optional[str] + code: Optional[str] + + +class ErrorResponseFormatters: + """Collection of formatters for error responses. + + Args: + request (Union[CreateCompletionRequest, CreateChatCompletionRequest]): + Request body + match (Match[str]): Match object from regex pattern + + Returns: + Tuple[int, ErrorResponse]: Status code and error response + """ + + @staticmethod + def context_length_exceeded( + request: Union["CreateCompletionRequest", "CreateChatCompletionRequest"], + match, # type: Match[str] # type: ignore + ) -> Tuple[int, ErrorResponse]: + """Formatter for context length exceeded error""" + + context_window = int(match.group(2)) + prompt_tokens = int(match.group(1)) + completion_tokens = request.max_tokens + if hasattr(request, "messages"): + # Chat completion + message = ( + "This model's maximum context length is {} tokens. " + "However, you requested {} tokens " + "({} in the messages, {} in the completion). " + "Please reduce the length of the messages or completion." + ) + else: + # Text completion + message = ( + "This model's maximum context length is {} tokens, " + "however you requested {} tokens " + "({} in your prompt; {} for the completion). " + "Please reduce your prompt; or completion length." + ) + return 400, ErrorResponse( + message=message.format( + context_window, + (completion_tokens or 0) + prompt_tokens, + prompt_tokens, + completion_tokens, + ), # type: ignore + type="invalid_request_error", + param="messages", + code="context_length_exceeded", + ) + + @staticmethod + def model_not_found( + request: Union["CreateCompletionRequest", "CreateChatCompletionRequest"], + match, # type: Match[str] # type: ignore + ) -> Tuple[int, ErrorResponse]: + """Formatter for model_not_found error""" + + model_path = str(match.group(1)) + message = f"The model `{model_path}` does not exist" + return 400, ErrorResponse( + message=message, + type="invalid_request_error", + param=None, + code="model_not_found", + ) + + +class RouteErrorHandler(APIRoute): + """Custom APIRoute that handles application errors and exceptions""" + + # key: regex pattern for original error message from llama_cpp + # value: formatter function + pattern_and_formatters: Dict[ + "Pattern[str]", + Callable[ + [ + Union["CreateCompletionRequest", "CreateChatCompletionRequest"], + "Match[str]", + ], + Tuple[int, ErrorResponse], + ], + ] = { + compile( + r"Requested tokens \((\d+)\) exceed context window of (\d+)" + ): ErrorResponseFormatters.context_length_exceeded, + compile( + r"Model path does not exist: (.+)" + ): ErrorResponseFormatters.model_not_found, + } + + def error_message_wrapper( + self, + error: Exception, + body: Optional[ + Union[ + "CreateChatCompletionRequest", + "CreateCompletionRequest", + "CreateEmbeddingRequest", + ] + ] = None, + ) -> Tuple[int, ErrorResponse]: + """Wraps error message in OpenAI style error response""" + if body is not None and isinstance( + body, + ( + CreateCompletionRequest, + CreateChatCompletionRequest, + ), + ): + # When text completion or chat completion + for pattern, callback in self.pattern_and_formatters.items(): + match = pattern.search(str(error)) + if match is not None: + return callback(body, match) + + # Only print the trace on unexpected exceptions + print(f"Exception: {str(error)}", file=sys.stderr) + traceback.print_exc(file=sys.stderr) + + # Wrap other errors as internal server error + return 500, ErrorResponse( + message=str(error), + type="internal_server_error", + param=None, + code=None, + ) + + def get_route_handler( + self, + ) -> Callable[[Request], Coroutine[None, None, Response]]: + """Defines custom route handler that catches exceptions and formats + in OpenAI style error response""" + + original_route_handler = super().get_route_handler() + + async def custom_route_handler(request: Request) -> Response: + try: + start_sec = time.perf_counter() + response = await original_route_handler(request) + elapsed_time_ms = int((time.perf_counter() - start_sec) * 1000) + response.headers["openai-processing-ms"] = f"{elapsed_time_ms}" + return response + except HTTPException as unauthorized: + # api key check failed + raise unauthorized + except Exception as exc: + json_body = await request.json() + try: + if "messages" in json_body: + # Chat completion + body: Optional[ + Union[ + CreateChatCompletionRequest, + CreateCompletionRequest, + CreateEmbeddingRequest, + ] + ] = CreateChatCompletionRequest(**json_body) + elif "prompt" in json_body: + # Text completion + body = CreateCompletionRequest(**json_body) + else: + # Embedding + body = CreateEmbeddingRequest(**json_body) + except Exception: + # Invalid request body + body = None + + # Get proper error message from the exception + ( + status_code, + error_message, + ) = self.error_message_wrapper(error=exc, body=body) + return JSONResponse( + {"error": error_message}, + status_code=status_code, + ) + + return custom_route_handler diff --git a/llama_cpp/server/model.py b/llama_cpp/server/model.py new file mode 100644 index 0000000000..8aa929202c --- /dev/null +++ b/llama_cpp/server/model.py @@ -0,0 +1,344 @@ +from __future__ import annotations + +import json + +from typing import Dict, Optional, Union, List + +import llama_cpp +import llama_cpp.llama_speculative as llama_speculative +import llama_cpp.llama_tokenizer as llama_tokenizer + +from llama_cpp.server.settings import ModelSettings + + +class LlamaProxy: + def __init__(self, models: List[ModelSettings]) -> None: + assert len(models) > 0, "No models provided!" + + self._model_settings_dict: dict[str, ModelSettings] = {} + for model in models: + if not model.model_alias: + model.model_alias = model.model + self._model_settings_dict[model.model_alias] = model + + self._current_model: Optional[llama_cpp.Llama] = None + self._current_model_alias: Optional[str] = None + + self._default_model_settings: ModelSettings = models[0] + self._default_model_alias: str = self._default_model_settings.model_alias # type: ignore + + # Load default model + self._current_model = self.load_llama_from_model_settings( + self._default_model_settings + ) + self._current_model_alias = self._default_model_alias + + def __call__(self, model: Optional[str] = None) -> llama_cpp.Llama: + if model is None: + model = self._default_model_alias + + if model not in self._model_settings_dict: + model = self._default_model_alias + + if model == self._current_model_alias: + if self._current_model is not None: + return self._current_model + + if self._current_model: + self._current_model.close() + self._current_model = None + + settings = self._model_settings_dict[model] + self._current_model = self.load_llama_from_model_settings(settings) + self._current_model_alias = model + return self._current_model + + def __getitem__(self, model: str): + return self._model_settings_dict[model].model_dump() + + def __setitem__(self, model: str, settings: Union[ModelSettings, str, bytes]): + if isinstance(settings, (bytes, str)): + settings = ModelSettings.model_validate_json(settings) + self._model_settings_dict[model] = settings + + def __iter__(self): + for model in self._model_settings_dict: + yield model + + def free(self): + if self._current_model: + self._current_model.close() + del self._current_model + + @staticmethod + def load_llama_from_model_settings(settings: ModelSettings) -> llama_cpp.Llama: + chat_handler = None + if settings.chat_format == "llava-1-5": + assert settings.clip_model_path is not None, "clip model not found" + if settings.hf_model_repo_id is not None: + chat_handler = ( + llama_cpp.llama_chat_format.Llava15ChatHandler.from_pretrained( + repo_id=settings.hf_model_repo_id, + filename=settings.clip_model_path, + verbose=settings.verbose, + ) + ) + else: + chat_handler = llama_cpp.llama_chat_format.Llava15ChatHandler( + clip_model_path=settings.clip_model_path, verbose=settings.verbose + ) + elif settings.chat_format == "obsidian": + assert settings.clip_model_path is not None, "clip model not found" + if settings.hf_model_repo_id is not None: + chat_handler = ( + llama_cpp.llama_chat_format.ObsidianChatHandler.from_pretrained( + repo_id=settings.hf_model_repo_id, + filename=settings.clip_model_path, + verbose=settings.verbose, + ) + ) + else: + chat_handler = llama_cpp.llama_chat_format.ObsidianChatHandler( + clip_model_path=settings.clip_model_path, verbose=settings.verbose + ) + elif settings.chat_format == "llava-1-6": + assert settings.clip_model_path is not None, "clip model not found" + if settings.hf_model_repo_id is not None: + chat_handler = ( + llama_cpp.llama_chat_format.Llava16ChatHandler.from_pretrained( + repo_id=settings.hf_model_repo_id, + filename=settings.clip_model_path, + verbose=settings.verbose, + ) + ) + else: + chat_handler = llama_cpp.llama_chat_format.Llava16ChatHandler( + clip_model_path=settings.clip_model_path, verbose=settings.verbose + ) + elif settings.chat_format in ("mtmd", "gemma4"): + assert settings.clip_model_path is not None, "clip model not found" + chat_handler_cls = ( + llama_cpp.llama_chat_format.MTMDChatHandler + if settings.chat_format == "mtmd" + else llama_cpp.llama_chat_format.Gemma4ChatHandler + ) + if settings.hf_model_repo_id is not None: + chat_handler = chat_handler_cls.from_pretrained( + repo_id=settings.hf_model_repo_id, + filename=settings.clip_model_path, + verbose=settings.verbose, + ) + else: + chat_handler = chat_handler_cls( + clip_model_path=settings.clip_model_path, verbose=settings.verbose + ) + elif settings.chat_format == "moondream": + assert settings.clip_model_path is not None, "clip model not found" + if settings.hf_model_repo_id is not None: + chat_handler = ( + llama_cpp.llama_chat_format.MoondreamChatHandler.from_pretrained( + repo_id=settings.hf_model_repo_id, + filename=settings.clip_model_path, + verbose=settings.verbose, + ) + ) + else: + chat_handler = llama_cpp.llama_chat_format.MoondreamChatHandler( + clip_model_path=settings.clip_model_path, verbose=settings.verbose + ) + elif settings.chat_format == "nanollava": + assert settings.clip_model_path is not None, "clip model not found" + if settings.hf_model_repo_id is not None: + chat_handler = ( + llama_cpp.llama_chat_format.NanoLlavaChatHandler.from_pretrained( + repo_id=settings.hf_model_repo_id, + filename=settings.clip_model_path, + verbose=settings.verbose, + ) + ) + else: + chat_handler = llama_cpp.llama_chat_format.NanoLlavaChatHandler( + clip_model_path=settings.clip_model_path, verbose=settings.verbose + ) + elif settings.chat_format == "llama-3-vision-alpha": + assert settings.clip_model_path is not None, "clip model not found" + if settings.hf_model_repo_id is not None: + chat_handler = ( + llama_cpp.llama_chat_format.Llama3VisionAlpha.from_pretrained( + repo_id=settings.hf_model_repo_id, + filename=settings.clip_model_path, + verbose=settings.verbose, + ) + ) + else: + chat_handler = llama_cpp.llama_chat_format.Llama3VisionAlpha( + clip_model_path=settings.clip_model_path, verbose=settings.verbose + ) + elif settings.chat_format == "minicpm-v-2.6": + assert settings.clip_model_path is not None, "clip model not found" + if settings.hf_model_repo_id is not None: + chat_handler = ( + llama_cpp.llama_chat_format.MiniCPMv26ChatHandler.from_pretrained( + repo_id=settings.hf_model_repo_id, + filename=settings.clip_model_path, + verbose=settings.verbose, + ) + ) + else: + chat_handler = llama_cpp.llama_chat_format.MiniCPMv26ChatHandler( + clip_model_path=settings.clip_model_path, verbose=settings.verbose + ) + elif settings.chat_format == "qwen2.5-vl": + assert settings.clip_model_path is not None, "clip model not found" + if settings.hf_model_repo_id is not None: + chat_handler = ( + llama_cpp.llama_chat_format.Qwen25VLChatHandler.from_pretrained( + repo_id=settings.hf_model_repo_id, + filename=settings.clip_model_path, + verbose=settings.verbose, + ) + ) + else: + chat_handler = llama_cpp.llama_chat_format.Qwen25VLChatHandler( + clip_model_path=settings.clip_model_path, verbose=settings.verbose + ) + elif settings.chat_format == "hf-autotokenizer": + assert settings.hf_pretrained_model_name_or_path is not None, ( + "hf_pretrained_model_name_or_path must be set for hf-autotokenizer" + ) + chat_handler = ( + llama_cpp.llama_chat_format.hf_autotokenizer_to_chat_completion_handler( + settings.hf_pretrained_model_name_or_path + ) + ) + elif settings.chat_format == "hf-tokenizer-config": + assert settings.hf_tokenizer_config_path is not None, ( + "hf_tokenizer_config_path must be set for hf-tokenizer-config" + ) + chat_handler = llama_cpp.llama_chat_format.hf_tokenizer_config_to_chat_completion_handler( + json.load(open(settings.hf_tokenizer_config_path)) + ) + + tokenizer: Optional[llama_cpp.BaseLlamaTokenizer] = None + if settings.hf_pretrained_model_name_or_path is not None: + tokenizer = llama_tokenizer.LlamaHFTokenizer.from_pretrained( + settings.hf_pretrained_model_name_or_path + ) + + draft_model = None + if settings.draft_model is not None: + draft_model = llama_speculative.LlamaPromptLookupDecoding( + num_pred_tokens=settings.draft_model_num_pred_tokens + ) + + kv_overrides: Optional[Dict[str, Union[bool, int, float, str]]] = None + if settings.kv_overrides is not None: + assert isinstance(settings.kv_overrides, list) + kv_overrides = {} + for kv in settings.kv_overrides: + key, value = kv.split("=") + if ":" in value: + value_type, value = value.split(":") + if value_type == "bool": + kv_overrides[key] = value.lower() in ["true", "1"] + elif value_type == "int": + kv_overrides[key] = int(value) + elif value_type == "float": + kv_overrides[key] = float(value) + elif value_type == "str": + kv_overrides[key] = value + else: + raise ValueError(f"Unknown value type {value_type}") + + import functools + + kwargs = {} + + if settings.hf_model_repo_id is not None: + create_fn = functools.partial( + llama_cpp.Llama.from_pretrained, + repo_id=settings.hf_model_repo_id, + filename=settings.model, + ) + else: + create_fn = llama_cpp.Llama + kwargs["model_path"] = settings.model + + _model = create_fn( + **kwargs, + # Model Params + n_gpu_layers=settings.n_gpu_layers, + split_mode=settings.split_mode, + main_gpu=settings.main_gpu, + tensor_split=settings.tensor_split, + vocab_only=settings.vocab_only, + use_mmap=settings.use_mmap, + use_mlock=settings.use_mlock, + kv_overrides=kv_overrides, + rpc_servers=settings.rpc_servers, + # Context Params + seed=settings.seed, + n_ctx=settings.n_ctx, + n_batch=settings.n_batch, + n_ubatch=settings.n_ubatch, + n_threads=settings.n_threads, + n_threads_batch=settings.n_threads_batch, + rope_scaling_type=settings.rope_scaling_type, + rope_freq_base=settings.rope_freq_base, + rope_freq_scale=settings.rope_freq_scale, + yarn_ext_factor=settings.yarn_ext_factor, + yarn_attn_factor=settings.yarn_attn_factor, + yarn_beta_fast=settings.yarn_beta_fast, + yarn_beta_slow=settings.yarn_beta_slow, + yarn_orig_ctx=settings.yarn_orig_ctx, + mul_mat_q=settings.mul_mat_q, + logits_all=settings.logits_all, + embedding=settings.embedding, + offload_kqv=settings.offload_kqv, + flash_attn=settings.flash_attn, + # Sampling Params + last_n_tokens_size=settings.last_n_tokens_size, + # LoRA Params + lora_base=settings.lora_base, + lora_path=settings.lora_path, + # Backend Params + numa=settings.numa, + # Chat Format Params + chat_format=settings.chat_format, + chat_handler=chat_handler, + # Speculative Decoding + draft_model=draft_model, + # KV Cache Quantization + type_k=settings.type_k, + type_v=settings.type_v, + # Tokenizer + tokenizer=tokenizer, + # Misc + verbose=settings.verbose, + ) + if settings.chat_template_kwargs: + base_chat_handler = ( + _model.chat_handler + or _model._chat_handlers.get(_model.chat_format) + or llama_cpp.llama_chat_format.get_chat_completion_handler( + _model.chat_format + ) + ) + + def chat_handler_with_kwargs(*args, **kwargs): + return base_chat_handler( + *args, **{**settings.chat_template_kwargs, **kwargs} + ) + + _model.chat_handler = chat_handler_with_kwargs + if settings.cache: + if settings.cache_type == "disk": + if settings.verbose: + print(f"Using disk cache with size {settings.cache_size}") + cache = llama_cpp.LlamaDiskCache(capacity_bytes=settings.cache_size) + else: + if settings.verbose: + print(f"Using ram cache with size {settings.cache_size}") + cache = llama_cpp.LlamaRAMCache(capacity_bytes=settings.cache_size) + _model.set_cache(cache) + return _model diff --git a/llama_cpp/server/settings.py b/llama_cpp/server/settings.py new file mode 100644 index 0000000000..78dd7cdeb8 --- /dev/null +++ b/llama_cpp/server/settings.py @@ -0,0 +1,244 @@ +from __future__ import annotations + +import multiprocessing + +from typing import Any, Optional, List, Literal, Union, Dict, cast +from typing_extensions import Self + +from pydantic import Field, model_validator +from pydantic_settings import BaseSettings + +import llama_cpp + +# Disable warning for model and model_alias settings +BaseSettings.model_config["protected_namespaces"] = () + + +class ModelSettings(BaseSettings): + """Model settings used to load a Llama model.""" + + model: str = Field( + description="The path to the model to use for generating completions." + ) + model_alias: Optional[str] = Field( + default=None, + description="The alias of the model to use for generating completions.", + ) + # Model Params + n_gpu_layers: int = Field( + default=0, + ge=-1, + description="The number of layers to put on the GPU. The rest will be on the CPU. Set -1 to move all to GPU.", + ) + split_mode: int = Field( + default=llama_cpp.LLAMA_SPLIT_MODE_LAYER, + description="The split mode to use.", + ) + main_gpu: int = Field( + default=0, + ge=0, + description="Main GPU to use.", + ) + tensor_split: Optional[List[float]] = Field( + default=None, + description="Split layers across multiple GPUs in proportion.", + ) + vocab_only: bool = Field( + default=False, description="Whether to only return the vocabulary." + ) + use_mmap: bool = Field( + default=llama_cpp.llama_supports_mmap(), + description="Use mmap.", + ) + use_mlock: bool = Field( + default=llama_cpp.llama_supports_mlock(), + description="Use mlock.", + ) + kv_overrides: Optional[List[str]] = Field( + default=None, + description="List of model kv overrides in the format key=type:value where type is one of (bool, int, float). Valid true values are (true, TRUE, 1), otherwise false.", + ) + rpc_servers: Optional[str] = Field( + default=None, + description="comma separated list of rpc servers for offloading", + ) + # Context Params + seed: int = Field( + default=llama_cpp.LLAMA_DEFAULT_SEED, description="Random seed. -1 for random." + ) + n_ctx: int = Field(default=2048, ge=0, description="The context size.") + n_batch: int = Field( + default=512, ge=1, description="The batch size to use per eval." + ) + n_ubatch: int = Field( + default=512, ge=1, description="The physical batch size used by llama.cpp" + ) + n_threads: int = Field( + default=max(multiprocessing.cpu_count() // 2, 1), + ge=1, + description="The number of threads to use. Use -1 for max cpu threads", + ) + n_threads_batch: int = Field( + default=max(multiprocessing.cpu_count(), 1), + ge=0, + description="The number of threads to use when batch processing. Use -1 for max cpu threads", + ) + rope_scaling_type: int = Field( + default=llama_cpp.LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED + ) + rope_freq_base: float = Field(default=0.0, description="RoPE base frequency") + rope_freq_scale: float = Field( + default=0.0, description="RoPE frequency scaling factor" + ) + yarn_ext_factor: float = Field(default=-1.0) + yarn_attn_factor: float = Field(default=1.0) + yarn_beta_fast: float = Field(default=32.0) + yarn_beta_slow: float = Field(default=1.0) + yarn_orig_ctx: int = Field(default=0) + mul_mat_q: bool = Field( + default=True, description="if true, use experimental mul_mat_q kernels" + ) + logits_all: bool = Field(default=True, description="Whether to return logits.") + embedding: bool = Field(default=False, description="Whether to use embeddings.") + offload_kqv: bool = Field( + default=True, description="Whether to offload kqv to the GPU." + ) + flash_attn: bool = Field( + default=False, description="Whether to use flash attention." + ) + # Sampling Params + last_n_tokens_size: int = Field( + default=64, + ge=0, + description="Last n tokens to keep for repeat penalty calculation.", + ) + # LoRA Params + lora_base: Optional[str] = Field( + default=None, + description="Optional path to base model, useful if using a quantized base model and you want to apply LoRA to an f16 model.", + ) + lora_path: Optional[str] = Field( + default=None, + description="Path to a LoRA file to apply to the model.", + ) + # Backend Params + numa: Union[bool, int] = Field( + default=False, + description="Enable NUMA support.", + ) + # Chat Format Params + chat_format: Optional[str] = Field( + default=None, + description="Chat format to use.", + ) + chat_template_kwargs: Optional[Dict[str, Any]] = Field( + default=None, + description="Extra keyword arguments forwarded to chat templates at model load time. Matches llama.cpp server `chat_template_kwargs`.", + ) + clip_model_path: Optional[str] = Field( + default=None, + description="Path to a CLIP model to use for multi-modal chat completion.", + ) + # Cache Params + cache: bool = Field( + default=False, + description="Use a cache to reduce processing times for evaluated prompts.", + ) + cache_type: Literal["ram", "disk"] = Field( + default="ram", + description="The type of cache to use. Only used if cache is True.", + ) + cache_size: int = Field( + default=2 << 30, + description="The size of the cache in bytes. Only used if cache is True.", + ) + # Tokenizer Options + hf_tokenizer_config_path: Optional[str] = Field( + default=None, + description="The path to a HuggingFace tokenizer_config.json file.", + ) + hf_pretrained_model_name_or_path: Optional[str] = Field( + default=None, + description="The model name or path to a pretrained HuggingFace tokenizer model. Same as you would pass to AutoTokenizer.from_pretrained().", + ) + # Loading from HuggingFace Model Hub + hf_model_repo_id: Optional[str] = Field( + default=None, + description="The model repo id to use for the HuggingFace tokenizer model.", + ) + # Speculative Decoding + draft_model: Optional[str] = Field( + default=None, + description="Method to use for speculative decoding. One of (prompt-lookup-decoding).", + ) + draft_model_num_pred_tokens: int = Field( + default=10, + description="Number of tokens to predict using the draft model.", + ) + # KV Cache Quantization + type_k: Optional[int] = Field( + default=None, + description="Type of the key cache quantization.", + ) + type_v: Optional[int] = Field( + default=None, + description="Type of the value cache quantization.", + ) + # Misc + verbose: bool = Field( + default=True, description="Whether to print debug information." + ) + + @model_validator( + mode="before" + ) # pre=True to ensure this runs before any other validation + def set_dynamic_defaults(self) -> Self: + # If n_threads or n_threads_batch is -1, set it to multiprocessing.cpu_count() + cpu_count = multiprocessing.cpu_count() + values = cast(Dict[str, int], self) + if values.get("n_threads", 0) == -1: + values["n_threads"] = cpu_count + if values.get("n_threads_batch", 0) == -1: + values["n_threads_batch"] = cpu_count + return self + + +class ServerSettings(BaseSettings): + """Server settings used to configure the FastAPI and Uvicorn server.""" + + # Uvicorn Settings + host: str = Field(default="localhost", description="Listen address") + port: int = Field(default=8000, description="Listen port") + ssl_keyfile: Optional[str] = Field( + default=None, description="SSL key file for HTTPS" + ) + ssl_certfile: Optional[str] = Field( + default=None, description="SSL certificate file for HTTPS" + ) + # FastAPI Settings + api_key: Optional[str] = Field( + default=None, + description="API key for authentication. If set all requests need to be authenticated.", + ) + interrupt_requests: bool = Field( + default=True, + description="Whether to interrupt requests when a new request is received.", + ) + disable_ping_events: bool = Field( + default=False, + description="Disable EventSource pings (may be needed for some clients).", + ) + root_path: str = Field( + default="", + description="The root path for the server. Useful when running behind a reverse proxy.", + ) + + +class Settings(ServerSettings, ModelSettings): + pass + + +class ConfigFileSettings(ServerSettings): + """Configuration file format settings.""" + + models: List[ModelSettings] = Field(default=[], description="Model configs") diff --git a/llama_cpp/server/types.py b/llama_cpp/server/types.py new file mode 100644 index 0000000000..eda2b92801 --- /dev/null +++ b/llama_cpp/server/types.py @@ -0,0 +1,316 @@ +from __future__ import annotations + +from typing import List, Optional, Union, Dict +from typing_extensions import TypedDict, Literal + +from pydantic import BaseModel, Field + +import llama_cpp + + +model_field = Field( + description="The model to use for generating completions.", default=None +) + +max_tokens_field = Field( + default=16, ge=1, description="The maximum number of tokens to generate." +) + +min_tokens_field = Field( + default=0, + ge=0, + description="The minimum number of tokens to generate. It may return fewer tokens if another condition is met (e.g. max_tokens, stop).", +) + +temperature_field = Field( + default=0.8, + description="Adjust the randomness of the generated text.\n\n" + + "Temperature is a hyperparameter that controls the randomness of the generated text. It affects the probability distribution of the model's output tokens. A higher temperature (e.g., 1.5) makes the output more random and creative, while a lower temperature (e.g., 0.5) makes the output more focused, deterministic, and conservative. The default value is 0.8, which provides a balance between randomness and determinism. At the extreme, a temperature of 0 will always pick the most likely next token, leading to identical outputs in each run.", +) + +top_p_field = Field( + default=0.95, + ge=0.0, + le=1.0, + description="Limit the next token selection to a subset of tokens with a cumulative probability above a threshold P.\n\n" + + "Top-p sampling, also known as nucleus sampling, is another text generation method that selects the next token from a subset of tokens that together have a cumulative probability of at least p. This method provides a balance between diversity and quality by considering both the probabilities of tokens and the number of tokens to sample from. A higher value for top_p (e.g., 0.95) will lead to more diverse text, while a lower value (e.g., 0.5) will generate more focused and conservative text.", +) + +min_p_field = Field( + default=0.05, + ge=0.0, + le=1.0, + description="Sets a minimum base probability threshold for token selection.\n\n" + + "The Min-P sampling method was designed as an alternative to Top-P, and aims to ensure a balance of quality and variety. The parameter min_p represents the minimum probability for a token to be considered, relative to the probability of the most likely token. For example, with min_p=0.05 and the most likely token having a probability of 0.9, logits with a value less than 0.045 are filtered out.", +) + +stop_field = Field( + default=None, + description="A list of tokens at which to stop generation. If None, no stop tokens are used.", +) + +stream_field = Field( + default=False, + description="Whether to stream the results as they are generated. Useful for chatbots.", +) + +top_k_field = Field( + default=40, + ge=0, + description="Limit the next token selection to the K most probable tokens.\n\n" + + "Top-k sampling is a text generation method that selects the next token only from the top k most likely tokens predicted by the model. It helps reduce the risk of generating low-probability or nonsensical tokens, but it may also limit the diversity of the output. A higher value for top_k (e.g., 100) will consider more tokens and lead to more diverse text, while a lower value (e.g., 10) will focus on the most probable tokens and generate more conservative text.", +) + +repeat_penalty_field = Field( + default=1.1, + ge=0.0, + description="A penalty applied to each token that is already generated. This helps prevent the model from repeating itself.\n\n" + + "Repeat penalty is a hyperparameter used to penalize the repetition of token sequences during text generation. It helps prevent the model from generating repetitive or monotonous text. A higher value (e.g., 1.5) will penalize repetitions more strongly, while a lower value (e.g., 0.9) will be more lenient.", +) + +presence_penalty_field = Field( + default=0.0, + ge=-2.0, + le=2.0, + description="Positive values penalize new tokens based on whether they appear in the text so far, increasing the model's likelihood to talk about new topics.", +) + +frequency_penalty_field = Field( + default=0.0, + ge=-2.0, + le=2.0, + description="Positive values penalize new tokens based on their existing frequency in the text so far, decreasing the model's likelihood to repeat the same line verbatim.", +) + +mirostat_mode_field = Field( + default=0, + ge=0, + le=2, + description="Enable Mirostat constant-perplexity algorithm of the specified version (1 or 2; 0 = disabled)", +) + +mirostat_tau_field = Field( + default=5.0, + ge=0.0, + le=10.0, + description="Mirostat target entropy, i.e. the target perplexity - lower values produce focused and coherent text, larger values produce more diverse and less coherent text", +) + +mirostat_eta_field = Field( + default=0.1, ge=0.001, le=1.0, description="Mirostat learning rate" +) + +grammar = Field( + default=None, + description="A CBNF grammar (as string) to be used for formatting the model's output.", +) + + +class CreateCompletionRequest(BaseModel): + model: Optional[str] = model_field + prompt: Union[str, List[str]] = Field( + default="", description="The prompt to generate completions for." + ) + suffix: Optional[str] = Field( + default=None, + description="A suffix to append to the generated text. If None, no suffix is appended. Useful for chatbots.", + ) + max_tokens: Optional[int] = Field( + default=16, ge=0, description="The maximum number of tokens to generate." + ) + min_tokens: int = min_tokens_field + temperature: float = temperature_field + top_p: float = top_p_field + min_p: float = min_p_field + echo: bool = Field( + default=False, + description="Whether to echo the prompt in the generated text. Useful for chatbots.", + ) + stop: Optional[Union[str, List[str]]] = stop_field + stream: bool = stream_field + logprobs: Optional[int] = Field( + default=None, + ge=0, + description="The number of logprobs to generate. If None, no logprobs are generated.", + ) + presence_penalty: Optional[float] = presence_penalty_field + frequency_penalty: Optional[float] = frequency_penalty_field + logit_bias: Optional[Dict[str, float]] = Field(None) + seed: Optional[int] = Field(None) + + # ignored or currently unsupported + n: Optional[int] = 1 + best_of: Optional[int] = 1 + user: Optional[str] = Field(default=None) + + # llama.cpp specific parameters + top_k: int = top_k_field + repeat_penalty: float = repeat_penalty_field + logit_bias_type: Optional[Literal["input_ids", "tokens"]] = Field(None) + mirostat_mode: int = mirostat_mode_field + mirostat_tau: float = mirostat_tau_field + mirostat_eta: float = mirostat_eta_field + grammar: Optional[str] = None + + model_config = { + "json_schema_extra": { + "examples": [ + { + "prompt": "\n\n### Instructions:\nWhat is the capital of France?\n\n### Response:\n", + "stop": ["\n", "###"], + } + ] + } + } + + +class CreateEmbeddingRequest(BaseModel): + model: Optional[str] = model_field + input: Union[str, List[str]] = Field(description="The input to embed.") + user: Optional[str] = Field(default=None) + + model_config = { + "json_schema_extra": { + "examples": [ + { + "input": "The food was delicious and the waiter...", + } + ] + } + } + + +class ChatCompletionRequestMessage(BaseModel): + role: Literal["system", "user", "assistant", "function"] = Field( + default="user", description="The role of the message." + ) + content: Optional[str] = Field( + default="", description="The content of the message." + ) + + +class CreateChatCompletionRequest(BaseModel): + model: Optional[str] = model_field + messages: List[llama_cpp.ChatCompletionRequestMessage] = Field( + default=[], description="A list of messages to generate completions for." + ) + functions: Optional[List[llama_cpp.ChatCompletionFunction]] = Field( + default=None, + description="A list of functions to apply to the generated completions.", + ) + function_call: Optional[llama_cpp.ChatCompletionRequestFunctionCall] = Field( + default=None, + description="A function to apply to the generated completions.", + ) + tools: Optional[List[llama_cpp.ChatCompletionTool]] = Field( + default=None, + description="A list of tools to apply to the generated completions.", + ) + tool_choice: Optional[llama_cpp.ChatCompletionToolChoiceOption] = Field( + default=None, + description="A tool to apply to the generated completions.", + ) # TODO: verify + max_tokens: Optional[int] = Field( + default=None, + description="The maximum number of tokens to generate. Defaults to inf", + ) + min_tokens: int = min_tokens_field + logprobs: Optional[bool] = Field( + default=False, + description="Whether to output the logprobs or not. Default is True", + ) + top_logprobs: Optional[int] = Field( + default=None, + ge=0, + description="The number of logprobs to generate. If None, no logprobs are generated. logprobs need to set to True.", + ) + temperature: float = temperature_field + top_p: float = top_p_field + min_p: float = min_p_field + stop: Optional[Union[str, List[str]]] = stop_field + stream: bool = stream_field + presence_penalty: Optional[float] = presence_penalty_field + frequency_penalty: Optional[float] = frequency_penalty_field + logit_bias: Optional[Dict[str, float]] = Field(None) + seed: Optional[int] = Field(None) + response_format: Optional[llama_cpp.ChatCompletionRequestResponseFormat] = Field( + default=None, + ) + + # ignored or currently unsupported + n: Optional[int] = 1 + user: Optional[str] = Field(None) + + # llama.cpp specific parameters + top_k: int = top_k_field + repeat_penalty: float = repeat_penalty_field + logit_bias_type: Optional[Literal["input_ids", "tokens"]] = Field(None) + mirostat_mode: int = mirostat_mode_field + mirostat_tau: float = mirostat_tau_field + mirostat_eta: float = mirostat_eta_field + grammar: Optional[str] = None + + model_config = { + "json_schema_extra": { + "examples": [ + { + "messages": [ + ChatCompletionRequestMessage( + role="system", content="You are a helpful assistant." + ).model_dump(), + ChatCompletionRequestMessage( + role="user", content="What is the capital of France?" + ).model_dump(), + ] + } + ] + } + } + + +class ModelData(TypedDict): + id: str + object: Literal["model"] + owned_by: str + permissions: List[str] + + +class ModelList(TypedDict): + object: Literal["list"] + data: List[ModelData] + + +class TokenizeInputRequest(BaseModel): + model: Optional[str] = model_field + input: str = Field(description="The input to tokenize.") + + model_config = { + "json_schema_extra": {"examples": [{"input": "How many tokens in this query?"}]} + } + + +class TokenizeInputResponse(BaseModel): + tokens: List[int] = Field(description="A list of tokens.") + + model_config = {"json_schema_extra": {"example": {"tokens": [123, 321, 222]}}} + + +class TokenizeInputCountResponse(BaseModel): + count: int = Field(description="The number of tokens in the input.") + + model_config = {"json_schema_extra": {"example": {"count": 5}}} + + +class DetokenizeInputRequest(BaseModel): + model: Optional[str] = model_field + tokens: List[int] = Field(description="A list of toekns to detokenize.") + + model_config = {"json_schema_extra": {"example": [{"tokens": [123, 321, 222]}]}} + + +class DetokenizeInputResponse(BaseModel): + text: str = Field(description="The detokenized text.") + + model_config = { + "json_schema_extra": {"example": {"text": "How many tokens in this query?"}} + } diff --git a/mkdocs.yml b/mkdocs.yml index 2865811765..37e1002e8d 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -2,20 +2,73 @@ site_name: llama-cpp-python repo_url: https://github.com/abetlen/llama-cpp-python theme: - name: "material" + name: material + palette: + + # Palette toggle for light mode + - scheme: default + primary: indigo + toggle: + icon: material/brightness-7 + name: Switch to dark mode + + # Palette toggle for dark mode + - scheme: slate + primary: indigo + toggle: + icon: material/brightness-4 + name: Switch to light mode plugins: - - mkdocstrings - search + - mkdocstrings: + handlers: + python: + options: + members_order: source + group_by_category: false + signature_crossrefs: true + show_signature: true + docstring_section_style: list + show_root_heading: true + heading_level: 3 + preload_modules: + - typing + - typing_extensions + - ctypes + inventories: + - https://docs.python.org/3/objects.inv + - https://numpy.org/doc/stable/objects.inv watch: - llama_cpp + - README.md + +nav: + - "Getting Started": "index.md" + - "Installation Guides": + - "macOS (Metal)": "install/macos.md" + - "API Reference": "api-reference.md" + - "OpenAI Compatible Web Server": "server.md" + - "Changelog": "changelog.md" markdown_extensions: + - attr_list + - pymdownx.emoji: + emoji_index: !!python/name:materialx.emoji.twemoji + emoji_generator: !!python/name:materialx.emoji.to_svg - pymdownx.highlight: anchor_linenums: true line_spans: __span pygments_lang_class: true - pymdownx.inlinehilite + - pymdownx.magiclink: + repo_url_shorthand: true + user: abetlen + repo: llama-cpp-python - pymdownx.snippets - - pymdownx.superfences \ No newline at end of file + - pymdownx.superfences + - pymdownx.tabbed: + alternate_style: true + - pymdownx.tilde + - tables diff --git a/poetry.lock b/poetry.lock deleted file mode 100644 index 9d129663ce..0000000000 --- a/poetry.lock +++ /dev/null @@ -1,1636 +0,0 @@ -# This file is automatically @generated by Poetry 1.5.1 and should not be changed by hand. - -[[package]] -name = "anyio" -version = "3.6.2" -description = "High level compatibility layer for multiple asynchronous event loop implementations" -optional = false -python-versions = ">=3.6.2" -files = [ - {file = "anyio-3.6.2-py3-none-any.whl", hash = "sha256:fbbe32bd270d2a2ef3ed1c5d45041250284e31fc0a4df4a5a6071842051a51e3"}, - {file = "anyio-3.6.2.tar.gz", hash = "sha256:25ea0d673ae30af41a0c442f81cf3b38c7e79fdc7b60335a4c14e05eb0947421"}, -] - -[package.dependencies] -idna = ">=2.8" -sniffio = ">=1.1" - -[package.extras] -doc = ["packaging", "sphinx-autodoc-typehints (>=1.2.0)", "sphinx-rtd-theme"] -test = ["contextlib2", "coverage[toml] (>=4.5)", "hypothesis (>=4.0)", "mock (>=4)", "pytest (>=7.0)", "pytest-mock (>=3.6.1)", "trustme", "uvloop (<0.15)", "uvloop (>=0.15)"] -trio = ["trio (>=0.16,<0.22)"] - -[[package]] -name = "black" -version = "23.3.0" -description = "The uncompromising code formatter." -optional = false -python-versions = ">=3.7" -files = [ - {file = "black-23.3.0-cp310-cp310-macosx_10_16_arm64.whl", hash = "sha256:0945e13506be58bf7db93ee5853243eb368ace1c08a24c65ce108986eac65915"}, - {file = "black-23.3.0-cp310-cp310-macosx_10_16_universal2.whl", hash = "sha256:67de8d0c209eb5b330cce2469503de11bca4085880d62f1628bd9972cc3366b9"}, - {file = "black-23.3.0-cp310-cp310-macosx_10_16_x86_64.whl", hash = "sha256:7c3eb7cea23904399866c55826b31c1f55bbcd3890ce22ff70466b907b6775c2"}, - {file = "black-23.3.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:32daa9783106c28815d05b724238e30718f34155653d4d6e125dc7daec8e260c"}, - {file = "black-23.3.0-cp310-cp310-win_amd64.whl", hash = "sha256:35d1381d7a22cc5b2be2f72c7dfdae4072a3336060635718cc7e1ede24221d6c"}, - {file = "black-23.3.0-cp311-cp311-macosx_10_16_arm64.whl", hash = "sha256:a8a968125d0a6a404842fa1bf0b349a568634f856aa08ffaff40ae0dfa52e7c6"}, - {file = "black-23.3.0-cp311-cp311-macosx_10_16_universal2.whl", hash = "sha256:c7ab5790333c448903c4b721b59c0d80b11fe5e9803d8703e84dcb8da56fec1b"}, - {file = "black-23.3.0-cp311-cp311-macosx_10_16_x86_64.whl", hash = "sha256:a6f6886c9869d4daae2d1715ce34a19bbc4b95006d20ed785ca00fa03cba312d"}, - {file = "black-23.3.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6f3c333ea1dd6771b2d3777482429864f8e258899f6ff05826c3a4fcc5ce3f70"}, - {file = "black-23.3.0-cp311-cp311-win_amd64.whl", hash = "sha256:11c410f71b876f961d1de77b9699ad19f939094c3a677323f43d7a29855fe326"}, - {file = "black-23.3.0-cp37-cp37m-macosx_10_16_x86_64.whl", hash = "sha256:1d06691f1eb8de91cd1b322f21e3bfc9efe0c7ca1f0e1eb1db44ea367dff656b"}, - {file = "black-23.3.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:50cb33cac881766a5cd9913e10ff75b1e8eb71babf4c7104f2e9c52da1fb7de2"}, - {file = "black-23.3.0-cp37-cp37m-win_amd64.whl", hash = "sha256:e114420bf26b90d4b9daa597351337762b63039752bdf72bf361364c1aa05925"}, - {file = "black-23.3.0-cp38-cp38-macosx_10_16_arm64.whl", hash = "sha256:48f9d345675bb7fbc3dd85821b12487e1b9a75242028adad0333ce36ed2a6d27"}, - {file = "black-23.3.0-cp38-cp38-macosx_10_16_universal2.whl", hash = "sha256:714290490c18fb0126baa0fca0a54ee795f7502b44177e1ce7624ba1c00f2331"}, - {file = "black-23.3.0-cp38-cp38-macosx_10_16_x86_64.whl", hash = "sha256:064101748afa12ad2291c2b91c960be28b817c0c7eaa35bec09cc63aa56493c5"}, - {file = "black-23.3.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:562bd3a70495facf56814293149e51aa1be9931567474993c7942ff7d3533961"}, - {file = "black-23.3.0-cp38-cp38-win_amd64.whl", hash = "sha256:e198cf27888ad6f4ff331ca1c48ffc038848ea9f031a3b40ba36aced7e22f2c8"}, - {file = "black-23.3.0-cp39-cp39-macosx_10_16_arm64.whl", hash = "sha256:3238f2aacf827d18d26db07524e44741233ae09a584273aa059066d644ca7b30"}, - {file = "black-23.3.0-cp39-cp39-macosx_10_16_universal2.whl", hash = "sha256:f0bd2f4a58d6666500542b26354978218a9babcdc972722f4bf90779524515f3"}, - {file = "black-23.3.0-cp39-cp39-macosx_10_16_x86_64.whl", hash = "sha256:92c543f6854c28a3c7f39f4d9b7694f9a6eb9d3c5e2ece488c327b6e7ea9b266"}, - {file = "black-23.3.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3a150542a204124ed00683f0db1f5cf1c2aaaa9cc3495b7a3b5976fb136090ab"}, - {file = "black-23.3.0-cp39-cp39-win_amd64.whl", hash = "sha256:6b39abdfb402002b8a7d030ccc85cf5afff64ee90fa4c5aebc531e3ad0175ddb"}, - {file = "black-23.3.0-py3-none-any.whl", hash = "sha256:ec751418022185b0c1bb7d7736e6933d40bbb14c14a0abcf9123d1b159f98dd4"}, - {file = "black-23.3.0.tar.gz", hash = "sha256:1c7b8d606e728a41ea1ccbd7264677e494e87cf630e399262ced92d4a8dac940"}, -] - -[package.dependencies] -click = ">=8.0.0" -mypy-extensions = ">=0.4.3" -packaging = ">=22.0" -pathspec = ">=0.9.0" -platformdirs = ">=2" -tomli = {version = ">=1.1.0", markers = "python_version < \"3.11\""} -typing-extensions = {version = ">=3.10.0.0", markers = "python_version < \"3.10\""} - -[package.extras] -colorama = ["colorama (>=0.4.3)"] -d = ["aiohttp (>=3.7.4)"] -jupyter = ["ipython (>=7.8.0)", "tokenize-rt (>=3.2.0)"] -uvloop = ["uvloop (>=0.15.2)"] - -[[package]] -name = "bleach" -version = "6.0.0" -description = "An easy safelist-based HTML-sanitizing tool." -optional = false -python-versions = ">=3.7" -files = [ - {file = "bleach-6.0.0-py3-none-any.whl", hash = "sha256:33c16e3353dbd13028ab4799a0f89a83f113405c766e9c122df8a06f5b85b3f4"}, - {file = "bleach-6.0.0.tar.gz", hash = "sha256:1a1a85c1595e07d8db14c5f09f09e6433502c51c595970edc090551f0db99414"}, -] - -[package.dependencies] -six = ">=1.9.0" -webencodings = "*" - -[package.extras] -css = ["tinycss2 (>=1.1.0,<1.2)"] - -[[package]] -name = "certifi" -version = "2023.5.7" -description = "Python package for providing Mozilla's CA Bundle." -optional = false -python-versions = ">=3.6" -files = [ - {file = "certifi-2023.5.7-py3-none-any.whl", hash = "sha256:c6c2e98f5c7869efca1f8916fed228dd91539f9f1b444c314c06eef02980c716"}, - {file = "certifi-2023.5.7.tar.gz", hash = "sha256:0f0d56dc5a6ad56fd4ba36484d6cc34451e1c6548c61daad8c320169f91eddc7"}, -] - -[[package]] -name = "cffi" -version = "1.15.1" -description = "Foreign Function Interface for Python calling C code." -optional = false -python-versions = "*" -files = [ - {file = "cffi-1.15.1-cp27-cp27m-macosx_10_9_x86_64.whl", hash = "sha256:a66d3508133af6e8548451b25058d5812812ec3798c886bf38ed24a98216fab2"}, - {file = "cffi-1.15.1-cp27-cp27m-manylinux1_i686.whl", hash = "sha256:470c103ae716238bbe698d67ad020e1db9d9dba34fa5a899b5e21577e6d52ed2"}, - {file = "cffi-1.15.1-cp27-cp27m-manylinux1_x86_64.whl", hash = "sha256:9ad5db27f9cabae298d151c85cf2bad1d359a1b9c686a275df03385758e2f914"}, - {file = "cffi-1.15.1-cp27-cp27m-win32.whl", hash = "sha256:b3bbeb01c2b273cca1e1e0c5df57f12dce9a4dd331b4fa1635b8bec26350bde3"}, - {file = "cffi-1.15.1-cp27-cp27m-win_amd64.whl", hash = "sha256:e00b098126fd45523dd056d2efba6c5a63b71ffe9f2bbe1a4fe1716e1d0c331e"}, - {file = "cffi-1.15.1-cp27-cp27mu-manylinux1_i686.whl", hash = "sha256:d61f4695e6c866a23a21acab0509af1cdfd2c013cf256bbf5b6b5e2695827162"}, - {file = "cffi-1.15.1-cp27-cp27mu-manylinux1_x86_64.whl", hash = "sha256:ed9cb427ba5504c1dc15ede7d516b84757c3e3d7868ccc85121d9310d27eed0b"}, - {file = "cffi-1.15.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:39d39875251ca8f612b6f33e6b1195af86d1b3e60086068be9cc053aa4376e21"}, - {file = "cffi-1.15.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:285d29981935eb726a4399badae8f0ffdff4f5050eaa6d0cfc3f64b857b77185"}, - {file = "cffi-1.15.1-cp310-cp310-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:3eb6971dcff08619f8d91607cfc726518b6fa2a9eba42856be181c6d0d9515fd"}, - {file = "cffi-1.15.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:21157295583fe8943475029ed5abdcf71eb3911894724e360acff1d61c1d54bc"}, - {file = "cffi-1.15.1-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:5635bd9cb9731e6d4a1132a498dd34f764034a8ce60cef4f5319c0541159392f"}, - {file = "cffi-1.15.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:2012c72d854c2d03e45d06ae57f40d78e5770d252f195b93f581acf3ba44496e"}, - {file = "cffi-1.15.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:dd86c085fae2efd48ac91dd7ccffcfc0571387fe1193d33b6394db7ef31fe2a4"}, - {file = "cffi-1.15.1-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:fa6693661a4c91757f4412306191b6dc88c1703f780c8234035eac011922bc01"}, - {file = "cffi-1.15.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:59c0b02d0a6c384d453fece7566d1c7e6b7bae4fc5874ef2ef46d56776d61c9e"}, - {file = "cffi-1.15.1-cp310-cp310-win32.whl", hash = "sha256:cba9d6b9a7d64d4bd46167096fc9d2f835e25d7e4c121fb2ddfc6528fb0413b2"}, - {file = "cffi-1.15.1-cp310-cp310-win_amd64.whl", hash = "sha256:ce4bcc037df4fc5e3d184794f27bdaab018943698f4ca31630bc7f84a7b69c6d"}, - {file = "cffi-1.15.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:3d08afd128ddaa624a48cf2b859afef385b720bb4b43df214f85616922e6a5ac"}, - {file = "cffi-1.15.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:3799aecf2e17cf585d977b780ce79ff0dc9b78d799fc694221ce814c2c19db83"}, - {file = "cffi-1.15.1-cp311-cp311-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:a591fe9e525846e4d154205572a029f653ada1a78b93697f3b5a8f1f2bc055b9"}, - {file = "cffi-1.15.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3548db281cd7d2561c9ad9984681c95f7b0e38881201e157833a2342c30d5e8c"}, - {file = "cffi-1.15.1-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:91fc98adde3d7881af9b59ed0294046f3806221863722ba7d8d120c575314325"}, - {file = "cffi-1.15.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:94411f22c3985acaec6f83c6df553f2dbe17b698cc7f8ae751ff2237d96b9e3c"}, - {file = "cffi-1.15.1-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:03425bdae262c76aad70202debd780501fabeaca237cdfddc008987c0e0f59ef"}, - {file = "cffi-1.15.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:cc4d65aeeaa04136a12677d3dd0b1c0c94dc43abac5860ab33cceb42b801c1e8"}, - {file = "cffi-1.15.1-cp311-cp311-win32.whl", hash = "sha256:a0f100c8912c114ff53e1202d0078b425bee3649ae34d7b070e9697f93c5d52d"}, - {file = "cffi-1.15.1-cp311-cp311-win_amd64.whl", hash = "sha256:04ed324bda3cda42b9b695d51bb7d54b680b9719cfab04227cdd1e04e5de3104"}, - {file = "cffi-1.15.1-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:50a74364d85fd319352182ef59c5c790484a336f6db772c1a9231f1c3ed0cbd7"}, - {file = "cffi-1.15.1-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e263d77ee3dd201c3a142934a086a4450861778baaeeb45db4591ef65550b0a6"}, - {file = "cffi-1.15.1-cp36-cp36m-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:cec7d9412a9102bdc577382c3929b337320c4c4c4849f2c5cdd14d7368c5562d"}, - {file = "cffi-1.15.1-cp36-cp36m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:4289fc34b2f5316fbb762d75362931e351941fa95fa18789191b33fc4cf9504a"}, - {file = "cffi-1.15.1-cp36-cp36m-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:173379135477dc8cac4bc58f45db08ab45d228b3363adb7af79436135d028405"}, - {file = "cffi-1.15.1-cp36-cp36m-manylinux_2_5_x86_64.manylinux1_x86_64.whl", hash = "sha256:6975a3fac6bc83c4a65c9f9fcab9e47019a11d3d2cf7f3c0d03431bf145a941e"}, - {file = "cffi-1.15.1-cp36-cp36m-win32.whl", hash = "sha256:2470043b93ff09bf8fb1d46d1cb756ce6132c54826661a32d4e4d132e1977adf"}, - {file = "cffi-1.15.1-cp36-cp36m-win_amd64.whl", hash = "sha256:30d78fbc8ebf9c92c9b7823ee18eb92f2e6ef79b45ac84db507f52fbe3ec4497"}, - {file = "cffi-1.15.1-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:198caafb44239b60e252492445da556afafc7d1e3ab7a1fb3f0584ef6d742375"}, - {file = "cffi-1.15.1-cp37-cp37m-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:5ef34d190326c3b1f822a5b7a45f6c4535e2f47ed06fec77d3d799c450b2651e"}, - {file = "cffi-1.15.1-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8102eaf27e1e448db915d08afa8b41d6c7ca7a04b7d73af6514df10a3e74bd82"}, - {file = "cffi-1.15.1-cp37-cp37m-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:5df2768244d19ab7f60546d0c7c63ce1581f7af8b5de3eb3004b9b6fc8a9f84b"}, - {file = "cffi-1.15.1-cp37-cp37m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a8c4917bd7ad33e8eb21e9a5bbba979b49d9a97acb3a803092cbc1133e20343c"}, - {file = "cffi-1.15.1-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0e2642fe3142e4cc4af0799748233ad6da94c62a8bec3a6648bf8ee68b1c7426"}, - {file = "cffi-1.15.1-cp37-cp37m-win32.whl", hash = "sha256:e229a521186c75c8ad9490854fd8bbdd9a0c9aa3a524326b55be83b54d4e0ad9"}, - {file = "cffi-1.15.1-cp37-cp37m-win_amd64.whl", hash = "sha256:a0b71b1b8fbf2b96e41c4d990244165e2c9be83d54962a9a1d118fd8657d2045"}, - {file = "cffi-1.15.1-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:320dab6e7cb2eacdf0e658569d2575c4dad258c0fcc794f46215e1e39f90f2c3"}, - {file = "cffi-1.15.1-cp38-cp38-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1e74c6b51a9ed6589199c787bf5f9875612ca4a8a0785fb2d4a84429badaf22a"}, - {file = "cffi-1.15.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a5c84c68147988265e60416b57fc83425a78058853509c1b0629c180094904a5"}, - {file = "cffi-1.15.1-cp38-cp38-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:3b926aa83d1edb5aa5b427b4053dc420ec295a08e40911296b9eb1b6170f6cca"}, - {file = "cffi-1.15.1-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:87c450779d0914f2861b8526e035c5e6da0a3199d8f1add1a665e1cbc6fc6d02"}, - {file = "cffi-1.15.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4f2c9f67e9821cad2e5f480bc8d83b8742896f1242dba247911072d4fa94c192"}, - {file = "cffi-1.15.1-cp38-cp38-win32.whl", hash = "sha256:8b7ee99e510d7b66cdb6c593f21c043c248537a32e0bedf02e01e9553a172314"}, - {file = "cffi-1.15.1-cp38-cp38-win_amd64.whl", hash = "sha256:00a9ed42e88df81ffae7a8ab6d9356b371399b91dbdf0c3cb1e84c03a13aceb5"}, - {file = "cffi-1.15.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:54a2db7b78338edd780e7ef7f9f6c442500fb0d41a5a4ea24fff1c929d5af585"}, - {file = "cffi-1.15.1-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:fcd131dd944808b5bdb38e6f5b53013c5aa4f334c5cad0c72742f6eba4b73db0"}, - {file = "cffi-1.15.1-cp39-cp39-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:7473e861101c9e72452f9bf8acb984947aa1661a7704553a9f6e4baa5ba64415"}, - {file = "cffi-1.15.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6c9a799e985904922a4d207a94eae35c78ebae90e128f0c4e521ce339396be9d"}, - {file = "cffi-1.15.1-cp39-cp39-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:3bcde07039e586f91b45c88f8583ea7cf7a0770df3a1649627bf598332cb6984"}, - {file = "cffi-1.15.1-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:33ab79603146aace82c2427da5ca6e58f2b3f2fb5da893ceac0c42218a40be35"}, - {file = "cffi-1.15.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5d598b938678ebf3c67377cdd45e09d431369c3b1a5b331058c338e201f12b27"}, - {file = "cffi-1.15.1-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:db0fbb9c62743ce59a9ff687eb5f4afbe77e5e8403d6697f7446e5f609976f76"}, - {file = "cffi-1.15.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:98d85c6a2bef81588d9227dde12db8a7f47f639f4a17c9ae08e773aa9c697bf3"}, - {file = "cffi-1.15.1-cp39-cp39-win32.whl", hash = "sha256:40f4774f5a9d4f5e344f31a32b5096977b5d48560c5592e2f3d2c4374bd543ee"}, - {file = "cffi-1.15.1-cp39-cp39-win_amd64.whl", hash = "sha256:70df4e3b545a17496c9b3f41f5115e69a4f2e77e94e1d2a8e1070bc0c38c8a3c"}, - {file = "cffi-1.15.1.tar.gz", hash = "sha256:d400bfb9a37b1351253cb402671cea7e89bdecc294e8016a707f6d1d8ac934f9"}, -] - -[package.dependencies] -pycparser = "*" - -[[package]] -name = "charset-normalizer" -version = "3.1.0" -description = "The Real First Universal Charset Detector. Open, modern and actively maintained alternative to Chardet." -optional = false -python-versions = ">=3.7.0" -files = [ - {file = "charset-normalizer-3.1.0.tar.gz", hash = "sha256:34e0a2f9c370eb95597aae63bf85eb5e96826d81e3dcf88b8886012906f509b5"}, - {file = "charset_normalizer-3.1.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:e0ac8959c929593fee38da1c2b64ee9778733cdf03c482c9ff1d508b6b593b2b"}, - {file = "charset_normalizer-3.1.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:d7fc3fca01da18fbabe4625d64bb612b533533ed10045a2ac3dd194bfa656b60"}, - {file = "charset_normalizer-3.1.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:04eefcee095f58eaabe6dc3cc2262f3bcd776d2c67005880894f447b3f2cb9c1"}, - {file = "charset_normalizer-3.1.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:20064ead0717cf9a73a6d1e779b23d149b53daf971169289ed2ed43a71e8d3b0"}, - {file = "charset_normalizer-3.1.0-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1435ae15108b1cb6fffbcea2af3d468683b7afed0169ad718451f8db5d1aff6f"}, - {file = "charset_normalizer-3.1.0-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:c84132a54c750fda57729d1e2599bb598f5fa0344085dbde5003ba429a4798c0"}, - {file = "charset_normalizer-3.1.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:75f2568b4189dda1c567339b48cba4ac7384accb9c2a7ed655cd86b04055c795"}, - {file = "charset_normalizer-3.1.0-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:11d3bcb7be35e7b1bba2c23beedac81ee893ac9871d0ba79effc7fc01167db6c"}, - {file = "charset_normalizer-3.1.0-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:891cf9b48776b5c61c700b55a598621fdb7b1e301a550365571e9624f270c203"}, - {file = "charset_normalizer-3.1.0-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:5f008525e02908b20e04707a4f704cd286d94718f48bb33edddc7d7b584dddc1"}, - {file = "charset_normalizer-3.1.0-cp310-cp310-musllinux_1_1_ppc64le.whl", hash = "sha256:b06f0d3bf045158d2fb8837c5785fe9ff9b8c93358be64461a1089f5da983137"}, - {file = "charset_normalizer-3.1.0-cp310-cp310-musllinux_1_1_s390x.whl", hash = "sha256:49919f8400b5e49e961f320c735388ee686a62327e773fa5b3ce6721f7e785ce"}, - {file = "charset_normalizer-3.1.0-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:22908891a380d50738e1f978667536f6c6b526a2064156203d418f4856d6e86a"}, - {file = "charset_normalizer-3.1.0-cp310-cp310-win32.whl", hash = "sha256:12d1a39aa6b8c6f6248bb54550efcc1c38ce0d8096a146638fd4738e42284448"}, - {file = "charset_normalizer-3.1.0-cp310-cp310-win_amd64.whl", hash = "sha256:65ed923f84a6844de5fd29726b888e58c62820e0769b76565480e1fdc3d062f8"}, - {file = "charset_normalizer-3.1.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:9a3267620866c9d17b959a84dd0bd2d45719b817245e49371ead79ed4f710d19"}, - {file = "charset_normalizer-3.1.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:6734e606355834f13445b6adc38b53c0fd45f1a56a9ba06c2058f86893ae8017"}, - {file = "charset_normalizer-3.1.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:f8303414c7b03f794347ad062c0516cee0e15f7a612abd0ce1e25caf6ceb47df"}, - {file = "charset_normalizer-3.1.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:aaf53a6cebad0eae578f062c7d462155eada9c172bd8c4d250b8c1d8eb7f916a"}, - {file = "charset_normalizer-3.1.0-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:3dc5b6a8ecfdc5748a7e429782598e4f17ef378e3e272eeb1340ea57c9109f41"}, - {file = "charset_normalizer-3.1.0-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:e1b25e3ad6c909f398df8921780d6a3d120d8c09466720226fc621605b6f92b1"}, - {file = "charset_normalizer-3.1.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0ca564606d2caafb0abe6d1b5311c2649e8071eb241b2d64e75a0d0065107e62"}, - {file = "charset_normalizer-3.1.0-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:b82fab78e0b1329e183a65260581de4375f619167478dddab510c6c6fb04d9b6"}, - {file = "charset_normalizer-3.1.0-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:bd7163182133c0c7701b25e604cf1611c0d87712e56e88e7ee5d72deab3e76b5"}, - {file = "charset_normalizer-3.1.0-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:11d117e6c63e8f495412d37e7dc2e2fff09c34b2d09dbe2bee3c6229577818be"}, - {file = "charset_normalizer-3.1.0-cp311-cp311-musllinux_1_1_ppc64le.whl", hash = "sha256:cf6511efa4801b9b38dc5546d7547d5b5c6ef4b081c60b23e4d941d0eba9cbeb"}, - {file = "charset_normalizer-3.1.0-cp311-cp311-musllinux_1_1_s390x.whl", hash = "sha256:abc1185d79f47c0a7aaf7e2412a0eb2c03b724581139193d2d82b3ad8cbb00ac"}, - {file = "charset_normalizer-3.1.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:cb7b2ab0188829593b9de646545175547a70d9a6e2b63bf2cd87a0a391599324"}, - {file = "charset_normalizer-3.1.0-cp311-cp311-win32.whl", hash = "sha256:c36bcbc0d5174a80d6cccf43a0ecaca44e81d25be4b7f90f0ed7bcfbb5a00909"}, - {file = "charset_normalizer-3.1.0-cp311-cp311-win_amd64.whl", hash = "sha256:cca4def576f47a09a943666b8f829606bcb17e2bc2d5911a46c8f8da45f56755"}, - {file = "charset_normalizer-3.1.0-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:0c95f12b74681e9ae127728f7e5409cbbef9cd914d5896ef238cc779b8152373"}, - {file = "charset_normalizer-3.1.0-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:fca62a8301b605b954ad2e9c3666f9d97f63872aa4efcae5492baca2056b74ab"}, - {file = "charset_normalizer-3.1.0-cp37-cp37m-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:ac0aa6cd53ab9a31d397f8303f92c42f534693528fafbdb997c82bae6e477ad9"}, - {file = "charset_normalizer-3.1.0-cp37-cp37m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:c3af8e0f07399d3176b179f2e2634c3ce9c1301379a6b8c9c9aeecd481da494f"}, - {file = "charset_normalizer-3.1.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3a5fc78f9e3f501a1614a98f7c54d3969f3ad9bba8ba3d9b438c3bc5d047dd28"}, - {file = "charset_normalizer-3.1.0-cp37-cp37m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:628c985afb2c7d27a4800bfb609e03985aaecb42f955049957814e0491d4006d"}, - {file = "charset_normalizer-3.1.0-cp37-cp37m-musllinux_1_1_aarch64.whl", hash = "sha256:74db0052d985cf37fa111828d0dd230776ac99c740e1a758ad99094be4f1803d"}, - {file = "charset_normalizer-3.1.0-cp37-cp37m-musllinux_1_1_i686.whl", hash = "sha256:1e8fcdd8f672a1c4fc8d0bd3a2b576b152d2a349782d1eb0f6b8e52e9954731d"}, - {file = "charset_normalizer-3.1.0-cp37-cp37m-musllinux_1_1_ppc64le.whl", hash = "sha256:04afa6387e2b282cf78ff3dbce20f0cc071c12dc8f685bd40960cc68644cfea6"}, - {file = "charset_normalizer-3.1.0-cp37-cp37m-musllinux_1_1_s390x.whl", hash = "sha256:dd5653e67b149503c68c4018bf07e42eeed6b4e956b24c00ccdf93ac79cdff84"}, - {file = "charset_normalizer-3.1.0-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:d2686f91611f9e17f4548dbf050e75b079bbc2a82be565832bc8ea9047b61c8c"}, - {file = "charset_normalizer-3.1.0-cp37-cp37m-win32.whl", hash = "sha256:4155b51ae05ed47199dc5b2a4e62abccb274cee6b01da5b895099b61b1982974"}, - {file = "charset_normalizer-3.1.0-cp37-cp37m-win_amd64.whl", hash = "sha256:322102cdf1ab682ecc7d9b1c5eed4ec59657a65e1c146a0da342b78f4112db23"}, - {file = "charset_normalizer-3.1.0-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:e633940f28c1e913615fd624fcdd72fdba807bf53ea6925d6a588e84e1151531"}, - {file = "charset_normalizer-3.1.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:3a06f32c9634a8705f4ca9946d667609f52cf130d5548881401f1eb2c39b1e2c"}, - {file = "charset_normalizer-3.1.0-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:7381c66e0561c5757ffe616af869b916c8b4e42b367ab29fedc98481d1e74e14"}, - {file = "charset_normalizer-3.1.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3573d376454d956553c356df45bb824262c397c6e26ce43e8203c4c540ee0acb"}, - {file = "charset_normalizer-3.1.0-cp38-cp38-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:e89df2958e5159b811af9ff0f92614dabf4ff617c03a4c1c6ff53bf1c399e0e1"}, - {file = "charset_normalizer-3.1.0-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:78cacd03e79d009d95635e7d6ff12c21eb89b894c354bd2b2ed0b4763373693b"}, - {file = "charset_normalizer-3.1.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:de5695a6f1d8340b12a5d6d4484290ee74d61e467c39ff03b39e30df62cf83a0"}, - {file = "charset_normalizer-3.1.0-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1c60b9c202d00052183c9be85e5eaf18a4ada0a47d188a83c8f5c5b23252f649"}, - {file = "charset_normalizer-3.1.0-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:f645caaf0008bacf349875a974220f1f1da349c5dbe7c4ec93048cdc785a3326"}, - {file = "charset_normalizer-3.1.0-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:ea9f9c6034ea2d93d9147818f17c2a0860d41b71c38b9ce4d55f21b6f9165a11"}, - {file = "charset_normalizer-3.1.0-cp38-cp38-musllinux_1_1_ppc64le.whl", hash = "sha256:80d1543d58bd3d6c271b66abf454d437a438dff01c3e62fdbcd68f2a11310d4b"}, - {file = "charset_normalizer-3.1.0-cp38-cp38-musllinux_1_1_s390x.whl", hash = "sha256:73dc03a6a7e30b7edc5b01b601e53e7fc924b04e1835e8e407c12c037e81adbd"}, - {file = "charset_normalizer-3.1.0-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:6f5c2e7bc8a4bf7c426599765b1bd33217ec84023033672c1e9a8b35eaeaaaf8"}, - {file = "charset_normalizer-3.1.0-cp38-cp38-win32.whl", hash = "sha256:12a2b561af122e3d94cdb97fe6fb2bb2b82cef0cdca131646fdb940a1eda04f0"}, - {file = "charset_normalizer-3.1.0-cp38-cp38-win_amd64.whl", hash = "sha256:3160a0fd9754aab7d47f95a6b63ab355388d890163eb03b2d2b87ab0a30cfa59"}, - {file = "charset_normalizer-3.1.0-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:38e812a197bf8e71a59fe55b757a84c1f946d0ac114acafaafaf21667a7e169e"}, - {file = "charset_normalizer-3.1.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:6baf0baf0d5d265fa7944feb9f7451cc316bfe30e8df1a61b1bb08577c554f31"}, - {file = "charset_normalizer-3.1.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:8f25e17ab3039b05f762b0a55ae0b3632b2e073d9c8fc88e89aca31a6198e88f"}, - {file = "charset_normalizer-3.1.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3747443b6a904001473370d7810aa19c3a180ccd52a7157aacc264a5ac79265e"}, - {file = "charset_normalizer-3.1.0-cp39-cp39-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b116502087ce8a6b7a5f1814568ccbd0e9f6cfd99948aa59b0e241dc57cf739f"}, - {file = "charset_normalizer-3.1.0-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:d16fd5252f883eb074ca55cb622bc0bee49b979ae4e8639fff6ca3ff44f9f854"}, - {file = "charset_normalizer-3.1.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:21fa558996782fc226b529fdd2ed7866c2c6ec91cee82735c98a197fae39f706"}, - {file = "charset_normalizer-3.1.0-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:6f6c7a8a57e9405cad7485f4c9d3172ae486cfef1344b5ddd8e5239582d7355e"}, - {file = "charset_normalizer-3.1.0-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:ac3775e3311661d4adace3697a52ac0bab17edd166087d493b52d4f4f553f9f0"}, - {file = "charset_normalizer-3.1.0-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:10c93628d7497c81686e8e5e557aafa78f230cd9e77dd0c40032ef90c18f2230"}, - {file = "charset_normalizer-3.1.0-cp39-cp39-musllinux_1_1_ppc64le.whl", hash = "sha256:6f4f4668e1831850ebcc2fd0b1cd11721947b6dc7c00bf1c6bd3c929ae14f2c7"}, - {file = "charset_normalizer-3.1.0-cp39-cp39-musllinux_1_1_s390x.whl", hash = "sha256:0be65ccf618c1e7ac9b849c315cc2e8a8751d9cfdaa43027d4f6624bd587ab7e"}, - {file = "charset_normalizer-3.1.0-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:53d0a3fa5f8af98a1e261de6a3943ca631c526635eb5817a87a59d9a57ebf48f"}, - {file = "charset_normalizer-3.1.0-cp39-cp39-win32.whl", hash = "sha256:a04f86f41a8916fe45ac5024ec477f41f886b3c435da2d4e3d2709b22ab02af1"}, - {file = "charset_normalizer-3.1.0-cp39-cp39-win_amd64.whl", hash = "sha256:830d2948a5ec37c386d3170c483063798d7879037492540f10a475e3fd6f244b"}, - {file = "charset_normalizer-3.1.0-py3-none-any.whl", hash = "sha256:3d9098b479e78c85080c98e1e35ff40b4a31d8953102bb0fd7d1b6f8a2111a3d"}, -] - -[[package]] -name = "click" -version = "8.1.3" -description = "Composable command line interface toolkit" -optional = false -python-versions = ">=3.7" -files = [ - {file = "click-8.1.3-py3-none-any.whl", hash = "sha256:bb4d8133cb15a609f44e8213d9b391b0809795062913b383c62be0ee95b1db48"}, - {file = "click-8.1.3.tar.gz", hash = "sha256:7682dc8afb30297001674575ea00d1814d808d6a36af415a82bd481d37ba7b8e"}, -] - -[package.dependencies] -colorama = {version = "*", markers = "platform_system == \"Windows\""} - -[[package]] -name = "colorama" -version = "0.4.6" -description = "Cross-platform colored terminal text." -optional = false -python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,>=2.7" -files = [ - {file = "colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6"}, - {file = "colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44"}, -] - -[[package]] -name = "cryptography" -version = "40.0.2" -description = "cryptography is a package which provides cryptographic recipes and primitives to Python developers." -optional = false -python-versions = ">=3.6" -files = [ - {file = "cryptography-40.0.2-cp36-abi3-macosx_10_12_universal2.whl", hash = "sha256:8f79b5ff5ad9d3218afb1e7e20ea74da5f76943ee5edb7f76e56ec5161ec782b"}, - {file = "cryptography-40.0.2-cp36-abi3-macosx_10_12_x86_64.whl", hash = "sha256:05dc219433b14046c476f6f09d7636b92a1c3e5808b9a6536adf4932b3b2c440"}, - {file = "cryptography-40.0.2-cp36-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4df2af28d7bedc84fe45bd49bc35d710aede676e2a4cb7fc6d103a2adc8afe4d"}, - {file = "cryptography-40.0.2-cp36-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0dcca15d3a19a66e63662dc8d30f8036b07be851a8680eda92d079868f106288"}, - {file = "cryptography-40.0.2-cp36-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:a04386fb7bc85fab9cd51b6308633a3c271e3d0d3eae917eebab2fac6219b6d2"}, - {file = "cryptography-40.0.2-cp36-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:adc0d980fd2760c9e5de537c28935cc32b9353baaf28e0814df417619c6c8c3b"}, - {file = "cryptography-40.0.2-cp36-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:d5a1bd0e9e2031465761dfa920c16b0065ad77321d8a8c1f5ee331021fda65e9"}, - {file = "cryptography-40.0.2-cp36-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:a95f4802d49faa6a674242e25bfeea6fc2acd915b5e5e29ac90a32b1139cae1c"}, - {file = "cryptography-40.0.2-cp36-abi3-win32.whl", hash = "sha256:aecbb1592b0188e030cb01f82d12556cf72e218280f621deed7d806afd2113f9"}, - {file = "cryptography-40.0.2-cp36-abi3-win_amd64.whl", hash = "sha256:b12794f01d4cacfbd3177b9042198f3af1c856eedd0a98f10f141385c809a14b"}, - {file = "cryptography-40.0.2-pp38-pypy38_pp73-macosx_10_12_x86_64.whl", hash = "sha256:142bae539ef28a1c76794cca7f49729e7c54423f615cfd9b0b1fa90ebe53244b"}, - {file = "cryptography-40.0.2-pp38-pypy38_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:956ba8701b4ffe91ba59665ed170a2ebbdc6fc0e40de5f6059195d9f2b33ca0e"}, - {file = "cryptography-40.0.2-pp38-pypy38_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:4f01c9863da784558165f5d4d916093737a75203a5c5286fde60e503e4276c7a"}, - {file = "cryptography-40.0.2-pp38-pypy38_pp73-win_amd64.whl", hash = "sha256:3daf9b114213f8ba460b829a02896789751626a2a4e7a43a28ee77c04b5e4958"}, - {file = "cryptography-40.0.2-pp39-pypy39_pp73-macosx_10_12_x86_64.whl", hash = "sha256:48f388d0d153350f378c7f7b41497a54ff1513c816bcbbcafe5b829e59b9ce5b"}, - {file = "cryptography-40.0.2-pp39-pypy39_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:c0764e72b36a3dc065c155e5b22f93df465da9c39af65516fe04ed3c68c92636"}, - {file = "cryptography-40.0.2-pp39-pypy39_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:cbaba590180cba88cb99a5f76f90808a624f18b169b90a4abb40c1fd8c19420e"}, - {file = "cryptography-40.0.2-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:7a38250f433cd41df7fcb763caa3ee9362777fdb4dc642b9a349721d2bf47404"}, - {file = "cryptography-40.0.2.tar.gz", hash = "sha256:c33c0d32b8594fa647d2e01dbccc303478e16fdd7cf98652d5b3ed11aa5e5c99"}, -] - -[package.dependencies] -cffi = ">=1.12" - -[package.extras] -docs = ["sphinx (>=5.3.0)", "sphinx-rtd-theme (>=1.1.1)"] -docstest = ["pyenchant (>=1.6.11)", "sphinxcontrib-spelling (>=4.0.1)", "twine (>=1.12.0)"] -pep8test = ["black", "check-manifest", "mypy", "ruff"] -sdist = ["setuptools-rust (>=0.11.4)"] -ssh = ["bcrypt (>=3.1.5)"] -test = ["iso8601", "pretend", "pytest (>=6.2.0)", "pytest-benchmark", "pytest-cov", "pytest-shard (>=0.1.2)", "pytest-subtests", "pytest-xdist"] -test-randomorder = ["pytest-randomly"] -tox = ["tox"] - -[[package]] -name = "diskcache" -version = "5.6.1" -description = "Disk Cache -- Disk and file backed persistent cache." -optional = false -python-versions = ">=3" -files = [ - {file = "diskcache-5.6.1-py3-none-any.whl", hash = "sha256:558c6a2d5d7c721bb00e40711803d6804850c9f76c426ed81ecc627fe9d2ce2d"}, - {file = "diskcache-5.6.1.tar.gz", hash = "sha256:e4c978532feff5814c4cc00fe1e11e40501985946643d73220d41ee7737c72c3"}, -] - -[[package]] -name = "distro" -version = "1.8.0" -description = "Distro - an OS platform information API" -optional = false -python-versions = ">=3.6" -files = [ - {file = "distro-1.8.0-py3-none-any.whl", hash = "sha256:99522ca3e365cac527b44bde033f64c6945d90eb9f769703caaec52b09bbd3ff"}, - {file = "distro-1.8.0.tar.gz", hash = "sha256:02e111d1dc6a50abb8eed6bf31c3e48ed8b0830d1ea2a1b78c61765c2513fdd8"}, -] - -[[package]] -name = "docutils" -version = "0.20" -description = "Docutils -- Python Documentation Utilities" -optional = false -python-versions = ">=3.7" -files = [ - {file = "docutils-0.20-py3-none-any.whl", hash = "sha256:a428f10de4de4774389734c986a01b4af2d802d26717108b0f1b9356862937c5"}, - {file = "docutils-0.20.tar.gz", hash = "sha256:f75a5a52fbcacd81b47e42888ad2b380748aaccfb3f13af0fe69deb759f01eb6"}, -] - -[[package]] -name = "exceptiongroup" -version = "1.1.1" -description = "Backport of PEP 654 (exception groups)" -optional = false -python-versions = ">=3.7" -files = [ - {file = "exceptiongroup-1.1.1-py3-none-any.whl", hash = "sha256:232c37c63e4f682982c8b6459f33a8981039e5fb8756b2074364e5055c498c9e"}, - {file = "exceptiongroup-1.1.1.tar.gz", hash = "sha256:d484c3090ba2889ae2928419117447a14daf3c1231d5e30d0aae34f354f01785"}, -] - -[package.extras] -test = ["pytest (>=6)"] - -[[package]] -name = "fastapi" -version = "0.99.1" -description = "FastAPI framework, high performance, easy to learn, fast to code, ready for production" -optional = true -python-versions = ">=3.7" -files = [ - {file = "fastapi-0.99.1-py3-none-any.whl", hash = "sha256:976df7bab51ac7beda9f68c4513b8c4490b5c1135c72aafd0a5ee4023ec5282e"}, - {file = "fastapi-0.99.1.tar.gz", hash = "sha256:ac78f717cd80d657bd183f94d33b9bda84aa376a46a9dab513586b8eef1dc6fc"}, -] - -[package.dependencies] -pydantic = ">=1.7.4,<1.8 || >1.8,<1.8.1 || >1.8.1,<2.0.0" -starlette = ">=0.27.0,<0.28.0" -typing-extensions = ">=4.5.0" - -[package.extras] -all = ["email-validator (>=1.1.1)", "httpx (>=0.23.0)", "itsdangerous (>=1.1.0)", "jinja2 (>=2.11.2)", "orjson (>=3.2.1)", "python-multipart (>=0.0.5)", "pyyaml (>=5.3.1)", "ujson (>=4.0.1,!=4.0.2,!=4.1.0,!=4.2.0,!=4.3.0,!=5.0.0,!=5.1.0)", "uvicorn[standard] (>=0.12.0)"] - -[[package]] -name = "ghp-import" -version = "2.1.0" -description = "Copy your docs directly to the gh-pages branch." -optional = false -python-versions = "*" -files = [ - {file = "ghp-import-2.1.0.tar.gz", hash = "sha256:9c535c4c61193c2df8871222567d7fd7e5014d835f97dc7b7439069e2413d343"}, - {file = "ghp_import-2.1.0-py3-none-any.whl", hash = "sha256:8337dd7b50877f163d4c0289bc1f1c7f127550241988d568c1db512c4324a619"}, -] - -[package.dependencies] -python-dateutil = ">=2.8.1" - -[package.extras] -dev = ["flake8", "markdown", "twine", "wheel"] - -[[package]] -name = "griffe" -version = "0.27.3" -description = "Signatures for entire Python programs. Extract the structure, the frame, the skeleton of your project, to generate API documentation or find breaking changes in your API." -optional = false -python-versions = ">=3.7" -files = [ - {file = "griffe-0.27.3-py3-none-any.whl", hash = "sha256:094513b209d4acd4b2680c2415d3af5f8ed925714795380c2a7d070e222e0b27"}, - {file = "griffe-0.27.3.tar.gz", hash = "sha256:a3d0f75aa76b80f181f818cf605f658a69fccf287aaeeeafc7a6cf4e6a2ca27e"}, -] - -[package.dependencies] -colorama = ">=0.4" - -[[package]] -name = "h11" -version = "0.14.0" -description = "A pure-Python, bring-your-own-I/O implementation of HTTP/1.1" -optional = false -python-versions = ">=3.7" -files = [ - {file = "h11-0.14.0-py3-none-any.whl", hash = "sha256:e3fe4ac4b851c468cc8363d500db52c2ead036020723024a109d37346efaa761"}, - {file = "h11-0.14.0.tar.gz", hash = "sha256:8f19fbbe99e72420ff35c00b27a34cb9937e902a8b810e2c88300c6f0a3b699d"}, -] - -[[package]] -name = "httpcore" -version = "0.17.0" -description = "A minimal low-level HTTP client." -optional = false -python-versions = ">=3.7" -files = [ - {file = "httpcore-0.17.0-py3-none-any.whl", hash = "sha256:0fdfea45e94f0c9fd96eab9286077f9ff788dd186635ae61b312693e4d943599"}, - {file = "httpcore-0.17.0.tar.gz", hash = "sha256:cc045a3241afbf60ce056202301b4d8b6af08845e3294055eb26b09913ef903c"}, -] - -[package.dependencies] -anyio = ">=3.0,<5.0" -certifi = "*" -h11 = ">=0.13,<0.15" -sniffio = "==1.*" - -[package.extras] -http2 = ["h2 (>=3,<5)"] -socks = ["socksio (==1.*)"] - -[[package]] -name = "httpx" -version = "0.24.1" -description = "The next generation HTTP client." -optional = false -python-versions = ">=3.7" -files = [ - {file = "httpx-0.24.1-py3-none-any.whl", hash = "sha256:06781eb9ac53cde990577af654bd990a4949de37a28bdb4a230d434f3a30b9bd"}, - {file = "httpx-0.24.1.tar.gz", hash = "sha256:5853a43053df830c20f8110c5e69fe44d035d850b2dfe795e196f00fdb774bdd"}, -] - -[package.dependencies] -certifi = "*" -httpcore = ">=0.15.0,<0.18.0" -idna = "*" -sniffio = "*" - -[package.extras] -brotli = ["brotli", "brotlicffi"] -cli = ["click (==8.*)", "pygments (==2.*)", "rich (>=10,<14)"] -http2 = ["h2 (>=3,<5)"] -socks = ["socksio (==1.*)"] - -[[package]] -name = "idna" -version = "3.4" -description = "Internationalized Domain Names in Applications (IDNA)" -optional = false -python-versions = ">=3.5" -files = [ - {file = "idna-3.4-py3-none-any.whl", hash = "sha256:90b77e79eaa3eba6de819a0c442c0b4ceefc341a7a2ab77d7562bf49f425c5c2"}, - {file = "idna-3.4.tar.gz", hash = "sha256:814f528e8dead7d329833b91c5faa87d60bf71824cd12a7530b5526063d02cb4"}, -] - -[[package]] -name = "importlib-metadata" -version = "6.6.0" -description = "Read metadata from Python packages" -optional = false -python-versions = ">=3.7" -files = [ - {file = "importlib_metadata-6.6.0-py3-none-any.whl", hash = "sha256:43dd286a2cd8995d5eaef7fee2066340423b818ed3fd70adf0bad5f1fac53fed"}, - {file = "importlib_metadata-6.6.0.tar.gz", hash = "sha256:92501cdf9cc66ebd3e612f1b4f0c0765dfa42f0fa38ffb319b6bd84dd675d705"}, -] - -[package.dependencies] -zipp = ">=0.5" - -[package.extras] -docs = ["furo", "jaraco.packaging (>=9)", "jaraco.tidelift (>=1.4)", "rst.linker (>=1.9)", "sphinx (>=3.5)", "sphinx-lint"] -perf = ["ipython"] -testing = ["flake8 (<5)", "flufl.flake8", "importlib-resources (>=1.3)", "packaging", "pyfakefs", "pytest (>=6)", "pytest-black (>=0.3.7)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=1.3)", "pytest-flake8", "pytest-mypy (>=0.9.1)", "pytest-perf (>=0.9.2)"] - -[[package]] -name = "importlib-resources" -version = "5.12.0" -description = "Read resources from Python packages" -optional = false -python-versions = ">=3.7" -files = [ - {file = "importlib_resources-5.12.0-py3-none-any.whl", hash = "sha256:7b1deeebbf351c7578e09bf2f63fa2ce8b5ffec296e0d349139d43cca061a81a"}, - {file = "importlib_resources-5.12.0.tar.gz", hash = "sha256:4be82589bf5c1d7999aedf2a45159d10cb3ca4f19b2271f8792bc8e6da7b22f6"}, -] - -[package.dependencies] -zipp = {version = ">=3.1.0", markers = "python_version < \"3.10\""} - -[package.extras] -docs = ["furo", "jaraco.packaging (>=9)", "jaraco.tidelift (>=1.4)", "rst.linker (>=1.9)", "sphinx (>=3.5)", "sphinx-lint"] -testing = ["flake8 (<5)", "pytest (>=6)", "pytest-black (>=0.3.7)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=1.3)", "pytest-flake8", "pytest-mypy (>=0.9.1)"] - -[[package]] -name = "iniconfig" -version = "2.0.0" -description = "brain-dead simple config-ini parsing" -optional = false -python-versions = ">=3.7" -files = [ - {file = "iniconfig-2.0.0-py3-none-any.whl", hash = "sha256:b6a85871a79d2e3b22d2d1b94ac2824226a63c6b741c88f7ae975f18b6778374"}, - {file = "iniconfig-2.0.0.tar.gz", hash = "sha256:2d91e135bf72d31a410b17c16da610a82cb55f6b0477d1a902134b24a455b8b3"}, -] - -[[package]] -name = "jaraco-classes" -version = "3.2.3" -description = "Utility functions for Python class constructs" -optional = false -python-versions = ">=3.7" -files = [ - {file = "jaraco.classes-3.2.3-py3-none-any.whl", hash = "sha256:2353de3288bc6b82120752201c6b1c1a14b058267fa424ed5ce5984e3b922158"}, - {file = "jaraco.classes-3.2.3.tar.gz", hash = "sha256:89559fa5c1d3c34eff6f631ad80bb21f378dbcbb35dd161fd2c6b93f5be2f98a"}, -] - -[package.dependencies] -more-itertools = "*" - -[package.extras] -docs = ["jaraco.packaging (>=9)", "jaraco.tidelift (>=1.4)", "rst.linker (>=1.9)", "sphinx (>=3.5)"] -testing = ["flake8 (<5)", "pytest (>=6)", "pytest-black (>=0.3.7)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=1.3)", "pytest-flake8", "pytest-mypy (>=0.9.1)"] - -[[package]] -name = "jeepney" -version = "0.8.0" -description = "Low-level, pure Python DBus protocol wrapper." -optional = false -python-versions = ">=3.7" -files = [ - {file = "jeepney-0.8.0-py3-none-any.whl", hash = "sha256:c0a454ad016ca575060802ee4d590dd912e35c122fa04e70306de3d076cce755"}, - {file = "jeepney-0.8.0.tar.gz", hash = "sha256:5efe48d255973902f6badc3ce55e2aa6c5c3b3bc642059ef3a91247bcfcc5806"}, -] - -[package.extras] -test = ["async-timeout", "pytest", "pytest-asyncio (>=0.17)", "pytest-trio", "testpath", "trio"] -trio = ["async_generator", "trio"] - -[[package]] -name = "jinja2" -version = "3.1.2" -description = "A very fast and expressive template engine." -optional = false -python-versions = ">=3.7" -files = [ - {file = "Jinja2-3.1.2-py3-none-any.whl", hash = "sha256:6088930bfe239f0e6710546ab9c19c9ef35e29792895fed6e6e31a023a182a61"}, - {file = "Jinja2-3.1.2.tar.gz", hash = "sha256:31351a702a408a9e7595a8fc6150fc3f43bb6bf7e319770cbc0db9df9437e852"}, -] - -[package.dependencies] -MarkupSafe = ">=2.0" - -[package.extras] -i18n = ["Babel (>=2.7)"] - -[[package]] -name = "keyring" -version = "23.13.1" -description = "Store and access your passwords safely." -optional = false -python-versions = ">=3.7" -files = [ - {file = "keyring-23.13.1-py3-none-any.whl", hash = "sha256:771ed2a91909389ed6148631de678f82ddc73737d85a927f382a8a1b157898cd"}, - {file = "keyring-23.13.1.tar.gz", hash = "sha256:ba2e15a9b35e21908d0aaf4e0a47acc52d6ae33444df0da2b49d41a46ef6d678"}, -] - -[package.dependencies] -importlib-metadata = {version = ">=4.11.4", markers = "python_version < \"3.12\""} -importlib-resources = {version = "*", markers = "python_version < \"3.9\""} -"jaraco.classes" = "*" -jeepney = {version = ">=0.4.2", markers = "sys_platform == \"linux\""} -pywin32-ctypes = {version = ">=0.2.0", markers = "sys_platform == \"win32\""} -SecretStorage = {version = ">=3.2", markers = "sys_platform == \"linux\""} - -[package.extras] -completion = ["shtab"] -docs = ["furo", "jaraco.packaging (>=9)", "jaraco.tidelift (>=1.4)", "rst.linker (>=1.9)", "sphinx (>=3.5)"] -testing = ["flake8 (<5)", "pytest (>=6)", "pytest-black (>=0.3.7)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=1.3)", "pytest-flake8", "pytest-mypy (>=0.9.1)"] - -[[package]] -name = "markdown" -version = "3.3.7" -description = "Python implementation of Markdown." -optional = false -python-versions = ">=3.6" -files = [ - {file = "Markdown-3.3.7-py3-none-any.whl", hash = "sha256:f5da449a6e1c989a4cea2631aa8ee67caa5a2ef855d551c88f9e309f4634c621"}, - {file = "Markdown-3.3.7.tar.gz", hash = "sha256:cbb516f16218e643d8e0a95b309f77eb118cb138d39a4f27851e6a63581db874"}, -] - -[package.dependencies] -importlib-metadata = {version = ">=4.4", markers = "python_version < \"3.10\""} - -[package.extras] -testing = ["coverage", "pyyaml"] - -[[package]] -name = "markdown-it-py" -version = "2.2.0" -description = "Python port of markdown-it. Markdown parsing, done right!" -optional = false -python-versions = ">=3.7" -files = [ - {file = "markdown-it-py-2.2.0.tar.gz", hash = "sha256:7c9a5e412688bc771c67432cbfebcdd686c93ce6484913dccf06cb5a0bea35a1"}, - {file = "markdown_it_py-2.2.0-py3-none-any.whl", hash = "sha256:5a35f8d1870171d9acc47b99612dc146129b631baf04970128b568f190d0cc30"}, -] - -[package.dependencies] -mdurl = ">=0.1,<1.0" - -[package.extras] -benchmarking = ["psutil", "pytest", "pytest-benchmark"] -code-style = ["pre-commit (>=3.0,<4.0)"] -compare = ["commonmark (>=0.9,<1.0)", "markdown (>=3.4,<4.0)", "mistletoe (>=1.0,<2.0)", "mistune (>=2.0,<3.0)", "panflute (>=2.3,<3.0)"] -linkify = ["linkify-it-py (>=1,<3)"] -plugins = ["mdit-py-plugins"] -profiling = ["gprof2dot"] -rtd = ["attrs", "myst-parser", "pyyaml", "sphinx", "sphinx-copybutton", "sphinx-design", "sphinx_book_theme"] -testing = ["coverage", "pytest", "pytest-cov", "pytest-regressions"] - -[[package]] -name = "markupsafe" -version = "2.1.2" -description = "Safely add untrusted strings to HTML/XML markup." -optional = false -python-versions = ">=3.7" -files = [ - {file = "MarkupSafe-2.1.2-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:665a36ae6f8f20a4676b53224e33d456a6f5a72657d9c83c2aa00765072f31f7"}, - {file = "MarkupSafe-2.1.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:340bea174e9761308703ae988e982005aedf427de816d1afe98147668cc03036"}, - {file = "MarkupSafe-2.1.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:22152d00bf4a9c7c83960521fc558f55a1adbc0631fbb00a9471e097b19d72e1"}, - {file = "MarkupSafe-2.1.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:28057e985dace2f478e042eaa15606c7efccb700797660629da387eb289b9323"}, - {file = "MarkupSafe-2.1.2-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ca244fa73f50a800cf8c3ebf7fd93149ec37f5cb9596aa8873ae2c1d23498601"}, - {file = "MarkupSafe-2.1.2-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:d9d971ec1e79906046aa3ca266de79eac42f1dbf3612a05dc9368125952bd1a1"}, - {file = "MarkupSafe-2.1.2-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:7e007132af78ea9df29495dbf7b5824cb71648d7133cf7848a2a5dd00d36f9ff"}, - {file = "MarkupSafe-2.1.2-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:7313ce6a199651c4ed9d7e4cfb4aa56fe923b1adf9af3b420ee14e6d9a73df65"}, - {file = "MarkupSafe-2.1.2-cp310-cp310-win32.whl", hash = "sha256:c4a549890a45f57f1ebf99c067a4ad0cb423a05544accaf2b065246827ed9603"}, - {file = "MarkupSafe-2.1.2-cp310-cp310-win_amd64.whl", hash = "sha256:835fb5e38fd89328e9c81067fd642b3593c33e1e17e2fdbf77f5676abb14a156"}, - {file = "MarkupSafe-2.1.2-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:2ec4f2d48ae59bbb9d1f9d7efb9236ab81429a764dedca114f5fdabbc3788013"}, - {file = "MarkupSafe-2.1.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:608e7073dfa9e38a85d38474c082d4281f4ce276ac0010224eaba11e929dd53a"}, - {file = "MarkupSafe-2.1.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:65608c35bfb8a76763f37036547f7adfd09270fbdbf96608be2bead319728fcd"}, - {file = "MarkupSafe-2.1.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f2bfb563d0211ce16b63c7cb9395d2c682a23187f54c3d79bfec33e6705473c6"}, - {file = "MarkupSafe-2.1.2-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:da25303d91526aac3672ee6d49a2f3db2d9502a4a60b55519feb1a4c7714e07d"}, - {file = "MarkupSafe-2.1.2-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:9cad97ab29dfc3f0249b483412c85c8ef4766d96cdf9dcf5a1e3caa3f3661cf1"}, - {file = "MarkupSafe-2.1.2-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:085fd3201e7b12809f9e6e9bc1e5c96a368c8523fad5afb02afe3c051ae4afcc"}, - {file = "MarkupSafe-2.1.2-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:1bea30e9bf331f3fef67e0a3877b2288593c98a21ccb2cf29b74c581a4eb3af0"}, - {file = "MarkupSafe-2.1.2-cp311-cp311-win32.whl", hash = "sha256:7df70907e00c970c60b9ef2938d894a9381f38e6b9db73c5be35e59d92e06625"}, - {file = "MarkupSafe-2.1.2-cp311-cp311-win_amd64.whl", hash = "sha256:e55e40ff0cc8cc5c07996915ad367fa47da6b3fc091fdadca7f5403239c5fec3"}, - {file = "MarkupSafe-2.1.2-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:a6e40afa7f45939ca356f348c8e23048e02cb109ced1eb8420961b2f40fb373a"}, - {file = "MarkupSafe-2.1.2-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:cf877ab4ed6e302ec1d04952ca358b381a882fbd9d1b07cccbfd61783561f98a"}, - {file = "MarkupSafe-2.1.2-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:63ba06c9941e46fa389d389644e2d8225e0e3e5ebcc4ff1ea8506dce646f8c8a"}, - {file = "MarkupSafe-2.1.2-cp37-cp37m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f1cd098434e83e656abf198f103a8207a8187c0fc110306691a2e94a78d0abb2"}, - {file = "MarkupSafe-2.1.2-cp37-cp37m-musllinux_1_1_aarch64.whl", hash = "sha256:55f44b440d491028addb3b88f72207d71eeebfb7b5dbf0643f7c023ae1fba619"}, - {file = "MarkupSafe-2.1.2-cp37-cp37m-musllinux_1_1_i686.whl", hash = "sha256:a6f2fcca746e8d5910e18782f976489939d54a91f9411c32051b4aab2bd7c513"}, - {file = "MarkupSafe-2.1.2-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:0b462104ba25f1ac006fdab8b6a01ebbfbce9ed37fd37fd4acd70c67c973e460"}, - {file = "MarkupSafe-2.1.2-cp37-cp37m-win32.whl", hash = "sha256:7668b52e102d0ed87cb082380a7e2e1e78737ddecdde129acadb0eccc5423859"}, - {file = "MarkupSafe-2.1.2-cp37-cp37m-win_amd64.whl", hash = "sha256:6d6607f98fcf17e534162f0709aaad3ab7a96032723d8ac8750ffe17ae5a0666"}, - {file = "MarkupSafe-2.1.2-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:a806db027852538d2ad7555b203300173dd1b77ba116de92da9afbc3a3be3eed"}, - {file = "MarkupSafe-2.1.2-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:a4abaec6ca3ad8660690236d11bfe28dfd707778e2442b45addd2f086d6ef094"}, - {file = "MarkupSafe-2.1.2-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f03a532d7dee1bed20bc4884194a16160a2de9ffc6354b3878ec9682bb623c54"}, - {file = "MarkupSafe-2.1.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4cf06cdc1dda95223e9d2d3c58d3b178aa5dacb35ee7e3bbac10e4e1faacb419"}, - {file = "MarkupSafe-2.1.2-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:22731d79ed2eb25059ae3df1dfc9cb1546691cc41f4e3130fe6bfbc3ecbbecfa"}, - {file = "MarkupSafe-2.1.2-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:f8ffb705ffcf5ddd0e80b65ddf7bed7ee4f5a441ea7d3419e861a12eaf41af58"}, - {file = "MarkupSafe-2.1.2-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:8db032bf0ce9022a8e41a22598eefc802314e81b879ae093f36ce9ddf39ab1ba"}, - {file = "MarkupSafe-2.1.2-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:2298c859cfc5463f1b64bd55cb3e602528db6fa0f3cfd568d3605c50678f8f03"}, - {file = "MarkupSafe-2.1.2-cp38-cp38-win32.whl", hash = "sha256:50c42830a633fa0cf9e7d27664637532791bfc31c731a87b202d2d8ac40c3ea2"}, - {file = "MarkupSafe-2.1.2-cp38-cp38-win_amd64.whl", hash = "sha256:bb06feb762bade6bf3c8b844462274db0c76acc95c52abe8dbed28ae3d44a147"}, - {file = "MarkupSafe-2.1.2-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:99625a92da8229df6d44335e6fcc558a5037dd0a760e11d84be2260e6f37002f"}, - {file = "MarkupSafe-2.1.2-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:8bca7e26c1dd751236cfb0c6c72d4ad61d986e9a41bbf76cb445f69488b2a2bd"}, - {file = "MarkupSafe-2.1.2-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:40627dcf047dadb22cd25ea7ecfe9cbf3bbbad0482ee5920b582f3809c97654f"}, - {file = "MarkupSafe-2.1.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:40dfd3fefbef579ee058f139733ac336312663c6706d1163b82b3003fb1925c4"}, - {file = "MarkupSafe-2.1.2-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:090376d812fb6ac5f171e5938e82e7f2d7adc2b629101cec0db8b267815c85e2"}, - {file = "MarkupSafe-2.1.2-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:2e7821bffe00aa6bd07a23913b7f4e01328c3d5cc0b40b36c0bd81d362faeb65"}, - {file = "MarkupSafe-2.1.2-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:c0a33bc9f02c2b17c3ea382f91b4db0e6cde90b63b296422a939886a7a80de1c"}, - {file = "MarkupSafe-2.1.2-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:b8526c6d437855442cdd3d87eede9c425c4445ea011ca38d937db299382e6fa3"}, - {file = "MarkupSafe-2.1.2-cp39-cp39-win32.whl", hash = "sha256:137678c63c977754abe9086a3ec011e8fd985ab90631145dfb9294ad09c102a7"}, - {file = "MarkupSafe-2.1.2-cp39-cp39-win_amd64.whl", hash = "sha256:0576fe974b40a400449768941d5d0858cc624e3249dfd1e0c33674e5c7ca7aed"}, - {file = "MarkupSafe-2.1.2.tar.gz", hash = "sha256:abcabc8c2b26036d62d4c746381a6f7cf60aafcc653198ad678306986b09450d"}, -] - -[[package]] -name = "mdurl" -version = "0.1.2" -description = "Markdown URL utilities" -optional = false -python-versions = ">=3.7" -files = [ - {file = "mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8"}, - {file = "mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba"}, -] - -[[package]] -name = "mergedeep" -version = "1.3.4" -description = "A deep merge function for 🐍." -optional = false -python-versions = ">=3.6" -files = [ - {file = "mergedeep-1.3.4-py3-none-any.whl", hash = "sha256:70775750742b25c0d8f36c55aed03d24c3384d17c951b3175d898bd778ef0307"}, - {file = "mergedeep-1.3.4.tar.gz", hash = "sha256:0096d52e9dad9939c3d975a774666af186eda617e6ca84df4c94dec30004f2a8"}, -] - -[[package]] -name = "mkdocs" -version = "1.4.3" -description = "Project documentation with Markdown." -optional = false -python-versions = ">=3.7" -files = [ - {file = "mkdocs-1.4.3-py3-none-any.whl", hash = "sha256:6ee46d309bda331aac915cd24aab882c179a933bd9e77b80ce7d2eaaa3f689dd"}, - {file = "mkdocs-1.4.3.tar.gz", hash = "sha256:5955093bbd4dd2e9403c5afaf57324ad8b04f16886512a3ee6ef828956481c57"}, -] - -[package.dependencies] -click = ">=7.0" -colorama = {version = ">=0.4", markers = "platform_system == \"Windows\""} -ghp-import = ">=1.0" -importlib-metadata = {version = ">=4.3", markers = "python_version < \"3.10\""} -jinja2 = ">=2.11.1" -markdown = ">=3.2.1,<3.4" -mergedeep = ">=1.3.4" -packaging = ">=20.5" -pyyaml = ">=5.1" -pyyaml-env-tag = ">=0.1" -watchdog = ">=2.0" - -[package.extras] -i18n = ["babel (>=2.9.0)"] -min-versions = ["babel (==2.9.0)", "click (==7.0)", "colorama (==0.4)", "ghp-import (==1.0)", "importlib-metadata (==4.3)", "jinja2 (==2.11.1)", "markdown (==3.2.1)", "markupsafe (==2.0.1)", "mergedeep (==1.3.4)", "packaging (==20.5)", "pyyaml (==5.1)", "pyyaml-env-tag (==0.1)", "typing-extensions (==3.10)", "watchdog (==2.0)"] - -[[package]] -name = "mkdocs-autorefs" -version = "0.4.1" -description = "Automatically link across pages in MkDocs." -optional = false -python-versions = ">=3.7" -files = [ - {file = "mkdocs-autorefs-0.4.1.tar.gz", hash = "sha256:70748a7bd025f9ecd6d6feeba8ba63f8e891a1af55f48e366d6d6e78493aba84"}, - {file = "mkdocs_autorefs-0.4.1-py3-none-any.whl", hash = "sha256:a2248a9501b29dc0cc8ba4c09f4f47ff121945f6ce33d760f145d6f89d313f5b"}, -] - -[package.dependencies] -Markdown = ">=3.3" -mkdocs = ">=1.1" - -[[package]] -name = "mkdocs-material" -version = "9.1.18" -description = "Documentation that simply works" -optional = false -python-versions = ">=3.7" -files = [ - {file = "mkdocs_material-9.1.18-py3-none-any.whl", hash = "sha256:5bcf8fb79ac2f253c0ffe93fa181cba87718c6438f459dc4180ac7418cc9a450"}, - {file = "mkdocs_material-9.1.18.tar.gz", hash = "sha256:981dd39979723d4cda7cfc77bbbe5e54922d5761a7af23fb8ba9edb52f114b13"}, -] - -[package.dependencies] -colorama = ">=0.4" -jinja2 = ">=3.0" -markdown = ">=3.2" -mkdocs = ">=1.4.2" -mkdocs-material-extensions = ">=1.1" -pygments = ">=2.14" -pymdown-extensions = ">=9.9.1" -regex = ">=2022.4.24" -requests = ">=2.26" - -[[package]] -name = "mkdocs-material-extensions" -version = "1.1.1" -description = "Extension pack for Python Markdown and MkDocs Material." -optional = false -python-versions = ">=3.7" -files = [ - {file = "mkdocs_material_extensions-1.1.1-py3-none-any.whl", hash = "sha256:e41d9f38e4798b6617ad98ca8f7f1157b1e4385ac1459ca1e4ea219b556df945"}, - {file = "mkdocs_material_extensions-1.1.1.tar.gz", hash = "sha256:9c003da71e2cc2493d910237448c672e00cefc800d3d6ae93d2fc69979e3bd93"}, -] - -[[package]] -name = "mkdocstrings" -version = "0.22.0" -description = "Automatic documentation from sources, for MkDocs." -optional = false -python-versions = ">=3.7" -files = [ - {file = "mkdocstrings-0.22.0-py3-none-any.whl", hash = "sha256:2d4095d461554ff6a778fdabdca3c00c468c2f1459d469f7a7f622a2b23212ba"}, - {file = "mkdocstrings-0.22.0.tar.gz", hash = "sha256:82a33b94150ebb3d4b5c73bab4598c3e21468c79ec072eff6931c8f3bfc38256"}, -] - -[package.dependencies] -importlib-metadata = {version = ">=4.6", markers = "python_version < \"3.10\""} -Jinja2 = ">=2.11.1" -Markdown = ">=3.3" -MarkupSafe = ">=1.1" -mkdocs = ">=1.2" -mkdocs-autorefs = ">=0.3.1" -mkdocstrings-python = {version = ">=0.5.2", optional = true, markers = "extra == \"python\""} -pymdown-extensions = ">=6.3" -typing-extensions = {version = ">=4.1", markers = "python_version < \"3.10\""} - -[package.extras] -crystal = ["mkdocstrings-crystal (>=0.3.4)"] -python = ["mkdocstrings-python (>=0.5.2)"] -python-legacy = ["mkdocstrings-python-legacy (>=0.2.1)"] - -[[package]] -name = "mkdocstrings-python" -version = "0.10.1" -description = "A Python handler for mkdocstrings." -optional = false -python-versions = ">=3.7" -files = [ - {file = "mkdocstrings_python-0.10.1-py3-none-any.whl", hash = "sha256:ef239cee2c688e2b949a0a47e42a141d744dd12b7007311b3309dc70e3bafc5c"}, - {file = "mkdocstrings_python-0.10.1.tar.gz", hash = "sha256:b72301fff739070ec517b5b36bf2f7c49d1360a275896a64efb97fc17d3f3968"}, -] - -[package.dependencies] -griffe = ">=0.24" -mkdocstrings = ">=0.20" - -[[package]] -name = "more-itertools" -version = "9.1.0" -description = "More routines for operating on iterables, beyond itertools" -optional = false -python-versions = ">=3.7" -files = [ - {file = "more-itertools-9.1.0.tar.gz", hash = "sha256:cabaa341ad0389ea83c17a94566a53ae4c9d07349861ecb14dc6d0345cf9ac5d"}, - {file = "more_itertools-9.1.0-py3-none-any.whl", hash = "sha256:d2bc7f02446e86a68911e58ded76d6561eea00cddfb2a91e7019bbb586c799f3"}, -] - -[[package]] -name = "mypy-extensions" -version = "1.0.0" -description = "Type system extensions for programs checked with the mypy type checker." -optional = false -python-versions = ">=3.5" -files = [ - {file = "mypy_extensions-1.0.0-py3-none-any.whl", hash = "sha256:4392f6c0eb8a5668a69e23d168ffa70f0be9ccfd32b5cc2d26a34ae5b844552d"}, - {file = "mypy_extensions-1.0.0.tar.gz", hash = "sha256:75dbf8955dc00442a438fc4d0666508a9a97b6bd41aa2f0ffe9d2f2725af0782"}, -] - -[[package]] -name = "numpy" -version = "1.24.4" -description = "Fundamental package for array computing in Python" -optional = false -python-versions = ">=3.8" -files = [ - {file = "numpy-1.24.4-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:c0bfb52d2169d58c1cdb8cc1f16989101639b34c7d3ce60ed70b19c63eba0b64"}, - {file = "numpy-1.24.4-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:ed094d4f0c177b1b8e7aa9cba7d6ceed51c0e569a5318ac0ca9a090680a6a1b1"}, - {file = "numpy-1.24.4-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:79fc682a374c4a8ed08b331bef9c5f582585d1048fa6d80bc6c35bc384eee9b4"}, - {file = "numpy-1.24.4-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7ffe43c74893dbf38c2b0a1f5428760a1a9c98285553c89e12d70a96a7f3a4d6"}, - {file = "numpy-1.24.4-cp310-cp310-win32.whl", hash = "sha256:4c21decb6ea94057331e111a5bed9a79d335658c27ce2adb580fb4d54f2ad9bc"}, - {file = "numpy-1.24.4-cp310-cp310-win_amd64.whl", hash = "sha256:b4bea75e47d9586d31e892a7401f76e909712a0fd510f58f5337bea9572c571e"}, - {file = "numpy-1.24.4-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:f136bab9c2cfd8da131132c2cf6cc27331dd6fae65f95f69dcd4ae3c3639c810"}, - {file = "numpy-1.24.4-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:e2926dac25b313635e4d6cf4dc4e51c8c0ebfed60b801c799ffc4c32bf3d1254"}, - {file = "numpy-1.24.4-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:222e40d0e2548690405b0b3c7b21d1169117391c2e82c378467ef9ab4c8f0da7"}, - {file = "numpy-1.24.4-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7215847ce88a85ce39baf9e89070cb860c98fdddacbaa6c0da3ffb31b3350bd5"}, - {file = "numpy-1.24.4-cp311-cp311-win32.whl", hash = "sha256:4979217d7de511a8d57f4b4b5b2b965f707768440c17cb70fbf254c4b225238d"}, - {file = "numpy-1.24.4-cp311-cp311-win_amd64.whl", hash = "sha256:b7b1fc9864d7d39e28f41d089bfd6353cb5f27ecd9905348c24187a768c79694"}, - {file = "numpy-1.24.4-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:1452241c290f3e2a312c137a9999cdbf63f78864d63c79039bda65ee86943f61"}, - {file = "numpy-1.24.4-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:04640dab83f7c6c85abf9cd729c5b65f1ebd0ccf9de90b270cd61935eef0197f"}, - {file = "numpy-1.24.4-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a5425b114831d1e77e4b5d812b69d11d962e104095a5b9c3b641a218abcc050e"}, - {file = "numpy-1.24.4-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:dd80e219fd4c71fc3699fc1dadac5dcf4fd882bfc6f7ec53d30fa197b8ee22dc"}, - {file = "numpy-1.24.4-cp38-cp38-win32.whl", hash = "sha256:4602244f345453db537be5314d3983dbf5834a9701b7723ec28923e2889e0bb2"}, - {file = "numpy-1.24.4-cp38-cp38-win_amd64.whl", hash = "sha256:692f2e0f55794943c5bfff12b3f56f99af76f902fc47487bdfe97856de51a706"}, - {file = "numpy-1.24.4-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:2541312fbf09977f3b3ad449c4e5f4bb55d0dbf79226d7724211acc905049400"}, - {file = "numpy-1.24.4-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:9667575fb6d13c95f1b36aca12c5ee3356bf001b714fc354eb5465ce1609e62f"}, - {file = "numpy-1.24.4-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f3a86ed21e4f87050382c7bc96571755193c4c1392490744ac73d660e8f564a9"}, - {file = "numpy-1.24.4-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d11efb4dbecbdf22508d55e48d9c8384db795e1b7b51ea735289ff96613ff74d"}, - {file = "numpy-1.24.4-cp39-cp39-win32.whl", hash = "sha256:6620c0acd41dbcb368610bb2f4d83145674040025e5536954782467100aa8835"}, - {file = "numpy-1.24.4-cp39-cp39-win_amd64.whl", hash = "sha256:befe2bf740fd8373cf56149a5c23a0f601e82869598d41f8e188a0e9869926f8"}, - {file = "numpy-1.24.4-pp38-pypy38_pp73-macosx_10_9_x86_64.whl", hash = "sha256:31f13e25b4e304632a4619d0e0777662c2ffea99fcae2029556b17d8ff958aef"}, - {file = "numpy-1.24.4-pp38-pypy38_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:95f7ac6540e95bc440ad77f56e520da5bf877f87dca58bd095288dce8940532a"}, - {file = "numpy-1.24.4-pp38-pypy38_pp73-win_amd64.whl", hash = "sha256:e98f220aa76ca2a977fe435f5b04d7b3470c0a2e6312907b37ba6068f26787f2"}, - {file = "numpy-1.24.4.tar.gz", hash = "sha256:80f5e3a4e498641401868df4208b74581206afbee7cf7b8329daae82676d9463"}, -] - -[[package]] -name = "packaging" -version = "23.1" -description = "Core utilities for Python packages" -optional = false -python-versions = ">=3.7" -files = [ - {file = "packaging-23.1-py3-none-any.whl", hash = "sha256:994793af429502c4ea2ebf6bf664629d07c1a9fe974af92966e4b8d2df7edc61"}, - {file = "packaging-23.1.tar.gz", hash = "sha256:a392980d2b6cffa644431898be54b0045151319d1e7ec34f0cfed48767dd334f"}, -] - -[[package]] -name = "pathspec" -version = "0.11.1" -description = "Utility library for gitignore style pattern matching of file paths." -optional = false -python-versions = ">=3.7" -files = [ - {file = "pathspec-0.11.1-py3-none-any.whl", hash = "sha256:d8af70af76652554bd134c22b3e8a1cc46ed7d91edcdd721ef1a0c51a84a5293"}, - {file = "pathspec-0.11.1.tar.gz", hash = "sha256:2798de800fa92780e33acca925945e9a19a133b715067cf165b8866c15a31687"}, -] - -[[package]] -name = "pkginfo" -version = "1.9.6" -description = "Query metadata from sdists / bdists / installed packages." -optional = false -python-versions = ">=3.6" -files = [ - {file = "pkginfo-1.9.6-py3-none-any.whl", hash = "sha256:4b7a555a6d5a22169fcc9cf7bfd78d296b0361adad412a346c1226849af5e546"}, - {file = "pkginfo-1.9.6.tar.gz", hash = "sha256:8fd5896e8718a4372f0ea9cc9d96f6417c9b986e23a4d116dda26b62cc29d046"}, -] - -[package.extras] -testing = ["pytest", "pytest-cov"] - -[[package]] -name = "platformdirs" -version = "3.5.0" -description = "A small Python package for determining appropriate platform-specific dirs, e.g. a \"user data dir\"." -optional = false -python-versions = ">=3.7" -files = [ - {file = "platformdirs-3.5.0-py3-none-any.whl", hash = "sha256:47692bc24c1958e8b0f13dd727307cff1db103fca36399f457da8e05f222fdc4"}, - {file = "platformdirs-3.5.0.tar.gz", hash = "sha256:7954a68d0ba23558d753f73437c55f89027cf8f5108c19844d4b82e5af396335"}, -] - -[package.extras] -docs = ["furo (>=2023.3.27)", "proselint (>=0.13)", "sphinx (>=6.1.3)", "sphinx-autodoc-typehints (>=1.23,!=1.23.4)"] -test = ["appdirs (==1.4.4)", "covdefaults (>=2.3)", "pytest (>=7.3.1)", "pytest-cov (>=4)", "pytest-mock (>=3.10)"] - -[[package]] -name = "pluggy" -version = "1.0.0" -description = "plugin and hook calling mechanisms for python" -optional = false -python-versions = ">=3.6" -files = [ - {file = "pluggy-1.0.0-py2.py3-none-any.whl", hash = "sha256:74134bbf457f031a36d68416e1509f34bd5ccc019f0bcc952c7b909d06b37bd3"}, - {file = "pluggy-1.0.0.tar.gz", hash = "sha256:4224373bacce55f955a878bf9cfa763c1e360858e330072059e10bad68531159"}, -] - -[package.extras] -dev = ["pre-commit", "tox"] -testing = ["pytest", "pytest-benchmark"] - -[[package]] -name = "pycparser" -version = "2.21" -description = "C parser in Python" -optional = false -python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*" -files = [ - {file = "pycparser-2.21-py2.py3-none-any.whl", hash = "sha256:8ee45429555515e1f6b185e78100aea234072576aa43ab53aefcae078162fca9"}, - {file = "pycparser-2.21.tar.gz", hash = "sha256:e644fdec12f7872f86c58ff790da456218b10f863970249516d60a5eaca77206"}, -] - -[[package]] -name = "pydantic" -version = "1.10.7" -description = "Data validation and settings management using python type hints" -optional = true -python-versions = ">=3.7" -files = [ - {file = "pydantic-1.10.7-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:e79e999e539872e903767c417c897e729e015872040e56b96e67968c3b918b2d"}, - {file = "pydantic-1.10.7-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:01aea3a42c13f2602b7ecbbea484a98169fb568ebd9e247593ea05f01b884b2e"}, - {file = "pydantic-1.10.7-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:516f1ed9bc2406a0467dd777afc636c7091d71f214d5e413d64fef45174cfc7a"}, - {file = "pydantic-1.10.7-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ae150a63564929c675d7f2303008d88426a0add46efd76c3fc797cd71cb1b46f"}, - {file = "pydantic-1.10.7-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:ecbbc51391248116c0a055899e6c3e7ffbb11fb5e2a4cd6f2d0b93272118a209"}, - {file = "pydantic-1.10.7-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:f4a2b50e2b03d5776e7f21af73e2070e1b5c0d0df255a827e7c632962f8315af"}, - {file = "pydantic-1.10.7-cp310-cp310-win_amd64.whl", hash = "sha256:a7cd2251439988b413cb0a985c4ed82b6c6aac382dbaff53ae03c4b23a70e80a"}, - {file = "pydantic-1.10.7-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:68792151e174a4aa9e9fc1b4e653e65a354a2fa0fed169f7b3d09902ad2cb6f1"}, - {file = "pydantic-1.10.7-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:dfe2507b8ef209da71b6fb5f4e597b50c5a34b78d7e857c4f8f3115effaef5fe"}, - {file = "pydantic-1.10.7-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:10a86d8c8db68086f1e30a530f7d5f83eb0685e632e411dbbcf2d5c0150e8dcd"}, - {file = "pydantic-1.10.7-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:d75ae19d2a3dbb146b6f324031c24f8a3f52ff5d6a9f22f0683694b3afcb16fb"}, - {file = "pydantic-1.10.7-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:464855a7ff7f2cc2cf537ecc421291b9132aa9c79aef44e917ad711b4a93163b"}, - {file = "pydantic-1.10.7-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:193924c563fae6ddcb71d3f06fa153866423ac1b793a47936656e806b64e24ca"}, - {file = "pydantic-1.10.7-cp311-cp311-win_amd64.whl", hash = "sha256:b4a849d10f211389502059c33332e91327bc154acc1845f375a99eca3afa802d"}, - {file = "pydantic-1.10.7-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:cc1dde4e50a5fc1336ee0581c1612215bc64ed6d28d2c7c6f25d2fe3e7c3e918"}, - {file = "pydantic-1.10.7-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e0cfe895a504c060e5d36b287ee696e2fdad02d89e0d895f83037245218a87fe"}, - {file = "pydantic-1.10.7-cp37-cp37m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:670bb4683ad1e48b0ecb06f0cfe2178dcf74ff27921cdf1606e527d2617a81ee"}, - {file = "pydantic-1.10.7-cp37-cp37m-musllinux_1_1_i686.whl", hash = "sha256:950ce33857841f9a337ce07ddf46bc84e1c4946d2a3bba18f8280297157a3fd1"}, - {file = "pydantic-1.10.7-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:c15582f9055fbc1bfe50266a19771bbbef33dd28c45e78afbe1996fd70966c2a"}, - {file = "pydantic-1.10.7-cp37-cp37m-win_amd64.whl", hash = "sha256:82dffb306dd20bd5268fd6379bc4bfe75242a9c2b79fec58e1041fbbdb1f7914"}, - {file = "pydantic-1.10.7-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:8c7f51861d73e8b9ddcb9916ae7ac39fb52761d9ea0df41128e81e2ba42886cd"}, - {file = "pydantic-1.10.7-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:6434b49c0b03a51021ade5c4daa7d70c98f7a79e95b551201fff682fc1661245"}, - {file = "pydantic-1.10.7-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:64d34ab766fa056df49013bb6e79921a0265204c071984e75a09cbceacbbdd5d"}, - {file = "pydantic-1.10.7-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:701daea9ffe9d26f97b52f1d157e0d4121644f0fcf80b443248434958fd03dc3"}, - {file = "pydantic-1.10.7-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:cf135c46099ff3f919d2150a948ce94b9ce545598ef2c6c7bf55dca98a304b52"}, - {file = "pydantic-1.10.7-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:b0f85904f73161817b80781cc150f8b906d521fa11e3cdabae19a581c3606209"}, - {file = "pydantic-1.10.7-cp38-cp38-win_amd64.whl", hash = "sha256:9f6f0fd68d73257ad6685419478c5aece46432f4bdd8d32c7345f1986496171e"}, - {file = "pydantic-1.10.7-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:c230c0d8a322276d6e7b88c3f7ce885f9ed16e0910354510e0bae84d54991143"}, - {file = "pydantic-1.10.7-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:976cae77ba6a49d80f461fd8bba183ff7ba79f44aa5cfa82f1346b5626542f8e"}, - {file = "pydantic-1.10.7-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7d45fc99d64af9aaf7e308054a0067fdcd87ffe974f2442312372dfa66e1001d"}, - {file = "pydantic-1.10.7-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:d2a5ebb48958754d386195fe9e9c5106f11275867051bf017a8059410e9abf1f"}, - {file = "pydantic-1.10.7-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:abfb7d4a7cd5cc4e1d1887c43503a7c5dd608eadf8bc615413fc498d3e4645cd"}, - {file = "pydantic-1.10.7-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:80b1fab4deb08a8292d15e43a6edccdffa5377a36a4597bb545b93e79c5ff0a5"}, - {file = "pydantic-1.10.7-cp39-cp39-win_amd64.whl", hash = "sha256:d71e69699498b020ea198468e2480a2f1e7433e32a3a99760058c6520e2bea7e"}, - {file = "pydantic-1.10.7-py3-none-any.whl", hash = "sha256:0cd181f1d0b1d00e2b705f1bf1ac7799a2d938cce3376b8007df62b29be3c2c6"}, - {file = "pydantic-1.10.7.tar.gz", hash = "sha256:cfc83c0678b6ba51b0532bea66860617c4cd4251ecf76e9846fa5a9f3454e97e"}, -] - -[package.dependencies] -typing-extensions = ">=4.2.0" - -[package.extras] -dotenv = ["python-dotenv (>=0.10.4)"] -email = ["email-validator (>=1.0.3)"] - -[[package]] -name = "pygments" -version = "2.15.1" -description = "Pygments is a syntax highlighting package written in Python." -optional = false -python-versions = ">=3.7" -files = [ - {file = "Pygments-2.15.1-py3-none-any.whl", hash = "sha256:db2db3deb4b4179f399a09054b023b6a586b76499d36965813c71aa8ed7b5fd1"}, - {file = "Pygments-2.15.1.tar.gz", hash = "sha256:8ace4d3c1dd481894b2005f560ead0f9f19ee64fe983366be1a21e171d12775c"}, -] - -[package.extras] -plugins = ["importlib-metadata"] - -[[package]] -name = "pymdown-extensions" -version = "9.11" -description = "Extension pack for Python Markdown." -optional = false -python-versions = ">=3.7" -files = [ - {file = "pymdown_extensions-9.11-py3-none-any.whl", hash = "sha256:a499191d8d869f30339de86fcf072a787e86c42b6f16f280f5c2cf174182b7f3"}, - {file = "pymdown_extensions-9.11.tar.gz", hash = "sha256:f7e86c1d3981f23d9dc43294488ecb54abadd05b0be4bf8f0e15efc90f7853ff"}, -] - -[package.dependencies] -markdown = ">=3.2" -pyyaml = "*" - -[[package]] -name = "pytest" -version = "7.4.0" -description = "pytest: simple powerful testing with Python" -optional = false -python-versions = ">=3.7" -files = [ - {file = "pytest-7.4.0-py3-none-any.whl", hash = "sha256:78bf16451a2eb8c7a2ea98e32dc119fd2aa758f1d5d66dbf0a59d69a3969df32"}, - {file = "pytest-7.4.0.tar.gz", hash = "sha256:b4bf8c45bd59934ed84001ad51e11b4ee40d40a1229d2c79f9c592b0a3f6bd8a"}, -] - -[package.dependencies] -colorama = {version = "*", markers = "sys_platform == \"win32\""} -exceptiongroup = {version = ">=1.0.0rc8", markers = "python_version < \"3.11\""} -iniconfig = "*" -packaging = "*" -pluggy = ">=0.12,<2.0" -tomli = {version = ">=1.0.0", markers = "python_version < \"3.11\""} - -[package.extras] -testing = ["argcomplete", "attrs (>=19.2.0)", "hypothesis (>=3.56)", "mock", "nose", "pygments (>=2.7.2)", "requests", "setuptools", "xmlschema"] - -[[package]] -name = "python-dateutil" -version = "2.8.2" -description = "Extensions to the standard Python datetime module" -optional = false -python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,>=2.7" -files = [ - {file = "python-dateutil-2.8.2.tar.gz", hash = "sha256:0123cacc1627ae19ddf3c27a5de5bd67ee4586fbdd6440d9748f8abb483d3e86"}, - {file = "python_dateutil-2.8.2-py2.py3-none-any.whl", hash = "sha256:961d03dc3453ebbc59dbdea9e4e11c5651520a876d0f4db161e8674aae935da9"}, -] - -[package.dependencies] -six = ">=1.5" - -[[package]] -name = "pywin32-ctypes" -version = "0.2.0" -description = "" -optional = false -python-versions = "*" -files = [ - {file = "pywin32-ctypes-0.2.0.tar.gz", hash = "sha256:24ffc3b341d457d48e8922352130cf2644024a4ff09762a2261fd34c36ee5942"}, - {file = "pywin32_ctypes-0.2.0-py2.py3-none-any.whl", hash = "sha256:9dc2d991b3479cc2df15930958b674a48a227d5361d413827a4cfd0b5876fc98"}, -] - -[[package]] -name = "pyyaml" -version = "6.0" -description = "YAML parser and emitter for Python" -optional = false -python-versions = ">=3.6" -files = [ - {file = "PyYAML-6.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:d4db7c7aef085872ef65a8fd7d6d09a14ae91f691dec3e87ee5ee0539d516f53"}, - {file = "PyYAML-6.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:9df7ed3b3d2e0ecfe09e14741b857df43adb5a3ddadc919a2d94fbdf78fea53c"}, - {file = "PyYAML-6.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:77f396e6ef4c73fdc33a9157446466f1cff553d979bd00ecb64385760c6babdc"}, - {file = "PyYAML-6.0-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a80a78046a72361de73f8f395f1f1e49f956c6be882eed58505a15f3e430962b"}, - {file = "PyYAML-6.0-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:f84fbc98b019fef2ee9a1cb3ce93e3187a6df0b2538a651bfb890254ba9f90b5"}, - {file = "PyYAML-6.0-cp310-cp310-win32.whl", hash = "sha256:2cd5df3de48857ed0544b34e2d40e9fac445930039f3cfe4bcc592a1f836d513"}, - {file = "PyYAML-6.0-cp310-cp310-win_amd64.whl", hash = "sha256:daf496c58a8c52083df09b80c860005194014c3698698d1a57cbcfa182142a3a"}, - {file = "PyYAML-6.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:d4b0ba9512519522b118090257be113b9468d804b19d63c71dbcf4a48fa32358"}, - {file = "PyYAML-6.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:81957921f441d50af23654aa6c5e5eaf9b06aba7f0a19c18a538dc7ef291c5a1"}, - {file = "PyYAML-6.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:afa17f5bc4d1b10afd4466fd3a44dc0e245382deca5b3c353d8b757f9e3ecb8d"}, - {file = "PyYAML-6.0-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:dbad0e9d368bb989f4515da330b88a057617d16b6a8245084f1b05400f24609f"}, - {file = "PyYAML-6.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:432557aa2c09802be39460360ddffd48156e30721f5e8d917f01d31694216782"}, - {file = "PyYAML-6.0-cp311-cp311-win32.whl", hash = "sha256:bfaef573a63ba8923503d27530362590ff4f576c626d86a9fed95822a8255fd7"}, - {file = "PyYAML-6.0-cp311-cp311-win_amd64.whl", hash = "sha256:01b45c0191e6d66c470b6cf1b9531a771a83c1c4208272ead47a3ae4f2f603bf"}, - {file = "PyYAML-6.0-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:897b80890765f037df3403d22bab41627ca8811ae55e9a722fd0392850ec4d86"}, - {file = "PyYAML-6.0-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:50602afada6d6cbfad699b0c7bb50d5ccffa7e46a3d738092afddc1f9758427f"}, - {file = "PyYAML-6.0-cp36-cp36m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:48c346915c114f5fdb3ead70312bd042a953a8ce5c7106d5bfb1a5254e47da92"}, - {file = "PyYAML-6.0-cp36-cp36m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:98c4d36e99714e55cfbaaee6dd5badbc9a1ec339ebfc3b1f52e293aee6bb71a4"}, - {file = "PyYAML-6.0-cp36-cp36m-win32.whl", hash = "sha256:0283c35a6a9fbf047493e3a0ce8d79ef5030852c51e9d911a27badfde0605293"}, - {file = "PyYAML-6.0-cp36-cp36m-win_amd64.whl", hash = "sha256:07751360502caac1c067a8132d150cf3d61339af5691fe9e87803040dbc5db57"}, - {file = "PyYAML-6.0-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:819b3830a1543db06c4d4b865e70ded25be52a2e0631ccd2f6a47a2822f2fd7c"}, - {file = "PyYAML-6.0-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:473f9edb243cb1935ab5a084eb238d842fb8f404ed2193a915d1784b5a6b5fc0"}, - {file = "PyYAML-6.0-cp37-cp37m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:0ce82d761c532fe4ec3f87fc45688bdd3a4c1dc5e0b4a19814b9009a29baefd4"}, - {file = "PyYAML-6.0-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:231710d57adfd809ef5d34183b8ed1eeae3f76459c18fb4a0b373ad56bedcdd9"}, - {file = "PyYAML-6.0-cp37-cp37m-win32.whl", hash = "sha256:c5687b8d43cf58545ade1fe3e055f70eac7a5a1a0bf42824308d868289a95737"}, - {file = "PyYAML-6.0-cp37-cp37m-win_amd64.whl", hash = "sha256:d15a181d1ecd0d4270dc32edb46f7cb7733c7c508857278d3d378d14d606db2d"}, - {file = "PyYAML-6.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:0b4624f379dab24d3725ffde76559cff63d9ec94e1736b556dacdfebe5ab6d4b"}, - {file = "PyYAML-6.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:213c60cd50106436cc818accf5baa1aba61c0189ff610f64f4a3e8c6726218ba"}, - {file = "PyYAML-6.0-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:9fa600030013c4de8165339db93d182b9431076eb98eb40ee068700c9c813e34"}, - {file = "PyYAML-6.0-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:277a0ef2981ca40581a47093e9e2d13b3f1fbbeffae064c1d21bfceba2030287"}, - {file = "PyYAML-6.0-cp38-cp38-win32.whl", hash = "sha256:d4eccecf9adf6fbcc6861a38015c2a64f38b9d94838ac1810a9023a0609e1b78"}, - {file = "PyYAML-6.0-cp38-cp38-win_amd64.whl", hash = "sha256:1e4747bc279b4f613a09eb64bba2ba602d8a6664c6ce6396a4d0cd413a50ce07"}, - {file = "PyYAML-6.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:055d937d65826939cb044fc8c9b08889e8c743fdc6a32b33e2390f66013e449b"}, - {file = "PyYAML-6.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:e61ceaab6f49fb8bdfaa0f92c4b57bcfbea54c09277b1b4f7ac376bfb7a7c174"}, - {file = "PyYAML-6.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d67d839ede4ed1b28a4e8909735fc992a923cdb84e618544973d7dfc71540803"}, - {file = "PyYAML-6.0-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:cba8c411ef271aa037d7357a2bc8f9ee8b58b9965831d9e51baf703280dc73d3"}, - {file = "PyYAML-6.0-cp39-cp39-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:40527857252b61eacd1d9af500c3337ba8deb8fc298940291486c465c8b46ec0"}, - {file = "PyYAML-6.0-cp39-cp39-win32.whl", hash = "sha256:b5b9eccad747aabaaffbc6064800670f0c297e52c12754eb1d976c57e4f74dcb"}, - {file = "PyYAML-6.0-cp39-cp39-win_amd64.whl", hash = "sha256:b3d267842bf12586ba6c734f89d1f5b871df0273157918b0ccefa29deb05c21c"}, - {file = "PyYAML-6.0.tar.gz", hash = "sha256:68fb519c14306fec9720a2a5b45bc9f0c8d1b9c72adf45c37baedfcd949c35a2"}, -] - -[[package]] -name = "pyyaml-env-tag" -version = "0.1" -description = "A custom YAML tag for referencing environment variables in YAML files. " -optional = false -python-versions = ">=3.6" -files = [ - {file = "pyyaml_env_tag-0.1-py3-none-any.whl", hash = "sha256:af31106dec8a4d68c60207c1886031cbf839b68aa7abccdb19868200532c2069"}, - {file = "pyyaml_env_tag-0.1.tar.gz", hash = "sha256:70092675bda14fdec33b31ba77e7543de9ddc88f2e5b99160396572d11525bdb"}, -] - -[package.dependencies] -pyyaml = "*" - -[[package]] -name = "readme-renderer" -version = "37.3" -description = "readme_renderer is a library for rendering \"readme\" descriptions for Warehouse" -optional = false -python-versions = ">=3.7" -files = [ - {file = "readme_renderer-37.3-py3-none-any.whl", hash = "sha256:f67a16caedfa71eef48a31b39708637a6f4664c4394801a7b0d6432d13907343"}, - {file = "readme_renderer-37.3.tar.gz", hash = "sha256:cd653186dfc73055656f090f227f5cb22a046d7f71a841dfa305f55c9a513273"}, -] - -[package.dependencies] -bleach = ">=2.1.0" -docutils = ">=0.13.1" -Pygments = ">=2.5.1" - -[package.extras] -md = ["cmarkgfm (>=0.8.0)"] - -[[package]] -name = "regex" -version = "2023.5.5" -description = "Alternative regular expression module, to replace re." -optional = false -python-versions = ">=3.6" -files = [ - {file = "regex-2023.5.5-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:48c9ec56579d4ba1c88f42302194b8ae2350265cb60c64b7b9a88dcb7fbde309"}, - {file = "regex-2023.5.5-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:02f4541550459c08fdd6f97aa4e24c6f1932eec780d58a2faa2068253df7d6ff"}, - {file = "regex-2023.5.5-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:53e22e4460f0245b468ee645156a4f84d0fc35a12d9ba79bd7d79bdcd2f9629d"}, - {file = "regex-2023.5.5-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:4b870b6f632fc74941cadc2a0f3064ed8409e6f8ee226cdfd2a85ae50473aa94"}, - {file = "regex-2023.5.5-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:171c52e320fe29260da550d81c6b99f6f8402450dc7777ef5ced2e848f3b6f8f"}, - {file = "regex-2023.5.5-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:aad5524c2aedaf9aa14ef1bc9327f8abd915699dea457d339bebbe2f0d218f86"}, - {file = "regex-2023.5.5-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:5a0f874ee8c0bc820e649c900243c6d1e6dc435b81da1492046716f14f1a2a96"}, - {file = "regex-2023.5.5-cp310-cp310-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:e645c757183ee0e13f0bbe56508598e2d9cd42b8abc6c0599d53b0d0b8dd1479"}, - {file = "regex-2023.5.5-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:a4c5da39bca4f7979eefcbb36efea04471cd68db2d38fcbb4ee2c6d440699833"}, - {file = "regex-2023.5.5-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:5e3f4468b8c6fd2fd33c218bbd0a1559e6a6fcf185af8bb0cc43f3b5bfb7d636"}, - {file = "regex-2023.5.5-cp310-cp310-musllinux_1_1_ppc64le.whl", hash = "sha256:59e4b729eae1a0919f9e4c0fc635fbcc9db59c74ad98d684f4877be3d2607dd6"}, - {file = "regex-2023.5.5-cp310-cp310-musllinux_1_1_s390x.whl", hash = "sha256:ba73a14e9c8f9ac409863543cde3290dba39098fc261f717dc337ea72d3ebad2"}, - {file = "regex-2023.5.5-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:0bbd5dcb19603ab8d2781fac60114fb89aee8494f4505ae7ad141a3314abb1f9"}, - {file = "regex-2023.5.5-cp310-cp310-win32.whl", hash = "sha256:40005cbd383438aecf715a7b47fe1e3dcbc889a36461ed416bdec07e0ef1db66"}, - {file = "regex-2023.5.5-cp310-cp310-win_amd64.whl", hash = "sha256:59597cd6315d3439ed4b074febe84a439c33928dd34396941b4d377692eca810"}, - {file = "regex-2023.5.5-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:8f08276466fedb9e36e5193a96cb944928301152879ec20c2d723d1031cd4ddd"}, - {file = "regex-2023.5.5-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:cd46f30e758629c3ee91713529cfbe107ac50d27110fdcc326a42ce2acf4dafc"}, - {file = "regex-2023.5.5-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f2910502f718828cecc8beff004917dcf577fc5f8f5dd40ffb1ea7612124547b"}, - {file = "regex-2023.5.5-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:445d6f4fc3bd9fc2bf0416164454f90acab8858cd5a041403d7a11e3356980e8"}, - {file = "regex-2023.5.5-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:18196c16a584619c7c1d843497c069955d7629ad4a3fdee240eb347f4a2c9dbe"}, - {file = "regex-2023.5.5-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:33d430a23b661629661f1fe8395be2004006bc792bb9fc7c53911d661b69dd7e"}, - {file = "regex-2023.5.5-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:72a28979cc667e5f82ef433db009184e7ac277844eea0f7f4d254b789517941d"}, - {file = "regex-2023.5.5-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:f764e4dfafa288e2eba21231f455d209f4709436baeebb05bdecfb5d8ddc3d35"}, - {file = "regex-2023.5.5-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:23d86ad2121b3c4fc78c58f95e19173790e22ac05996df69b84e12da5816cb17"}, - {file = "regex-2023.5.5-cp311-cp311-musllinux_1_1_ppc64le.whl", hash = "sha256:690a17db524ee6ac4a27efc5406530dd90e7a7a69d8360235323d0e5dafb8f5b"}, - {file = "regex-2023.5.5-cp311-cp311-musllinux_1_1_s390x.whl", hash = "sha256:1ecf3dcff71f0c0fe3e555201cbe749fa66aae8d18f80d2cc4de8e66df37390a"}, - {file = "regex-2023.5.5-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:811040d7f3dd9c55eb0d8b00b5dcb7fd9ae1761c454f444fd9f37fe5ec57143a"}, - {file = "regex-2023.5.5-cp311-cp311-win32.whl", hash = "sha256:c8c143a65ce3ca42e54d8e6fcaf465b6b672ed1c6c90022794a802fb93105d22"}, - {file = "regex-2023.5.5-cp311-cp311-win_amd64.whl", hash = "sha256:586a011f77f8a2da4b888774174cd266e69e917a67ba072c7fc0e91878178a80"}, - {file = "regex-2023.5.5-cp36-cp36m-macosx_10_9_x86_64.whl", hash = "sha256:b6365703e8cf1644b82104cdd05270d1a9f043119a168d66c55684b1b557d008"}, - {file = "regex-2023.5.5-cp36-cp36m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a56c18f21ac98209da9c54ae3ebb3b6f6e772038681d6cb43b8d53da3b09ee81"}, - {file = "regex-2023.5.5-cp36-cp36m-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b8b942d8b3ce765dbc3b1dad0a944712a89b5de290ce8f72681e22b3c55f3cc8"}, - {file = "regex-2023.5.5-cp36-cp36m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:844671c9c1150fcdac46d43198364034b961bd520f2c4fdaabfc7c7d7138a2dd"}, - {file = "regex-2023.5.5-cp36-cp36m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c2ce65bdeaf0a386bb3b533a28de3994e8e13b464ac15e1e67e4603dd88787fa"}, - {file = "regex-2023.5.5-cp36-cp36m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:fee0016cc35a8a91e8cc9312ab26a6fe638d484131a7afa79e1ce6165328a135"}, - {file = "regex-2023.5.5-cp36-cp36m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:18f05d14f14a812fe9723f13afafefe6b74ca042d99f8884e62dbd34dcccf3e2"}, - {file = "regex-2023.5.5-cp36-cp36m-musllinux_1_1_aarch64.whl", hash = "sha256:941b3f1b2392f0bcd6abf1bc7a322787d6db4e7457be6d1ffd3a693426a755f2"}, - {file = "regex-2023.5.5-cp36-cp36m-musllinux_1_1_i686.whl", hash = "sha256:921473a93bcea4d00295799ab929522fc650e85c6b9f27ae1e6bb32a790ea7d3"}, - {file = "regex-2023.5.5-cp36-cp36m-musllinux_1_1_ppc64le.whl", hash = "sha256:e2205a81f815b5bb17e46e74cc946c575b484e5f0acfcb805fb252d67e22938d"}, - {file = "regex-2023.5.5-cp36-cp36m-musllinux_1_1_s390x.whl", hash = "sha256:385992d5ecf1a93cb85adff2f73e0402dd9ac29b71b7006d342cc920816e6f32"}, - {file = "regex-2023.5.5-cp36-cp36m-musllinux_1_1_x86_64.whl", hash = "sha256:890a09cb0a62198bff92eda98b2b507305dd3abf974778bae3287f98b48907d3"}, - {file = "regex-2023.5.5-cp36-cp36m-win32.whl", hash = "sha256:821a88b878b6589c5068f4cc2cfeb2c64e343a196bc9d7ac68ea8c2a776acd46"}, - {file = "regex-2023.5.5-cp36-cp36m-win_amd64.whl", hash = "sha256:7918a1b83dd70dc04ab5ed24c78ae833ae8ea228cef84e08597c408286edc926"}, - {file = "regex-2023.5.5-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:338994d3d4ca4cf12f09822e025731a5bdd3a37aaa571fa52659e85ca793fb67"}, - {file = "regex-2023.5.5-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0a69cf0c00c4d4a929c6c7717fd918414cab0d6132a49a6d8fc3ded1988ed2ea"}, - {file = "regex-2023.5.5-cp37-cp37m-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:8f5e06df94fff8c4c85f98c6487f6636848e1dc85ce17ab7d1931df4a081f657"}, - {file = "regex-2023.5.5-cp37-cp37m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a8906669b03c63266b6a7693d1f487b02647beb12adea20f8840c1a087e2dfb5"}, - {file = "regex-2023.5.5-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9fda3e50abad8d0f48df621cf75adc73c63f7243cbe0e3b2171392b445401550"}, - {file = "regex-2023.5.5-cp37-cp37m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:5ac2b7d341dc1bd102be849d6dd33b09701223a851105b2754339e390be0627a"}, - {file = "regex-2023.5.5-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:fb2b495dd94b02de8215625948132cc2ea360ae84fe6634cd19b6567709c8ae2"}, - {file = "regex-2023.5.5-cp37-cp37m-musllinux_1_1_aarch64.whl", hash = "sha256:aa7d032c1d84726aa9edeb6accf079b4caa87151ca9fabacef31fa028186c66d"}, - {file = "regex-2023.5.5-cp37-cp37m-musllinux_1_1_i686.whl", hash = "sha256:3d45864693351c15531f7e76f545ec35000d50848daa833cead96edae1665559"}, - {file = "regex-2023.5.5-cp37-cp37m-musllinux_1_1_ppc64le.whl", hash = "sha256:21e90a288e6ba4bf44c25c6a946cb9b0f00b73044d74308b5e0afd190338297c"}, - {file = "regex-2023.5.5-cp37-cp37m-musllinux_1_1_s390x.whl", hash = "sha256:10250a093741ec7bf74bcd2039e697f519b028518f605ff2aa7ac1e9c9f97423"}, - {file = "regex-2023.5.5-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:6b8d0c153f07a953636b9cdb3011b733cadd4178123ef728ccc4d5969e67f3c2"}, - {file = "regex-2023.5.5-cp37-cp37m-win32.whl", hash = "sha256:10374c84ee58c44575b667310d5bbfa89fb2e64e52349720a0182c0017512f6c"}, - {file = "regex-2023.5.5-cp37-cp37m-win_amd64.whl", hash = "sha256:9b320677521aabf666cdd6e99baee4fb5ac3996349c3b7f8e7c4eee1c00dfe3a"}, - {file = "regex-2023.5.5-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:afb1c70ec1e594a547f38ad6bf5e3d60304ce7539e677c1429eebab115bce56e"}, - {file = "regex-2023.5.5-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:cf123225945aa58b3057d0fba67e8061c62d14cc8a4202630f8057df70189051"}, - {file = "regex-2023.5.5-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a99757ad7fe5c8a2bb44829fc57ced11253e10f462233c1255fe03888e06bc19"}, - {file = "regex-2023.5.5-cp38-cp38-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:a623564d810e7a953ff1357f7799c14bc9beeab699aacc8b7ab7822da1e952b8"}, - {file = "regex-2023.5.5-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:ced02e3bd55e16e89c08bbc8128cff0884d96e7f7a5633d3dc366b6d95fcd1d6"}, - {file = "regex-2023.5.5-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d1cbe6b5be3b9b698d8cc4ee4dee7e017ad655e83361cd0ea8e653d65e469468"}, - {file = "regex-2023.5.5-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:4a6e4b0e0531223f53bad07ddf733af490ba2b8367f62342b92b39b29f72735a"}, - {file = "regex-2023.5.5-cp38-cp38-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:2e9c4f778514a560a9c9aa8e5538bee759b55f6c1dcd35613ad72523fd9175b8"}, - {file = "regex-2023.5.5-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:256f7f4c6ba145f62f7a441a003c94b8b1af78cee2cccacfc1e835f93bc09426"}, - {file = "regex-2023.5.5-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:bd7b68fd2e79d59d86dcbc1ccd6e2ca09c505343445daaa4e07f43c8a9cc34da"}, - {file = "regex-2023.5.5-cp38-cp38-musllinux_1_1_ppc64le.whl", hash = "sha256:4a5059bd585e9e9504ef9c07e4bc15b0a621ba20504388875d66b8b30a5c4d18"}, - {file = "regex-2023.5.5-cp38-cp38-musllinux_1_1_s390x.whl", hash = "sha256:6893544e06bae009916a5658ce7207e26ed17385149f35a3125f5259951f1bbe"}, - {file = "regex-2023.5.5-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:c64d5abe91a3dfe5ff250c6bb267ef00dbc01501518225b45a5f9def458f31fb"}, - {file = "regex-2023.5.5-cp38-cp38-win32.whl", hash = "sha256:7923470d6056a9590247ff729c05e8e0f06bbd4efa6569c916943cb2d9b68b91"}, - {file = "regex-2023.5.5-cp38-cp38-win_amd64.whl", hash = "sha256:4035d6945cb961c90c3e1c1ca2feb526175bcfed44dfb1cc77db4fdced060d3e"}, - {file = "regex-2023.5.5-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:50fd2d9b36938d4dcecbd684777dd12a407add4f9f934f235c66372e630772b0"}, - {file = "regex-2023.5.5-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:d19e57f888b00cd04fc38f5e18d0efbd91ccba2d45039453ab2236e6eec48d4d"}, - {file = "regex-2023.5.5-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bd966475e963122ee0a7118ec9024388c602d12ac72860f6eea119a3928be053"}, - {file = "regex-2023.5.5-cp39-cp39-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:db09e6c18977a33fea26fe67b7a842f706c67cf8bda1450974d0ae0dd63570df"}, - {file = "regex-2023.5.5-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:6164d4e2a82f9ebd7752a06bd6c504791bedc6418c0196cd0a23afb7f3e12b2d"}, - {file = "regex-2023.5.5-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:84397d3f750d153ebd7f958efaa92b45fea170200e2df5e0e1fd4d85b7e3f58a"}, - {file = "regex-2023.5.5-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:9c3efee9bb53cbe7b285760c81f28ac80dc15fa48b5fe7e58b52752e642553f1"}, - {file = "regex-2023.5.5-cp39-cp39-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:144b5b017646b5a9392a5554a1e5db0000ae637be4971c9747566775fc96e1b2"}, - {file = "regex-2023.5.5-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:1189fbbb21e2c117fda5303653b61905aeeeea23de4a94d400b0487eb16d2d60"}, - {file = "regex-2023.5.5-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:f83fe9e10f9d0b6cf580564d4d23845b9d692e4c91bd8be57733958e4c602956"}, - {file = "regex-2023.5.5-cp39-cp39-musllinux_1_1_ppc64le.whl", hash = "sha256:72aa4746993a28c841e05889f3f1b1e5d14df8d3daa157d6001a34c98102b393"}, - {file = "regex-2023.5.5-cp39-cp39-musllinux_1_1_s390x.whl", hash = "sha256:de2f780c3242ea114dd01f84848655356af4dd561501896c751d7b885ea6d3a1"}, - {file = "regex-2023.5.5-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:290fd35219486dfbc00b0de72f455ecdd63e59b528991a6aec9fdfc0ce85672e"}, - {file = "regex-2023.5.5-cp39-cp39-win32.whl", hash = "sha256:732176f5427e72fa2325b05c58ad0b45af341c459910d766f814b0584ac1f9ac"}, - {file = "regex-2023.5.5-cp39-cp39-win_amd64.whl", hash = "sha256:1307aa4daa1cbb23823d8238e1f61292fd07e4e5d8d38a6efff00b67a7cdb764"}, - {file = "regex-2023.5.5.tar.gz", hash = "sha256:7d76a8a1fc9da08296462a18f16620ba73bcbf5909e42383b253ef34d9d5141e"}, -] - -[[package]] -name = "requests" -version = "2.30.0" -description = "Python HTTP for Humans." -optional = false -python-versions = ">=3.7" -files = [ - {file = "requests-2.30.0-py3-none-any.whl", hash = "sha256:10e94cc4f3121ee6da529d358cdaeaff2f1c409cd377dbc72b825852f2f7e294"}, - {file = "requests-2.30.0.tar.gz", hash = "sha256:239d7d4458afcb28a692cdd298d87542235f4ca8d36d03a15bfc128a6559a2f4"}, -] - -[package.dependencies] -certifi = ">=2017.4.17" -charset-normalizer = ">=2,<4" -idna = ">=2.5,<4" -urllib3 = ">=1.21.1,<3" - -[package.extras] -socks = ["PySocks (>=1.5.6,!=1.5.7)"] -use-chardet-on-py3 = ["chardet (>=3.0.2,<6)"] - -[[package]] -name = "requests-toolbelt" -version = "1.0.0" -description = "A utility belt for advanced users of python-requests" -optional = false -python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*" -files = [ - {file = "requests-toolbelt-1.0.0.tar.gz", hash = "sha256:7681a0a3d047012b5bdc0ee37d7f8f07ebe76ab08caeccfc3921ce23c88d5bc6"}, - {file = "requests_toolbelt-1.0.0-py2.py3-none-any.whl", hash = "sha256:cccfdd665f0a24fcf4726e690f65639d272bb0637b9b92dfd91a5568ccf6bd06"}, -] - -[package.dependencies] -requests = ">=2.0.1,<3.0.0" - -[[package]] -name = "rfc3986" -version = "2.0.0" -description = "Validating URI References per RFC 3986" -optional = false -python-versions = ">=3.7" -files = [ - {file = "rfc3986-2.0.0-py2.py3-none-any.whl", hash = "sha256:50b1502b60e289cb37883f3dfd34532b8873c7de9f49bb546641ce9cbd256ebd"}, - {file = "rfc3986-2.0.0.tar.gz", hash = "sha256:97aacf9dbd4bfd829baad6e6309fa6573aaf1be3f6fa735c8ab05e46cecb261c"}, -] - -[package.extras] -idna2008 = ["idna"] - -[[package]] -name = "rich" -version = "13.3.5" -description = "Render rich text, tables, progress bars, syntax highlighting, markdown and more to the terminal" -optional = false -python-versions = ">=3.7.0" -files = [ - {file = "rich-13.3.5-py3-none-any.whl", hash = "sha256:69cdf53799e63f38b95b9bf9c875f8c90e78dd62b2f00c13a911c7a3b9fa4704"}, - {file = "rich-13.3.5.tar.gz", hash = "sha256:2d11b9b8dd03868f09b4fffadc84a6a8cda574e40dc90821bd845720ebb8e89c"}, -] - -[package.dependencies] -markdown-it-py = ">=2.2.0,<3.0.0" -pygments = ">=2.13.0,<3.0.0" -typing-extensions = {version = ">=4.0.0,<5.0", markers = "python_version < \"3.9\""} - -[package.extras] -jupyter = ["ipywidgets (>=7.5.1,<9)"] - -[[package]] -name = "scikit-build" -version = "0.17.6" -description = "Improved build system generator for Python C/C++/Fortran/Cython extensions" -optional = false -python-versions = ">=3.7" -files = [ - {file = "scikit_build-0.17.6-py3-none-any.whl", hash = "sha256:18bd55e81841106eec93f30a297df4f301003791c41be46ef6428d58bd42d6b3"}, - {file = "scikit_build-0.17.6.tar.gz", hash = "sha256:b51a51a36b37c42650994b5047912f59b22e3210b23e321f287611f9ef6e5c9d"}, -] - -[package.dependencies] -distro = "*" -packaging = "*" -setuptools = ">=42.0.0" -tomli = {version = "*", markers = "python_version < \"3.11\""} -wheel = ">=0.32.0" - -[package.extras] -cov = ["coverage[toml] (>=4.2)", "pytest-cov (>=2.7.1)"] -docs = ["pygments", "sphinx (>=4)", "sphinx-issues", "sphinx-rtd-theme (>=1.0)", "sphinxcontrib-moderncmakedomain (>=3.19)"] -doctest = ["ubelt (>=0.8.2)", "xdoctest (>=0.10.0)"] -test = ["build (>=0.7)", "cython (>=0.25.1)", "importlib-metadata", "pytest (>=6.0.0)", "pytest-mock (>=1.10.4)", "requests", "virtualenv"] - -[[package]] -name = "secretstorage" -version = "3.3.3" -description = "Python bindings to FreeDesktop.org Secret Service API" -optional = false -python-versions = ">=3.6" -files = [ - {file = "SecretStorage-3.3.3-py3-none-any.whl", hash = "sha256:f356e6628222568e3af06f2eba8df495efa13b3b63081dafd4f7d9a7b7bc9f99"}, - {file = "SecretStorage-3.3.3.tar.gz", hash = "sha256:2403533ef369eca6d2ba81718576c5e0f564d5cca1b58f73a8b23e7d4eeebd77"}, -] - -[package.dependencies] -cryptography = ">=2.0" -jeepney = ">=0.6" - -[[package]] -name = "setuptools" -version = "67.7.2" -description = "Easily download, build, install, upgrade, and uninstall Python packages" -optional = false -python-versions = ">=3.7" -files = [ - {file = "setuptools-67.7.2-py3-none-any.whl", hash = "sha256:23aaf86b85ca52ceb801d32703f12d77517b2556af839621c641fca11287952b"}, - {file = "setuptools-67.7.2.tar.gz", hash = "sha256:f104fa03692a2602fa0fec6c6a9e63b6c8a968de13e17c026957dd1f53d80990"}, -] - -[package.extras] -docs = ["furo", "jaraco.packaging (>=9)", "jaraco.tidelift (>=1.4)", "pygments-github-lexers (==0.0.5)", "rst.linker (>=1.9)", "sphinx (>=3.5)", "sphinx-favicon", "sphinx-hoverxref (<2)", "sphinx-inline-tabs", "sphinx-lint", "sphinx-notfound-page (==0.8.3)", "sphinx-reredirects", "sphinxcontrib-towncrier"] -testing = ["build[virtualenv]", "filelock (>=3.4.0)", "flake8 (<5)", "flake8-2020", "ini2toml[lite] (>=0.9)", "jaraco.envs (>=2.2)", "jaraco.path (>=3.2.0)", "pip (>=19.1)", "pip-run (>=8.8)", "pytest (>=6)", "pytest-black (>=0.3.7)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=1.3)", "pytest-flake8", "pytest-mypy (>=0.9.1)", "pytest-perf", "pytest-timeout", "pytest-xdist", "tomli-w (>=1.0.0)", "virtualenv (>=13.0.0)", "wheel"] -testing-integration = ["build[virtualenv]", "filelock (>=3.4.0)", "jaraco.envs (>=2.2)", "jaraco.path (>=3.2.0)", "pytest", "pytest-enabler", "pytest-xdist", "tomli", "virtualenv (>=13.0.0)", "wheel"] - -[[package]] -name = "six" -version = "1.16.0" -description = "Python 2 and 3 compatibility utilities" -optional = false -python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*" -files = [ - {file = "six-1.16.0-py2.py3-none-any.whl", hash = "sha256:8abb2f1d86890a2dfb989f9a77cfcfd3e47c2a354b01111771326f8aa26e0254"}, - {file = "six-1.16.0.tar.gz", hash = "sha256:1e61c37477a1626458e36f7b1d82aa5c9b094fa4802892072e49de9c60c4c926"}, -] - -[[package]] -name = "sniffio" -version = "1.3.0" -description = "Sniff out which async library your code is running under" -optional = false -python-versions = ">=3.7" -files = [ - {file = "sniffio-1.3.0-py3-none-any.whl", hash = "sha256:eecefdce1e5bbfb7ad2eeaabf7c1eeb404d7757c379bd1f7e5cce9d8bf425384"}, - {file = "sniffio-1.3.0.tar.gz", hash = "sha256:e60305c5e5d314f5389259b7f22aaa33d8f7dee49763119234af3755c55b9101"}, -] - -[[package]] -name = "sse-starlette" -version = "1.6.1" -description = "\"SSE plugin for Starlette\"" -optional = true -python-versions = ">=3.8" -files = [ - {file = "sse-starlette-1.6.1.tar.gz", hash = "sha256:6208af2bd7d0887c92f1379da14bd1f4db56bd1274cc5d36670c683d2aa1de6a"}, - {file = "sse_starlette-1.6.1-py3-none-any.whl", hash = "sha256:d8f18f1c633e355afe61cc5e9c92eea85badcb8b2d56ec8cfb0a006994aa55da"}, -] - -[package.dependencies] -starlette = "*" - -[[package]] -name = "starlette" -version = "0.27.0" -description = "The little ASGI library that shines." -optional = true -python-versions = ">=3.7" -files = [ - {file = "starlette-0.27.0-py3-none-any.whl", hash = "sha256:918416370e846586541235ccd38a474c08b80443ed31c578a418e2209b3eef91"}, - {file = "starlette-0.27.0.tar.gz", hash = "sha256:6a6b0d042acb8d469a01eba54e9cda6cbd24ac602c4cd016723117d6a7e73b75"}, -] - -[package.dependencies] -anyio = ">=3.4.0,<5" -typing-extensions = {version = ">=3.10.0", markers = "python_version < \"3.10\""} - -[package.extras] -full = ["httpx (>=0.22.0)", "itsdangerous", "jinja2", "python-multipart", "pyyaml"] - -[[package]] -name = "tomli" -version = "2.0.1" -description = "A lil' TOML parser" -optional = false -python-versions = ">=3.7" -files = [ - {file = "tomli-2.0.1-py3-none-any.whl", hash = "sha256:939de3e7a6161af0c887ef91b7d41a53e7c5a1ca976325f429cb46ea9bc30ecc"}, - {file = "tomli-2.0.1.tar.gz", hash = "sha256:de526c12914f0c550d15924c62d72abc48d6fe7364aa87328337a31007fe8a4f"}, -] - -[[package]] -name = "twine" -version = "4.0.2" -description = "Collection of utilities for publishing packages on PyPI" -optional = false -python-versions = ">=3.7" -files = [ - {file = "twine-4.0.2-py3-none-any.whl", hash = "sha256:929bc3c280033347a00f847236564d1c52a3e61b1ac2516c97c48f3ceab756d8"}, - {file = "twine-4.0.2.tar.gz", hash = "sha256:9e102ef5fdd5a20661eb88fad46338806c3bd32cf1db729603fe3697b1bc83c8"}, -] - -[package.dependencies] -importlib-metadata = ">=3.6" -keyring = ">=15.1" -pkginfo = ">=1.8.1" -readme-renderer = ">=35.0" -requests = ">=2.20" -requests-toolbelt = ">=0.8.0,<0.9.0 || >0.9.0" -rfc3986 = ">=1.4.0" -rich = ">=12.0.0" -urllib3 = ">=1.26.0" - -[[package]] -name = "typing-extensions" -version = "4.7.1" -description = "Backported and Experimental Type Hints for Python 3.7+" -optional = false -python-versions = ">=3.7" -files = [ - {file = "typing_extensions-4.7.1-py3-none-any.whl", hash = "sha256:440d5dd3af93b060174bf433bccd69b0babc3b15b1a8dca43789fd7f61514b36"}, - {file = "typing_extensions-4.7.1.tar.gz", hash = "sha256:b75ddc264f0ba5615db7ba217daeb99701ad295353c45f9e95963337ceeeffb2"}, -] - -[[package]] -name = "urllib3" -version = "2.0.2" -description = "HTTP library with thread-safe connection pooling, file post, and more." -optional = false -python-versions = ">=3.7" -files = [ - {file = "urllib3-2.0.2-py3-none-any.whl", hash = "sha256:d055c2f9d38dc53c808f6fdc8eab7360b6fdbbde02340ed25cfbcd817c62469e"}, - {file = "urllib3-2.0.2.tar.gz", hash = "sha256:61717a1095d7e155cdb737ac7bb2f4324a858a1e2e6466f6d03ff630ca68d3cc"}, -] - -[package.extras] -brotli = ["brotli (>=1.0.9)", "brotlicffi (>=0.8.0)"] -secure = ["certifi", "cryptography (>=1.9)", "idna (>=2.0.0)", "pyopenssl (>=17.1.0)", "urllib3-secure-extra"] -socks = ["pysocks (>=1.5.6,!=1.5.7,<2.0)"] -zstd = ["zstandard (>=0.18.0)"] - -[[package]] -name = "uvicorn" -version = "0.22.0" -description = "The lightning-fast ASGI server." -optional = true -python-versions = ">=3.7" -files = [ - {file = "uvicorn-0.22.0-py3-none-any.whl", hash = "sha256:e9434d3bbf05f310e762147f769c9f21235ee118ba2d2bf1155a7196448bd996"}, - {file = "uvicorn-0.22.0.tar.gz", hash = "sha256:79277ae03db57ce7d9aa0567830bbb51d7a612f54d6e1e3e92da3ef24c2c8ed8"}, -] - -[package.dependencies] -click = ">=7.0" -h11 = ">=0.8" - -[package.extras] -standard = ["colorama (>=0.4)", "httptools (>=0.5.0)", "python-dotenv (>=0.13)", "pyyaml (>=5.1)", "uvloop (>=0.14.0,!=0.15.0,!=0.15.1)", "watchfiles (>=0.13)", "websockets (>=10.4)"] - -[[package]] -name = "watchdog" -version = "3.0.0" -description = "Filesystem events monitoring" -optional = false -python-versions = ">=3.7" -files = [ - {file = "watchdog-3.0.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:336adfc6f5cc4e037d52db31194f7581ff744b67382eb6021c868322e32eef41"}, - {file = "watchdog-3.0.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:a70a8dcde91be523c35b2bf96196edc5730edb347e374c7de7cd20c43ed95397"}, - {file = "watchdog-3.0.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:adfdeab2da79ea2f76f87eb42a3ab1966a5313e5a69a0213a3cc06ef692b0e96"}, - {file = "watchdog-3.0.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:2b57a1e730af3156d13b7fdddfc23dea6487fceca29fc75c5a868beed29177ae"}, - {file = "watchdog-3.0.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:7ade88d0d778b1b222adebcc0927428f883db07017618a5e684fd03b83342bd9"}, - {file = "watchdog-3.0.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:7e447d172af52ad204d19982739aa2346245cc5ba6f579d16dac4bfec226d2e7"}, - {file = "watchdog-3.0.0-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:9fac43a7466eb73e64a9940ac9ed6369baa39b3bf221ae23493a9ec4d0022674"}, - {file = "watchdog-3.0.0-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:8ae9cda41fa114e28faf86cb137d751a17ffd0316d1c34ccf2235e8a84365c7f"}, - {file = "watchdog-3.0.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:25f70b4aa53bd743729c7475d7ec41093a580528b100e9a8c5b5efe8899592fc"}, - {file = "watchdog-3.0.0-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:4f94069eb16657d2c6faada4624c39464f65c05606af50bb7902e036e3219be3"}, - {file = "watchdog-3.0.0-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:7c5f84b5194c24dd573fa6472685b2a27cc5a17fe5f7b6fd40345378ca6812e3"}, - {file = "watchdog-3.0.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:3aa7f6a12e831ddfe78cdd4f8996af9cf334fd6346531b16cec61c3b3c0d8da0"}, - {file = "watchdog-3.0.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:233b5817932685d39a7896b1090353fc8efc1ef99c9c054e46c8002561252fb8"}, - {file = "watchdog-3.0.0-pp37-pypy37_pp73-macosx_10_9_x86_64.whl", hash = "sha256:13bbbb462ee42ec3c5723e1205be8ced776f05b100e4737518c67c8325cf6100"}, - {file = "watchdog-3.0.0-pp38-pypy38_pp73-macosx_10_9_x86_64.whl", hash = "sha256:8f3ceecd20d71067c7fd4c9e832d4e22584318983cabc013dbf3f70ea95de346"}, - {file = "watchdog-3.0.0-pp39-pypy39_pp73-macosx_10_9_x86_64.whl", hash = "sha256:c9d8c8ec7efb887333cf71e328e39cffbf771d8f8f95d308ea4125bf5f90ba64"}, - {file = "watchdog-3.0.0-py3-none-manylinux2014_aarch64.whl", hash = "sha256:0e06ab8858a76e1219e68c7573dfeba9dd1c0219476c5a44d5333b01d7e1743a"}, - {file = "watchdog-3.0.0-py3-none-manylinux2014_armv7l.whl", hash = "sha256:d00e6be486affb5781468457b21a6cbe848c33ef43f9ea4a73b4882e5f188a44"}, - {file = "watchdog-3.0.0-py3-none-manylinux2014_i686.whl", hash = "sha256:c07253088265c363d1ddf4b3cdb808d59a0468ecd017770ed716991620b8f77a"}, - {file = "watchdog-3.0.0-py3-none-manylinux2014_ppc64.whl", hash = "sha256:5113334cf8cf0ac8cd45e1f8309a603291b614191c9add34d33075727a967709"}, - {file = "watchdog-3.0.0-py3-none-manylinux2014_ppc64le.whl", hash = "sha256:51f90f73b4697bac9c9a78394c3acbbd331ccd3655c11be1a15ae6fe289a8c83"}, - {file = "watchdog-3.0.0-py3-none-manylinux2014_s390x.whl", hash = "sha256:ba07e92756c97e3aca0912b5cbc4e5ad802f4557212788e72a72a47ff376950d"}, - {file = "watchdog-3.0.0-py3-none-manylinux2014_x86_64.whl", hash = "sha256:d429c2430c93b7903914e4db9a966c7f2b068dd2ebdd2fa9b9ce094c7d459f33"}, - {file = "watchdog-3.0.0-py3-none-win32.whl", hash = "sha256:3ed7c71a9dccfe838c2f0b6314ed0d9b22e77d268c67e015450a29036a81f60f"}, - {file = "watchdog-3.0.0-py3-none-win_amd64.whl", hash = "sha256:4c9956d27be0bb08fc5f30d9d0179a855436e655f046d288e2bcc11adfae893c"}, - {file = "watchdog-3.0.0-py3-none-win_ia64.whl", hash = "sha256:5d9f3a10e02d7371cd929b5d8f11e87d4bad890212ed3901f9b4d68767bee759"}, - {file = "watchdog-3.0.0.tar.gz", hash = "sha256:4d98a320595da7a7c5a18fc48cb633c2e73cda78f93cac2ef42d42bf609a33f9"}, -] - -[package.extras] -watchmedo = ["PyYAML (>=3.10)"] - -[[package]] -name = "webencodings" -version = "0.5.1" -description = "Character encoding aliases for legacy web content" -optional = false -python-versions = "*" -files = [ - {file = "webencodings-0.5.1-py2.py3-none-any.whl", hash = "sha256:a0af1213f3c2226497a97e2b3aa01a7e4bee4f403f95be16fc9acd2947514a78"}, - {file = "webencodings-0.5.1.tar.gz", hash = "sha256:b36a1c245f2d304965eb4e0a82848379241dc04b865afcc4aab16748587e1923"}, -] - -[[package]] -name = "wheel" -version = "0.40.0" -description = "A built-package format for Python" -optional = false -python-versions = ">=3.7" -files = [ - {file = "wheel-0.40.0-py3-none-any.whl", hash = "sha256:d236b20e7cb522daf2390fa84c55eea81c5c30190f90f29ae2ca1ad8355bf247"}, - {file = "wheel-0.40.0.tar.gz", hash = "sha256:cd1196f3faee2b31968d626e1731c94f99cbdb67cf5a46e4f5656cbee7738873"}, -] - -[package.extras] -test = ["pytest (>=6.0.0)"] - -[[package]] -name = "zipp" -version = "3.15.0" -description = "Backport of pathlib-compatible object wrapper for zip files" -optional = false -python-versions = ">=3.7" -files = [ - {file = "zipp-3.15.0-py3-none-any.whl", hash = "sha256:48904fc76a60e542af151aded95726c1a5c34ed43ab4134b597665c86d7ad556"}, - {file = "zipp-3.15.0.tar.gz", hash = "sha256:112929ad649da941c23de50f356a2b5570c954b65150642bccdd66bf194d224b"}, -] - -[package.extras] -docs = ["furo", "jaraco.packaging (>=9)", "jaraco.tidelift (>=1.4)", "rst.linker (>=1.9)", "sphinx (>=3.5)", "sphinx-lint"] -testing = ["big-O", "flake8 (<5)", "jaraco.functools", "jaraco.itertools", "more-itertools", "pytest (>=6)", "pytest-black (>=0.3.7)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=1.3)", "pytest-flake8", "pytest-mypy (>=0.9.1)"] - -[extras] -server = ["fastapi", "sse-starlette", "uvicorn"] - -[metadata] -lock-version = "2.0" -python-versions = "^3.8.1" -content-hash = "da42c48a426b64ce393b4febca1be0e2ea0fe9d48cedb2392b390d4a49276474" diff --git a/poetry.toml b/poetry.toml deleted file mode 100644 index be97f1ef2f..0000000000 --- a/poetry.toml +++ /dev/null @@ -1,3 +0,0 @@ -[virtualenvs] -in-project = true -prefer-active-python = true \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index a9e012e6b7..6bfacf279c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,44 +1,95 @@ -[tool.poetry] +[build-system] +requires = ["scikit-build-core[pyproject]>=0.9.2"] +build-backend = "scikit_build_core.build" + +[project] name = "llama_cpp_python" -version = "0.1.70" +dynamic = ["version"] description = "Python bindings for the llama.cpp library" -authors = ["Andrei Betlen <abetlen@gmail.com>"] -license = "MIT" readme = "README.md" -homepage = "https://github.com/abetlen/llama-cpp-python" -repository = "https://github.com/abetlen/llama-cpp-python" -packages = [{include = "llama_cpp"}] -include = [ - "LICENSE.md", +license = { text = "MIT" } +authors = [ + { name = "Andrei Betlen", email = "abetlen@gmail.com" }, +] +dependencies = [ + "typing-extensions>=4.5.0", + "numpy>=1.20.0", + "diskcache>=5.6.1", + "jinja2>=2.11.3", +] +requires-python = ">=3.8" +classifiers = [ + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Python :: 3.14", ] -[tool.poetry.dependencies] -python = "^3.8.1" -typing-extensions = "^4.7.1" -numpy = "^1.24.4" -diskcache = "^5.6.1" -uvicorn = { version = "^0.22.0", optional = true } -fastapi = { version = "^0.99.1", optional = true } -sse-starlette = { version = "^1.6.1", optional = true } - -[tool.poetry.group.dev.dependencies] -black = "^23.3.0" -twine = "^4.0.2" -mkdocs = "^1.4.3" -mkdocstrings = {extras = ["python"], version = "^0.22.0"} -mkdocs-material = "^9.1.18" -pytest = "^7.4.0" -httpx = "^0.24.1" -scikit-build = "0.17.6" - -[tool.poetry.extras] -server = ["uvicorn>=0.22.0", "fastapi>=0.100.0", "pydantic-settings>=2.0.1", "sse-starlette>=1.6.1"] -[build-system] -requires = [ - "setuptools>=42", - "scikit-build>=0.13", - "cmake>=3.18", - "ninja", +[project.optional-dependencies] +server = [ + "uvicorn>=0.22.0", + "fastapi>=0.100.0", + "pydantic-settings>=2.0.1", + "sse-starlette>=1.6.1", + "starlette-context>=0.3.6,<0.4", + "PyYAML>=5.1", +] +test = [ + "pytest>=7.4.0", + "httpx>=0.24.1", + "scipy>=1.10", + "fastapi>=0.100.0", + "sse-starlette>=1.6.1", + "starlette-context>=0.3.6,<0.4", + "pydantic-settings>=2.0.1", + "huggingface-hub>=0.23.0" +] +dev = [ + "ruff>=0.15.7", + "twine>=4.0.2", + "mkdocs>=1.4.3", + "mkdocstrings[python]>=0.22.0", + "mkdocs-material>=9.1.18", + "pytest>=7.4.0", + "httpx>=0.24.1", ] -build-backend = "setuptools.build_meta" \ No newline at end of file +all = [ + "llama_cpp_python[server,test,dev]", +] + +[tool.scikit-build] +wheel.packages = ["llama_cpp"] +wheel.py-api = "py3" +cmake.verbose = true +cmake.minimum-version = "3.21" +minimum-version = "0.5.1" +sdist.include = [".git", "vendor/llama.cpp/*"] + +[tool.scikit-build.metadata.version] +provider = "scikit_build_core.metadata.regex" +input = "llama_cpp/__init__.py" + +[project.urls] +Homepage = "https://github.com/abetlen/llama-cpp-python" +Issues = "https://github.com/abetlen/llama-cpp-python/issues" +Documentation = "https://llama-cpp-python.readthedocs.io/en/latest/" +Changelog = "https://llama-cpp-python.readthedocs.io/en/latest/changelog/" + +[tool.ruff] +target-version = "py38" +line-length = 88 +required-version = ">=0.15.7" +src = ["llama_cpp", "tests"] +extend-exclude = ["vendor", "examples/notebooks"] + +[tool.ruff.lint] +select = ["E4", "E7", "E9"] +ignore = ["E712"] + +[tool.pytest.ini_options] +testpaths = "tests" diff --git a/scripts/get-releases.sh b/scripts/get-releases.sh new file mode 100755 index 0000000000..4c904da78c --- /dev/null +++ b/scripts/get-releases.sh @@ -0,0 +1,47 @@ +#!/bin/bash + +# Function to get all releases +get_all_releases() { + local page=1 + local per_page=100 + local releases="" + local new_releases + + # Prepare headers + local headers=(-H "Accept: application/vnd.github.v3+json") + if [ -n "$GITHUB_TOKEN" ]; then + headers+=(-H "Authorization: Bearer $GITHUB_TOKEN") + fi + + while true; do + response=$(curl -s "${headers[@]}" \ + "https://api.github.com/repos/abetlen/llama-cpp-python/releases?page=$page&per_page=$per_page") + + # Check if the response is valid JSON + if ! echo "$response" | jq empty > /dev/null 2>&1; then + echo "Error: Invalid response from GitHub API" >&2 + echo "Response: $response" >&2 + return 1 + fi + + new_releases=$(echo "$response" | jq -r '.[].tag_name') + if [ -z "$new_releases" ]; then + break + fi + releases="$releases $new_releases" + ((page++)) + done + + echo $releases +} + +# Get all releases and save to file +releases=$(get_all_releases) +if [ $? -ne 0 ]; then + echo "Failed to fetch releases. Please check your internet connection and try again later." >&2 + exit 1 +fi + +echo "$releases" | tr ' ' '\n' > all_releases.txt + +echo "All releases have been saved to all_releases.txt" diff --git a/scripts/releases-to-pep-503.sh b/scripts/releases-to-pep-503.sh new file mode 100755 index 0000000000..8359624492 --- /dev/null +++ b/scripts/releases-to-pep-503.sh @@ -0,0 +1,108 @@ +#!/bin/bash + +# Enable exit on error +set -e + +# Function for logging +log_error() { + echo "ERROR: $1" >&2 +} + +log_info() { + echo "INFO: $1" +} + +# Get output directory or default to index/whl/cpu +output_dir=${1:-"index/whl/cpu"} + +# Get pattern from second arg or default to valid python package version pattern +pattern=${2:-"^[v]?[0-9]+\.[0-9]+\.[0-9]+$"} + +# Get the current directory (where the script is run from) +current_dir="$(pwd)" + +# Check if all_releases.txt exists +if [ ! -f "$current_dir/all_releases.txt" ]; then + log_error "all_releases.txt not found in the current directory." + exit 1 +fi + +# Create output directory +mkdir -p "$output_dir" + +# Create an index html file +cat << EOF > "$output_dir/index.html" +<!DOCTYPE html> +<html> + <head></head> + <body> + <a href="llama-cpp-python/">llama-cpp-python</a> + <br> + </body> +</html> + +EOF + +# Create llama-cpp-python directory +mkdir -p "$output_dir/llama-cpp-python" + +# Create an index html file in llama-cpp-python directory +cat << EOF > "$output_dir/llama-cpp-python/index.html" +<!DOCTYPE html> +<html> + <body> + <h1>Links for llama-cpp-python</h1> +EOF + +# Filter releases by pattern. Some backend indexes are valid even when there +# are no matching releases yet. +releases=$(grep -E "$pattern" "$current_dir/all_releases.txt" || true) +if [ -z "$releases" ]; then + log_info "No releases found matching pattern: $pattern" +fi + +# Prepare curl headers +headers=('--header' 'Accept: application/vnd.github.v3+json') +if [ -n "$GITHUB_TOKEN" ]; then + headers+=('--header' "authorization: Bearer $GITHUB_TOKEN") +fi +headers+=('--header' 'content-type: application/json') + +# For each release, get all assets +for release in $releases; do + log_info "Processing release: $release" + response=$(curl -s "${headers[@]}" \ + "https://api.github.com/repos/abetlen/llama-cpp-python/releases/tags/$release") + + if [ -z "$response" ]; then + log_error "Empty response from GitHub API for release $release" + continue + fi + + if ! echo "$response" | jq -e '.assets' > /dev/null 2>&1; then + log_error "Invalid or unexpected response from GitHub API for release $release" + log_error "Response: $response" + continue + fi + + wheel_urls=$(echo "$response" | jq -r '.assets[] | select(.name | endswith(".whl")) | .browser_download_url') + if [ -z "$wheel_urls" ]; then + log_error "No wheel files found for release $release" + continue + fi + + # Get release version from release ie v0.1.0-cu121 -> v0.1.0 + release_version=$(echo "$release" | grep -oE "^[v]?[0-9]+\.[0-9]+\.[0-9]+") + echo " <h2>$release_version</h2>" >> "$output_dir/llama-cpp-python/index.html" + + echo "$wheel_urls" | while read -r asset; do + echo " <a href=\"$asset\">$asset</a>" >> "$output_dir/llama-cpp-python/index.html" + echo " <br>" >> "$output_dir/llama-cpp-python/index.html" + done +done + +echo " </body>" >> "$output_dir/llama-cpp-python/index.html" +echo "</html>" >> "$output_dir/llama-cpp-python/index.html" +echo "" >> "$output_dir/llama-cpp-python/index.html" + +log_info "Index generation complete. Output directory: $output_dir" diff --git a/setup.py b/setup.py deleted file mode 100644 index b8acedb5a6..0000000000 --- a/setup.py +++ /dev/null @@ -1,32 +0,0 @@ -from skbuild import setup - -from pathlib import Path - -this_directory = Path(__file__).parent -long_description = (this_directory / "README.md").read_text(encoding="utf-8") - -setup( - name="llama_cpp_python", - description="A Python wrapper for llama.cpp", - long_description=long_description, - long_description_content_type="text/markdown", - version="0.1.70", - author="Andrei Betlen", - author_email="abetlen@gmail.com", - license="MIT", - package_dir={"llama_cpp": "llama_cpp", "llama_cpp.server": "llama_cpp/server"}, - packages=["llama_cpp", "llama_cpp.server"], - install_requires=["typing-extensions>=4.5.0", "numpy>=1.20.0", "diskcache>=5.6.1"], - extras_require={ - "server": ["uvicorn>=0.22.1", "fastapi>=0.100.0", "pydantic-settings>=2.0.1", "sse-starlette>=1.6.1"], - }, - python_requires=">=3.7", - classifiers=[ - "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.7", - "Programming Language :: Python :: 3.8", - "Programming Language :: Python :: 3.9", - "Programming Language :: Python :: 3.10", - "Programming Language :: Python :: 3.11", - ], -) diff --git a/tests/test_llama.py b/tests/test_llama.py index 941287de62..336d6a6122 100644 --- a/tests/test_llama.py +++ b/tests/test_llama.py @@ -1,171 +1,367 @@ +import ctypes +import multiprocessing + +import numpy as np +from scipy.special import log_softmax + +from huggingface_hub import hf_hub_download + +import pytest + import llama_cpp +import llama_cpp._internals as internals + + +MODEL = "./vendor/llama.cpp/models/ggml-vocab-llama-spm.gguf" -MODEL = "./vendor/llama.cpp/models/ggml-vocab.bin" +def test_llama_cpp_version(): + assert llama_cpp.__version__ -def test_llama(): - llama = llama_cpp.Llama(model_path=MODEL, vocab_only=True) + +def test_llama_cpp_tokenization(): + llama = llama_cpp.Llama(model_path=MODEL, vocab_only=True, verbose=False) assert llama - assert llama.ctx is not None + assert llama._ctx.ctx is not None text = b"Hello World" - assert llama.detokenize(llama.tokenize(text)) == text + tokens = llama.tokenize(text) + assert tokens[0] == llama.token_bos() + assert tokens == [1, 15043, 2787] + detokenized = llama.detokenize(tokens) + assert detokenized == text + + tokens = llama.tokenize(text, add_bos=False) + assert tokens[0] != llama.token_bos() + assert tokens == [15043, 2787] + + detokenized = llama.detokenize(tokens) + assert detokenized != text + + text = b"Hello World</s>" + tokens = llama.tokenize(text) + assert tokens[-1] != llama.token_eos() + assert tokens == [1, 15043, 2787, 829, 29879, 29958] + + tokens = llama.tokenize(text, special=True) + assert tokens[-1] == llama.token_eos() + assert tokens == [1, 15043, 2787, 2] + + text = b"" + tokens = llama.tokenize(text, add_bos=True, special=True) + assert tokens[-1] != llama.token_eos() + assert tokens == [llama.token_bos()] + assert text == llama.detokenize(tokens) + + +@pytest.fixture +def llama_cpp_model_path(): + repo_id = "lmstudio-community/Qwen3.5-0.8B-GGUF" + filename = "Qwen3.5-0.8B-Q8_0.gguf" + model_path = hf_hub_download(repo_id, filename) + return model_path + + +@pytest.fixture +def llama_cpp_embedding_model_path(): + repo_id = "CompendiumLabs/bge-small-en-v1.5-gguf" + filename = "bge-small-en-v1.5-q4_k_m.gguf" + model_path = hf_hub_download(repo_id, filename) + return model_path + + +@pytest.fixture +def llama_cpp_recurrent_model_path(): + repo_id = "QuantFactory/mamba-130m-hf-GGUF" + filename = "mamba-130m-hf.Q2_K.gguf" + model_path = hf_hub_download(repo_id, filename) + return model_path + + +@pytest.fixture +def llama_cpp_hybrid_model_path(): + repo_id = "tiiuae/Falcon-H1-Tiny-90M-Instruct-GGUF" + filename = "Falcon-H1-Tiny-90M-Instruct-Q2_K.gguf" + model_path = hf_hub_download(repo_id, filename) + return model_path + + +def test_real_model(llama_cpp_model_path): + import os + + assert os.path.exists(llama_cpp_model_path) + params = llama_cpp.llama_model_default_params() + params.use_mmap = llama_cpp.llama_supports_mmap() + params.use_mlock = llama_cpp.llama_supports_mlock() + params.check_tensors = False -# @pytest.mark.skip(reason="need to update sample mocking") -def test_llama_patch(monkeypatch): - llama = llama_cpp.Llama(model_path=MODEL, vocab_only=True) - n_vocab = llama_cpp.llama_n_vocab(llama.ctx) + model = internals.LlamaModel(path_model=llama_cpp_model_path, params=params) - ## Set up mock function - def mock_eval(*args, **kwargs): - return 0 + cparams = llama_cpp.llama_context_default_params() + cparams.n_ctx = 16 + cparams.n_batch = 16 + cparams.n_ubatch = 16 + cparams.n_threads = multiprocessing.cpu_count() + cparams.n_threads_batch = multiprocessing.cpu_count() + cparams.logits_all = False + cparams.flash_attn_type = llama_cpp.LLAMA_FLASH_ATTN_TYPE_ENABLED - def mock_get_logits(*args, **kwargs): - return (llama_cpp.c_float * n_vocab)( - *[llama_cpp.c_float(0) for _ in range(n_vocab)] - ) + context = internals.LlamaContext(model=model, params=cparams) + tokens = model.tokenize(b"Hello, world!", add_bos=True, special=True) - monkeypatch.setattr("llama_cpp.llama_cpp.llama_eval", mock_eval) - monkeypatch.setattr("llama_cpp.llama_cpp.llama_get_logits", mock_get_logits) + assert tokens == [9419, 11, 1814, 0] - output_text = " jumps over the lazy dog." - output_tokens = llama.tokenize(output_text.encode("utf-8")) - token_eos = llama.token_eos() - n = 0 + tokens = model.tokenize( + b"The quick brown fox jumps over the lazy dog. The quick brown fox jumps ", + add_bos=True, + special=True, + ) + prompt_token_count = len(tokens) + + batch = internals.LlamaBatch(n_tokens=len(tokens), embd=0, n_seq_max=1) + + seed = 1337 + sampler = internals.LlamaSampler() + sampler.add_top_k(50) + sampler.add_top_p(0.9, 1) + sampler.add_temp(0.8) + sampler.add_dist(seed) + + result = tokens + n_eval = 0 + for _ in range(4): + batch.set_batch(tokens, n_past=n_eval, logits_all=False) + context.decode(batch) + n_eval += len(tokens) + token_id = sampler.sample(context, -1) + tokens = [token_id] + result += tokens + + output = result[prompt_token_count:] + output_text = model.detokenize(output, special=True) + # Low-level sampling output varies across CPU and Metal backends. + assert len(output) == 4 + assert output_text + + +def test_real_llama(llama_cpp_model_path): + model = llama_cpp.Llama( + llama_cpp_model_path, + n_ctx=32, + n_batch=32, + n_ubatch=32, + n_threads=multiprocessing.cpu_count(), + n_threads_batch=multiprocessing.cpu_count(), + logits_all=False, + flash_attn=True, + ) - def mock_sample(*args, **kwargs): - nonlocal n - if n < len(output_tokens): - n += 1 - return output_tokens[n - 1] - else: - return token_eos + output = model.create_completion( + "The quick brown fox jumps over the lazy dog. The quick brown fox", + max_tokens=6, + top_k=50, + top_p=0.9, + temperature=0.0, + seed=1337, + ) + assert output["choices"][0]["text"] == " jumps over the lazy dog." + + output = model.create_completion( + "The capital of france is paris, 'true' or 'false'?:\n", + max_tokens=4, + top_k=50, + top_p=0.9, + temperature=0.8, + seed=1337, + grammar=llama_cpp.LlamaGrammar.from_string(""" +root ::= "true" | "false" +"""), + ) + assert output["choices"][0]["text"] == "true" + + suffix = b"rot" + tokens = model.tokenize(suffix, add_bos=True, special=True) + + def logit_processor_func(input_ids, logits): + for token in tokens: + logits[token] *= 1000 + return logits + + logit_processors = llama_cpp.LogitsProcessorList([logit_processor_func]) + + output = model.create_completion( + "The capital of france is par", + max_tokens=4, + top_k=50, + top_p=0.9, + temperature=0.8, + seed=1337, + logits_processor=logit_processors, + ) + assert output["choices"][0]["text"].lower().startswith("rot") - monkeypatch.setattr("llama_cpp.llama_cpp.llama_sample_token", mock_sample) + model.set_seed(1337) - text = "The quick brown fox" + state = model.save_state() - ## Test basic completion until eos - n = 0 # reset - completion = llama.create_completion(text, max_tokens=20) - assert completion["choices"][0]["text"] == output_text - assert completion["choices"][0]["finish_reason"] == "stop" + output = model.create_completion( + "Pick a number from 1 to 10?:\n", + max_tokens=4, + top_k=50, + top_p=0.9, + temperature=1.0, + grammar=llama_cpp.LlamaGrammar.from_string(""" +root ::= "1" | "2" | "3" | "4" | "5" | "6" | "7" | "8" | "9" | "10" +"""), + ) + number_1 = output["choices"][0]["text"] + + output = model.create_completion( + "Pick a number from 1 to 10?:\n", + max_tokens=4, + top_k=50, + top_p=0.9, + temperature=1.0, + grammar=llama_cpp.LlamaGrammar.from_string(""" +root ::= "1" | "2" | "3" | "4" | "5" | "6" | "7" | "8" | "9" | "10" +"""), + ) + number_2 = output["choices"][0]["text"] + + model.load_state(state) + + output = model.create_completion( + "Pick a number from 1 to 10?:\n", + max_tokens=4, + top_k=50, + top_p=0.9, + temperature=1.0, + grammar=llama_cpp.LlamaGrammar.from_string(""" +root ::= "1" | "2" | "3" | "4" | "5" | "6" | "7" | "8" | "9" | "10" +"""), + ) + number_3 = output["choices"][0]["text"] + + assert number_1 != number_2 + assert number_1 == number_3 + + +def test_real_llama_repeated_prompt_cache(llama_cpp_model_path): + model = llama_cpp.Llama( + llama_cpp_model_path, + n_ctx=32, + n_batch=32, + n_ubatch=32, + n_threads=multiprocessing.cpu_count(), + n_threads_batch=multiprocessing.cpu_count(), + logits_all=False, + flash_attn=True, + verbose=False, + ) + prompt = "The quick brown fox jumps over the lazy dog. The quick brown fox" - ## Test streaming completion until eos - n = 0 # reset - chunks = llama.create_completion(text, max_tokens=20, stream=True) - assert "".join(chunk["choices"][0]["text"] for chunk in chunks) == output_text - assert completion["choices"][0]["finish_reason"] == "stop" + output_1 = model.create_completion( + prompt, + max_tokens=6, + temperature=0.0, + seed=1337, + ) + output_2 = model.create_completion( + prompt, + max_tokens=6, + temperature=0.0, + seed=1337, + ) - ## Test basic completion until stop sequence - n = 0 # reset - completion = llama.create_completion(text, max_tokens=20, stop=["lazy"]) - assert completion["choices"][0]["text"] == " jumps over the " - assert completion["choices"][0]["finish_reason"] == "stop" + assert output_1["choices"][0]["text"] == " jumps over the lazy dog." + assert output_2["choices"][0]["text"] == output_1["choices"][0]["text"] + + +def _assert_prompt_cache_reset_handles_history_edit( + model_path, + *, + is_recurrent: bool, + is_hybrid: bool, +): + model = llama_cpp.Llama( + model_path, + n_ctx=32, + n_batch=32, + n_ubatch=32, + n_threads=multiprocessing.cpu_count(), + n_threads_batch=multiprocessing.cpu_count(), + logits_all=False, + verbose=False, + ) - ## Test streaming completion until stop sequence - n = 0 # reset - chunks = llama.create_completion(text, max_tokens=20, stream=True, stop=["lazy"]) - assert ( - "".join(chunk["choices"][0]["text"] for chunk in chunks) == " jumps over the " + assert model._is_recurrent is is_recurrent + assert model._is_hybrid is is_hybrid + + first_prompt = "The quick brown fox" + second_prompt = "The slow brown fox" + first_tokens = model.tokenize(first_prompt.encode(), add_bos=True, special=True) + second_tokens = model.tokenize(second_prompt.encode(), add_bos=True, special=True) + + assert first_tokens != second_tokens + assert first_tokens[0] == second_tokens[0] + + first_output = model.create_completion( + first_prompt, + max_tokens=1, + temperature=0.0, ) - assert completion["choices"][0]["finish_reason"] == "stop" + assert isinstance(first_output["choices"][0]["text"], str) - ## Test basic completion until length - n = 0 # reset - completion = llama.create_completion(text, max_tokens=2) - assert completion["choices"][0]["text"] == " j" - assert completion["choices"][0]["finish_reason"] == "length" + second_output = model.create_completion( + second_prompt, + max_tokens=1, + temperature=0.0, + ) + assert isinstance(second_output["choices"][0]["text"], str) - ## Test streaming completion until length - n = 0 # reset - chunks = llama.create_completion(text, max_tokens=2, stream=True) - assert "".join(chunk["choices"][0]["text"] for chunk in chunks) == " j" - assert completion["choices"][0]["finish_reason"] == "length" +def test_recurrent_model_prompt_cache_reset(llama_cpp_recurrent_model_path): + _assert_prompt_cache_reset_handles_history_edit( + llama_cpp_recurrent_model_path, + is_recurrent=True, + is_hybrid=False, + ) -def test_llama_pickle(): - import pickle - import tempfile - fp = tempfile.TemporaryFile() - llama = llama_cpp.Llama(model_path=MODEL, vocab_only=True) - pickle.dump(llama, fp) - fp.seek(0) - llama = pickle.load(fp) +def test_hybrid_model_prompt_cache_reset(llama_cpp_hybrid_model_path): + _assert_prompt_cache_reset_handles_history_edit( + llama_cpp_hybrid_model_path, + is_recurrent=False, + is_hybrid=True, + ) - assert llama - assert llama.ctx is not None - text = b"Hello World" +def test_real_llama_embeddings(llama_cpp_embedding_model_path): + model = llama_cpp.Llama( + llama_cpp_embedding_model_path, + n_ctx=32, + n_batch=32, + n_ubatch=32, + n_threads=multiprocessing.cpu_count(), + n_threads_batch=multiprocessing.cpu_count(), + logits_all=False, + flash_attn=True, + embedding=True, + ) + embedding = model.embed("Hello World") + assert len(embedding) > 0 + + prompts = ["Hello World", "A different prompt"] + individual_embeddings = [model.embed(prompt) for prompt in prompts] + batched_embeddings = model.embed(prompts) + + assert len(batched_embeddings) == len(prompts) + for individual, batched in zip(individual_embeddings, batched_embeddings): + np.testing.assert_allclose(batched, individual, rtol=1e-4, atol=1e-4) - assert llama.detokenize(llama.tokenize(text)) == text - - -def test_utf8(monkeypatch): - llama = llama_cpp.Llama(model_path=MODEL, vocab_only=True) - n_vocab = llama_cpp.llama_n_vocab(llama.ctx) - - ## Set up mock function - def mock_eval(*args, **kwargs): - return 0 - - def mock_get_logits(*args, **kwargs): - return (llama_cpp.c_float * n_vocab)( - *[llama_cpp.c_float(0) for _ in range(n_vocab)] - ) - - monkeypatch.setattr("llama_cpp.llama_cpp.llama_eval", mock_eval) - monkeypatch.setattr("llama_cpp.llama_cpp.llama_get_logits", mock_get_logits) - - output_text = "😀" - output_tokens = llama.tokenize(output_text.encode("utf-8")) - token_eos = llama.token_eos() - n = 0 - - def mock_sample(*args, **kwargs): - nonlocal n - if n < len(output_tokens): - n += 1 - return output_tokens[n - 1] - else: - return token_eos - - monkeypatch.setattr("llama_cpp.llama_cpp.llama_sample_token", mock_sample) - - ## Test basic completion with utf8 multibyte - n = 0 # reset - completion = llama.create_completion("", max_tokens=4) - assert completion["choices"][0]["text"] == output_text - - ## Test basic completion with incomplete utf8 multibyte - n = 0 # reset - completion = llama.create_completion("", max_tokens=1) - assert completion["choices"][0]["text"] == "" - - -def test_llama_server(): - from fastapi.testclient import TestClient - from llama_cpp.server.app import create_app, Settings - - settings = Settings( - model=MODEL, - vocab_only=True, - ) - app = create_app(settings) - client = TestClient(app) - response = client.get("/v1/models") - assert response.json() == { - "object": "list", - "data": [ - { - "id": MODEL, - "object": "model", - "owned_by": "me", - "permissions": [], - } - ], - } + repeated_embeddings = model.embed(list(reversed(prompts))) + assert len(repeated_embeddings) == len(prompts) + assert all(len(repeated) == len(embedding) for repeated in repeated_embeddings) diff --git a/tests/test_llama_chat_format.py b/tests/test_llama_chat_format.py new file mode 100644 index 0000000000..18c7279cf0 --- /dev/null +++ b/tests/test_llama_chat_format.py @@ -0,0 +1,94 @@ +import json + +import jinja2 + +from llama_cpp import ( + ChatCompletionRequestUserMessage, +) +import llama_cpp.llama_types as llama_types +import llama_cpp.llama_chat_format as llama_chat_format + +from llama_cpp.llama_chat_format import hf_tokenizer_config_to_chat_formatter + + +def test_mistral_instruct(): + chat_template = "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ '[INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ message['content'] + eos_token}}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}" + chat_formatter = jinja2.Template(chat_template) + messages = [ + llama_types.ChatCompletionRequestUserMessage( + role="user", content="Instruction" + ), + llama_types.ChatCompletionRequestAssistantMessage( + role="assistant", content="Model answer" + ), + llama_types.ChatCompletionRequestUserMessage( + role="user", content="Follow-up instruction" + ), + ] + response = llama_chat_format.format_mistral_instruct( + messages=messages, + ) + prompt = ("" if response.added_special else "<s>") + response.prompt + reference = chat_formatter.render( + messages=messages, + bos_token="<s>", + eos_token="</s>", + ) + assert prompt == reference + + +mistral_7b_tokenizer_config = """{ + "add_bos_token": true, + "add_eos_token": false, + "added_tokens_decoder": { + "0": { + "content": "<unk>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "1": { + "content": "<s>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + }, + "2": { + "content": "</s>", + "lstrip": false, + "normalized": false, + "rstrip": false, + "single_word": false, + "special": true + } + }, + "additional_special_tokens": [], + "bos_token": "<s>", + "clean_up_tokenization_spaces": false, + "eos_token": "</s>", + "legacy": true, + "model_max_length": 1000000000000000019884624838656, + "pad_token": null, + "sp_model_kwargs": {}, + "spaces_between_special_tokens": false, + "tokenizer_class": "LlamaTokenizer", + "unk_token": "<unk>", + "use_default_system_prompt": false, + "chat_template": "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ '[INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ message['content'] + eos_token}}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}" +}""" + + +def test_hf_tokenizer_config_str_to_chat_formatter(): + tokenizer_config = json.loads(mistral_7b_tokenizer_config) + chat_formatter = hf_tokenizer_config_to_chat_formatter(tokenizer_config) + chat_formatter_respoonse = chat_formatter( + messages=[ + ChatCompletionRequestUserMessage(role="user", content="Hello, world!"), + ] + ) + + assert chat_formatter_respoonse.prompt == ("<s>[INST] Hello, world! [/INST]</s>") diff --git a/tests/test_llama_grammar.py b/tests/test_llama_grammar.py new file mode 100644 index 0000000000..34ef2874df --- /dev/null +++ b/tests/test_llama_grammar.py @@ -0,0 +1,78 @@ +import llama_cpp +import json + +tree = """ +leaf ::= "." +node ::= leaf | "(" node node ")" +root ::= node +""" + + +def test_grammar_from_string(): + grammar = llama_cpp.LlamaGrammar.from_string(tree) + # assert grammar._n_rules == 3 + # assert grammar._start_rule_index == 2 + # assert grammar.grammar is not None + + +def test_composed_pydantic_grammar(): + """ + from pydantic import BaseModel + + class A(BaseModel): + a: int + + class B(BaseModel): + a: A + b: int + """ + + # This schema corresponds to the grammar in the comment above. + # We don't use the pydantic models directly to avoid the dependency. + schema = { + "$defs": { + "A": { + "properties": {"a": {"title": "A", "type": "integer"}}, + "required": ["a"], + "title": "A", + "type": "object", + } + }, + "properties": { + "a": {"$ref": "#/$defs/A"}, + "b": {"title": "B", "type": "integer"}, + }, + "required": ["a", "b"], + "title": "B", + "type": "object", + } + + grammar = llama_cpp.LlamaGrammar.from_json_schema(json.dumps(schema)) + + # assert grammar.grammar is not None + + +def test_grammar_anyof(): + sch = { + "properties": { + "temperature": { + "description": "The temperature mentioned", + "type": "number", + }, + "unit": { + "anyOf": [ + { + "description": "Unit for temperature", + "enum": ["celsius", "fahrenheit"], + "type": "string", + }, + {"type": "null"}, + ], + }, + }, + "type": "object", + } + + grammar = llama_cpp.LlamaGrammar.from_json_schema(json.dumps(sch)) + + # assert grammar.grammar is not None diff --git a/tests/test_llama_speculative.py b/tests/test_llama_speculative.py new file mode 100644 index 0000000000..d28c9ca9c6 --- /dev/null +++ b/tests/test_llama_speculative.py @@ -0,0 +1,21 @@ +import numpy as np + +from llama_cpp.llama_speculative import LlamaPromptLookupDecoding + + +def test_find_candidate_pred_tokens(): + find_candidate_pred_tokens = LlamaPromptLookupDecoding.find_candidate_pred_tokens + + # Test Case 1: Matching ngram is found + input_ids1 = np.array([1, 2, 3, 1, 2, 3, 1, 2, 3]) + result1 = find_candidate_pred_tokens( + input_ids1, max_ngram_size=3, num_pred_tokens=2 + ) + assert np.array_equal(result1, np.array([1, 2])) + + # Test Case 2: Matching ngram is not found + input_ids2 = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9]) + result2 = find_candidate_pred_tokens( + input_ids2, max_ngram_size=3, num_pred_tokens=2 + ) + assert np.array_equal(result2, np.array([])) diff --git a/vendor/llama.cpp b/vendor/llama.cpp index 1d16309969..e3471b3e73 160000 --- a/vendor/llama.cpp +++ b/vendor/llama.cpp @@ -1 +1 @@ -Subproject commit 1d1630996920f889cdc08de26cebf2415958540e +Subproject commit e3471b3e7306fe120dc8f38a2263c1293fc2add7