diff --git a/.dockerignore b/.dockerignore index fd64c09b3..e7993c69d 100644 --- a/.dockerignore +++ b/.dockerignore @@ -164,3 +164,6 @@ cython_debug/ # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. .idea/ + +.git/ +*.tgz diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 91abb11fd..6bf90273a 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 63c81f1e3..785f38847 100644 --- a/.github/workflows/build-and-release.yaml +++ b/.github/workflows/build-and-release.yaml @@ -11,34 +11,64 @@ jobs: runs-on: ${{ matrix.os }} strategy: matrix: - os: [ubuntu-latest, windows-latest, macOS-latest] + os: [ubuntu-20.04, windows-2019, macos-12] steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v4 with: submodules: "recursive" # Used to host cibuildwheel - - uses: actions/setup-python@v3 + - uses: actions/setup-python@v5 with: python-version: "3.8" - - name: Install cibuildwheel - run: python -m pip install cibuildwheel==2.12.1 - - name: Install dependencies run: | python -m pip install --upgrade pip python -m pip install -e .[all] - name: Build wheels - run: python -m cibuildwheel --output-dir wheelhouse + uses: pypa/cibuildwheel@v2.20.0 env: # disable repair CIBW_REPAIR_WHEEL_COMMAND: "" + with: + package-dir: . + output-dir: wheelhouse + + - uses: actions/upload-artifact@v4 + with: + name: wheels-${{ matrix.os }} + path: ./wheelhouse/*.whl + + build_wheels_arm64: + name: Build arm64 wheels + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + submodules: "recursive" + + - name: Set up QEMU + uses: docker/setup-qemu-action@v3 + with: + platforms: linux/arm64 + + - name: Build wheels + uses: pypa/cibuildwheel@v2.20.0 + env: + CIBW_SKIP: "*musllinux* pp*" + CIBW_REPAIR_WHEEL_COMMAND: "" + CIBW_ARCHS: "aarch64" + CIBW_BUILD: "cp38-* cp39-* cp310-* cp311-* cp312-*" + with: + output-dir: wheelhouse - - uses: actions/upload-artifact@v3 + - name: Upload wheels as artifacts + uses: actions/upload-artifact@v4 with: + name: wheels_arm64 path: ./wheelhouse/*.whl build_sdist: @@ -46,10 +76,10 @@ jobs: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v4 with: submodules: "recursive" - - uses: actions/setup-python@v3 + - uses: actions/setup-python@v5 with: python-version: "3.8" - name: Install dependencies @@ -59,22 +89,24 @@ jobs: - name: Build source distribution run: | python -m build --sdist - - uses: actions/upload-artifact@v3 + - uses: actions/upload-artifact@v4 with: + name: sdist path: ./dist/*.tar.gz release: name: Release - needs: [build_wheels, build_sdist] + needs: [build_wheels, build_wheels_arm64, build_sdist] runs-on: ubuntu-latest steps: - - uses: actions/download-artifact@v3 + - uses: actions/download-artifact@v4 with: - name: artifact + merge-multiple: true path: dist - - uses: softprops/action-gh-release@v1 + + - uses: softprops/action-gh-release@v2 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 750b91e1f..b5c7346db 100644 --- a/.github/workflows/build-docker.yaml +++ b/.github/workflows/build-docker.yaml @@ -12,18 +12,18 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout - uses: actions/checkout@v3 + uses: actions/checkout@v4 with: submodules: "recursive" - name: Set up QEMU - uses: docker/setup-qemu-action@v2 + uses: docker/setup-qemu-action@v3 - name: Set up Docker Buildx - uses: docker/setup-buildx-action@v2 + uses: docker/setup-buildx-action@v3 - name: Login to GitHub Container Registry - uses: docker/login-action@v2 + uses: docker/login-action@v3 with: registry: ghcr.io username: ${{ github.repository_owner }} @@ -31,7 +31,7 @@ jobs: - name: Build and push id: docker_build - uses: docker/build-push-action@v4 + uses: docker/build-push-action@v6 with: context: . file: "docker/simple/Dockerfile" diff --git a/.github/workflows/build-wheels-cuda.yaml b/.github/workflows/build-wheels-cuda.yaml new file mode 100644 index 000000000..0733a68c5 --- /dev/null +++ b/.github/workflows/build-wheels-cuda.yaml @@ -0,0 +1,138 @@ +name: Build Wheels (CUDA) + +on: workflow_dispatch + +permissions: + contents: write + +jobs: + define_matrix: + name: Define Build Matrix + runs-on: ubuntu-latest + outputs: + matrix: ${{ steps.set-matrix.outputs.matrix }} + defaults: + run: + shell: pwsh + + steps: + - name: Define Job Output + id: set-matrix + run: | + $matrix = @{ + 'os' = @('ubuntu-latest', 'windows-2019') + 'pyver' = @("3.9", "3.10", "3.11", "3.12") + 'cuda' = @("12.1.1", "12.2.2", "12.3.2", "12.4.1") + 'releasetag' = @("basic") + } + + $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: Add MSBuild to PATH + if: runner.os == 'Windows' + uses: microsoft/setup-msbuild@v2 + with: + vs-version: '[16.11,16.12)' + + - uses: actions/checkout@v4 + with: + submodules: "recursive" + + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.pyver }} + cache: 'pip' + + - name: Setup Mamba + uses: conda-incubator/setup-miniconda@v3.0.4 + with: + activate-environment: "build" + python-version: ${{ matrix.pyver }} + miniforge-variant: Mambaforge + miniforge-version: latest + use-mamba: true + add-pip-as-python-dependency: true + auto-activate-base: false + + - name: VS Integration Cache + id: vs-integration-cache + if: runner.os == 'Windows' + uses: actions/cache@v4.0.2 + with: + path: ./MSBuildExtensions + key: cuda-${{ matrix.cuda }}-vs-integration + + - name: Get Visual Studio Integration + if: runner.os == 'Windows' && steps.vs-integration-cache.outputs.cache-hit != 'true' + run: | + if ($env:CUDAVER -eq '12.1.1') {$x = '12.1.0'} else {$x = $env:CUDAVER} + $links = (Invoke-RestMethod 'https://raw.githubusercontent.com/Jimver/cuda-toolkit/master/src/links/windows-links.ts').Trim().split().where({$_ -ne ''}) + for ($i=$q=0;$i -lt $links.count -and $q -lt 2;$i++) {if ($links[$i] -eq "'$x',") {$q++}} + Invoke-RestMethod $links[$i].Trim("'") -OutFile 'cudainstaller.zip' + & 'C:\Program Files\7-Zip\7z.exe' e cudainstaller.zip -oMSBuildExtensions -r *\MSBuildExtensions\* > $null + Remove-Item 'cudainstaller.zip' + + - name: Install Visual Studio Integration + if: runner.os == 'Windows' + run: | + $y = (gi '.\MSBuildExtensions').fullname + '\*' + (gi 'C:\Program Files (x86)\Microsoft Visual Studio\2019\Enterprise\MSBuild\Microsoft\VC\*\BuildCustomizations').fullname.foreach({cp $y $_}) + $cupath = 'CUDA_PATH_V' + $env:CUDAVER.Remove($env:CUDAVER.LastIndexOf('.')).Replace('.','_') + echo "$cupath=$env:CONDA_PREFIX" >> $env:GITHUB_ENV + + - name: Install Dependencies + env: + MAMBA_DOWNLOAD_FAILFAST: "0" + MAMBA_NO_LOW_SPEED_LIMIT: "1" + run: | + $cudaVersion = $env:CUDAVER + mamba install -y 'cuda' -c nvidia/label/cuda-$cudaVersion + python -m pip install build wheel + + - name: Build Wheel + run: | + $cudaVersion = $env:CUDAVER.Remove($env:CUDAVER.LastIndexOf('.')).Replace('.','') + $env:CUDA_PATH = $env:CONDA_PREFIX + $env:CUDA_HOME = $env:CONDA_PREFIX + $env:CUDA_TOOLKIT_ROOT_DIR = $env:CONDA_PREFIX + if ($IsLinux) { + $env:LD_LIBRARY_PATH = $env:CONDA_PREFIX + '/lib:' + $env:LD_LIBRARY_PATH + } + $env:VERBOSE = '1' + $env:CMAKE_ARGS = '-DGGML_CUDA=on -DCMAKE_CUDA_ARCHITECTURES=all' + $env:CMAKE_ARGS = "-DGGML_CUDA_FORCE_MMQ=ON $env:CMAKE_ARGS" + # if ($env:AVXVER -eq 'AVX') { + $env:CMAKE_ARGS = $env:CMAKE_ARGS + ' -DGGML_AVX2=off -DGGML_FMA=off -DGGML_F16C=off' + # } + # if ($env:AVXVER -eq 'AVX512') { + # $env:CMAKE_ARGS = $env:CMAKE_ARGS + ' -DGGML_AVX512=on' + # } + # if ($env:AVXVER -eq 'basic') { + # $env:CMAKE_ARGS = $env:CMAKE_ARGS + ' -DGGML_AVX=off -DGGML_AVX2=off -DGGML_FMA=off -DGGML_F16C=off' + # } + python -m build --wheel + # write the build tag to the output + Write-Output "CUDA_VERSION=$cudaVersion" >> $env:GITHUB_ENV + + - uses: softprops/action-gh-release@v2 + with: + files: dist/* + # Set tag_name to -cu + tag_name: ${{ 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 000000000..dd8b9c8fc --- /dev/null +++ b/.github/workflows/build-wheels-metal.yaml @@ -0,0 +1,66 @@ +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-12, macos-13, macos-14] + + steps: + - uses: actions/checkout@v4 + with: + submodules: "recursive" + + # Used to host cibuildwheel + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + cache: 'pip' + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + python -m pip install -e .[all] + + - name: Build wheels + uses: pypa/cibuildwheel@v2.20.0 + 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" + CIBW_BUILD: "cp39-* cp310-* cp311-* cp312-*" + with: + package-dir: . + output-dir: wheelhouse2 + + - uses: actions/upload-artifact@v4 + with: + name: wheels-mac_${{ matrix.os }} + path: ./wheelhouse2/*.whl + + release: + name: Release + needs: [build_wheels] + runs-on: ubuntu-latest + + steps: + - uses: actions/download-artifact@v4 + with: + merge-multiple: true + path: dist2 + + - uses: softprops/action-gh-release@v2 + 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/generate-index-from-release.yaml b/.github/workflows/generate-index-from-release.yaml new file mode 100644 index 000000000..662cb6f65 --- /dev/null +++ b/.github/workflows/generate-index-from-release.yaml @@ -0,0 +1,52 @@ +name: Wheels Index + +on: + # Trigger on new release + workflow_run: + workflows: ["Release", "Build Wheels (CUDA)", "Build Wheels (Metal)"] + 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@v4 + - name: Setup Pages + uses: actions/configure-pages@v5 + - name: Build + run: | + ./scripts/releases-to-pep-503.sh index/whl/cpu '^[v]?[0-9]+\.[0-9]+\.[0-9]+$' + ./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/metal '^[v]?[0-9]+\.[0-9]+\.[0-9]+-metal$' + - name: Upload artifact + uses: actions/upload-pages-artifact@v3 + with: + # Upload entire repository + path: 'index' + - name: Deploy to GitHub Pages + id: deployment + uses: actions/deploy-pages@v4 diff --git a/.github/workflows/publish-to-test.yaml b/.github/workflows/publish-to-test.yaml index 47e7c40b1..19613233b 100644 --- a/.github/workflows/publish-to-test.yaml +++ b/.github/workflows/publish-to-test.yaml @@ -16,13 +16,14 @@ jobs: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v4 with: submodules: "recursive" - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: - python-version: "3.8" + python-version: "3.11" + cache: 'pip' - name: Append Dev Version to __version__ run: | DEV_VERSION=${{ github.event.inputs.dev_version }} @@ -31,13 +32,13 @@ jobs: sed -i 's/__version__ = \".*\"/__version__ = \"'"${NEW_VERSION}"'\"/' llama_cpp/__init__.py - name: Install dependencies run: | - python3 -m pip install --upgrade pip build - python3 -m pip install -e .[all] + python -m pip install --upgrade pip build + python -m pip install -e .[all] - name: Build source distribution run: | - python3 -m build --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 1afdd667d..c6abb43b3 100644 --- a/.github/workflows/publish.yaml +++ b/.github/workflows/publish.yaml @@ -10,20 +10,20 @@ jobs: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v4 with: submodules: "recursive" - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: - python-version: "3.8" + python-version: "3.9" - name: Install dependencies run: | - python3 -m pip install --upgrade pip build - python3 -m pip install -e .[all] + python -m pip install --upgrade pip build + python -m pip install -e .[all] - name: Build source distribution run: | - python3 -m build --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 cc6a3a725..d7131956d 100644 --- a/.github/workflows/test-pypi.yaml +++ b/.github/workflows/test-pypi.yaml @@ -8,57 +8,60 @@ 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@v5 with: python-version: ${{ matrix.python-version }} + cache: 'pip' - name: Install dependencies run: | - python3 -m pip install --upgrade pip - python3 -m pip install --verbose llama-cpp-python[all] + python -m pip install --upgrade pip + python -m pip install --verbose llama-cpp-python[all] - 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@v5 with: python-version: ${{ matrix.python-version }} + cache: 'pip' - name: Install dependencies run: | - python3 -m pip install --upgrade pip - python3 -m pip install --verbose llama-cpp-python[all] + python -m pip install --upgrade pip + python -m pip install --verbose llama-cpp-python[all] - 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@v5 with: python-version: ${{ matrix.python-version }} + cache: 'pip' - name: Install dependencies run: | - python3 -m pip install --upgrade pip - python3 -m pip install --verbose llama-cpp-python[all] + python -m pip install --upgrade pip + python -m pip install --verbose llama-cpp-python[all] - 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 77df54697..78f0b4983 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -14,95 +14,72 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"] + python-version: ["3.9", "3.10", "3.11", "3.12"] steps: - uses: actions/checkout@v4 with: submodules: "recursive" - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} + cache: 'pip' - name: Install dependencies run: | - python3 -m pip install --upgrade pip - python3 -m pip install .[all] -v + python -m pip install --upgrade pip + python -m pip install .[all] -v - name: Test with pytest run: | - python3 -m pytest + python -m pytest build-windows: runs-on: windows-latest strategy: matrix: - python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"] + python-version: ["3.9", "3.10", "3.11", "3.12"] steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v4 with: submodules: "recursive" - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} + cache: 'pip' - name: Install dependencies run: | - python3 -m pip install --upgrade pip - python3 -m pip install .[all] -v + python -m pip install --upgrade pip + python -m pip install .[all] -v - name: Test with pytest run: | - python3 -m pytest + python -m pytest build-macos: runs-on: macos-latest strategy: matrix: - python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"] + python-version: ["3.9", "3.10", "3.11", "3.12"] steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v4 with: submodules: "recursive" - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} + cache: 'pip' - name: Install dependencies run: | - python3 -m pip install --upgrade pip - python3 -m pip install .[all] --verbose + python -m pip install --upgrade pip + python -m pip install .[all] --verbose - name: Test with pytest run: | - python3 -m pytest - - # build-linux-opencl: - - # runs-on: ubuntu-latest - - # steps: - # - uses: actions/checkout@v3 - # with: - # submodules: "recursive" - # - name: Set up Python 3.8 - # uses: actions/setup-python@v4 - # with: - # python-version: "3.8" - # - name: Set up OpenCL & CLBlast - # run: | - # wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | sudo tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null - # echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | sudo tee /etc/apt/sources.list.d/oneAPI.list - # sudo apt-get update - # sudo apt-get install -y --no-install-recommends llvm intel-oneapi-runtime-opencl intel-oneapi-runtime-compilers libclblast-dev - # - name: Install dependencies - # run: | - # python3 -m pip install --upgrade pip - # CMAKE_ARGS="-DLLAMA_CLBLAST=on" python3 -m pip install .[all] --verbose - # - name: Test with pytest - # run: | - # python3 -m pytest + python -m pytest build-macos-metal: @@ -110,17 +87,17 @@ jobs: runs-on: macos-latest steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v4 with: submodules: "recursive" - - name: Set up Python 3.8 - uses: actions/setup-python@v4 + - name: Set up Python 3.9 + uses: actions/setup-python@v5 with: - python-version: "3.8" + python-version: "3.9" - name: Install dependencies run: | - python3 -m pip install --upgrade pip - CMAKE_ARGS="-DLLAMA_METAL=on" python3 -m pip install .[all] --verbose + python -m pip install --upgrade pip + CMAKE_ARGS="-DGGML_METAL=on" python -m pip install .[all] --verbose - name: Test with pytest run: | - python3 -m pytest + python -m pytest diff --git a/CHANGELOG.md b/CHANGELOG.md index 90dd1e690..063fe722b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,240 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +## [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 @@ -17,7 +251,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [0.2.55] -- feat: Update llama.cpp to ggerganov/9731134296af3a6839cd682e51d9c2109a871de5 +- 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] diff --git a/CMakeLists.txt b/CMakeLists.txt index b4df8ef45..c6b35ed6c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -5,49 +5,98 @@ project(llama_cpp) 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) +function(llama_cpp_python_install_target target) + 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( + 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 + ) + 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) + # Building llama if (APPLE AND NOT CMAKE_SYSTEM_PROCESSOR MATCHES "arm64") # Need to disable these llama.cpp flags on Apple x86_64, # otherwise users may encounter invalid instruction errors - set(LLAMA_AVX "Off" CACHE BOOL "llama: enable AVX" FORCE) - set(LLAMA_AVX2 "Off" CACHE BOOL "llama: enable AVX2" FORCE) - set(LLAMA_FMA "Off" CACHE BOOL "llama: enable FMA" FORCE) - set(LLAMA_F16C "Off" CACHE BOOL "llama: enable F16C" FORCE) + 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 "gml: enable FMA" FORCE) + set(GGML_F16C "Off" CACHE BOOL "gml: enable F16C" FORCE) + endif() + + if (APPLE) + set(GGML_METAL_EMBED_LIBRARY "On" CACHE BOOL "llama: embed metal library" FORCE) endif() + add_subdirectory(vendor/llama.cpp) - install( - TARGETS llama - LIBRARY DESTINATION ${SKBUILD_PLATLIB_DIR}/llama_cpp - RUNTIME DESTINATION ${SKBUILD_PLATLIB_DIR}/llama_cpp - ARCHIVE DESTINATION ${SKBUILD_PLATLIB_DIR}/llama_cpp - FRAMEWORK DESTINATION ${SKBUILD_PLATLIB_DIR}/llama_cpp - RESOURCE DESTINATION ${SKBUILD_PLATLIB_DIR}/llama_cpp - ) - # Temporary fix for https://github.com/scikit-build/scikit-build-core/issues/374 - install( - TARGETS llama - LIBRARY DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/llama_cpp - RUNTIME DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/llama_cpp - ARCHIVE DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/llama_cpp - FRAMEWORK DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/llama_cpp - RESOURCE DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/llama_cpp - ) + llama_cpp_python_install_target(llama) + llama_cpp_python_install_target(ggml) + # Workaround for Windows + CUDA https://github.com/abetlen/llama-cpp-python/issues/563 - install( - FILES $ - DESTINATION ${SKBUILD_PLATLIB_DIR}/llama_cpp - ) - install( - FILES $ - DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/llama_cpp - ) + 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) + if (LLAMA_CUBLAS OR LLAMA_CUDA) add_compile_definitions(GGML_USE_CUBLAS) + add_compile_definitions(GGML_USE_CUDA) endif() if (LLAMA_METAL) @@ -61,22 +110,16 @@ if (LLAMA_BUILD) if (WIN32) set_target_properties(llava_shared PROPERTIES CUDA_ARCHITECTURES OFF) endif() - install( - TARGETS llava_shared - LIBRARY DESTINATION ${SKBUILD_PLATLIB_DIR}/llama_cpp - RUNTIME DESTINATION ${SKBUILD_PLATLIB_DIR}/llama_cpp - ARCHIVE DESTINATION ${SKBUILD_PLATLIB_DIR}/llama_cpp - FRAMEWORK DESTINATION ${SKBUILD_PLATLIB_DIR}/llama_cpp - RESOURCE DESTINATION ${SKBUILD_PLATLIB_DIR}/llama_cpp - ) - # Temporary fix for https://github.com/scikit-build/scikit-build-core/issues/374 - install( - TARGETS llava_shared - LIBRARY DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/llama_cpp - RUNTIME DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/llama_cpp - ARCHIVE DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/llama_cpp - FRAMEWORK DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/llama_cpp - RESOURCE DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/llama_cpp - ) + llama_cpp_python_install_target(llava_shared) + if (WIN32) + install( + FILES $ + DESTINATION ${CMAKE_CURRENT_SOURCE_DIR}/llama_cpp/lib + ) + install( + FILES $ + DESTINATION ${SKBUILD_PLATLIB_DIR}/llama_cpp/lib + ) + endif() endif() endif() diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 000000000..d54f6b430 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,46 @@ +FROM ubuntu:20.04 + +# We need to set the host to 0.0.0.0 to allow outside access +ENV HOST=0.0.0.0 + +# Needs to be <= 2.31 to work with older Linux +RUN echo "Glibc version:\n" && /lib/aarch64-linux-gnu/libc.so.6 + +RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y \ + --no-install-recommends ninja-build pkg-config python3.9 \ + python3.9-dev python3-pip git + +# Install toolchain (GCC 11) +RUN DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends software-properties-common && \ + add-apt-repository -y ppa:ubuntu-toolchain-r/test && \ + apt-get update && \ + apt-get install -y --no-install-recommends gcc-11 g++-11 make && \ + update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-11 110 \ + --slave /usr/bin/g++ g++ /usr/bin/g++-11 --slave /usr/bin/gcov gcov /usr/bin/gcov-11 \ + --slave /usr/bin/gcc-ar gcc-ar /usr/bin/gcc-ar-11 \ + --slave /usr/bin/gcc-ranlib gcc-ranlib /usr/bin/gcc-ranlib-11 \ + --slave /usr/bin/cpp cpp /usr/bin/cpp-11 + +RUN python3.9 -m pip install --upgrade pip cmake scikit-build-core[pyproject] setuptools pyinstaller + +# Install the repo at current state +COPY . . + +# Set HW specific flags for us +ENV march=armv8.2-a+crypto+fp16+rcpc+dotprod +ENV mcpu=cortex-a78c+crypto+noprofile+nossbs+noflagm+nopauth +ENV mtune=cortex-a78c + +ENV compiler_flags="-march=${march} -mcpu=${mcpu} -mtune=${mtune}" + +# This is a release build that works (have to disable GGML_LLAMAFILE for Q4_0_4_4 quantization) +RUN CC=gcc-11 CXX=g++-11 CMAKE_BUILD_TYPE=Release \ + CMAKE_ARGS="-DGGML_LLAMAFILE=OFF -DCMAKE_C_FLAGS='${compiler_flags}' -DCMAKE_CXX_FLAGS='${compiler_flags}'" \ + python3.9 -m pip install -v -e .[server] 2>&1 | tee buildlog.txt + +# TODO: Export buildlog.txt in `make deploy.docker` step for review after +# building. +RUN cd /root && pyinstaller -DF /llama_cpp/server/__main__.py \ + --add-data /llama_cpp/lib/libllama.so:llama_cpp/lib \ + --add-data /llama_cpp/lib/libggml.so:llama_cpp/lib \ + -n llama-cpp-py-server diff --git a/Makefile b/Makefile index 4ae011074..c3027ddb4 100644 --- a/Makefile +++ b/Makefile @@ -13,31 +13,37 @@ build: python3 -m pip install --verbose -e . build.debug: - CMAKE_ARGS="-DCMAKE_BUILD_TYPE=Debug" python3 -m pip install --verbose --config-settings=cmake.verbose=true --config-settings=logging.level=INFO --config-settings=install.strip=false --editable . + 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.cuda: - CMAKE_ARGS="-DLLAMA_CUBLAS=on" python3 -m pip install --verbose -e . - -build.opencl: - CMAKE_ARGS="-DLLAMA_CLBLAST=on" python3 -m pip install --verbose -e . + CMAKE_ARGS="-DGGML_CUDA=on" python3 -m pip install --verbose -e . build.openblas: - CMAKE_ARGS="-DLLAMA_BLAS=ON -DLLAMA_BLAS_VENDOR=OpenBLAS" python3 -m pip install --verbose -e . + CMAKE_ARGS="-DGGML_BLAS=ON -DGGML_BLAS_VENDOR=OpenBLAS" python3 -m pip install --verbose -e . build.blis: - CMAKE_ARGS="-DLLAMA_BLAS=on -DLLAMA_BLAS_VENDOR=FLAME" python3 -m pip install --verbose -e . + CMAKE_ARGS="-DGGML_BLAS=on -DGGML_BLAS_VENDOR=FLAME" python3 -m pip install --verbose -e . build.metal: - CMAKE_ARGS="-DLLAMA_METAL=on" python3 -m pip install --verbose -e . + CMAKE_ARGS="-DGGML_METAL=on" python3 -m pip install --verbose -e . build.vulkan: - CMAKE_ARGS="-DLLAMA_VULKAN=on" python3 -m pip install --verbose -e . + CMAKE_ARGS="-DGGML_VULKAN=on" python3 -m pip install --verbose -e . build.kompute: - CMAKE_ARGS="-DLLAMA_KOMPUTE=on" python3 -m pip install --verbose -e . + CMAKE_ARGS="-DGGML_KOMPUTE=on" python3 -m pip install --verbose -e . build.sycl: - CMAKE_ARGS="-DLLAMA_SYCL=on" python3 -m pip install --verbose -e . + 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 -m build --sdist @@ -49,6 +55,43 @@ deploy.gh-docs: mkdocs build mkdocs gh-deploy +COMMIT := $(shell git rev-parse --short HEAD) + +deploy.docker: + # Make image with commit in name + docker build -t openblas_server_$(COMMIT) . + + # Run image and immediately exit (just want to create the container) + docker run openblas_server_$(COMMIT) bash + + # Get container ID, copy server tarball + libllama.so tarball, and delete + # temp container + CONTAINER_ID=$$(docker ps -lq --filter ancestor=openblas_server_$(COMMIT)) ; \ + echo Container ID: $$CONTAINER_ID ; \ + docker cp $$CONTAINER_ID:/root/dist/llama-cpp-py-server - | pigz -9 > llama-cpp-py-server.tgz ; \ + docker rm $$CONTAINER_ID + + # More cleanup + yes | docker image prune + +# Build standalone server, may want to do in fresh venv to avoid bloat +deploy.pyinstaller.mac: + # CPU must be aarch64 and OS is MacOS + @if [ `uname -m` != "arm64" ]; then echo "Must be on aarch64"; exit 1; fi + @if [ `uname` != "Darwin" ]; then echo "Must be on MacOS"; exit 1; fi + @echo "Building and installing with proper env vars for aarch64-specific ops" + CMAKE_ARGS="-DGGML_METAL=off -DGGML_LLAMAFILE=OFF -DGGML_BLAS=OFF -DCMAKE_BUILD_TYPE=Release" python3 -m pip install -v -e .[server,dev] + @server_path=$$(python -c 'import llama_cpp.server; print(llama_cpp.server.__file__)' | sed s/init/main/) ; \ + echo "Server path: $$server_path" ; \ + libllama_path=$$(python -c 'import llama_cpp.llama_cpp; print(llama_cpp.llama_cpp._load_shared_library("llama")._name)') ; \ + libggml_path=$$(python -c 'import llama_cpp.llama_cpp; print(llama_cpp.llama_cpp._load_shared_library("ggml")._name)') ; \ + echo "libllama path: $$libllama_path" ; \ + echo "libggml path: $$libggml_path" ; \ + pyinstaller -DF $$server_path \ + --add-data $$libllama_path:llama_cpp/lib \ + --add-data $$libggml_path:llama_cpp/lib \ + -n llama-cpp-py-server + test: python3 -m pytest @@ -78,5 +121,7 @@ clean: build.sdist \ deploy.pypi \ deploy.gh-docs \ + deploy.docker \ + deploy.pyinstaller.mac \ docker \ - clean \ No newline at end of file + clean diff --git a/README.md b/README.md index 3323f3899..b0dfdd5b5 100644 --- a/README.md +++ b/README.md @@ -6,6 +6,7 @@ [![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/) +[![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: @@ -43,6 +44,15 @@ This will also build `llama.cpp` from source and install it alongside this pytho If this fails, add `--verbose` to the `pip install` see the full cmake build log. +**Pre-built Wheel (New)** + +It is also possible to install a pre-built wheel with basic CPU support. + +```bash +pip install llama-cpp-python \ + --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cpu +``` + ### 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/ggerganov/llama.cpp#build) for a full list. @@ -54,13 +64,13 @@ All `llama.cpp` cmake build options can be set via the `CMAKE_ARGS` environment ```bash # Linux and Mac -CMAKE_ARGS="-DLLAMA_BLAS=ON -DLLAMA_BLAS_VENDOR=OpenBLAS" \ +CMAKE_ARGS="-DGGML_BLAS=ON -DGGML_BLAS_VENDOR=OpenBLAS" \ pip install llama-cpp-python ``` ```powershell # Windows -$env:CMAKE_ARGS = "-DLLAMA_BLAS=ON -DLLAMA_BLAS_VENDOR=OpenBLAS" +$env:CMAKE_ARGS = "-DGGML_BLAS=ON -DGGML_BLAS_VENDOR=OpenBLAS" pip install llama-cpp-python ``` @@ -73,13 +83,13 @@ They can also be set via `pip install -C / --config-settings` command and saved ```bash pip install --upgrade pip # ensure pip is up to date pip install llama-cpp-python \ - -C cmake.args="-DLLAMA_BLAS=ON;-DLLAMA_BLAS_VENDOR=OpenBLAS" + -C cmake.args="-DGGML_BLAS=ON;-DGGML_BLAS_VENDOR=OpenBLAS" ``` ```txt # requirements.txt -llama-cpp-python -C cmake.args="-DLLAMA_BLAS=ON;-DLLAMA_BLAS_VENDOR=OpenBLAS" +llama-cpp-python -C cmake.args="-DGGML_BLAS=ON;-DGGML_BLAS_VENDOR=OpenBLAS" ``` @@ -91,20 +101,45 @@ Below are some common backends, their build commands and any additional environm
OpenBLAS (CPU) -To install with OpenBLAS, set the `LLAMA_BLAS` and `LLAMA_BLAS_VENDOR` environment variables before installing: +To install with OpenBLAS, set the `GGML_BLAS` and `GGML_BLAS_VENDOR` environment variables before installing: ```bash -CMAKE_ARGS="-DLLAMA_BLAS=ON -DLLAMA_BLAS_VENDOR=OpenBLAS" pip install llama-cpp-python +CMAKE_ARGS="-DGGML_BLAS=ON -DGGML_BLAS_VENDOR=OpenBLAS" pip install llama-cpp-python ```
-cuBLAS (CUDA) +CUDA -To install with cuBLAS, set the `LLAMA_CUBLAS=on` environment variable before installing: +To install with CUDA support, set the `GGML_CUDA=on` environment variable before installing: ```bash -CMAKE_ARGS="-DLLAMA_CUBLAS=on" pip install llama-cpp-python +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 12.1, 12.2, 12.3, or 12.4 +- 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/ +``` + +Where `` is one of the following: +- `cu121`: CUDA 12.1 +- `cu122`: CUDA 12.2 +- `cu123`: CUDA 12.3 +- `cu124`: CUDA 12.4 + +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 ```
@@ -112,21 +147,22 @@ CMAKE_ARGS="-DLLAMA_CUBLAS=on" pip install llama-cpp-python
Metal -To install with Metal (MPS), set the `LLAMA_METAL=on` environment variable before installing: +To install with Metal (MPS), set the `GGML_METAL=on` environment variable before installing: ```bash -CMAKE_ARGS="-DLLAMA_METAL=on" pip install llama-cpp-python +CMAKE_ARGS="-DGGML_METAL=on" pip install llama-cpp-python ``` -
-
+**Pre-built Wheel (New)** -CLBlast (OpenCL) +It is also possible to install a pre-built wheel with Metal support. As long as your system meets some requirements: -To install with CLBlast, set the `LLAMA_CLBLAST=on` environment variable before installing: +- MacOS Version is 11.0 or later +- Python Version is 3.10, 3.11 or 3.12 ```bash -CMAKE_ARGS="-DLLAMA_CLBLAST=on" pip install llama-cpp-python +pip install llama-cpp-python \ + --extra-index-url https://abetlen.github.io/llama-cpp-python/whl/metal ```
@@ -134,10 +170,10 @@ CMAKE_ARGS="-DLLAMA_CLBLAST=on" pip install llama-cpp-python
hipBLAS (ROCm) -To install with hipBLAS / ROCm support for AMD cards, set the `LLAMA_HIPBLAS=on` environment variable before installing: +To install with hipBLAS / ROCm support for AMD cards, set the `GGML_HIPBLAS=on` environment variable before installing: ```bash -CMAKE_ARGS="-DLLAMA_HIPBLAS=on" pip install llama-cpp-python +CMAKE_ARGS="-DGGML_HIPBLAS=on" pip install llama-cpp-python ```
@@ -145,32 +181,33 @@ CMAKE_ARGS="-DLLAMA_HIPBLAS=on" pip install llama-cpp-python
Vulkan -To install with Vulkan support, set the `LLAMA_VULKAN=on` environment variable before installing: +To install with Vulkan support, set the `GGML_VULKAN=on` environment variable before installing: ```bash -CMAKE_ARGS="-DLLAMA_VULKAN=on" pip install llama-cpp-python +CMAKE_ARGS="-DGGML_VULKAN=on" pip install llama-cpp-python ```
-Kompute +SYCL -To install with Kompute support, set the `LLAMA_KOMPUTE=on` environment variable before installing: +To install with SYCL support, set the `GGML_SYCL=on` environment variable before installing: ```bash -CMAKE_ARGS="-DLLAMA_KOMPUTE=on" 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 ```
-SYCL +RPC -To install with SYCL support, set the `LLAMA_SYCL=on` environment variable before installing: +To install with RPC support, set the `GGML_RPC=on` environment variable before installing: ```bash source /opt/intel/oneapi/setvars.sh -CMAKE_ARGS="-DLLAMA_SYCL=on -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx" pip install llama-cpp-python +CMAKE_ARGS="-DGGML_RPC=on" pip install llama-cpp-python ```
@@ -184,7 +221,7 @@ If you run into issues where it complains it can't find `'nmake'` `'?'` or CMAKE ```ps $env:CMAKE_GENERATOR = "MinGW Makefiles" -$env:CMAKE_ARGS = "-DLLAMA_OPENBLAS=on -DCMAKE_C_COMPILER=C:/w64devkit/bin/gcc.exe -DCMAKE_CXX_COMPILER=C:/w64devkit/bin/g++.exe" +$env:CMAKE_ARGS = "-DGGML_OPENBLAS=on -DCMAKE_C_COMPILER=C:/w64devkit/bin/gcc.exe -DCMAKE_CXX_COMPILER=C:/w64devkit/bin/g++.exe" ``` See the above instructions and set `CMAKE_ARGS` to the BLAS backend you want to use. @@ -213,7 +250,7 @@ Otherwise, while installing it will build the llama.cpp x86 version which will b Try installing with ```bash -CMAKE_ARGS="-DCMAKE_OSX_ARCHITECTURES=arm64 -DCMAKE_APPLE_SILICON_PROCESSOR=arm64 -DLLAMA_METAL=on" pip install --upgrade --verbose --force-reinstall --no-cache-dir llama-cpp-python +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 ``` @@ -230,20 +267,26 @@ The high-level API provides a simple managed interface through the [`Llama`](htt Below is a short example demonstrating how to use the high-level API to for basic text completion: ```python ->>> from llama_cpp import Llama ->>> llm = Llama( +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( +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) +print(output) +``` + +By default `llama-cpp-python` generates completions in an OpenAI compatible format: + +```python { "id": "cmpl-xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx", "object": "text_completion", @@ -274,7 +317,7 @@ You'll need to install the `huggingface-hub` package to use this feature (`pip i ```python llm = Llama.from_pretrained( - repo_id="Qwen/Qwen1.5-0.5B-Chat-GGUF", + repo_id="Qwen/Qwen2-0.5B-Instruct-GGUF", filename="*q8_0.gguf", verbose=False ) @@ -298,12 +341,12 @@ The model will will format the messages into a single prompt using the following Set `verbose=True` to see the selected chat format. ```python ->>> from llama_cpp import Llama ->>> llm = Llama( +from llama_cpp import Llama +llm = Llama( model_path="path/to/llama-2/llama-model.gguf", chat_format="llama-2" ) ->>> llm.create_chat_completion( +llm.create_chat_completion( messages = [ {"role": "system", "content": "You are an assistant who perfectly describes images."}, { @@ -321,16 +364,16 @@ For OpenAI API v1 compatibility, you use the [`create_chat_completion_openai_v1` ### 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`](http://localhost:8000/api-reference/#llama_cpp.Llama.create_chat_completion). +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( +from llama_cpp import Llama +llm = Llama(model_path="path/to/model.gguf", chat_format="chatml") +llm.create_chat_completion( messages=[ { "role": "system", @@ -350,9 +393,9 @@ The following example will constrain the response to valid JSON strings only. 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( +from llama_cpp import Llama +llm = Llama(model_path="path/to/model.gguf", chat_format="chatml") +llm.create_chat_completion( messages=[ { "role": "system", @@ -377,9 +420,9 @@ To constrain the response further to a specific JSON Schema add the schema to th 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( +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", @@ -412,12 +455,12 @@ The high-level API supports OpenAI compatible function and tool calling. This is } } }], - tool_choice=[{ + tool_choice={ "type": "function", "function": { "name": "UserDetail" } - }] + } ) ``` @@ -429,54 +472,94 @@ The various gguf-converted files for this set of models can be found [here](http 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( +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.). ### Multi-modal 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. +`llama-cpp-python` supports such as llava1.5 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: +Below are the supported multi-modal models and their respective chat handlers (Python API) and chat formats (Server API). -- [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) +| 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` | 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( +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 accomodate the image embedding - logits_all=True,# needed to make llava work + n_ctx=2048, # n_ctx should be increased to accommodate the image embedding ) ->>> llm.create_chat_completion( +llm.create_chat_completion( messages = [ {"role": "system", "content": "You are an assistant who perfectly describes images."}, { "role": "user", "content": [ - {"type": "image_url", "image_url": {"url": "https://.../image.png"}}, - {"type" : "text", "text": "Describe this image in detail please."} + {"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. +
Loading a Local Image @@ -529,7 +612,7 @@ llama = Llama( ### Embeddings -To generate text embeddings use [`create_embedding`](http://localhost:8000/api-reference/#llama_cpp.Llama.create_embedding). +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 @@ -543,6 +626,12 @@ embeddings = llm.create_embedding("Hello, world!") 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. @@ -568,7 +657,7 @@ 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="-DLLAMA_CUBLAS=on" FORCE_CMAKE=1 pip install 'llama-cpp-python[server]' +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 ``` @@ -589,7 +678,7 @@ For possible options, see [llama_cpp/llama_chat_format.py](llama_cpp/llama_chat_ 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 Qwen/Qwen1.5-0.5B-Chat-GGUF --model '*q8_0.gguf' +python3 -m llama_cpp.server --hf_model_repo_id Qwen/Qwen2-0.5B-Instruct-GGUF --model '*q8_0.gguf' ``` ### Web Server Features @@ -619,18 +708,18 @@ The entire low-level API can be found in [llama_cpp/llama_cpp.py](https://github Below is a short example demonstrating how to use the low-level API to tokenize a prompt: ```python ->>> import llama_cpp ->>> import ctypes ->>> llama_cpp.llama_backend_init(False) # Must be called once at the start of each program ->>> params = llama_cpp.llama_context_default_params() +import llama_cpp +import ctypes +llama_cpp.llama_backend_init(False) # Must be called once at the start of each program +params = llama_cpp.llama_context_default_params() # use bytes for char * params ->>> model = llama_cpp.llama_load_model_from_file(b"./models/7b/llama-model.gguf", params) ->>> ctx = llama_cpp.llama_new_context_with_model(model, params) ->>> max_tokens = params.n_ctx +model = llama_cpp.llama_load_model_from_file(b"./models/7b/llama-model.gguf", params) +ctx = llama_cpp.llama_new_context_with_model(model, params) +max_tokens = 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, llama_cpp.c_bool(True)) ->>> llama_cpp.llama_free(ctx) +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, llama_cpp.c_bool(True)) +llama_cpp.llama_free(ctx) ``` Check out the [examples folder](examples/low_level_api) for more examples of using the low-level API. @@ -666,7 +755,7 @@ pip install -e .[all] make clean ``` -You can also test out specific commits of `lama.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 +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 diff --git a/docker/cuda_simple/Dockerfile b/docker/cuda_simple/Dockerfile index a9e51cdc1..0bbf20ffe 100644 --- a/docker/cuda_simple/Dockerfile +++ b/docker/cuda_simple/Dockerfile @@ -1,4 +1,4 @@ -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 @@ -15,13 +15,13 @@ COPY . . # setting build related env vars ENV CUDA_DOCKER_ARCH=all -ENV LLAMA_CUBLAS=1 +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 # Install llama-cpp-python (build with cuda) -RUN CMAKE_ARGS="-DLLAMA_CUBLAS=on" pip install llama-cpp-python +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 23e37d4d4..f05e66ef2 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 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 e213518b9..5ee667dc0 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 pkg-config +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 CMAKE_ARGS="-DLLAMA_BLAS=ON -DLLAMA_BLAS_VENDOR=OpenBLAS" 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 41cc68ea2..3594df1a5 100644 --- a/docker/simple/Dockerfile +++ b/docker/simple/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} @@ -13,7 +13,9 @@ RUN apt-get update && apt-get upgrade -y && apt-get install -y --no-install-reco 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 @@ -21,7 +23,9 @@ COPY . /app RUN python3 -m pip install --upgrade pip -RUN make deps && 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 pip install llama-cpp-python --verbose; # Set environment variable for the host ENV HOST=0.0.0.0 diff --git a/docs/install/macos.md b/docs/install/macos.md index 240422832..e006fc0a3 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" 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 diff --git a/docs/server.md b/docs/server.md index c66c9cce0..cd6f86c51 100644 --- a/docs/server.md +++ b/docs/server.md @@ -98,6 +98,8 @@ You'll first need to download one of the available multi-modal models in GGUF fo - [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 diff --git a/examples/batch-processing/server.py b/examples/batch-processing/server.py new file mode 100644 index 000000000..0b36746f9 --- /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_load_model_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_new_context_with_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/gradio_chat/local.py b/examples/gradio_chat/local.py index a7de8e842..e16bf234a 100644 --- a/examples/gradio_chat/local.py +++ b/examples/gradio_chat/local.py @@ -6,25 +6,26 @@ llama = llama_cpp.Llama.from_pretrained( repo_id="Qwen/Qwen1.5-0.5B-Chat-GGUF", filename="*q8_0.gguf", - tokenizer=llama_cpp.llama_tokenizer.LlamaHFTokenizer.from_pretrained("Qwen/Qwen1.5-0.5B"), - verbose=False + tokenizer=llama_cpp.llama_tokenizer.LlamaHFTokenizer.from_pretrained( + "Qwen/Qwen1.5-0.5B" + ), + 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 + model=model, messages=messages, stream=True ) text = "" @@ -52,7 +53,14 @@ def predict(message, history): """ 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?"]) + gr.ChatInterface( + predict, + fill_height=True, + examples=[ + "What is the capital of France?", + "Who was the first person on the moon?", + ], + ) if __name__ == "__main__": diff --git a/examples/gradio_chat/server.py b/examples/gradio_chat/server.py index 36fa43fbd..52061bea6 100644 --- a/examples/gradio_chat/server.py +++ b/examples/gradio_chat/server.py @@ -2,26 +2,22 @@ from openai import OpenAI -client = OpenAI( - base_url="http://localhost:8000/v1", - api_key="llama.cpp" -) +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 + model=model, messages=messages, stream=True ) text = "" @@ -49,7 +45,14 @@ def predict(message, history): """ 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?"]) + gr.ChatInterface( + predict, + fill_height=True, + examples=[ + "What is the capital of France?", + "Who was the first person on the moon?", + ], + ) if __name__ == "__main__": diff --git a/examples/hf_pull/main.py b/examples/hf_pull/main.py index d3eb11c39..dfed17516 100644 --- a/examples/hf_pull/main.py +++ b/examples/hf_pull/main.py @@ -5,29 +5,26 @@ llama = llama_cpp.Llama.from_pretrained( repo_id="Qwen/Qwen1.5-0.5B-Chat-GGUF", filename="*q8_0.gguf", - tokenizer=llama_cpp.llama_tokenizer.LlamaHFTokenizer.from_pretrained("Qwen/Qwen1.5-0.5B"), - verbose=False + tokenizer=llama_cpp.llama_tokenizer.LlamaHFTokenizer.from_pretrained( + "Qwen/Qwen1.5-0.5B" + ), + verbose=False, ) response = llama.create_chat_completion( - messages=[ - { - "role": "user", - "content": "What is the capital of France?" - } - ], + 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"} + "capital": {"type": "string"}, }, "required": ["country", "capital"], - } + }, }, - stream=True + stream=True, ) for chunk in response: @@ -36,4 +33,4 @@ continue print(delta["content"], end="", flush=True) -print() \ No newline at end of file +print() diff --git a/examples/high_level_api/create_disk_cache.py b/examples/high_level_api/create_disk_cache.py new file mode 100644 index 000000000..8e2eaa8aa --- /dev/null +++ b/examples/high_level_api/create_disk_cache.py @@ -0,0 +1,227 @@ +""" +Creates static disk cache give a dataset of snippets to use as contexts for a RAG application. + +Background: Want to embed the fixed system prompt + first context and then store on disk. +That way, when a question is asked and the first context is provided, can look up the +KV cache to find the prompt that matches the prompt tokens (including first context). + +This should save on prompt ingestion time, decreasing time to first token. +""" + +import argparse +import os +import pathlib + +import pandas as pd +import tqdm + +from llama_cpp.llama import Llama +from llama_cpp.llama_cache import LlamaStaticDiskCache +from llama_cpp.llama_chat_format import format_nekomata + +# Add additional formatters here as desired so that can swap out models. +CHAT_FORMATTER_MAP = { + "rinna/nekomata-7b-instruction-gguf": format_nekomata, +} + + +def combine_question_ctx_nekomata(question, contexts): + """ + Formats question and contexts for nekomata-7b. + """ + output = "" + for i, context in enumerate(contexts): + output += f"- 資料{i+1}: '{context}'\n" + + output += "\n" + + output += question + + return output + + +# How to combine contexts + user question when creating a *full* prompt +CONTEXT_QUESTION_FORMATTER_MAP = { + "rinna/nekomata-7b-instruction-gguf": combine_question_ctx_nekomata, +} + +DEFAULT_SYSTEM_PROMPT = """ +You are a virtual assistant. You will be provided with contexts and a user +question. Your job is to answer a user's question faithfully and concisely. + +If the context provided does not contain enough information to answer the question, +respond with "I don't know" - do not make up information. If you are helpful and provide +accurate information, you will be provided with a $10,000 bonus. If you provide inaccurate +information, unhelpful responses, or information not grounded in the context +provided, you will be penalized $10,000 and fired - into the Sun. +""".strip() + + +def _create_nekomata_prompt_prefix( + context: str, system_prompt=DEFAULT_SYSTEM_PROMPT +) -> str: + """ + Override this if using a different model. + + This provides a partially formatted prompt for the Nekomata model. + It passes in the system prompt and the first context to the model, + but not the question or prompt for assistant. + """ + + return """ +以下は、タスクを説明する指示と、文脈のある入力の組み合わせです。要求を適切に満たす応答を書きなさい。 + +### 指示: +{system_prompt} + +### 入力: +{input}""".format( + system_prompt=system_prompt, input=f"- 資料1: '{context.strip()}'\n" + ).lstrip( + "\n" + ) + + +PARTIAL_PROMPT_MODEL_MAP = { + "rinna/nekomata-7b-instruction-gguf": _create_nekomata_prompt_prefix, +} + + +def main(args: argparse.Namespace): + dataset_path: pathlib.Path = args.dataset + assert dataset_path.exists(), f"Dataset path {dataset_path} does not exist" + + dataset = pd.read_csv(str(dataset_path)) + + snippets = dataset.loc[:, args.column_name].tolist() + + model = Llama.from_pretrained( + args.model, + filename=args.model_filename, + n_ctx=args.n_ctx, + n_gpu_layers=-1, + n_batch=1, + n_threads=args.n_threads, + n_threads_batch=args.n_threads, + verbose=False, + ) + + prompt_formatter_func = PARTIAL_PROMPT_MODEL_MAP[args.model] + + # Have to format such that includes system prompt and the context + snippets = [prompt_formatter_func(context) for context in snippets] + + cache = LlamaStaticDiskCache.build_cache(args.output, tqdm.tqdm(snippets), model) + snippet_tokens = model.tokenize( + snippets[0].encode("utf-8"), add_bos=True, special=True + ) + assert snippet_tokens in cache, "First snippet not in cache" + + # pylint: disable=protected-access + cache_prefix_tokens = cache._find_longest_prefix_key(snippet_tokens) + + assert cache_prefix_tokens == tuple( + snippet_tokens + ), "Expected all snippet tokens to be in cache" + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "-d", + "--dataset", + type=pathlib.Path, + required=True, + help="Path to serialized dataframe with snippets to use", + ) + + parser.add_argument( + "-m", + "--model", + type=str, + default="rinna/nekomata-7b-instruction-gguf", + help="Hugging Face model name", + ) + + parser.add_argument( + "--model-filename", + type=str, + default="*Q4_K_M.gguf", + help="Name of model weights file to load from repo - may contain wildcards (like '*Q4_K_M.gguf')", + ) + + parser.add_argument( + "--n-ctx", + type=int, + required=True, + help="Context size (in tokens) - must be fixed for KV cache", + ) + + parser.add_argument( + "--n-threads", + type=int, + default=os.cpu_count(), + help="Number of threads to use for inference + batch processing", + ) + + parser.add_argument( + "--column-name", + type=str, + default="snippets", + help="Column name identifying snippets to use as contexts", + ) + + parser.add_argument( + "-o", + "--output", + type=str, + default="static_cache", + help="Output directory for static cache", + ) + + args = parser.parse_args() + + chat_formatter = CHAT_FORMATTER_MAP[args.model] + question_ctx_combiner = CONTEXT_QUESTION_FORMATTER_MAP[args.model] + + DUMMY_CONTEXTS = [ + "The air speed of an unladen swallow is 24 miles per hour.", + "Red pandas are not actually pandas, but are more closely related to raccoons.", + "Directly observing a quantum system can change its state.", + "The mitochondria is the powerhouse of the cell.", + "The least common multiple of 6 and 8 is 24.", + ] + + # Just a quick-and-dirty test so that can verify that a full prompt will contain + # the partial prompt (and so prefix matching should work) + def _generate_full_prompt(user_question: str): + user_msg = question_ctx_combiner(user_question, DUMMY_CONTEXTS) + msgs = [ + { + "role": "system", + "content": DEFAULT_SYSTEM_PROMPT, + }, + { + "role": "user", + "content": user_msg, + }, + ] + + full_prompt = chat_formatter(msgs).prompt + + return full_prompt + + question = "What is the velocity of an unladen swallow?" + + complete_prompt = _generate_full_prompt(question) + partial_context = _create_nekomata_prompt_prefix(DUMMY_CONTEXTS[0]) + + if not partial_context in complete_prompt: + print("Partial context:\n") + print(partial_context + "\n") + print("not found in complete prompt:\n") + print(complete_prompt) + assert False, "Sanity check failed" + + main(args) diff --git a/examples/high_level_api/fastapi_server.py b/examples/high_level_api/fastapi_server.py index 9421db57b..ee59767d6 100644 --- a/examples/high_level_api/fastapi_server.py +++ b/examples/high_level_api/fastapi_server.py @@ -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 000000000..282333e5a --- /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 fcef8cd80..a755089b2 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 eb9a2cfa9..e072ab1b1 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 82e5c4487..1f2c59017 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 1a5152530..a0212ff0d 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=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") + 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 02c09afb0..39081be17 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,198 +20,257 @@ 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.path_session is None: - self.params.path_session = "" - if self.params.antiprompt is None: - self.params.antiprompt = "" - - 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_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_load_model_from_file( - self.params.model.encode("utf8"), self.lparams) - - # Context Params. - self.cparams = llama_cpp.llama_context_default_params() - - self.ctx = llama_cpp.llama_new_context_with_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_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_load_model_from_file( + self.params.model.encode("utf8"), self.lparams + ) + + # Context Params. + self.cparams = llama_cpp.llama_context_default_params() + + self.ctx = llama_cpp.llama_new_context_with_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_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""" 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]} -> '{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},\ +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},\ @@ -228,77 +288,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.model, 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): +""", + 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.model, + 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: @@ -307,278 +386,356 @@ 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 - ) != 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.model) - - # 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(self.ctx)] - 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_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(self.ctx) 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(self.ctx): - 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.model, llama_cpp.llama_token(token_id), buffer, size) - 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 + 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_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.model) + + # 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(self.ctx)] + 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_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(self.ctx) + 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( + self.ctx + ): + 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.model, llama_cpp.llama_token(token_id), buffer, size + ) + 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) + # 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. 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. @@ -595,10 +752,10 @@ def interact(self): {USER_NAME}: Name a color. {AI_NAME}: Blue {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() + params = gpt_params_parse() + if params.prompt is None and params.file is None: + params.prompt = prompt + + 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 ef1b2c016..ba3545771 100644 --- a/examples/low_level_api/low_level_api_llama_cpp.py +++ b/examples/low_level_api/low_level_api_llama_cpp.py @@ -7,23 +7,20 @@ llama_cpp.llama_backend_init(numa=False) N_THREADS = multiprocessing.cpu_count() -MODEL_PATH = os.environ.get('MODEL', "../models/7B/ggml-model.bin") +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_model_default_params() cparams = llama_cpp.llama_context_default_params() -model = llama_cpp.llama_load_model_from_file(MODEL_PATH.encode('utf-8'), lparams) +model = llama_cpp.llama_load_model_from_file(MODEL_PATH.encode("utf-8"), lparams) ctx = llama_cpp.llama_new_context_with_model(model, cparams) # determine the required inference memory per token: tmp = [0, 1, 2, 3] llama_cpp.llama_eval( - ctx = ctx, - tokens=(llama_cpp.c_int * len(tmp))(*tmp), - n_tokens=len(tmp), - n_past=0 - )# Deprecated + ctx=ctx, tokens=(llama_cpp.c_int * len(tmp))(*tmp), n_tokens=len(tmp), n_past=0 +) # Deprecated n_past = 0 @@ -32,12 +29,12 @@ embd_inp = (llama_cpp.llama_token * (len(prompt) + 1))() n_of_tok = llama_cpp.llama_tokenize( model=model, - text=bytes(str(prompt),'utf-8'), - text_len=len(embd_inp), + text=bytes(str(prompt), "utf-8"), + text_len=len(embd_inp), tokens=embd_inp, n_max_tokens=len(embd_inp), add_bos=False, - special=False + special=False, ) embd_inp = embd_inp[:n_of_tok] @@ -63,11 +60,11 @@ while remaining_tokens > 0: if len(embd) > 0: llama_cpp.llama_eval( - ctx = ctx, + ctx=ctx, tokens=(llama_cpp.c_int * len(embd))(*embd), n_tokens=len(embd), - n_past=n_past - )# Deprecated + n_past=n_past, + ) # Deprecated n_past += len(embd) embd = [] @@ -75,20 +72,26 @@ logits = llama_cpp.llama_get_logits(ctx) n_vocab = llama_cpp.llama_n_vocab(model) - _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) - ]) + _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)) + 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_penalties(ctx, candidates_p, + llama_cpp.llama_sample_repetition_penalties( + ctx, + candidates_p, _arr, penalty_last_n=last_n_repeat, penalty_repeat=repeat_penalty, penalty_freq=frequency_penalty, - penalty_present=presence_penalty) + penalty_present=presence_penalty, + ) 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) @@ -111,10 +114,11 @@ size = 32 buffer = (ctypes.c_char * size)() n = llama_cpp.llama_token_to_piece( - model, llama_cpp.llama_token(id), buffer, size) + model, llama_cpp.llama_token(id), buffer, size + ) assert n <= size print( - buffer[:n].decode('utf-8'), + buffer[:n].decode("utf-8"), end="", flush=True, ) diff --git a/examples/low_level_api/quantize.py b/examples/low_level_api/quantize.py index 8bd03f88a..057ac389e 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/util.py b/examples/low_level_api/util.py index 9d0ec2f70..ef8b1c1ee 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 index 687316b32..73b28c744 100644 --- a/examples/notebooks/Batching.ipynb +++ b/examples/notebooks/Batching.ipynb @@ -392,7 +392,9 @@ "source": [ "params = llama_cpp.llama_model_default_params()\n", "params.n_gpu_layers = 35\n", - "model = llama_cpp.llama_load_model_from_file(b\"../../models/mistral-7b-v0.1-GGUF/ggml-model-Q4_K.gguf\", params=params) # Update this to whatever" + "model = llama_cpp.llama_load_model_from_file(\n", + " b\"../../models/mistral-7b-v0.1-GGUF/ggml-model-Q4_K.gguf\", params=params\n", + ") # Update this to whatever" ] }, { @@ -416,7 +418,9 @@ "prompt = b\"The quick brown fox\"\n", "\n", "tokens = (llama_cpp.llama_token * n_ctx)()\n", - "tokens_len = llama_cpp.llama_tokenize(model, prompt, len(prompt), tokens, len(tokens), True, True)\n", + "tokens_len = llama_cpp.llama_tokenize(\n", + " model, 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", @@ -447,7 +451,6 @@ } ], "source": [ - "\n", "ctx_params = llama_cpp.llama_context_default_params()\n", "ctx_params.seed = 1234\n", "ctx_params.n_ctx = n_kv_req\n", @@ -577,7 +580,7 @@ " for i in range(n_parallel):\n", " if i_batch[i] < 0:\n", " continue\n", - " \n", + "\n", " n_vocab = llama_cpp.llama_n_vocab(model)\n", " logits = llama_cpp.llama_get_logits_ith(ctx, i_batch[i])\n", "\n", @@ -588,7 +591,9 @@ " candidates[token_id].logit = logits[token_id]\n", " candidates[token_id].p = 0.0\n", "\n", - " candidates_p = llama_cpp.llama_token_data_array(candidates, len(candidates), False)\n", + " candidates_p = llama_cpp.llama_token_data_array(\n", + " candidates, len(candidates), False\n", + " )\n", "\n", " top_k = 40\n", " top_p = 0.9\n", @@ -596,8 +601,8 @@ "\n", " llama_cpp.llama_sample_top_k(ctx, ctypes.byref(candidates_p), top_k, 1)\n", " llama_cpp.llama_sample_top_p(ctx, ctypes.byref(candidates_p), top_p, 1)\n", - " llama_cpp.llama_sample_temp (ctx, ctypes.byref(candidates_p), temp)\n", - " \n", + " llama_cpp.llama_sample_temp(ctx, ctypes.byref(candidates_p), temp)\n", + "\n", " new_token_id = llama_cpp.llama_sample_token(ctx, ctypes.byref(candidates_p))\n", "\n", " if new_token_id == llama_cpp.llama_token_eos(ctx) or n_cur == n_len:\n", @@ -617,7 +622,7 @@ " i_batch[i] = batch.n_tokens\n", " batch.n_tokens += 1\n", " n_decode += 1\n", - " \n", + "\n", " if batch.n_tokens == 0:\n", " break\n", "\n", @@ -627,7 +632,7 @@ " print(\"Error decoding\", flush=True)\n", " break\n", " print(n_cur)\n", - " print(streams)\n" + " print(streams)" ] }, { diff --git a/examples/notebooks/Clients.ipynb b/examples/notebooks/Clients.ipynb index caebbb67f..fab82673e 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 index f1e5e9a1d..1f4138165 100644 --- a/examples/notebooks/Functions.ipynb +++ b/examples/notebooks/Functions.ipynb @@ -45,10 +45,11 @@ "\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", + " 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", @@ -56,15 +57,23 @@ " 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({\"location\": \"San Francisco\", \"temperature\": \"72\", \"unit\": \"fahrenheit\"})\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 = [{\"role\": \"user\", \"content\": \"What's the weather like in San Francisco, Tokyo, and Paris?\"}]\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", @@ -123,6 +132,8 @@ " 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())" ] }, @@ -169,16 +180,18 @@ "# 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", "\n", "assert isinstance(user, UserDetail)\n", @@ -213,17 +226,22 @@ "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", @@ -237,6 +255,7 @@ " ],\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)" @@ -265,19 +284,23 @@ "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", @@ -291,6 +314,7 @@ " ],\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", @@ -331,10 +355,12 @@ "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", @@ -351,6 +377,7 @@ ")\n", "print(qa)\n", "\n", + "\n", "class QuestionAnswerNoEvil(BaseModel):\n", " question: str\n", " answer: Annotated[\n", @@ -360,6 +387,7 @@ " ),\n", " ]\n", "\n", + "\n", "try:\n", " qa: QuestionAnswerNoEvil = client.chat.completions.create(\n", " model=\"gpt-3.5-turbo\",\n", @@ -405,6 +433,7 @@ "\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", @@ -424,6 +453,7 @@ " 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", @@ -440,13 +470,17 @@ " temperature=0.0,\n", " response_model=QuestionAnswer,\n", " messages=[\n", - " {\"role\": \"system\", \"content\": \"You are a world class algorithm to answer questions with correct and exact citations.\"},\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", + " {\"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", diff --git a/examples/notebooks/Guidance.ipynb b/examples/notebooks/Guidance.ipynb index 045856ea2..c52559853 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 index def68df00..8448ac1f7 100644 --- a/examples/notebooks/Multimodal.ipynb +++ b/examples/notebooks/Multimodal.ipynb @@ -38,23 +38,20 @@ " \"url\": \"https://user-images.githubusercontent.com/1991296/230134379-7181e485-c521-4d23-a0d6-f7b3b61ba524.png\",\n", " },\n", " },\n", - " {\"type\": \"text\", \"text\": \"What does the image say. Format your response as a json object with a single 'text' key.\"},\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", + " response_format={\n", " \"type\": \"json_object\",\n", - " \"schema\": {\n", - " \"type\": \"object\",\n", - " \"properties\": {\n", - " \"text\": {\n", - " \"type\": \"string\"\n", - " }\n", - " }\n", - " }\n", - " }\n", + " \"schema\": {\"type\": \"object\", \"properties\": {\"text\": {\"type\": \"string\"}}},\n", + " },\n", ")\n", "import json\n", + "\n", "print(json.loads(response.choices[0].message.content))" ] }, diff --git a/examples/notebooks/OpenHermesFunctionCalling.ipynb b/examples/notebooks/OpenHermesFunctionCalling.ipynb index c0de3fdc2..13128be04 100644 --- a/examples/notebooks/OpenHermesFunctionCalling.ipynb +++ b/examples/notebooks/OpenHermesFunctionCalling.ipynb @@ -42,42 +42,58 @@ "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", - "def calculate_mortgage_payment(loan_amount: int, interest_rate: float, loan_term: int) -> float:\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", + "\n", " # TODO: you must implement this to actually call it later\n", " pass\n", "\n", - "def get_article_details(title: str, authors: list[str], short_summary: str, date_published: str, tags: list[str]) -> Article:\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", + " 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", + "\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", + " 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", @@ -85,6 +101,7 @@ " 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", @@ -92,11 +109,8 @@ " function_info = {\n", " \"name\": func.__name__,\n", " \"description\": func.__doc__,\n", - " \"parameters\": {\n", - " \"type\": \"object\",\n", - " \"properties\": {}\n", - " },\n", - " \"returns\": type_hints.get('return', 'void').__name__\n", + " \"parameters\": {\"type\": \"object\", \"properties\": {}},\n", + " \"returns\": type_hints.get(\"return\", \"void\").__name__,\n", " }\n", "\n", " for name, _ in signature.parameters.items():\n", @@ -105,6 +119,7 @@ "\n", " return json.dumps(function_info, indent=2)\n", "\n", + "\n", "print(serialize_function_to_json(get_article_details))" ] }, @@ -117,13 +132,14 @@ "import xml.etree.ElementTree as ET\n", "import re\n", "\n", + "\n", "def extract_function_calls(completion):\n", " completion = completion.strip()\n", " pattern = r\"((.*?))\"\n", " match = re.search(pattern, completion, re.DOTALL)\n", " if not match:\n", " return None\n", - " \n", + "\n", " multiplefn = match.group(1)\n", " root = ET.fromstring(multiplefn)\n", " functions = root.findall(\"functioncall\")\n", @@ -399,7 +415,7 @@ "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", + " \"What's the current exchange rate for USD to EUR?\",\n", "]\n", "functions = [get_weather, calculate_mortgage_payment, get_article_details]\n", "\n", @@ -793,7 +809,12 @@ "source": [ "import llama_cpp\n", "\n", - "llama = llama_cpp.Llama(model_path=\"../../models/OpenHermes-2.5-Mistral-7B-GGUF/openhermes-2.5-mistral-7b.Q4_K_M.gguf\", n_gpu_layers=-1, n_ctx=2048, verbose=False)" + "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", + ")" ] }, { @@ -818,7 +839,7 @@ "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", + " \"What's the current exchange rate for USD to EUR?\",\n", "]\n", "functions = [get_weather, calculate_mortgage_payment, get_article_details]\n", "\n", @@ -830,7 +851,7 @@ " print(function_calls)\n", " else:\n", " print(completion.strip())\n", - " print(\"=\"*100)" + " print(\"=\" * 100)" ] }, { @@ -861,11 +882,13 @@ "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", + " \"What's the current exchange rate for USD to EUR?\",\n", "]\n", "\n", "for prompt in prompts:\n", - " completion = llama.create_completion(generate_hermes_prompt(prompt, functions), max_tokens=-1)[\"choices\"][0][\"text\"]\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", @@ -875,7 +898,7 @@ " else:\n", " print(completion.strip())\n", "\n", - " print(\"=\"*100)" + " print(\"=\" * 100)" ] }, { diff --git a/examples/notebooks/PerformanceTuning.ipynb b/examples/notebooks/PerformanceTuning.ipynb index 76e26fbd1..ba74e4a41 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 000000000..6e338ba17 --- /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 000000000..2325dd303 --- /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 000000000..a409fb882 --- /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/llama_cpp/__init__.py b/llama_cpp/__init__.py index fcbc7150c..e3cece2f6 100644 --- a/llama_cpp/__init__.py +++ b/llama_cpp/__init__.py @@ -1,4 +1,4 @@ from .llama_cpp import * from .llama import * -__version__ = "0.2.56" \ No newline at end of file +__version__ = "0.2.89" diff --git a/llama_cpp/_internals.py b/llama_cpp/_internals.py index 22d0bef8a..6dae88c8f 100644 --- a/llama_cpp/_internals.py +++ b/llama_cpp/_internals.py @@ -4,17 +4,20 @@ import ctypes from typing import ( + Dict, List, Optional, Sequence, ) 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 @@ -26,9 +29,6 @@ class _LlamaModel: """Intermediate Python wrapper for a llama.cpp llama_model. NOTE: For stability it's recommended you use the Llama class instead.""" - _llama_free_model = None - # NOTE: this must be "saved" here to avoid exceptions when calling __del__ - def __init__( self, *, @@ -39,26 +39,35 @@ def __init__( self.path_model = path_model self.params = params self.verbose = verbose - - self._llama_free_model = llama_cpp._lib.llama_free_model # type: ignore + self._exit_stack = ExitStack() self.model = None if not os.path.exists(path_model): raise ValueError(f"Model path does not exist: {path_model}") - self.model = llama_cpp.llama_load_model_from_file( - self.path_model.encode("utf-8"), self.params - ) + with suppress_stdout_stderr(disable=verbose): + self.model = llama_cpp.llama_load_model_from_file( + self.path_model.encode("utf-8"), self.params + ) if self.model is None: raise ValueError(f"Failed to load model from file: {path_model}") - def __del__(self): - if self.model is not None and self._llama_free_model is not None: - self._llama_free_model(self.model) + def free_model(): + if self.model is None: + return + llama_cpp.llama_free_model(self.model) self.model = None + self._exit_stack.callback(free_model) + + def close(self): + self._exit_stack.close() + + def __del__(self): + self.close() + def vocab_type(self) -> int: assert self.model is not None return llama_cpp.llama_vocab_type(self.model) @@ -109,9 +118,11 @@ def apply_lora_from_file( self.model, lora_path.encode("utf-8"), scale, - path_base_model.encode("utf-8") - if path_base_model is not None - else ctypes.c_char_p(0), + ( + path_base_model.encode("utf-8") + if path_base_model is not None + else ctypes.c_char_p(0) + ), n_threads, ) @@ -126,9 +137,9 @@ def token_get_score(self, token: int) -> float: assert self.model is not None return llama_cpp.llama_token_get_score(self.model, token) - def token_get_type(self, token: int) -> int: + def token_get_attr(self, token: int) -> int: assert self.model is not None - return llama_cpp.llama_token_get_type(self.model, token) + return llama_cpp.llama_token_get_attr(self.model, token) # Special tokens @@ -140,6 +151,14 @@ def token_eos(self) -> int: assert self.model is not None return llama_cpp.llama_token_eos(self.model) + def token_cls(self) -> int: + assert self.model is not None + return llama_cpp.llama_token_cls(self.model) + + def token_sep(self) -> int: + assert self.model is not None + return llama_cpp.llama_token_sep(self.model) + def token_nl(self) -> int: assert self.model is not None return llama_cpp.llama_token_nl(self.model) @@ -160,6 +179,14 @@ def token_eot(self) -> int: assert self.model is not None return llama_cpp.llama_token_eot(self.model) + def add_bos_token(self) -> bool: + assert self.model is not None + return llama_cpp.llama_add_bos_token(self.model) + + def add_eos_token(self) -> bool: + assert self.model is not None + return llama_cpp.llama_add_eos_token(self.model) + # Tokenization def tokenize(self, text: bytes, add_bos: bool, special: bool): @@ -181,27 +208,29 @@ def tokenize(self, text: bytes, add_bos: bool, special: bool): ) return list(tokens[:n_tokens]) - def token_to_piece(self, token: int) -> bytes: + def token_to_piece(self, token: int, special: bool = False) -> bytes: assert self.model is not None buf = ctypes.create_string_buffer(32) - llama_cpp.llama_token_to_piece(self.model, token, buf, 32) + llama_cpp.llama_token_to_piece(self.model, token, buf, 32, 0, special) return bytes(buf) - def detokenize(self, tokens: List[int]) -> bytes: + def detokenize(self, tokens: List[int], special: bool = False) -> bytes: assert self.model is not None output = b"" size = 32 buffer = (ctypes.c_char * size)() for token in tokens: n = llama_cpp.llama_token_to_piece( - self.model, llama_cpp.llama_token(token), buffer, size + self.model, 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() else output + output[1:] + if len(tokens) > 0 and tokens[0] == self.token_bos() and output[0:1] == b" " + else output ) # Extra @@ -211,20 +240,28 @@ def metadata(self) -> Dict[str, str]: buffer_size = 1024 buffer = ctypes.create_string_buffer(buffer_size) # zero the buffer - buffer.value = b'\0' * buffer_size + 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) + 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) + 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) + 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) + 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 @@ -239,8 +276,6 @@ class _LlamaContext: """Intermediate Python wrapper for a llama.cpp llama_context. NOTE: For stability it's recommended you use the Llama class instead.""" - _llama_free = None - def __init__( self, *, @@ -251,28 +286,39 @@ def __init__( self.model = model self.params = params self.verbose = verbose + self._exit_stack = ExitStack() - self._llama_free = llama_cpp._lib.llama_free # type: ignore self.ctx = None assert self.model.model is not None - self.ctx = llama_cpp.llama_new_context_with_model( - self.model.model, self.params - ) + self.ctx = llama_cpp.llama_new_context_with_model(self.model.model, self.params) if self.ctx is None: raise ValueError("Failed to create llama_context") - def __del__(self): - if self.ctx is not None and self._llama_free is not None: - self._llama_free(self.ctx) + 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: assert self.ctx is not None return llama_cpp.llama_n_ctx(self.ctx) + def pooling_type(self) -> int: + assert self.ctx is not None + return llama_cpp.llama_pooling_type(self.ctx) + def kv_cache_clear(self): assert self.ctx is not None llama_cpp.llama_kv_cache_clear(self.ctx) @@ -432,7 +478,11 @@ def sample_token_mirostat( ) def sample_token_mirostat_v2( - self, candidates: "_LlamaTokenDataArray", tau: float, eta: float, mu: llama_cpp.CtypesPointerOrRef[ctypes.c_float] + self, + candidates: "_LlamaTokenDataArray", + tau: float, + eta: float, + mu: llama_cpp.CtypesPointerOrRef[ctypes.c_float], ) -> int: assert self.ctx is not None return llama_cpp.llama_sample_token_mirostat_v2( @@ -461,7 +511,7 @@ def sample_token(self, candidates: "_LlamaTokenDataArray") -> int: def grammar_accept_token(self, grammar: LlamaGrammar, token: int): assert self.ctx is not None assert grammar.grammar is not None - llama_cpp.llama_grammar_accept_token(self.ctx, grammar.grammar, token) + llama_cpp.llama_grammar_accept_token(grammar.grammar, self.ctx, token) def reset_timings(self): assert self.ctx is not None @@ -479,8 +529,6 @@ def default_params(): class _LlamaBatch: - _llama_batch_free = None - def __init__( self, *, n_tokens: int, embd: int, n_seq_max: int, verbose: bool = True ): @@ -488,19 +536,27 @@ def __init__( self.embd = embd self.n_seq_max = n_seq_max self.verbose = verbose - - self._llama_batch_free = llama_cpp._lib.llama_batch_free # type: ignore + self._exit_stack = ExitStack() self.batch = None self.batch = llama_cpp.llama_batch_init( self._n_tokens, self.embd, self.n_seq_max ) - def __del__(self): - if self.batch is not None and self._llama_batch_free is not None: - self._llama_batch_free(self.batch) + 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: assert self.batch is not None return self.batch.n_tokens @@ -539,30 +595,26 @@ def add_sequence(self, batch: Sequence[int], seq_id: int, logits_all: bool): class _LlamaTokenDataArray: def __init__(self, *, n_vocab: int): self.n_vocab = n_vocab - self.candidates_data = np.array( - [], + self.candidates_data = np.recarray( + (self.n_vocab,), dtype=np.dtype( [("id", np.intc), ("logit", np.single), ("p", np.single)], align=True ), ) - self.candidates_data.resize(3, self.n_vocab, refcheck=False) 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_id = np.arange(self.n_vocab, dtype=np.intc) # type: ignore self.default_candidates_data_p = np.zeros(self.n_vocab, dtype=np.single) 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.data = self.candidates_data.ctypes.data_as( - llama_cpp.llama_token_data_p - ) - self.candidates.sorted = ctypes.c_bool(False) - self.candidates.size = ctypes.c_size_t(self.n_vocab) + 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 # Python wrappers over common/common @@ -597,13 +649,17 @@ def _tokenize(model: _LlamaModel, text: str, add_bos: bool, special: bool) -> li return list(result) -def _token_to_piece(model: _LlamaModel, token: int) -> str: +def _token_to_piece(model: _LlamaModel, token: int, special: bool = False) -> str: assert model.model is not None result = (ctypes.c_char * 8)(0) - n_tokens = llama_cpp.llama_token_to_piece(model.model, token, result, len(result)) + n_tokens = llama_cpp.llama_token_to_piece( + model.model, token, result, 0, len(result), special + ) if n_tokens < 0: result = (ctypes.c_char * -n_tokens)(0) - check = llama_cpp.llama_token_to_piece(model.model, token, result, len(result)) + check = llama_cpp.llama_token_to_piece( + model.model, token, result, 0, len(result), special + ) if check != -n_tokens: raise RuntimeError(f"Failed to get piece: token={token}") else: @@ -635,12 +691,22 @@ def _detokenize_bpe(model: _LlamaModel, tokens: List[int]) -> str: def _should_add_bos(model: _LlamaModel) -> bool: assert model.model is not None add_bos = llama_cpp.llama_add_bos_token(model.model) - if add_bos != -1: - return add_bos != 0 + if add_bos: + return add_bos else: return llama_cpp.llama_vocab_type(model.model) == llama_cpp.LLAMA_VOCAB_TYPE_SPM +# 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 @@ -655,7 +721,7 @@ class _LlamaSamplingParams: typical_p: float = 1.00 temp: float = 0.80 penalty_last_n: int = 64 - penalty_repeat: float = 1.10 + penalty_repeat: float = 1.0 penalty_freq: float = 0.00 penalty_present: float = 0.00 mirostat: int = 0 @@ -705,7 +771,10 @@ 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 + self, + ctx_main: _LlamaContext, + idx: int = 0, + logits_array: Optional[npt.NDArray[np.single]] = None, ): n_vocab = ctx_main.model.n_vocab() id: int = 0 @@ -730,25 +799,27 @@ def sample( if len(self.prev) > 0: nl_token = ctx_main.model.token_nl() nl_logit = logits_array[nl_token] - if self.params.penalty_last_n > 0: + last_tokens = self.prev[-self.params.penalty_last_n :] + last_tokens_size = min(len(last_tokens), self.params.penalty_last_n) + if last_tokens_size > 0: + last_tokens_p = (llama_cpp.llama_token * len(last_tokens))(*last_tokens) ctx_main.sample_repetition_penalties( token_data_array, - # TODO: Only create this once - (llama_cpp.llama_token * len(self.prev))(*self.prev), - self.params.penalty_last_n, + last_tokens_p, + last_tokens_size, self.params.penalty_repeat, self.params.penalty_freq, self.params.penalty_present, ) if not self.params.penalize_nl: - token_data_array.candidates_data["logit"][nl_token] = nl_logit + token_data_array.candidates_data.logit[nl_token] = nl_logit if self.grammar is not None: ctx_main.sample_grammar(token_data_array, self.grammar) if self.params.temp < 0: ctx_main.sample_softmax(token_data_array) - id = token_data_array.candidates_data["id"][0] + id = token_data_array.candidates_data.id[0] elif self.params.temp == 0: id = ctx_main.sample_token_greedy(token_data_array) else: @@ -794,4 +865,4 @@ def sample( def accept(self, ctx_main: _LlamaContext, id: int, apply_grammar: bool): if apply_grammar and self.grammar is not None: ctx_main.grammar_accept_token(self.grammar, id) - self.prev.append(id) \ No newline at end of file + self.prev.append(id) diff --git a/llama_cpp/_utils.py b/llama_cpp/_utils.py index 4a106470b..29628193b 100644 --- a/llama_cpp/_utils.py +++ b/llama_cpp/_utils.py @@ -1,13 +1,16 @@ import os import sys -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 @@ -22,12 +25,8 @@ def __enter__(self): if self.disable: return self - # Check if sys.stdout and sys.stderr have fileno method - if not hasattr(self.sys.stdout, 'fileno') or not hasattr(self.sys.stderr, 'fileno'): - return self # Return the instance without making changes - - self.old_stdout_fileno_undup = self.sys.stdout.fileno() - self.old_stderr_fileno_undup = self.sys.stderr.fileno() + 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) @@ -47,15 +46,14 @@ def __exit__(self, *_): return # Check if sys.stdout and sys.stderr have fileno method - if hasattr(self.sys.stdout, 'fileno') and hasattr(self.sys.stderr, 'fileno'): - self.sys.stdout = self.old_stdout - self.sys.stderr = self.old_stderr + 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.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) + self.os.close(self.old_stdout_fileno) + self.os.close(self.old_stderr_fileno) class MetaSingleton(type): diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py index fed84d579..5cc93c39b 100644 --- a/llama_cpp/llama.py +++ b/llama_cpp/llama.py @@ -1,16 +1,22 @@ from __future__ import annotations +import ctypes +import typing +import fnmatch +import json +import multiprocessing import os import sys -import uuid import time -import json -import ctypes -import fnmatch +import uuid +import warnings +import contextlib import multiprocessing from typing import ( + Any, List, + Literal, Optional, Union, Generator, @@ -18,30 +24,30 @@ Iterator, Deque, Callable, + Dict, ) from collections import deque from pathlib import Path -from llama_cpp.llama_types import List - from .llama_types import * from .llama_grammar import LlamaGrammar from .llama_cache import ( BaseLlamaCache, LlamaCache, # type: ignore LlamaDiskCache, # type: ignore + LlamaStaticDiskCache, # type: ignore LlamaRAMCache, # type: ignore + StateReloadError, # 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 +import llama_cpp.llama_chat_format as llama_chat_format +import llama_cpp.llama_cpp as llama_cpp +from llama_cpp.llama_speculative import LlamaDraftModel + from ._internals import ( _LlamaModel, # type: ignore _LlamaContext, # type: ignore @@ -49,9 +55,11 @@ _LlamaTokenDataArray, # type: ignore _LlamaSamplingParams, # type: ignore _LlamaSamplingContext, # type: ignore + _normalize_embedding, # type: ignore ) from ._logger import set_verbose from ._utils import suppress_stdout_stderr +from .llama_tokenizer import BaseLlamaTokenizer, LlamaTokenizer class Llama: @@ -68,17 +76,20 @@ def __init__( split_mode: int = llama_cpp.LLAMA_SPLIT_MODE_LAYER, main_gpu: int = 0, tensor_split: Optional[List[float]] = None, + rpc_servers: Optional[str] = None, vocab_only: bool = False, use_mmap: bool = True, use_mlock: bool = False, - kv_overrides: Optional[Dict[str, Union[bool, int, float]]] = 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_threads: Optional[int] = None, n_threads_batch: Optional[int] = None, - rope_scaling_type: Optional[int] = llama_cpp.LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED, + rope_scaling_type: Optional[ + int + ] = llama_cpp.LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED, pooling_type: int = llama_cpp.LLAMA_POOLING_TYPE_UNSPECIFIED, rope_freq_base: float = 0.0, rope_freq_scale: float = 0.0, @@ -90,6 +101,7 @@ def __init__( logits_all: bool = False, embedding: bool = False, offload_kqv: bool = True, + flash_attn: bool = False, # Sampling Params last_n_tokens_size: int = 64, # LoRA Params @@ -105,7 +117,11 @@ def __init__( 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 @@ -142,6 +158,7 @@ def __init__( 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_NONE: the GPU that is used for the entire model. LLAMA_SPLIT_ROW: the GPU that is used for small tensors and intermediate results. LLAMA_SPLIT_LAYER: ignored tensor_split: How split tensors should be distributed across GPUs. If None, the model is not split. + rpc_servers: Comma separated list of RPC servers to use for offloading 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. @@ -163,6 +180,7 @@ def __init__( logits_all: Return logits for all tokens, not just the last token. Must be True for completion to return logprobs. embedding: Embedding mode only. offload_kqv: Offload K, Q, V to GPU. + flash_attn: Use flash attention. 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. @@ -172,6 +190,9 @@ def __init__( 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. @@ -180,6 +201,7 @@ def __init__( A Llama instance. """ self.verbose = verbose + self._stack = contextlib.ExitStack() set_verbose(verbose) @@ -210,6 +232,11 @@ def __init__( ) # 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 + if rpc_servers is not None: + self.model_params.rpc_servers = rpc_servers.encode("utf-8") + self._rpc_servers = rpc_servers + else: + self._rpc_servers = None self.tensor_split = tensor_split self._c_tensor_split = None if self.tensor_split is not None: @@ -239,14 +266,40 @@ def __init__( 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.bool_value = v + 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.int_value = v + 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.float_value = v + 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}") @@ -257,9 +310,7 @@ def __init__( 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 max( - multiprocessing.cpu_count() // 2, 1 - ) + self.n_threads_batch = n_threads_batch or multiprocessing.cpu_count() # Context Params self.context_params = llama_cpp.llama_context_default_params() @@ -296,9 +347,14 @@ def __init__( self.context_params.logits_all = ( logits_all if draft_model is None else True ) # Must be set to True for speculative decoding - self.context_params.embeddings = embedding # TODO: Rename to embeddings + self.context_params.embeddings = embedding # TODO: Rename to embeddings self.context_params.offload_kqv = offload_kqv - + self.context_params.flash_attn = flash_attn + # 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.last_n_tokens_size = last_n_tokens_size @@ -308,11 +364,19 @@ def __init__( self.lora_scale = lora_scale self.lora_path = lora_path + self.spm_infill = spm_infill + if not os.path.exists(model_path): raise ValueError(f"Model path does not exist: {model_path}") - self._model = _LlamaModel( - path_model=self.model_path, params=self.model_params, verbose=self.verbose + self._model = self._stack.enter_context( + contextlib.closing( + _LlamaModel( + path_model=self.model_path, + params=self.model_params, + verbose=self.verbose, + ) + ) ) # Override tokenizer @@ -325,28 +389,54 @@ def __init__( self.context_params.n_ctx = self._model.n_ctx_train() self.context_params.n_batch = self.n_batch - self._ctx = _LlamaContext( - model=self._model, - params=self.context_params, - verbose=self.verbose, + self._ctx = self._stack.enter_context( + contextlib.closing( + _LlamaContext( + model=self._model, + params=self.context_params, + verbose=self.verbose, + ) + ) ) - self._batch = _LlamaBatch( - n_tokens=self.n_batch, - embd=0, - n_seq_max=self.context_params.n_ctx, - verbose=self.verbose, + self._batch = self._stack.enter_context( + contextlib.closing( + _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_lora_adapter_p] = None + if self.lora_path: - if self._model.apply_lora_from_file( - self.lora_path, - self.lora_scale, - self.lora_base, - self.n_threads, + assert self._model.model is not None + self._lora_adapter = llama_cpp.llama_lora_adapter_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_lora_adapter_free(self._lora_adapter) + self._lora_adapter = None + + self._stack.callback(free_lora_adapter) + + assert self._ctx.ctx is not None + if llama_cpp.llama_lora_adapter_set( + self._ctx.ctx, self._lora_adapter, self.lora_scale ): 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: @@ -354,6 +444,9 @@ def __init__( self.chat_format = chat_format self.chat_handler = chat_handler + self._chat_handlers: Dict[str, llama_chat_format.LlamaChatCompletionHandler] = ( + {} + ) self.draft_model = draft_model @@ -385,10 +478,46 @@ def __init__( 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 "tokenizer.chat_template" in self.metadata + and "chat_template.default" in template_choices ): chat_format = llama_chat_format.guess_chat_format_from_gguf_metadata( self.metadata @@ -399,32 +528,22 @@ def __init__( if self.verbose: print(f"Guessed chat format: {chat_format}", file=sys.stderr) else: - template = self.metadata["tokenizer.chat_template"] - try: - eos_token_id = int(self.metadata["tokenizer.ggml.eos_token_id"]) - except: - eos_token_id = self.token_eos() - try: - bos_token_id = int(self.metadata["tokenizer.ggml.bos_token_id"]) - except: - bos_token_id = self.token_bos() - - eos_token = self._model.token_get_text(eos_token_id) - bos_token = self._model.token_get_text(bos_token_id) - if self.verbose: - print(f"Using gguf chat template: {template}", file=sys.stderr) + 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_handler = llama_chat_format.Jinja2ChatFormatter( - template=template, eos_token=eos_token, bos_token=bos_token - ).to_chat_handler() + 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: {chat_format}", file=sys.stderr) + print( + f"Using fallback chat format: {self.chat_format}", file=sys.stderr + ) @property def ctx(self) -> llama_cpp.llama_context_p: @@ -526,14 +645,20 @@ def eval(self, tokens: Sequence[int]): # Save tokens self.input_ids[n_past : n_past + n_tokens] = batch # Save logits - rows = n_tokens - cols = self._n_vocab - offset = ( - 0 if self.context_params.logits_all else n_tokens - 1 - ) # NOTE: Only save the last token logits if logits_all is False - self.scores[n_past + offset : n_past + n_tokens, :].reshape(-1)[ - : - ] = self._ctx.get_logits()[offset * cols : rows * cols] + if self.context_params.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 # Update n_tokens self.n_tokens += n_tokens @@ -544,7 +669,7 @@ def sample( 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, @@ -619,7 +744,7 @@ def generate( 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, @@ -637,7 +762,7 @@ def generate( 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: @@ -663,11 +788,15 @@ def generate( else: break if longest_prefix > 0: - if self.verbose: - print("Llama.generate: prefix-match hit", file=sys.stderr) 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, + ) # Reset the model state if reset: @@ -747,7 +876,7 @@ def create_embedding( input = input if isinstance(input, list) else [input] # get numeric embeddings - embeds: List[List[float]] + embeds: Union[List[List[float]], List[List[List[float]]]] total_tokens: int embeds, total_tokens = self.embed(input, return_count=True) # type: ignore @@ -774,7 +903,7 @@ def create_embedding( def embed( self, input: Union[str, List[str]], - normalize: bool = True, + normalize: bool = False, truncate: bool = True, return_count: bool = False, ): @@ -790,7 +919,11 @@ def embed( n_embd = self.n_embd() n_batch = self.n_batch - if self.context_params.embeddings == False: + # get pooling information + pooling_type = self.pooling_type() + logits_all = pooling_type == llama_cpp.LLAMA_POOLING_TYPE_NONE + + if self.context_params.embeddings is False: raise RuntimeError( "Llama model must be created with embedding=True to call this method" ) @@ -807,29 +940,38 @@ def embed( self._batch.reset() # decode and fetch embeddings - data: List[List[float]] = [] + data: Union[List[List[float]], List[List[List[float]]]] = [] - def decode_batch(n_seq: int): + def decode_batch(seq_sizes: List[int]): assert self._ctx.ctx is not None llama_cpp.llama_kv_cache_clear(self._ctx.ctx) self._ctx.decode(self._batch) self._batch.reset() # store embeddings - for i in range(n_seq): - ptr = llama_cpp.llama_get_embeddings_seq( - self._ctx.ctx, i - ) - if not ptr: - raise RuntimeError("Failed to get embeddings from sequence pooling type is not set") - embedding: List[float] = ptr[:n_embd] - if normalize: - norm = float(np.linalg.norm(embedding)) - embedding = [v / norm for v in embedding] - data.append(embedding) + 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 = [_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 = _normalize_embedding(embedding) + data.append(embedding) # init state total_tokens = 0 + s_batch = [] t_batch = 0 p_batch = 0 @@ -850,17 +992,21 @@ def decode_batch(n_seq: int): # time to eval batch if t_batch + n_tokens > n_batch: - decode_batch(p_batch) + decode_batch(s_batch) + s_batch = [] t_batch = 0 p_batch = 0 # add to batch - self._batch.add_sequence(tokens, p_batch, False) + self._batch.add_sequence(tokens, p_batch, logits_all) + + # update batch stats + s_batch.append(n_tokens) t_batch += n_tokens p_batch += 1 # hanlde last batch - decode_batch(p_batch) + decode_batch(s_batch) if self.verbose: llama_cpp.llama_print_timings(self._ctx.ctx) @@ -889,7 +1035,7 @@ def _create_completion( 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, @@ -910,19 +1056,83 @@ def _create_completion( completion_id: str = f"cmpl-{str(uuid.uuid4())}" created: int = int(time.time()) + 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 = self._model.token_prefix() + middle_token_id: int = self._model.token_middle() + suffix_token_id: int = self._model.token_suffix() + 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 [self.token_bos()] + 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] = ( + prefix_tokens: List[int] = ( + [prefix_token_id] if prefix_token_id >= 0 and suffix is not None else [] + ) + ( ( - self.tokenize(prompt.encode("utf-8"), special=True) + self.tokenize( + prompt.encode("utf-8"), + add_bos=False, + special=(prefix_token_id < 0 or suffix is None), + ) if prompt != "" - else [self.token_bos()] + 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 = ( @@ -930,6 +1140,12 @@ 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: @@ -983,17 +1199,64 @@ def logit_bias_processor( if self.cache: try: - cache_item = self.cache[prompt_tokens] + # pylint: disable=protected-access + longest_cache_key = self.cache._find_longest_prefix_key(prompt_tokens) + + if longest_cache_key is None: + raise KeyError + + if self.verbose: + print("Item found in cache", file=sys.stderr) + cache_prefix_len = Llama.longest_token_prefix( - cache_item.input_ids.tolist(), prompt_tokens + longest_cache_key, prompt_tokens ) eval_prefix_len = Llama.longest_token_prefix( self._input_ids.tolist(), prompt_tokens ) if cache_prefix_len > eval_prefix_len: - self.load_state(cache_item) + # Debugging: print the portion of the prompt tokens that are + # found within prefix length For dumb reasons, have to skip + # the first token if it's a BOS token. + before = time.time() + cache_item = self.cache[prompt_tokens] + after = time.time() if self.verbose: - print("Llama._create_completion: cache hit", file=sys.stderr) + print("Cache lookup took", round((after - before) * 1_000, 4), "ms", file=sys.stderr) + if self.verbose and (prompt_tokens[0] == self.token_bos()): + print( + "Matching prompt tokens: " + f"{self.detokenize(prompt_tokens[1:cache_prefix_len - 1]).decode('utf-8')}", + file=sys.stderr, + ) + elif self.verbose: + print( + f"Matching prompt tokens: {self.detokenize(prompt_tokens[:cache_prefix_len]).decode('utf-8')}", + file=sys.stderr, + ) + + try: + before = time.time() + self.cache.reload_from_cache_state(self, cache_item) + after = time.time() + if self.verbose: + print("State loading took", round((after - before) * 1_000, 4), "ms", file=sys.stderr) + print( + f"Llama._create_completion: cache hit with len {cache_prefix_len} / {len(prompt_tokens)}", + file=sys.stderr, + ) + except StateReloadError as e: + if self.verbose: + print( + f"Llama._create_completion: cache hit with len {cache_prefix_len} / {len(prompt_tokens)}, but failed to reload state: {e}", + file=sys.stderr, + ) + print("Falling back to re-evaluating prompt", file=sys.stderr) + elif self.verbose: + print( + f"Llama._create_completion: not reloading from cache, cache prefix len {cache_prefix_len} < eval prefix len {eval_prefix_len}", + file=sys.stderr, + ) except KeyError: if self.verbose: print("Llama._create_completion: cache miss", file=sys.stderr) @@ -1021,7 +1284,8 @@ def logit_bias_processor( logits_processor=logits_processor, grammar=grammar, ): - if token == self._token_eos: + assert self._model.model is not None + if llama_cpp.llama_token_is_eog(self._model.model, token): text = self.detokenize(completion_tokens, prev_tokens=prompt_tokens) finish_reason = "stop" break @@ -1052,7 +1316,10 @@ def logit_bias_processor( if stream: remaining_tokens = completion_tokens[returned_tokens:] - remaining_text = self.detokenize(remaining_tokens, prev_tokens=prompt_tokens + completion_tokens[:returned_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 @@ -1072,21 +1339,31 @@ def logit_bias_processor( # not sure how to handle this branch when dealing # with CJK output, so keep it unchanged for token in remaining_tokens: - if token == self.token_bos(): + if token == bos_token_id: continue - token_end_position += len(self.detokenize([token], prev_tokens=prompt_tokens + completion_tokens[:returned_tokens])) + 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" - ) + 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], prev_tokens=prompt_tokens + completion_tokens[:returned_tokens]).decode( - "utf-8", errors="ignore" - ) + 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, :] @@ -1106,9 +1383,11 @@ def logit_bias_processor( top_logprob.update({token_str: current_logprobs[int(token)]}) logprobs_or_none = { "tokens": [ - self.detokenize([token], prev_tokens=prompt_tokens + completion_tokens[:returned_tokens]).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)]], @@ -1122,9 +1401,11 @@ def logit_bias_processor( "model": model_name, "choices": [ { - "text": self.detokenize([token], prev_tokens=prompt_tokens + completion_tokens[:returned_tokens]).decode( - "utf-8", errors="ignore" - ), + "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, @@ -1136,7 +1417,11 @@ def logit_bias_processor( 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]) + bs = self.detokenize( + remaining_tokens[:i], + prev_tokens=prompt_tokens + + completion_tokens[:returned_tokens], + ) ts = bs.decode("utf-8") decode_success = True break @@ -1186,7 +1471,10 @@ def logit_bias_processor( if stream: remaining_tokens = completion_tokens[returned_tokens:] - all_text = self.detokenize(remaining_tokens, prev_tokens=prompt_tokens + completion_tokens[:returned_tokens]) + all_text = self.detokenize( + remaining_tokens, + prev_tokens=prompt_tokens + completion_tokens[:returned_tokens], + ) any_stop = [s for s in stop_sequences if s in all_text] if len(any_stop) > 0: end = min(all_text.index(stop) for stop in any_stop) @@ -1195,17 +1483,26 @@ def logit_bias_processor( token_end_position = 0 for token in remaining_tokens: - token_end_position += len(self.detokenize([token], prev_tokens=prompt_tokens + completion_tokens[:returned_tokens])) + 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 == self.token_bos(): + 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], prev_tokens=prompt_tokens + 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, :] @@ -1283,14 +1580,15 @@ def logit_bias_processor( } ], } - if self.cache: + if self.cache and not self.cache.is_ro: 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: + if self.cache and not self.cache.is_ro: if self.verbose: print("Llama._create_completion: cache save", file=sys.stderr) self.cache[prompt_tokens + completion_tokens] = self.save_state() @@ -1300,7 +1598,7 @@ def logit_bias_processor( 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 @@ -1313,13 +1611,18 @@ def logit_bias_processor( 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], prev_tokens=all_tokens[:i]).decode("utf-8", errors="ignore") + 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(self._scores)[token_offset:] @@ -1327,7 +1630,7 @@ def logit_bias_processor( for idx, (token, token_str, logprobs_token) in enumerate( zip(all_tokens, all_token_strs, all_logprobs) ): - if token == self.token_bos(): + if token == bos_token_id: continue text_offsets.append( text_offset @@ -1345,7 +1648,9 @@ def logit_bias_processor( ) token_logprobs.append(logprobs_token[int(token)]) top_logprob: Optional[Dict[str, float]] = { - self.detokenize([i], prev_tokens=all_tokens[:idx]).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)]}) @@ -1397,7 +1702,7 @@ def create_completion( 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, @@ -1494,7 +1799,7 @@ def __call__( 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, @@ -1591,7 +1896,7 @@ def create_chat_completion( 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, @@ -1638,9 +1943,12 @@ def create_chat_completion( Returns: Generated chat completion or a stream of chat completion chunks. """ - handler = self.chat_handler or llama_chat_format.get_chat_completion_handler( - self.chat_format + handler = ( + self.chat_handler + or self._chat_handlers.get(self.chat_format) + or llama_chat_format.get_chat_completion_handler(self.chat_format) ) + print(f"Got seed: {seed}") return handler( llama=self, messages=messages, @@ -1653,6 +1961,8 @@ def create_chat_completion( top_k=top_k, min_p=min_p, typical_p=typical_p, + logprobs=logprobs, + top_logprobs=top_logprobs, stream=stream, stop=stop, seed=seed, @@ -1723,6 +2033,7 @@ def __getstate__(self): 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, 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, @@ -1732,6 +2043,8 @@ def __getstate__(self): yarn_orig_ctx=self.context_params.yarn_orig_ctx, logits_all=self.context_params.logits_all, embedding=self.context_params.embeddings, + offload_kqv=self.context_params.offload_kqv, + flash_attn=self.context_params.flash_attn, # Sampling Params last_n_tokens_size=self.last_n_tokens_size, # LoRA Params @@ -1743,51 +2056,18 @@ def __getstate__(self): # 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"], - # Model Params - n_gpu_layers=state["n_gpu_layers"], - split_mode=state["split_mode"], - main_gpu=state["main_gpu"], - tensor_split=state["tensor_split"], - vocab_only=state["vocab_only"], - use_mmap=state["use_mmap"], - use_mlock=state["use_mlock"], - kv_overrides=state["kv_overrides"], - # Context Params - seed=state["seed"], - n_ctx=state["n_ctx"], - n_batch=state["n_batch"], - n_threads=state["n_threads"], - n_threads_batch=state["n_threads_batch"], - rope_freq_base=state["rope_freq_base"], - rope_freq_scale=state["rope_freq_scale"], - rope_scaling_type=state["rope_scaling_type"], - yarn_ext_factor=state["yarn_ext_factor"], - yarn_attn_factor=state["yarn_attn_factor"], - yarn_beta_fast=state["yarn_beta_fast"], - yarn_beta_slow=state["yarn_beta_slow"], - yarn_orig_ctx=state["yarn_orig_ctx"], - logits_all=state["logits_all"], - embedding=state["embedding"], - # Sampling Params - last_n_tokens_size=state["last_n_tokens_size"], - # LoRA Params - lora_base=state["lora_base"], - lora_path=state["lora_path"], - # Backend Params - numa=state["numa"], - # Chat Format Params - chat_format=state["chat_format"], - chat_handler=state["chat_handler"], - # Misc - verbose=state["verbose"], - ) + self.__init__(**state) def save_state(self) -> LlamaState: assert self._ctx.ctx is not None @@ -1812,7 +2092,7 @@ 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), @@ -1821,7 +2101,9 @@ def save_state(self) -> LlamaState: def load_state(self, state: LlamaState) -> None: assert self._ctx.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() + self.scores[state.n_tokens :, :] = 0.0 self.input_ids = state.input_ids.copy() self.n_tokens = state.n_tokens state_size = state.llama_state_size @@ -1859,6 +2141,17 @@ def token_nl(self) -> int: """Return the newline token.""" 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: Union[npt.NDArray[np.single], List], axis: int = -1 @@ -1911,7 +2204,7 @@ def from_pretrained( Returns: A Llama model.""" try: - from huggingface_hub import hf_hub_download, HfFileSystem + from huggingface_hub import HfFileSystem, hf_hub_download from huggingface_hub.utils import validate_repo_id except ImportError: raise ImportError( @@ -1925,7 +2218,7 @@ def from_pretrained( files = [ file["name"] if isinstance(file, dict) else file - for file in hffs.ls(repo_id) + for file in hffs.ls(repo_id, recursive=True) ] # split each file into repo_id, subfolder, filename @@ -1972,7 +2265,6 @@ def from_pretrained( 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) @@ -2021,3 +2313,19 @@ 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 index 9e9870a52..4f754a8cc 100644 --- a/llama_cpp/llama_cache.py +++ b/llama_cpp/llama_cache.py @@ -1,19 +1,25 @@ +import ctypes +import pickle import sys from abc import ABC, abstractmethod -from typing import ( - Optional, - Sequence, - Tuple, -) from collections import OrderedDict +from typing import Optional, Sequence, Tuple import diskcache +import numpy as np +import pytrie import llama_cpp.llama from .llama_types import * +class StateReloadError(Exception): + """ + Error for when state from cache cannot be read by current model. + """ + + class BaseLlamaCache(ABC): """Base cache class for a llama.cpp model.""" @@ -25,6 +31,11 @@ def __init__(self, capacity_bytes: int = (2 << 30)): def cache_size(self) -> int: raise NotImplementedError + @property + @abstractmethod + def is_ro(self) -> bool: + raise NotImplementedError + def _find_longest_prefix_key( self, key: Tuple[int, ...], @@ -40,9 +51,25 @@ def __contains__(self, key: Sequence[int]) -> bool: raise NotImplementedError @abstractmethod - def __setitem__(self, key: Sequence[int], value: "llama_cpp.llama.LlamaState") -> None: + def __setitem__( + self, key: Sequence[int], value: "llama_cpp.llama.LlamaState" + ) -> None: raise NotImplementedError + @classmethod + def reload_from_cache_state( + cls, model: "llama_cpp.llama.Llama", state: "llama_cpp.llama.LlamaState" + ) -> None: + """ + Reload the state onto the model. Normally this is done with load_state + (as state is created with the corresponding `save_state`), but for some + caches may need special handling as an optimization. + + Throws a StateReloadError if the state is not compatible with the model + (for example, logits ) + """ + model.load_state(state) + class LlamaRAMCache(BaseLlamaCache): """Cache for a llama.cpp model using RAM.""" @@ -50,12 +77,18 @@ class LlamaRAMCache(BaseLlamaCache): 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() + 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()]) + @property + def is_ro(self) -> bool: + return False + def _find_longest_prefix_key( self, key: Tuple[int, ...], @@ -63,7 +96,8 @@ def _find_longest_prefix_key( min_len = 0 min_key = None keys = ( - (k, llama_cpp.llama.Llama.longest_token_prefix(k, key)) for k in self.cache_state.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: @@ -109,6 +143,10 @@ def __init__( def cache_size(self): return int(self.cache.volume()) # type: ignore + @property + def is_ro(self) -> bool: + return False + def _find_longest_prefix_key( self, key: Tuple[int, ...], @@ -148,3 +186,211 @@ def __setitem__(self, key: Sequence[int], value: "llama_cpp.llama.LlamaState"): key_to_remove = next(iter(self.cache)) del self.cache[key_to_remove] print("LlamaDiskCache.__setitem__: trim", file=sys.stderr) + + +class LlamaStaticDiskCache(BaseLlamaCache): + """ + Cache that only reads from the cache, doesn't store / overwrite items, and + doesn't pop from cache. + + Still using diskcache.Cache for underlying cache, but uses a trie to store + keys so that can more efficiently look for prefixes. + + Want to store C++ state as bytes (from `llama_copy_state_data`), but for now + still storing LlamaState, because need scores/input_ids/n_tokens so that Python + code can continue inference. + """ + + def __init__( + self, cache_dir: str = ".cache/llama_cache", capacity_bytes: int = (2 << 30) + ): + self.cache = diskcache.Cache( + cache_dir, size_limit=capacity_bytes, cull_limit=0, eviction_policy="none" + ) + self.capacity_bytes = capacity_bytes + # Don't want to have to iterate over all keys when doing longest matching prefix search + self.keys = pytrie.Trie.fromkeys(self.cache.iterkeys()) + + @property + def cache_size(self): + return int(self.cache.volume()) # type: ignore + + @property + def is_ro(self) -> bool: + return True + + def _private_setitem(self, key: Sequence[int], value: "llama_cpp.llama.LlamaState"): + if self.cache_size > self.capacity_bytes: + # I think it's okay to raise an error here, because only done when building cache anyway. + raise ValueError("Cache is full, refusing to set more") + + key = tuple(key) + if key in self.cache: + print( + "LlamaStaticDiskCache._private_setitem: delete (overwriting)", + file=sys.stderr, + ) + del self.cache[key] + + # This is what diskcache does anyway, eventually want this to be more compact + print("LlamaStaticDiskCache._private_setitem: set", file=sys.stderr) + self.cache[key] = pickle.dumps(value, pickle.HIGHEST_PROTOCOL) + + @staticmethod + def build_cache( + cache_dir: str, + prompts: Sequence[str], + model: "llama_cpp.Llama", + # Same default as LlamaDiskCache, 1 GB + capacity_bytes: int = 2 << 30, + seed: Optional[int] = None, + add_bos=True, + save_logits: bool = False, + ) -> "LlamaStaticDiskCache": + """ + Using model passed in, evaluates each prompt and stores LlamaState in cache. + + Returns a new LlamaStaticDiskCache instance with cache at cache_dir. + """ + cache = LlamaStaticDiskCache(cache_dir, capacity_bytes) + + for p in prompts: + if seed: + model.set_seed(seed) + # Special tokens == control characters like in ChatML + toks = model.tokenize(p.encode("utf-8"), add_bos=add_bos, special=True) + # Will always eval at least one token, same logic as in + # `Llama.generate` for prefix-match hit. + # pylint: disable=protected-access + shared_prefix_len = model.longest_token_prefix(toks[:-1], model._input_ids) + # Reset to shared prefix length so that don't have to re-eval system prompt + model.n_tokens = shared_prefix_len + eval_toks = toks[shared_prefix_len:] + print("LlamaStaticDiskCache.build_cache: eval", file=sys.stderr) + model.eval(eval_toks) + state = model.save_state() + + if not save_logits: + if ( + model.context_params.logits_all + or model.draft_model is not None + or model.context_params.embeddings + ): + # Erroring instead of falling back to just saving with scores + raise ValueError( + "Cannot save state without logits - model requires logits to sample." + ) + state.scores = None + + cache._private_setitem(toks, state) # pylint: disable=protected-access + + # Set up Trie for efficient prefix search + for key in cache.cache.iterkeys(): + cache.keys[key] = None + + return cache + + def _find_longest_prefix_key(self, key: Tuple[int]) -> Optional[Tuple[int, ...]]: + try: + longest_prefix = self.keys.longest_prefix(key) + return longest_prefix + except KeyError: + return None + + def __contains__(self, key: Sequence[int]) -> bool: + return self._find_longest_prefix_key(tuple(key)) is not None + + def __getitem__(self, key: Sequence[int]) -> "llama_cpp.llama.LlamaState": + """ + Only handling exact matches (not prefixes). Use case is that have some + prompt + context that want to match against. + """ + key = tuple(key) + # Don't worry about KeyError, that's handled by caller + longest_prefix = self._find_longest_prefix_key(key) + value: "llama_cpp.llama.LlamaState" = pickle.loads(self.cache[longest_prefix]) + return value + + def __setitem__(self, key: Sequence[int], value: "llama_cpp.llama.LlamaState"): + # Should this just be a warning? + raise ValueError("Cannot set items in a static cache") + + @classmethod + def reload_from_cache_state( + cls, model: "llama_cpp.llama.Llama", state: "llama_cpp.llama.LlamaState" + ) -> None: + """ + Skip reloading logits and set last logits from llama.cpp context struct + as the scores for last token of prompt. + """ + # pylint: disable=protected-access + + # Check if model needs logits (draft model, log probs required, etc.) + model_needs_scores_to_reload = ( + # May be overly pessimistic if don't want embeddings for prompt tokens. + model.context_params.embeddings + or model.context_params.logits_all + # Same: is this really a hard requirement? We need token IDs from + # draft model and all the logits from base model to do verification + # of candidate tokens, but not for prompt tokens. + or model.draft_model is not None + ) + + if model_needs_scores_to_reload: + if state.scores is None: + raise StateReloadError( + "Model requires logits to be reloaded, but static cache does not store logits" + ) + else: + model.load_state(state) + return + + # Case where don't need logits from numpy and can just get last-token + # logits from llama.cpp struct + model.n_tokens = state.n_tokens + model.input_ids = state.input_ids.copy() + model.scores[:] = 0.0 + + state_size = state.llama_state_size + + try: + llama_state_array_type = ctypes.c_uint8 * state_size + # Have to do from_buffer_copy since LlamaState.llama_state is + # non-mutable bytes, not mutable bytearray. + llama_state = llama_state_array_type.from_buffer_copy(state.llama_state) + reloaded_state_size = llama_cpp.llama_set_state_data( + model._ctx.ctx, llama_state + ) + + if reloaded_state_size != state_size: + raise StateReloadError( + "Failed to set llama state data - reloaded state size " + f"{reloaded_state_size} does not match original size {state_size}" + ) + + # cffi dtype, compatible w/ numpy through ducktyping :scared: + dtype = llama_cpp.llama_cpp.llama_get_logits_ith.restype._type_ + + # If model scores dtype doesn't match dtype from sig, then can't + # copy it. + if model.scores.dtype != dtype: + raise StateReloadError( + f"Expected scores to be {dtype} but got " + f"{model.scores.dtype} - are you running this in the future? Or the past?" + ) + + # Will have a ValueError for null pointers + last_position_logits = np.array( + ctypes.cast( + model._ctx.get_logits_ith(-1), + ctypes.POINTER(dtype * model.n_vocab()), + ).contents, + # Otherwise will be a view into C array on llama.cpp context + copy=True, + dtype=dtype, + ) + + model._scores[-1, :] = last_position_logits + + except ValueError as e: + raise StateReloadError from e diff --git a/llama_cpp/llama_chat_format.py b/llama_cpp/llama_chat_format.py index 81ca5520b..e60bb5ddb 100644 --- a/llama_cpp/llama_chat_format.py +++ b/llama_cpp/llama_chat_format.py @@ -1,14 +1,32 @@ from __future__ import annotations import os +import sys import json import ctypes import dataclasses import random import string -from typing import Any, Dict, Iterator, List, Literal, Optional, Tuple, Union, Protocol + +from contextlib import ExitStack +from typing import ( + Any, + Dict, + Iterator, + List, + Literal, + Optional, + Tuple, + Union, + Protocol, + cast, +) import jinja2 +from jinja2.sandbox import ImmutableSandboxedEnvironment + +import numpy as np +import numpy.typing as npt import llama_cpp.llama as llama import llama_cpp.llama_types as llama_types @@ -32,6 +50,9 @@ # 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 ### @@ -77,6 +98,8 @@ def __call__( 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, @@ -148,6 +171,8 @@ class ChatFormatterResponse: prompt: str stop: Optional[Union[str, List[str]]] = None + stopping_criteria: Optional[llama.StoppingCriteriaList] = None + added_special: bool = False class ChatFormatter(Protocol): @@ -171,14 +196,18 @@ def __init__( 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 + ) - self._environment = jinja2.Environment( + self._environment = ImmutableSandboxedEnvironment( loader=jinja2.BaseLoader(), trim_blocks=True, lstrip_blocks=True, @@ -188,6 +217,10 @@ 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): @@ -199,9 +232,28 @@ def raise_exception(message: str): 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, ) - return ChatFormatterResponse(prompt=prompt, stop=[self.eos_token]) + 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) @@ -223,6 +275,7 @@ def _convert_text_completion_to_chat( "role": "assistant", "content": completion["choices"][0]["text"], }, + "logprobs": completion["choices"][0]["logprobs"], "finish_reason": completion["choices"][0]["finish_reason"], } ], @@ -246,6 +299,7 @@ def _convert_text_completion_chunks_to_chat( "delta": { "role": "assistant", }, + "logprobs": None, "finish_reason": None, } ], @@ -265,6 +319,7 @@ def _convert_text_completion_chunks_to_chat( if chunk["choices"][0]["finish_reason"] is None else {} ), + "logprobs": chunk["choices"][0]["logprobs"], "finish_reason": chunk["choices"][0]["finish_reason"], } ], @@ -288,6 +343,183 @@ def _convert_completion_to_chat( 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": 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": 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": 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: @@ -322,6 +554,8 @@ def chat_completion_handler( 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, @@ -331,15 +565,74 @@ def chat_completion_handler( messages=messages, functions=functions, function_call=function_call, + tools=tools, + tool_choice=tool_choice, + ) + prompt = llama.tokenize( + result.prompt.encode("utf-8"), + add_bos=not result.added_special, + special=True, ) - prompt = result.prompt 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) + 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, @@ -348,6 +641,7 @@ def chat_completion_handler( 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, @@ -361,9 +655,15 @@ def chat_completion_handler( 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 @@ -387,7 +687,9 @@ def format_autotokenizer( prompt: str = tokenizer.apply_chat_template(messages, tokenize=False) # type: ignore assert isinstance(prompt, str) # Return formatted prompt and eos token by default - return ChatFormatterResponse(prompt=prompt, stop=tokenizer.eos_token) + return ChatFormatterResponse( + prompt=prompt, stop=tokenizer.eos_token, added_special=True + ) return format_autotokenizer @@ -417,8 +719,7 @@ def hf_tokenizer_config_to_chat_formatter( assert isinstance(tokenizer_config["eos_token"], str) eos_token = tokenizer_config["eos_token"] - env = jinja2.Environment( - loader=jinja2.BaseLoader(), + env = ImmutableSandboxedEnvironment( trim_blocks=True, lstrip_blocks=True, ).from_string(chat_template) @@ -441,7 +742,9 @@ def format_tokenizer_config( bos_token=bos_token, eos_token=eos_token, ) - return ChatFormatterResponse(prompt=prompt, stop=[eos_token, bos_token]) + return ChatFormatterResponse( + prompt=prompt, stop=[eos_token, bos_token], added_special=True + ) return format_tokenizer_config @@ -463,10 +766,15 @@ def guess_chat_format_from_gguf_metadata(metadata: Dict[str, str]) -> Optional[s 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): + 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 @@ -597,13 +905,15 @@ def _format_chatglm3( 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(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 + schema: str, verbose: bool = False, fallback_to_json: bool = True ): try: return llama_grammar.LlamaGrammar.from_json_schema(schema, verbose=verbose) @@ -613,9 +923,10 @@ def _grammar_for_json_schema( else: raise e + def _grammar_for_response_format( - response_format: llama_types.ChatCompletionRequestResponseFormat, - verbose: bool = False + response_format: llama_types.ChatCompletionRequestResponseFormat, + verbose: bool = False, ): if response_format["type"] != "json_object": return None @@ -627,6 +938,7 @@ def _grammar_for_response_format( else: return _grammar_for_json(verbose=verbose) + ### Chat Formats ### @@ -648,7 +960,7 @@ def format_llama2( messages: List[llama_types.ChatCompletionRequestMessage], **kwargs: Any, ) -> ChatFormatterResponse: - _system_template = "[INST] <>\n{system_message}\n<>" + _system_template = "[INST] <>\n{system_message}\n<>" _roles = dict(user="[INST]", assistant="[/INST]") _messages = _map_roles(messages, _roles) system_message = _get_system_message(messages) @@ -658,6 +970,25 @@ def format_llama2( 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], @@ -939,10 +1270,9 @@ def format_mistral_instruct( messages: List[llama_types.ChatCompletionRequestMessage], **kwargs: Any, ) -> ChatFormatterResponse: - bos = "" eos = "" stop = eos - prompt = bos + prompt = "" for message in messages: if ( message["role"] == "user" @@ -950,11 +1280,7 @@ def format_mistral_instruct( and isinstance(message["content"], str) ): prompt += "[INST] " + message["content"] - elif ( - message["role"] == "assistant" - and message["content"] is not None - and isinstance(message["content"], str) - ): + elif message["role"] == "assistant" and message["content"] is not None: prompt += " [/INST]" + message["content"] + eos prompt += " [/INST]" return ChatFormatterResponse(prompt=prompt, stop=stop) @@ -1025,7 +1351,7 @@ def format_gemma( **kwargs: Any, ) -> ChatFormatterResponse: system_message = _get_system_message(messages) - if system_message is not None and system_message != "": + if system_message != "": logger.debug( "`role='system'` messages are not allowed on Google's Gemma models." ) @@ -1037,6 +1363,55 @@ def format_gemma( return ChatFormatterResponse(prompt=_prompt, stop=_sep) +@register_chat_format("nekomata") +def format_nekomata( + messages: List[llama_types.ChatCompletionRequestMessage], + **kwargs: Any, +) -> ChatFormatterResponse: + # This is an Alpaca format model so don't support multi-turn + num_user_messages = sum(1 for m in messages if m["role"] == "user") + assert ( + num_user_messages == 1 + ), f"Only one user message allowed, got {num_user_messages}" + + user_msg = [m for m in messages if m["role"] == "user"][0] + try: + system_msg = [m for m in messages if m["role"] == "system"][0] + except IndexError: + system_msg = None + + if system_msg is not None: + prompt = """ +以下は、タスクを説明する指示と、文脈のある入力の組み合わせです。要求を適切に満たす応答を書きなさい。 + +### 指示: +{instruction} + +### 入力: +{input} + +### 応答: +""".format( + instruction=system_msg["content"].strip(), input=user_msg["content"].strip() + ) + + else: + prompt = """ +以下は、タスクを説明する指示と、文脈のある入力の組み合わせです。要求を適切に満たす応答を書きなさい。 + +### 入力: +{input} + +### 応答: +""".format( + input=user_msg["content"].strip() + ) + + prompt = prompt.lstrip("\n") + + return ChatFormatterResponse(prompt) + + # Tricky chat formats that require custom chat handlers @@ -1288,12 +1663,12 @@ def message_to_str(msg: llama_types.ChatCompletionRequestMessage): 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 + f":\n" + 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 + f":\n" + new_prompt = prompt + ":\n" function_body = None for function in functions or []: @@ -1390,6 +1765,7 @@ def message_to_str(msg: llama_types.ChatCompletionRequestMessage): } ], }, + "logprobs": completion["choices"][0]["logprobs"], "finish_reason": "tool_calls", } ], @@ -1540,27 +1916,35 @@ def prepare_messages_for_inference( 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 functions is not None: + if tool_choice == "none": all_messages.append( llama_types.ChatCompletionRequestSystemMessage( - role="system", content=generate_schema_from_functions(functions) + role="system", content=generate_schema_from_functions([]) ) ) - elif 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" - ] - ), + 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( @@ -1596,13 +1980,17 @@ def prepare_messages_for_inference( 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 + messages, tokenizer, version, functions, tools, function_call ) # If no tools/functions are provided - if function_call is None and (functions is None or len(functions) == 0): + if function_call == "none" or functions is None or len(functions) == 0: if version == "v1": stop = END_ASSISTANT_TOKEN else: @@ -1630,10 +2018,12 @@ def prepare_messages_for_inference( 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 - assert stream is False # TODO: support stream mode - def get_grammar(function_call): function_body = None for function in functions or []: @@ -1667,379 +2057,1304 @@ def get_grammar(function_call): return grammar - def create_completion(stop): - completion: 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, + 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 - if version == "v1": - # If no or "auto" tool_choice/function_call - if function_call is None or ( - isinstance(function_call, str) and function_call == "auto" - ): - stops = ["\n", END_ASSISTANT_TOKEN] - # If tool_choice/function_call is "none" - elif isinstance(function_call, str) and function_call == "none": - prompt = prepare_messages_for_inference( - messages, tokenizer, version, [], [] - ) - stops = 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] + def generate_streaming(tools, functions, function_call, prompt): + assert version == "v2", "Streaming for v1 is not supported" - completion = create_completion(stop=stops) - completion_text = completion["choices"][0]["text"] + chunk_id, chunk_created = None, None - # 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 - ): - 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" + # 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)] ) - function_calls.append( - completion_text.split(START_FUNCTION_CALL_TOKEN)[-1][:-1].strip() + 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": 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, + }, + } + ], ) - grammar = get_grammar(function_calls[-1]) - completion = create_completion(stop=END_FUNCTION_CALL_TOKEN) - 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: - # Loop until all parallel function calls are generated - while True: - # If no or "auto" tool_choice/function_call - if function_call is None or ( - isinstance(function_call, str) and function_call == "auto" - ): + # 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 - # If tool_choice/function_call is "none" - elif isinstance(function_call, str) and function_call == "none": - prompt = ( - prepare_messages_for_inference(messages, tokenizer, version, [], []) - + "all\n<|content|>" + completion = create_completion( + prompt=prompt, stop=stops, grammar=grammar ) - stops = STOP_TOKEN + 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": 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": 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": 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": 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"{function_call['name']}\n{CONTENT_TOKEN}" - stops = STOP_TOKEN + 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 = STOP_TOKEN + stops = ["\n", END_ASSISTANT_TOKEN] - completion = create_completion(stop=stops) + 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 prompt.endswith("all\n<|content|>") and not completion_text.startswith( - "all" + 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 - # Generate model response if the model decides not to call any function - elif prompt.endswith(RECIPIENT_TOKEN) and completion_text.startswith("all"): - prompt += completion_text + CONTENT_TOKEN - completion = create_completion(stop=STOP_TOKEN) - return _convert_completion_to_chat(completion, stream=stream) # type: ignore - # Generate parameters if model decides to call a function - elif prompt.endswith(RECIPIENT_TOKEN): - function_calls.append(completion_text[:-1]) + # 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(stop=[STOP_TOKEN, "\n"]) + 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()) - prompt += f"{function_calls[-1]}\n{CONTENT_TOKEN}{function_bodies[-1]}" - grammar = None - - # Try to generate the beginning of next turn - # If empty completion, break from loop - next_turn_completion_text = create_completion( - stop=[STOP_TOKEN, RECIPIENT_TOKEN] - )["choices"][0]["text"] - if len(next_turn_completion_text) > 0: - prompt += f"\n{FROM_TOKEN}assistant\n{RECIPIENT_TOKEN}" - else: - break - # Break from loop if tool_choice/function_call is provided as a dict + # If the prompt involves a function call, just append generated parameters to function_bodies else: function_bodies.append(completion_text.strip()) - break - - assert "usage" in completion - assert len(function_calls) > 0 - assert len(function_calls) == len(function_bodies) - - tool_calls = [] - 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, - }, - } - ) + 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 - 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": tool_calls[0]["function"]["name"], - "arguments": tool_calls[0]["function"]["arguments"], + # 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": completion["choices"][0]["logprobs"], + "message": { + "role": "assistant", + "content": None if content == "" else content, + **function_call_dict, }, - "tool_calls": tool_calls, - }, - "finish_reason": "tool_calls", - } - ], - usage=completion["usage"], - ) + "finish_reason": "tool_calls" if len(tool_calls) > 0 else "stop", + } + ], + usage=completion["usage"], + ) class Llava15ChatHandler: - _clip_free = None + 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 = False): + def __init__(self, clip_model_path: str, verbose: bool = True): import llama_cpp.llava_cpp as llava_cpp - self._llava_cpp = llava_cpp self.clip_model_path = clip_model_path self.verbose = verbose - self._clip_free = self._llava_cpp._libllava.clip_free # type: ignore - if not os.path.exists(clip_model_path): - raise ValueError(f"Clip model path does not exist: {clip_model_path}") + self._llava_cpp = llava_cpp # TODO: Fix + self._exit_stack = ExitStack() + self._last_image_embed: Optional[ + llava_cpp.CtypesPointer[llava_cpp.llava_image_embed] + ] = None + self._last_image_hash: Optional[int] = None + + if not os.path.exists(clip_model_path): + raise ValueError(f"Clip model path does not exist: {clip_model_path}") + + with suppress_stdout_stderr(disable=self.verbose): + clip_ctx = self._llava_cpp.clip_model_load(self.clip_model_path.encode(), 0) + + if clip_ctx is None: + raise ValueError(f"Failed to load clip model: {clip_model_path}") + + self.clip_ctx = clip_ctx + + def clip_free(): + with suppress_stdout_stderr(disable=self.verbose): + self._llava_cpp.clip_free(self.clip_ctx) + + self._exit_stack.callback(clip_free) + + def last_image_embed_free(): + with suppress_stdout_stderr(disable=self.verbose): + if self._last_image_embed is not None: + self._llava_cpp.llava_image_embed_free(self._last_image_embed) + self._last_image_embed = None + + self._exit_stack.callback(last_image_embed_free) + + def load_image(self, image_url: str) -> bytes: + return self._load_image(image_url) + + 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], + ]: + assert self.clip_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) + text = template.render( + messages=messages, + add_generation_prompt=True, + eos_token=llama.detokenize([llama.token_eos()]), + bos_token=llama.detokenize([llama.token_bos()]), + ) + split_text = self.split_text_on_image_urls(text, image_urls) + + def embed_image_bytes(image_bytes: bytes): + if ( + self._last_image_embed is not None + and self._last_image_hash is not None + and hash(image_bytes) == self._last_image_hash + ): + return self._last_image_embed + with suppress_stdout_stderr(disable=self.verbose): + # Free the previous image embed + if self._last_image_embed is not None: + self._llava_cpp.llava_image_embed_free(self._last_image_embed) + self._last_image_embed = None + self._last_image_hash = None + embed = self._llava_cpp.llava_image_embed_make_with_bytes( + self.clip_ctx, + llama.context_params.n_threads_batch, + (ctypes.c_uint8 * len(image_bytes)).from_buffer( + bytearray(image_bytes) + ), + len(image_bytes), + ) + self._last_image_embed = embed + self._last_image_hash = hash(image_bytes) + return embed + + # Evaluate prompt + llama.reset() + llama._ctx.kv_cache_clear() + for type_, value in split_text: + if type_ == "text": + tokens = llama.tokenize( + value.encode("utf8"), add_bos=False, special=True + ) + 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) + else: + image_bytes = self.load_image(value) + embed = embed_image_bytes(image_bytes) + if llama.n_tokens + embed.contents.n_image_pos > llama.n_ctx(): + raise ValueError( + f"Prompt exceeds n_ctx: {llama.n_tokens + embed.contents.n_image_pos} > {llama.n_ctx()}" + ) + n_past = ctypes.c_int(llama.n_tokens) + n_past_p = ctypes.pointer(n_past) + with suppress_stdout_stderr(disable=self.verbose): + self._llava_cpp.llava_eval_image_embed( + llama.ctx, + embed, + llama.n_batch, + n_past_p, + ) + # Required to avoid issues with hf tokenizer + llama.input_ids[llama.n_tokens : n_past.value] = -1 + llama.n_tokens = n_past.value + + # Get prompt tokens to avoid a cache miss + prompt = llama.input_ids[: llama.n_tokens].tolist() + + 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]): + 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 ObsidianChatHandler(Llava15ChatHandler): + # Prompt Format + # The model followed ChatML format. However, with ### as the seperator + + # <|im_start|>user + # What is this sign about?\n + # ### + # <|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"\n\n{chat_history}Question: {question}\n\nAnswer:" + CHAT_FORMAT = ( + "{% for message in messages %}" + "{% if message.role == 'user' %}" + "{% if message.content is iterable %}" + # + "{% 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: \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 %}" + # + "{% 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 %}" + ) + - with suppress_stdout_stderr(disable=self.verbose): - self.clip_ctx = self._llava_cpp.clip_model_load( - self.clip_model_path.encode(), 0 - ) +class NanoLlavaChatHandler(Llava15ChatHandler): + # Prompt Format + # The model follow the ChatML standard, however, without \n at the end of <|im_end|>: - def __del__(self): - with suppress_stdout_stderr(disable=self.verbose): - if self.clip_ctx is not None and self._clip_free is not None: - self._clip_free(self.clip_ctx) - self.clip_ctx = None + # <|im_start|>system + # Answer the question<|im_end|><|im_start|>user + # + # What is the picture about?<|im_end|><|im_start|>assistant + DEFAULT_SYSTEM_MESSAGE = "Answer the question" - def load_image(self, image_url: str) -> bytes: - if image_url.startswith("data:"): - import base64 + 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 %}" + ) - 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 +class Llama3VisionAlphaChatHandler(Llava15ChatHandler): + # question = "" + q - 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]]] = [], - 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.CreateChatCompletionResponse, - Iterator[llama_types.CreateChatCompletionStreamResponse], - ]: - assert ( - llama.context_params.logits_all is True - ) # BUG: logits_all=True is required for llava - assert self.clip_ctx is not None - system_prompt = _get_system_message(messages) - system_prompt = ( - system_prompt - if system_prompt != "" - else "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_role = "\nUSER:" - assistant_role = "\nASSISTANT:" - llama.reset() - llama.eval(llama.tokenize(system_prompt.encode("utf8"), add_bos=True)) - for message in messages: - if message["role"] == "user" and message["content"] is not None: - if isinstance(message["content"], str): - llama.eval( - llama.tokenize( - f"{user_role} {message['content']}".encode("utf8"), - add_bos=False, - ) - ) - else: - assert isinstance(message["content"], list) - llama.eval( - llama.tokenize(f"{user_role} ".encode("utf8"), add_bos=False) - ) - for content in message["content"]: - if content["type"] == "text": - llama.eval( - llama.tokenize( - f"{content['text']}".encode("utf8"), add_bos=False - ) - ) - if content["type"] == "image_url": - image_bytes = ( - self.load_image(content["image_url"]["url"]) - if isinstance(content["image_url"], dict) - else self.load_image(content["image_url"]) - ) - import array - - data_array = array.array("B", image_bytes) - c_ubyte_ptr = ( - ctypes.c_ubyte * len(data_array) - ).from_buffer(data_array) - with suppress_stdout_stderr(disable=self.verbose): - embed = ( - self._llava_cpp.llava_image_embed_make_with_bytes( - self.clip_ctx, - llama.context_params.n_threads, - c_ubyte_ptr, - len(image_bytes), - ) - ) - try: - n_past = ctypes.c_int(llama.n_tokens) - n_past_p = ctypes.pointer(n_past) - with suppress_stdout_stderr(disable=self.verbose): - self._llava_cpp.llava_eval_image_embed( - llama.ctx, - embed, - llama.n_batch, - n_past_p, - ) - assert llama.n_ctx() >= n_past.value - llama.n_tokens = n_past.value - finally: - with suppress_stdout_stderr(disable=self.verbose): - self._llava_cpp.llava_image_embed_free(embed) - if message["role"] == "assistant" and message["content"] is not None: - llama.eval( - llama.tokenize( - f"ASSISTANT: {message['content']}".encode("utf8"), add_bos=False - ) - ) - assert llama.n_ctx() >= llama.n_tokens - llama.eval(llama.tokenize(f"{assistant_role}".encode("utf8"), add_bos=False)) - assert llama.n_ctx() >= llama.n_tokens + # 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 - prompt = llama.input_ids[: llama.n_tokens].tolist() + CHAT_FORMAT = ( + "{% for message in messages %}" + "<|start_header_id|>" + "{% if message.role == 'user' %}" + "user<|end_header_id|>\n\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 %}" + # 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 %}" + ) - 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, - ), - stream=stream, - ) +# alias +Llama3VisionAlpha = Llama3VisionAlphaChatHandler @register_chat_completion_handler("chatml-function-calling") @@ -2069,6 +3384,8 @@ def chatml_function_calling( 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, @@ -2125,8 +3442,7 @@ def chatml_function_calling( "{% endfor %}" "{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}" ) - template_renderer = jinja2.Environment( - loader=jinja2.BaseLoader(), + template_renderer = ImmutableSandboxedEnvironment( autoescape=jinja2.select_autoescape(["html", "xml"]), undefined=jinja2.StrictUndefined, ).from_string(function_calling_template) @@ -2155,7 +3471,11 @@ def chatml_function_calling( }, } - stop = [stop, "<|im_end|>"] if isinstance(stop, str) else stop + ["<|im_end|>"] if stop else ["<|im_end|>"] + 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 ( @@ -2195,185 +3515,11 @@ def chatml_function_calling( model=model, logits_processor=logits_processor, grammar=grammar, + logprobs=top_logprobs if logprobs else None, ), stream=stream, ) - 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"], - }, - } - ], - }, - "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": None, - "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": "", - }, - } - ], - }, - } - ], - } - 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": None, - "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) - # Case 2: Tool choice by user if isinstance(tool_choice, dict): tool_name = tool_choice["function"]["name"] @@ -2482,6 +3628,7 @@ def _stream_response_to_function_stream( 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, @@ -2503,8 +3650,8 @@ def _stream_response_to_function_stream( tool_name = text[len("functions.") :] tool = next((tool for tool in tools if tool["function"]["name"] == tool_name), None) if not stream: - completions = [] - completions_tool_name = [] + completions: List[llama_types.CreateCompletionResponse] = [] + completions_tool_name: List[str] = [] while tool is not None: prompt += f"functions.{tool_name}:\n" try: @@ -2541,6 +3688,9 @@ def _stream_response_to_function_stream( 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"] @@ -2569,6 +3719,7 @@ def _stream_response_to_function_stream( follow_up_gbnf_tool_grammar, verbose=llama.verbose ), ) + response = cast(llama_types.CreateCompletionResponse, response) tool_name = response["choices"][0]["text"][len("functions.") :] tool = next( @@ -2576,12 +3727,22 @@ def _stream_response_to_function_stream( ) # Merge completions - function_call = { - "function_call": { - "name": tool_name, - "arguments": completions[0]["choices"][0]["text"], + 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 {} + if len(completions) == 1 + else {} + ) return { "id": "chat" + completion["id"], "object": "chat.completion", @@ -2591,6 +3752,7 @@ def _stream_response_to_function_stream( { "finish_reason": "tool_calls", "index": 0, + "logprobs": completion["choices"][0]["logprobs"], "message": { "role": "assistant", "content": None, @@ -2611,20 +3773,26 @@ def _stream_response_to_function_stream( zip(completions_tool_name, completions) ) ], - **function_call + **function_call_dict, }, } ], "usage": { "completion_tokens": sum( - completion["usage"]["completion_tokens"] + ( + completion["usage"]["completion_tokens"] + if "usage" in completion + else 0 + ) for completion in completions ), "prompt_tokens": sum( - completion["usage"]["prompt_tokens"] for completion in completions + completion["usage"]["prompt_tokens"] if "usage" in completion else 0 + for completion in completions ), "total_tokens": sum( - completion["usage"]["total_tokens"] for completion in completions + completion["usage"]["total_tokens"] if "usage" in completion else 0 + for completion in completions ), }, } diff --git a/llama_cpp/llama_cpp.py b/llama_cpp/llama_cpp.py index b9593cf7a..476713aee 100644 --- a/llama_cpp/llama_cpp.py +++ b/llama_cpp/llama_cpp.py @@ -23,12 +23,12 @@ # Load the library def _load_shared_library(lib_base_name: str): # Construct the paths to the possible shared library names - _base_path = pathlib.Path(os.path.abspath(os.path.dirname(__file__))) + _base_path = pathlib.Path(os.path.abspath(os.path.dirname(__file__))) / "lib" # 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"): + if sys.platform.startswith("linux") or sys.platform.startswith("freebsd"): _lib_paths += [ _base_path / f"lib{lib_base_name}.so", ] @@ -52,7 +52,12 @@ def _load_shared_library(lib_base_name: str): _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": + 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: @@ -141,6 +146,70 @@ def byref(obj: CtypesCData, offset: Optional[int] = None) -> CtypesRef[CtypesCDa byref = ctypes.byref # type: ignore +# 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_COUNT, +# }; +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_COUNT = 30 # from ggml-backend.h # typedef bool (*ggml_backend_sched_eval_callback)(struct ggml_tensor * t, bool ask, void * user_data); @@ -164,20 +233,24 @@ def byref(obj: CtypesCData, offset: Optional[int] = None) -> CtypesRef[CtypesCDa # define LLAMA_DEFAULT_SEED 0xFFFFFFFF LLAMA_DEFAULT_SEED = 0xFFFFFFFF -# define LLAMA_MAX_RNG_STATE (64*1024) -LLAMA_MAX_RNG_STATE = 64 * 1024 - # 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 -# define LLAMA_SESSION_VERSION 4 -LLAMA_SESSION_VERSION = 4 +# define LLAMA_SESSION_VERSION 8 +LLAMA_SESSION_VERSION = 8 +# 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_model; llama_model_p = NewType("llama_model_p", int) @@ -199,14 +272,78 @@ def byref(obj: CtypesCData, offset: Optional[int] = None) -> CtypesRef[CtypesCDa # enum llama_vocab_type { # LLAMA_VOCAB_TYPE_NONE = 0, // For models without vocab -# LLAMA_VOCAB_TYPE_SPM = 1, // SentencePiece -# LLAMA_VOCAB_TYPE_BPE = 2, // Byte Pair Encoding -# LLAMA_VOCAB_TYPE_WPM = 3, // WordPiece +# 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_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""" + + +# // 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_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_VOCAV_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 # // note: these values should be synchronized with ggml_rope @@ -214,16 +351,14 @@ def byref(obj: CtypesCData, offset: Optional[int] = None) -> CtypesRef[CtypesCDa # enum llama_rope_type { # LLAMA_ROPE_TYPE_NONE = -1, # LLAMA_ROPE_TYPE_NORM = 0, -# LLAMA_ROPE_TYPE_NEOX = 2, -# LLAMA_ROPE_TYPE_GLM = 4, +# LLAMA_ROPE_TYPE_NEOX = GGML_ROPE_TYPE_NEOX, # }; LLAMA_ROPE_TYPE_NONE = -1 LLAMA_ROPE_TYPE_NORM = 0 -LLAMA_ROPE_TYPE_NEOX = 2 -LLAMA_ROPE_TYPE_GLM = 4 +LLAMA_ROPE_TYPE_NEOX = GGML_ROPE_TYPE_NEOX = 2 -# enum llama_token_type { +# 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, @@ -241,13 +376,39 @@ def byref(obj: CtypesCData, offset: Optional[int] = None) -> CtypesRef[CtypesCDa 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_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 @@ -274,14 +435,18 @@ def byref(obj: CtypesCData, offset: Optional[int] = None) -> CtypesRef[CtypesCDa # 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, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_Q4_0_4_8 = 34, // except 1d tensors +# LLAMA_FTYPE_MOSTLY_Q4_0_8_8 = 35, // 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_Q4_1_SOME_F16 = 4 LLAMA_FTYPE_MOSTLY_Q8_0 = 7 LLAMA_FTYPE_MOSTLY_Q5_0 = 8 LLAMA_FTYPE_MOSTLY_Q5_1 = 9 @@ -306,6 +471,11 @@ def byref(obj: CtypesCData, offset: Optional[int] = None) -> CtypesRef[CtypesCDa 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_GUESSED = 1024 # enum llama_rope_scaling_type { @@ -326,11 +496,22 @@ def byref(obj: CtypesCData, offset: Optional[int] = None) -> CtypesRef[CtypesCDa # LLAMA_POOLING_TYPE_NONE = 0, # LLAMA_POOLING_TYPE_MEAN = 1, # LLAMA_POOLING_TYPE_CLS = 2, +# LLAMA_POOLING_TYPE_LAST = 3, # }; 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 + +# 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_split_mode { # LLAMA_SPLIT_MODE_NONE = 0, // single GPU @@ -355,6 +536,11 @@ class llama_token_data(ctypes.Structure): 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", ctypes.c_float), @@ -378,6 +564,11 @@ class llama_token_data_array(ctypes.Structure): size (int): size of the array sorted (bool): whether the array is sorted""" + if TYPE_CHECKING: + data: CtypesArray[llama_token_data] + size: int + sorted: bool + _fields_ = [ ("data", llama_token_data_p), ("size", ctypes.c_size_t), @@ -387,7 +578,7 @@ class llama_token_data_array(ctypes.Structure): llama_token_data_array_p = ctypes.POINTER(llama_token_data_array) -# typedef bool (*llama_progress_callback)(float progress, void *ctx); +# 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 ) @@ -439,6 +630,15 @@ class llama_batch(ctypes.Structure): 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] + _fields_ = [ ("n_tokens", ctypes.c_int32), ("token", ctypes.POINTER(llama_token)), @@ -457,36 +657,54 @@ class llama_batch(ctypes.Structure): # LLAMA_KV_OVERRIDE_TYPE_INT, # LLAMA_KV_OVERRIDE_TYPE_FLOAT, # LLAMA_KV_OVERRIDE_TYPE_BOOL, +# LLAMA_KV_OVERRIDE_TYPE_STR, # }; LLAMA_KV_OVERRIDE_TYPE_INT = 0 LLAMA_KV_OVERRIDE_TYPE_FLOAT = 1 LLAMA_KV_OVERRIDE_TYPE_BOOL = 2 +LLAMA_KV_OVERRIDE_TYPE_STR = 3 # struct llama_model_kv_override { -# char key[128]; # enum llama_model_kv_override_type tag; + +# char key[128]; + + # union { -# int64_t int_value; -# double float_value; -# bool bool_value; +# int64_t val_i64; +# double val_f64; +# bool val_bool; +# char val_str[128]; # }; # }; class llama_model_kv_override_value(ctypes.Union): _fields_ = [ - ("int_value", ctypes.c_int64), - ("float_value", ctypes.c_double), - ("bool_value", ctypes.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 + class llama_model_kv_override(ctypes.Structure): _fields_ = [ - ("key", ctypes.c_char * 128), ("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_params { # int32_t n_gpu_layers; // number of layers to store in VRAM @@ -501,6 +719,9 @@ class llama_model_kv_override(ctypes.Structure): # // proportion of the model (layers or rows) to offload to each GPU, size: llama_max_devices() # const float * tensor_split; +# // comma separated list of RPC servers to use for offloading +# const char * rpc_servers; + # // 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. @@ -514,9 +735,10 @@ class llama_model_kv_override(ctypes.Structure): # // 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_mlock; // force system to keep model in RAM +# 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 check_tensors; // validate model tensor data # }; class llama_model_params(ctypes.Structure): """Parameters for llama_model @@ -526,27 +748,47 @@ class llama_model_params(ctypes.Structure): split_mode (int): how to split the model across multiple GPUs main_gpu (int): the GPU that is used for the entire model. main_gpu interpretation depends on split_mode: LLAMA_SPLIT_NONE: the GPU that is used for the entire model LLAMA_SPLIT_ROW: the GPU that is used for small tensors and intermediate results LLAMA_SPLIT_LAYER: ignored tensor_split (ctypes.Array[ctypes.ctypes.c_float]): proportion of the model (layers or rows) to offload to each GPU, size: llama_max_devices() + rpc_servers (ctypes.c_char_p): comma separated list of RPC servers to use for offloading 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_mlock (bool): force system to keep model in RAM""" + use_mlock (bool): force system to keep model in RAM + check_tensors (bool): validate model tensor data""" + + if TYPE_CHECKING: + n_gpu_layers: int + split_mode: int + main_gpu: int + tensor_split: CtypesArray[ctypes.c_float] + rpc_servers: ctypes.c_char_p + 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_mlock: bool + check_tensors: bool _fields_ = [ ("n_gpu_layers", ctypes.c_int32), ("split_mode", ctypes.c_int), ("main_gpu", ctypes.c_int32), ("tensor_split", ctypes.POINTER(ctypes.c_float)), + ("rpc_servers", ctypes.c_char_p), ("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_mlock", ctypes.c_bool), + ("check_tensors", ctypes.c_bool), ] +# // NOTE: changing the default values of parameters marked as [EXPERIMENTAL] may cause crashes or incorrect results in certain configurations +# // https://github.com/ggerganov/llama.cpp/pull/7544 # struct llama_context_params { # uint32_t seed; // RNG seed, -1 for random # uint32_t n_ctx; // text context, 0 = from model @@ -558,7 +800,7 @@ class llama_model_params(ctypes.Structure): # 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 -# // (ignored if no pooling layer) +# enum llama_attention_type attention_type; // attention type to use for embeddings # // ref: https://github.com/ggerganov/llama.cpp/pull/2054 # float rope_freq_base; // RoPE base frequency, 0 = from model @@ -573,13 +815,15 @@ class llama_model_params(ctypes.Structure): # ggml_backend_sched_eval_callback cb_eval; # void * cb_eval_user_data; -# enum ggml_type type_k; // data type for K cache -# enum ggml_type type_v; // data type for V cache +# enum ggml_type type_k; // data type for K cache [EXPERIMENTAL] +# enum ggml_type type_v; // data type for V cache [EXPERIMENTAL] # // Keep the booleans together to avoid misalignment during copy-by-value. # bool logits_all; // the llama_decode() call computes all logits, not just the last one (DEPRECATED - set llama_batch.logits instead) # bool embeddings; // if true, extract embeddings (together with logits) # bool offload_kqv; // whether to offload the KQV ops (including the KV cache) to GPU +# bool flash_attn; // whether to use flash attention [EXPERIMENTAL] + # // Abort callback # // if it returns true, execution of llama_decode() will be aborted @@ -600,6 +844,7 @@ class llama_context_params(ctypes.Structure): n_threads_batch (int): number of threads to use for batch processing 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 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 @@ -615,10 +860,41 @@ class llama_context_params(ctypes.Structure): logits_all (bool): the llama_eval() call computes all logits, not just the last one (DEPRECATED - set llama_batch.logits instead) embeddings (bool): if true, extract embeddings (together with logits) offload_kqv (bool): whether to offload the KQV ops (including the KV cache) to GPU + flash_attn (bool): whether to use flash attention 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 """ + if TYPE_CHECKING: + seed: int + n_ctx: int + n_batch: int + n_ubatch: int + n_seq_max: int + n_threads: int + n_threads_batch: int + rope_scaling_type: int + pooling_type: int + attention_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 + logits_all: bool + embeddings: bool + offload_kqv: bool + flash_attn: bool + abort_callback: Callable[[ctypes.c_void_p], bool] + abort_callback_data: ctypes.c_void_p + _fields_ = [ ("seed", ctypes.c_uint32), ("n_ctx", ctypes.c_uint32), @@ -629,6 +905,7 @@ class llama_context_params(ctypes.Structure): ("n_threads_batch", ctypes.c_uint32), ("rope_scaling_type", ctypes.c_int), ("pooling_type", ctypes.c_int), + ("attention_type", ctypes.c_int), ("rope_freq_base", ctypes.c_float), ("rope_freq_scale", ctypes.c_float), ("yarn_ext_factor", ctypes.c_float), @@ -644,6 +921,7 @@ class llama_context_params(ctypes.Structure): ("logits_all", ctypes.c_bool), ("embeddings", ctypes.c_bool), ("offload_kqv", ctypes.c_bool), + ("flash_attn", ctypes.c_bool), ("abort_callback", ggml_abort_callback), ("abort_callback_data", ctypes.c_void_p), ] @@ -667,13 +945,17 @@ class llama_context_params(ctypes.Structure): # // 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 -# 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 -# void * imatrix; // pointer to importance matrix data +# 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 +# void * imatrix; // pointer to importance matrix data +# void * kv_overrides; // pointer to vector containing overrides # } llama_model_quantize_params; class llama_model_quantize_params(ctypes.Structure): """Parameters for llama_model_quantize @@ -681,21 +963,42 @@ class llama_model_quantize_params(ctypes.Structure): 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 - imatrix (ctypes.ctypes.c_void_p): pointer to importance matrix data + keep_split (bool): quantize to the same number of shards + imatrix (ctypes.c_void_p): pointer to importance matrix data + kv_overrides (ctypes.c_void_p): pointer to vector containing overrides """ + 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 + imatrix: ctypes.c_void_p + kv_overrides: ctypes.c_void_p + _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), ("imatrix", ctypes.c_void_p), + ("kv_overrides", ctypes.c_void_p), ] @@ -727,6 +1030,9 @@ class llama_model_quantize_params(ctypes.Structure): # // modifies a preceding LLAMA_GRETYPE_CHAR or # // LLAMA_GRETYPE_CHAR_RNG_UPPER to add an alternate char to match ([ab], [a-zA]) # LLAMA_GRETYPE_CHAR_ALT = 6, + +# // any character (.) +# LLAMA_GRETYPE_CHAR_ANY = 7, # }; LLAMA_GRETYPE_END = 0 LLAMA_GRETYPE_ALT = 1 @@ -735,6 +1041,7 @@ class llama_model_quantize_params(ctypes.Structure): LLAMA_GRETYPE_CHAR_NOT = 4 LLAMA_GRETYPE_CHAR_RNG_UPPER = 5 LLAMA_GRETYPE_CHAR_ALT = 6 +LLAMA_GRETYPE_CHAR_ANY = 7 # typedef struct llama_grammar_element { @@ -742,6 +1049,10 @@ class llama_model_quantize_params(ctypes.Structure): # uint32_t value; // Unicode code point or rule ID # } llama_grammar_element; class llama_grammar_element(ctypes.Structure): + if TYPE_CHECKING: + type: int + value: int + _fields_ = [ ("type", ctypes.c_int), ("value", ctypes.c_uint32), @@ -765,6 +1076,17 @@ class llama_grammar_element(ctypes.Structure): # int32_t n_eval; # }; class llama_timings(ctypes.Structure): + if TYPE_CHECKING: + t_start_ms: float + t_end_ms: float + t_load_ms: float + t_sample_ms: float + t_p_eval_ms: float + t_eval_ms: float + n_sample: int + n_p_eval: int + n_eval: int + _fields_ = [ ("t_start_ms", ctypes.c_double), ("t_end_ms", ctypes.c_double), @@ -790,6 +1112,12 @@ class llama_chat_message(ctypes.Structure): ] +# // lora adapter +# struct llama_lora_adapter; +llama_lora_adapter_p = ctypes.c_void_p +llama_lora_adapter_p_ctypes = ctypes.POINTER(ctypes.c_void_p) + + # // Helpers for getting default parameters # LLAMA_API struct llama_model_params llama_model_default_params(void); @ctypes_function( @@ -865,7 +1193,8 @@ def llama_backend_init(): [ctypes.c_int], None, ) -def llama_numa_init(numa: int, /): ... +def llama_numa_init(numa: int, /): + ... # // Call once at the end of the program - currently only used for MPI @@ -890,7 +1219,8 @@ def llama_backend_free(): ) def llama_load_model_from_file( path_model: bytes, params: llama_model_params, / -) -> Optional[llama_model_p]: ... +) -> Optional[llama_model_p]: + ... # LLAMA_API void llama_free_model(struct llama_model * model); @@ -899,7 +1229,8 @@ def llama_load_model_from_file( [llama_model_p_ctypes], None, ) -def llama_free_model(model: llama_model_p, /): ... +def llama_free_model(model: llama_model_p, /): + ... # LLAMA_API struct llama_context * llama_new_context_with_model( @@ -912,7 +1243,8 @@ def llama_free_model(model: llama_model_p, /): ... ) def llama_new_context_with_model( model: llama_model_p, params: llama_context_params, / -) -> Optional[llama_context_p]: ... +) -> Optional[llama_context_p]: + ... # // Frees all allocated memory @@ -933,77 +1265,104 @@ def llama_free(ctx: llama_context_p, /): [], ctypes.c_int64, ) -def llama_time_us() -> int: ... +def llama_time_us() -> int: + ... # LLAMA_API size_t llama_max_devices(void); @ctypes_function("llama_max_devices", [], ctypes.c_size_t) -def llama_max_devices() -> int: ... +def llama_max_devices() -> int: + ... # LLAMA_API bool llama_supports_mmap (void); @ctypes_function("llama_supports_mmap", [], ctypes.c_bool) -def llama_supports_mmap() -> bool: ... +def llama_supports_mmap() -> bool: + ... # LLAMA_API bool llama_supports_mlock (void); @ctypes_function("llama_supports_mlock", [], ctypes.c_bool) -def llama_supports_mlock() -> bool: ... +def llama_supports_mlock() -> bool: + ... # LLAMA_API bool llama_supports_gpu_offload(void); @ctypes_function("llama_supports_gpu_offload", [], ctypes.c_bool) -def llama_supports_gpu_offload() -> bool: ... +def llama_supports_gpu_offload() -> bool: + ... # 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]: ... +def llama_get_model(ctx: llama_context_p, /) -> Optional[llama_model_p]: + ... # 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: ... +def llama_n_ctx(ctx: llama_context_p, /) -> 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: ... +def llama_n_batch(ctx: llama_context_p, /) -> int: + ... # 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: ... +def llama_n_ubatch(ctx: llama_context_p, /) -> 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: ... +def llama_n_seq_max(ctx: llama_context_p, /) -> int: + ... + + +# 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: + ... -# LLAMA_API enum llama_vocab_type llama_vocab_type(const struct llama_model * model); +# LLAMA_API enum llama_vocab_type llama_vocab_type (const struct llama_model * model); @ctypes_function("llama_vocab_type", [llama_model_p_ctypes], ctypes.c_int) -def llama_vocab_type(model: llama_model_p, /) -> int: ... +def llama_vocab_type(model: llama_model_p, /) -> int: + ... -# LLAMA_API enum llama_rope_type llama_rope_type (const struct llama_model * model); +# LLAMA_API enum llama_rope_type llama_rope_type (const struct llama_model * model); @ctypes_function("llama_rope_type", [llama_model_p_ctypes], ctypes.c_int) -def llama_rope_type(model: llama_model_p, /) -> int: ... +def llama_rope_type(model: llama_model_p, /) -> int: + ... # LLAMA_API int32_t llama_n_vocab (const struct llama_model * model); @ctypes_function("llama_n_vocab", [llama_model_p_ctypes], ctypes.c_int32) -def llama_n_vocab(model: llama_model_p, /) -> int: ... +def llama_n_vocab(model: llama_model_p, /) -> int: + ... # LLAMA_API int32_t llama_n_ctx_train(const struct llama_model * model); @ctypes_function("llama_n_ctx_train", [llama_model_p_ctypes], ctypes.c_int32) -def llama_n_ctx_train(model: llama_model_p, /) -> int: ... +def llama_n_ctx_train(model: llama_model_p, /) -> int: + ... # LLAMA_API int32_t llama_n_embd (const struct llama_model * model); @ctypes_function("llama_n_embd", [llama_model_p_ctypes], ctypes.c_int32) -def llama_n_embd(model: llama_model_p, /) -> int: ... +def llama_n_embd(model: llama_model_p, /) -> int: + ... + + +# LLAMA_API int32_t llama_n_layer (const struct llama_model * model); +@ctypes_function("llama_n_layer", [llama_model_p_ctypes], ctypes.c_int32) +def llama_n_layer(model: llama_model_p, /) -> int: + ... # // Get the model's RoPE frequency scaling factor @@ -1142,6 +1501,35 @@ def llama_get_model_tensor( ... +# // 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 0 on success # LLAMA_API uint32_t llama_model_quantize( # const char * fname_inp, @@ -1166,31 +1554,130 @@ def llama_model_quantize( ... -# LLAMA_API int32_t llama_model_apply_lora_from_file( -# const struct llama_model * model, -# const char * path_lora, -# float scale, -# const char * path_base_model, -# int32_t n_threads); +# // Load a LoRA adapter from file +# // The loaded adapter will be associated to the given model, and will be free when the model is deleted +# LLAMA_API struct llama_lora_adapter * llama_lora_adapter_init( +# struct llama_model * model, +# const char * path_lora); +@ctypes_function( + "llama_lora_adapter_init", + [llama_model_p_ctypes, ctypes.c_char_p], + llama_lora_adapter_p_ctypes, +) +def llama_lora_adapter_init( + model: llama_model_p, path_lora: bytes, / +) -> Optional[llama_lora_adapter_p]: + """Load a LoRA adapter from file + The loaded adapter will be associated to the given model, and will be free when the model is deleted + """ + ... + + +# // Add a loaded LoRA adapter to given context +# // This will not modify model's weight +# LLAMA_API int32_t llama_lora_adapter_set( +# struct llama_context * ctx, +# struct llama_lora_adapter * adapter, +# float scale); +@ctypes_function( + "llama_lora_adapter_set", + [llama_context_p_ctypes, llama_lora_adapter_p_ctypes, ctypes.c_float], + ctypes.c_int32, +) +def llama_lora_adapter_set( + ctx: llama_context_p, adapter: llama_lora_adapter_p, scale: float, / +) -> int: + """Add a loaded LoRA adapter to given context + This will not modify model's weight""" + ... + + +# // Remove a specific LoRA adapter from given context +# // Return -1 if the adapter is not present in the context +# LLAMA_API int32_t llama_lora_adapter_remove( +# struct llama_context * ctx, +# struct llama_lora_adapter * adapter); +@ctypes_function( + "llama_lora_adapter_remove", + [llama_context_p_ctypes, llama_lora_adapter_p_ctypes], + ctypes.c_int32, +) +def llama_lora_adapter_remove( + ctx: llama_context_p, adapter: llama_lora_adapter_p, / +) -> int: + """Remove a LoRA adapter from given context + Return -1 if the adapter is not present in the context""" + ... + + +# // Remove all LoRA adapters from given context +# LLAMA_API void llama_lora_adapter_clear( +# struct llama_context * ctx); +@ctypes_function( + "llama_lora_adapter_clear", + [llama_context_p_ctypes], + None, +) +def llama_lora_adapter_clear(ctx: llama_context_p, /): + """Remove all LoRA adapters from given context""" + ... + + +# // Manually free a LoRA adapter +# // Note: loaded adapters will be free when the associated model is deleted +# LLAMA_API void llama_lora_adapter_free(struct llama_lora_adapter * adapter); +@ctypes_function( + "llama_lora_adapter_free", + [llama_lora_adapter_p_ctypes], + None, +) +def llama_lora_adapter_free(adapter: llama_lora_adapter_p, /): + """Manually free a LoRA adapter + Note: loaded adapters will be free when the associated model is deleted""" + ... + + +# // 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_control_vector_apply( +# struct llama_context * lctx, +# const float * data, +# size_t len, +# int32_t n_embd, +# int32_t il_start, +# int32_t il_end); @ctypes_function( - "llama_model_apply_lora_from_file", + "llama_control_vector_apply", [ - llama_model_p_ctypes, - ctypes.c_char_p, - ctypes.c_float, - ctypes.c_char_p, + 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_model_apply_lora_from_file( - model: llama_model_p, - path_lora: Union[ctypes.c_char_p, bytes], - scale: Union[ctypes.c_float, float], - path_base_model: Union[ctypes.c_char_p, bytes, None], - n_threads: Union[ctypes.c_int32, int], +def llama_control_vector_apply( + lctx: llama_context_p, + data: CtypesPointerOrRef[ctypes.c_float], + len: int, + n_embd: int, + il_start: int, + il_end: int, /, -) -> 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.""" + ... # // @@ -1205,6 +1692,15 @@ def llama_model_apply_lora_from_file( # llama_pos pos; # }; class llama_kv_cache_view_cell(ctypes.Structure): + """Information associated with an individual cell in the KV cache view. + + Attributes: + pos (llama_pos): The position for this cell. Takes KV cache shifts into account. + May be negative if the cell is not populated.""" + + if TYPE_CHECKING: + pos: llama_pos + _fields_ = [("pos", llama_pos)] @@ -1241,6 +1737,16 @@ class llama_kv_cache_view_cell(ctypes.Structure): # llama_seq_id * cells_sequences; # }; class llama_kv_cache_view(ctypes.Structure): + if TYPE_CHECKING: + n_cells: int + n_max_seq: int + token_count: int + used_cells: int + max_contiguous: int + max_contiguous_idx: int + cells: CtypesArray[llama_kv_cache_view_cell] + cells_sequences: CtypesArray[llama_seq_id] + _fields_ = [ ("n_cells", ctypes.c_int32), ("n_max_seq", ctypes.c_int32), @@ -1311,7 +1817,7 @@ def llama_get_kv_cache_used_cells(ctx: llama_context_p, /) -> int: ... -# // Clear the KV cache +# // Clear the KV cache - both cell info is erased and KV data is zeroed # LLAMA_API void llama_kv_cache_clear( # struct llama_context * ctx); @ctypes_function("llama_kv_cache_clear", [llama_context_p_ctypes], None) @@ -1321,6 +1827,7 @@ def llama_kv_cache_clear(ctx: llama_context_p, /): # // 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) @@ -1347,6 +1854,9 @@ def llama_kv_cache_seq_rm( /, ) -> 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)""" @@ -1504,9 +2014,18 @@ def llama_kv_cache_update(ctx: llama_context_p, /): # // -# 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); +# // Returns the *actual* size in bytes of the state +# // (rng, logits, embedding and kv_cache) +# // 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 (rng, logits, embedding and kv_cache) - will often be smaller after compacting tokens""" + ... + + +# 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 maximum size in bytes of the state (rng, logits, embedding @@ -1514,12 +2033,38 @@ def llama_get_state_size(ctx: llama_context_p, /) -> int: ... -# 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( +# // 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); +# uint8_t * dst), +# "use llama_state_get_data instead"); @ctypes_function( "llama_copy_state_data", [ @@ -1539,9 +2084,30 @@ def llama_copy_state_data( # // Set the state reading from the specified address # // Returns the number of bytes read -# LLAMA_API size_t llama_set_state_data( +# 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); +# 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)], @@ -1555,12 +2121,41 @@ def llama_set_state_data( # Save/load session file -# LLAMA_API bool llama_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, + 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", [ @@ -1579,14 +2174,41 @@ def llama_load_session_file( n_token_capacity: Union[ctypes.c_size_t, int], n_token_count_out: CtypesPointerOrRef[ctypes.c_size_t], /, -) -> int: ... +) -> int: + ... -# LLAMA_API bool llama_save_session_file( +# 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", [ @@ -1603,7 +2225,138 @@ def llama_save_session_file( tokens: CtypesArray[llama_token], n_token_count: Union[ctypes.c_size_t, int], /, -) -> int: ... +) -> int: + ... + + +# // Get the exact size needed to copy the KV cache 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 KV cache of a single sequence""" + ... + + +# // Copy the KV cache 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 KV cache 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 (originally copied with `llama_state_seq_get_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: + ... # // @@ -1682,6 +2435,22 @@ def llama_batch_free(batch: llama_batch, /): ... +# // Processes a batch of tokens with the ecoder part of the encoder-decoder model. +# // Stores the encoder output internally for later use by the decoder cross-attention layers. +# // 0 - success +# // < 0 - error +# 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: + """Processes a batch of tokens with the ecoder part of the encoder-decoder model. + Stores the encoder output internally for later use by the decoder cross-attention layers. + 0 - success + < 0 - error""" + ... + + # // Positive return values does not mean a fatal error, but rather a warning. # // 0 - success # // 1 - could not find a KV slot for the batch (try reducing the size of the batch or increase the context) @@ -1724,6 +2493,32 @@ def llama_set_n_threads( ... +# // Get the number of threads used for generation of a single token. +# LLAMA_API uint32_t llama_n_threads(struct llama_context * ctx); +@ctypes_function("llama_n_threads", [llama_context_p_ctypes], ctypes.c_uint32) +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 uint32_t llama_n_threads_batch(struct llama_context * ctx); +@ctypes_function("llama_n_threads_batch", [llama_context_p_ctypes], ctypes.c_uint32) +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 model is in embeddings mode or not +# // If true, embeddings will be returned but logits will not +# 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 model is in embeddings model or not + If true, embeddings will be returned but logits will 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); @@ -1764,9 +2559,9 @@ def llama_synchronize(ctx: llama_context_p, /): # // Token logits obtained from the last call to llama_decode() -# // The logits for the last token are stored in the last row -# // Logits for which llama_batch.logits[i] == 0 are undefined -# // Rows: n_tokens provided with llama_batch +# // 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 # LLAMA_API float * llama_get_logits(struct llama_context * ctx); @ctypes_function( @@ -1784,8 +2579,10 @@ def llama_get_logits(ctx: llama_context_p, /) -> CtypesArray[ctypes.c_float]: ... -# // Logits for the ith token. Equivalent to: -# // llama_get_logits(ctx) + i*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", @@ -1800,8 +2597,12 @@ def llama_get_logits_ith( ... -# // Get all output token embeddings -# // shape: [n_tokens*n_embd] (1-dimensional) +# // 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. # LLAMA_API float * llama_get_embeddings(struct llama_context * ctx); @ctypes_function( "llama_get_embeddings", [llama_context_p_ctypes], ctypes.POINTER(ctypes.c_float) @@ -1812,9 +2613,11 @@ def llama_get_embeddings(ctx: llama_context_p, /) -> CtypesArray[ctypes.c_float] ... -# // Get the embeddings for the ith token -# // llama_get_embeddings(ctx) + i*n_embd +# // 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", @@ -1858,7 +2661,8 @@ def llama_get_embeddings_seq( ) def llama_token_get_text( model: llama_model_p, token: Union[llama_token, int], / -) -> bytes: ... +) -> bytes: + ... # LLAMA_API float llama_token_get_score(const struct llama_model * model, llama_token token); @@ -1867,16 +2671,40 @@ def llama_token_get_text( ) def llama_token_get_score( model: llama_model_p, token: Union[llama_token, int], / -) -> float: ... +) -> float: + ... + + +# LLAMA_API enum llama_token_attr llama_token_get_attr(const struct llama_model * model, llama_token token); +@ctypes_function( + "llama_token_get_attr", [llama_model_p_ctypes, llama_token], ctypes.c_int +) +def llama_token_get_attr( + model: llama_model_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_token_is_eog(const struct llama_model * model, llama_token token); +@ctypes_function( + "llama_token_is_eog", [llama_model_p_ctypes, llama_token], ctypes.c_bool +) +def llama_token_is_eog(model: llama_model_p, token: Union[llama_token, int], /) -> bool: + """Check if the token is supposed to end generation (end-of-generation, eg. EOS, EOT, etc.)""" + ... -# LLAMA_API enum llama_token_type llama_token_get_type(const struct llama_model * model, llama_token token); +# // Identify if Token Id is a control token or a render-able token +# LLAMA_API bool llama_token_is_control(const struct llama_model * model, llama_token token); @ctypes_function( - "llama_token_get_type", [llama_model_p_ctypes, llama_token], ctypes.c_int + "llama_token_is_control", [llama_model_p_ctypes, llama_token], ctypes.c_bool ) -def llama_token_get_type( +def llama_token_is_control( model: llama_model_p, token: Union[llama_token, int], / -) -> int: ... +) -> bool: + """Identify if Token Id is a control token or a render-able token""" + ... # // Special tokens @@ -1896,6 +2724,20 @@ def llama_token_eos(model: llama_model_p, /) -> int: ... +# LLAMA_API llama_token llama_token_cls(const struct llama_model * model); // classification +@ctypes_function("llama_token_cls", [llama_model_p_ctypes], llama_token) +def llama_token_cls(model: llama_model_p, /) -> int: + """classification""" + ... + + +# LLAMA_API llama_token llama_token_sep(const struct llama_model * model); // sentence separator +@ctypes_function("llama_token_sep", [llama_model_p_ctypes], llama_token) +def llama_token_sep(model: llama_model_p, /) -> int: + """sentence separator""" + ... + + # LLAMA_API llama_token llama_token_nl (const struct llama_model * model); // next-line @ctypes_function("llama_token_nl", [llama_model_p_ctypes], llama_token) def llama_token_nl(model: llama_model_p, /) -> int: @@ -1903,23 +2745,19 @@ def llama_token_nl(model: llama_model_p, /) -> int: ... -# // Returns -1 if unknown, 1 for true or 0 for false. -# LLAMA_API int32_t llama_add_bos_token(const struct llama_model * model); -@ctypes_function("llama_add_bos_token", [llama_model_p_ctypes], ctypes.c_int32) -def llama_add_bos_token(model: llama_model_p, /) -> int: - """Returns -1 if unknown, 1 for true or 0 for false.""" +# LLAMA_API bool llama_add_bos_token(const struct llama_model * model); +@ctypes_function("llama_add_bos_token", [llama_model_p_ctypes], ctypes.c_bool) +def llama_add_bos_token(model: llama_model_p, /) -> bool: ... -# // Returns -1 if unknown, 1 for true or 0 for false. -# LLAMA_API int32_t llama_add_eos_token(const struct llama_model * model); -@ctypes_function("llama_add_eos_token", [llama_model_p_ctypes], ctypes.c_int32) -def llama_add_eos_token(model: llama_model_p, /) -> int: - """Returns -1 if unknown, 1 for true or 0 for false.""" +# LLAMA_API bool llama_add_eos_token(const struct llama_model * model); +@ctypes_function("llama_add_eos_token", [llama_model_p_ctypes], ctypes.c_bool) +def llama_add_eos_token(model: llama_model_p, /) -> bool: ... -# // codellama infill tokens +# // Codellama infill tokens # LLAMA_API llama_token llama_token_prefix(const struct llama_model * model); // Beginning of infill prefix @ctypes_function("llama_token_prefix", [llama_model_p_ctypes], llama_token) def llama_token_prefix(model: llama_model_p) -> int: @@ -1929,17 +2767,20 @@ def llama_token_prefix(model: llama_model_p) -> int: # LLAMA_API llama_token llama_token_middle(const struct llama_model * model); // Beginning of infill middle @ctypes_function("llama_token_middle", [llama_model_p_ctypes], llama_token) -def llama_token_middle(model: llama_model_p, /) -> int: ... +def llama_token_middle(model: llama_model_p, /) -> int: + ... # LLAMA_API llama_token llama_token_suffix(const struct llama_model * model); // Beginning of infill suffix @ctypes_function("llama_token_suffix", [llama_model_p_ctypes], llama_token) -def llama_token_suffix(model: llama_model_p, /) -> int: ... +def llama_token_suffix(model: llama_model_p, /) -> int: + ... # LLAMA_API llama_token llama_token_eot (const struct llama_model * model); // End of infill middle @ctypes_function("llama_token_eot", [llama_model_p_ctypes], llama_token) -def llama_token_eot(model: llama_model_p, /) -> int: ... +def llama_token_eot(model: llama_model_p, /) -> int: + ... # // @@ -1951,16 +2792,17 @@ def llama_token_eot(model: llama_model_p, /) -> int: ... # /// @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 -# /// @param special Allow tokenizing special and/or control tokens which otherwise are not exposed and treated as plaintext. -# /// Does not insert a leading space. +# /// @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_model * model, # const char * text, # int32_t text_len, # llama_token * tokens, # int32_t n_tokens_max, -# bool add_bos, -# bool special); +# bool add_special, +# bool parse_special); @ctypes_function( "llama_tokenize", [ @@ -1980,37 +2822,40 @@ def llama_tokenize( text_len: Union[ctypes.c_int, int], tokens: CtypesArray[llama_token], n_tokens_max: Union[ctypes.c_int, int], - add_bos: Union[ctypes.c_bool, bool], - special: Union[ctypes.c_bool, bool], + add_special: Union[ctypes.c_bool, bool], + parse_special: Union[ctypes.c_bool, bool], /, ) -> int: """Convert the provided text into tokens. - + Args: model: The model 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_bos: Whether to add a beginning-of-sentence token. - special: Allow tokenizing special and/or control tokens which otherwise are not exposed and treated as plaintext. - Does not insert a leading space. - + add_special: Allow adding special tokenns 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""" + 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 code is responsible to remove the leading whitespace of the first non-BOS token when decoding multiple tokens. +# // 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_model * model, # llama_token token, # char * buf, -# int32_t length); +# int32_t length, +# int32_t lstrip, +# bool special); @ctypes_function( "llama_token_to_piece", [ @@ -2018,6 +2863,8 @@ def llama_tokenize( llama_token, ctypes.c_char_p, ctypes.c_int32, + ctypes.c_int32, + ctypes.c_bool, ], ctypes.c_int32, ) @@ -2026,16 +2873,80 @@ def llama_token_to_piece( 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: + model: The model 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_model * model, +# 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_model_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( + model: llama_model_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: + """Convert the provided tokens into text (inverse of llama_tokenize()). + + Args: + model: The model 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/ggerganov/llama.cpp/wiki/Templates-supported-by-llama_chat_apply_template @@ -2070,7 +2981,8 @@ def llama_chat_apply_template( chat: CtypesArray[llama_chat_message], n_msg: int, /, -) -> int: ... +) -> int: + ... # // @@ -2098,7 +3010,7 @@ def llama_grammar_init( n_rules: Union[ctypes.c_size_t, int], start_rule_index: Union[ctypes.c_size_t, int], /, -) -> llama_grammar_p: +) -> Optional[llama_grammar_p]: """Initialize a grammar from a set of rules.""" ... @@ -2125,6 +3037,79 @@ def llama_grammar_copy(grammar: llama_grammar_p, /) -> llama_grammar_p: ... +# /// @details Apply constraints from grammar +# LLAMA_API void llama_grammar_sample( +# const struct llama_grammar * grammar, +# const struct llama_context * ctx, +# llama_token_data_array * candidates); +@ctypes_function( + "llama_grammar_sample", + [ + llama_grammar_p, + llama_context_p_ctypes, + llama_token_data_array_p, + ], + None, +) +def llama_grammar_sample( + grammar: llama_grammar_p, + ctx: llama_context_p, + candidates: Union[ + CtypesArray[llama_token_data_array], CtypesPointerOrRef[llama_token_data_array] + ], + /, +): + """Apply constraints from grammar""" + ... + + +# LLAMA_API DEPRECATED(void llama_sample_grammar( +# struct llama_context * ctx, +# llama_token_data_array * candidates, +# const struct llama_grammar * grammar), +# "use llama_grammar_sample instead"); +@ctypes_function( + "llama_sample_grammar", + [llama_context_p_ctypes, llama_token_data_array_p, llama_grammar_p], + None, +) +def llama_sample_grammar( + ctx: llama_context_p, + candidates: Union[ + CtypesArray[llama_token_data_array], CtypesPointerOrRef[llama_token_data_array] + ], + grammar, # type: llama_grammar_p + /, +): + """Apply constraints from grammar + + Parameters: + 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. + grammar: A grammar object containing the rules and constraints to apply to the generated text. + """ + ... + + +# /// @details Accepts the sampled token into the grammar +# LLAMA_API void llama_grammar_accept_token( +# struct llama_grammar * grammar, +# struct llama_context * ctx, +# llama_token token); +@ctypes_function( + "llama_grammar_accept_token", + [llama_grammar_p, llama_context_p_ctypes, llama_token], + None, +) +def llama_grammar_accept_token( + grammar: llama_grammar_p, + ctx: llama_context_p, + token: Union[llama_token, int], + /, +): + """Accepts the sampled token into the grammar""" + ... + + # // # // Sampling functions # // @@ -2411,33 +3396,6 @@ def llama_sample_temp( ... -# /// @details Apply constraints from grammar -# LLAMA_API void llama_sample_grammar( -# struct llama_context * ctx, -# llama_token_data_array * candidates, -# const struct llama_grammar * grammar); -@ctypes_function( - "llama_sample_grammar", - [llama_context_p_ctypes, llama_token_data_array_p, llama_grammar_p], - None, -) -def llama_sample_grammar( - ctx: llama_context_p, - candidates: Union[ - CtypesArray[llama_token_data_array], CtypesPointerOrRef[llama_token_data_array] - ], - grammar, # type: llama_grammar_p - /, -): - """Apply constraints from grammar - - Parameters: - 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. - grammar: A grammar object containing the rules and constraints to apply to the generated text. - """ - ... - - # /// @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. @@ -2550,7 +3508,7 @@ def llama_sample_token_greedy( ... -# /// @details Randomly selects a token from the candidates based on their probabilities. +# /// @details Randomly selects a token from the candidates based on their probabilities using the RNG of ctx. # LLAMA_API llama_token llama_sample_token( # struct llama_context * ctx, # llama_token_data_array * candidates); @@ -2570,108 +3528,51 @@ def llama_sample_token( ... -# /// @details Accepts the sampled token into the grammar -# LLAMA_API void llama_grammar_accept_token( -# struct llama_context * ctx, -# struct llama_grammar * grammar, -# llama_token token); -@ctypes_function( - "llama_grammar_accept_token", - [llama_context_p_ctypes, llama_grammar_p, llama_token], - None, -) -def llama_grammar_accept_token( - ctx: llama_context_p, grammar: llama_grammar_p, token: Union[llama_token, int], / -) -> None: - """Accepts the sampled token into the grammar""" - ... - - # // -# // Beam search +# // Model split # // -# struct llama_beam_view { -# const llama_token * tokens; - - -# size_t n_tokens; -# float p; // Cumulative beam probability (renormalized relative to all beams) -# bool eob; // Callback should set this to true when a beam is at end-of-beam. -# }; -class llama_beam_view(ctypes.Structure): - _fields_ = [ - ("tokens", llama_token_p), - ("n_tokens", ctypes.c_size_t), - ("p", ctypes.c_float), - ("eob", ctypes.c_bool), - ] - -# // Passed to beam_search_callback function. -# // Whenever 0 < common_prefix_length, this number of tokens should be copied from any of the beams -# // (e.g. beams[0]) as they will be removed (shifted) from all beams in all subsequent callbacks. -# // These pointers are valid only during the synchronous callback, so should not be saved. -# struct llama_beams_state { -# struct llama_beam_view * beam_views; -# size_t n_beams; // Number of elements in beam_views[]. -# size_t common_prefix_length; // Current max length of prefix tokens shared by all beams. -# bool last_call; // True iff this is the last callback invocation. -# }; -class llama_beams_state(ctypes.Structure): - _fields_ = [ - ("beam_views", ctypes.POINTER(llama_beam_view)), - ("n_beams", ctypes.c_size_t), - ("common_prefix_length", ctypes.c_size_t), - ("last_call", ctypes.c_bool), - ] +# /// @details Build a split GGUF final path for this chunk. +# /// llama_split_path(split_path, sizeof(split_path), "/models/ggml-model-q4_0", 2, 4) => split_path = "/models/ggml-model-q4_0-00002-of-00004.gguf" +# // Returns the split_path length. +# 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.""" + ... -# // Type of pointer to the beam_search_callback function. -# // void* callback_data is any custom data passed to llama_beam_search, that is subsequently -# // passed back to beam_search_callback. This avoids having to use global variables in the callback. -# typedef void (*llama_beam_search_callback_fn_t)(void * callback_data, struct llama_beams_state); -llama_beam_search_callback_fn_t = ctypes.CFUNCTYPE( - None, ctypes.c_void_p, llama_beams_state -) - - -# /// @details Deterministically returns entire sentence constructed by a beam search. -# /// @param ctx Pointer to the llama_context. -# /// @param callback Invoked for each iteration of the beam_search loop, passing in beams_state. -# /// @param callback_data A pointer that is simply passed back to callback. -# /// @param n_beams Number of beams to use. -# /// @param n_past Number of tokens already evaluated. -# /// @param n_predict Maximum number of tokens to predict. EOS may occur earlier. -# /// @param n_threads Number of threads as passed to llama_eval(). -# LLAMA_API void llama_beam_search( -# struct llama_context * ctx, -# llama_beam_search_callback_fn_t callback, -# void * callback_data, -# size_t n_beams, -# int32_t n_past, -# int32_t n_predict); -@ctypes_function( - "llama_beam_search", - [ - llama_context_p_ctypes, - llama_beam_search_callback_fn_t, - ctypes.c_void_p, - ctypes.c_size_t, - ctypes.c_int32, - ctypes.c_int32, - ], - None, +# /// @details Extract the path prefix from the split_path if and only if the split_no and split_count match. +# /// llama_split_prefix(split_prefix, 64, "/models/ggml-model-q4_0-00002-of-00004.gguf", 2, 4) => split_prefix = "/models/ggml-model-q4_0" +# // Returns the split_prefix length. +# 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_beam_search( - ctx: llama_context_p, - callback: CtypesFuncPointer, - callback_data: ctypes.c_void_p, - n_beams: Union[ctypes.c_size_t, int], - n_past: Union[ctypes.c_int, int], - n_predict: Union[ctypes.c_int, 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.""" + ... # Performance information @@ -2748,4 +3649,5 @@ def llama_log_set( [ctypes.c_void_p, llama_context_p_ctypes], None, ) -def llama_dump_timing_info_yaml(stream: ctypes.c_void_p, ctx: llama_context_p, /): ... +def llama_dump_timing_info_yaml(stream: ctypes.c_void_p, ctx: llama_context_p, /): + ... diff --git a/llama_cpp/llama_grammar.py b/llama_cpp/llama_grammar.py index 9cc48a93b..4cd52c2d5 100644 --- a/llama_cpp/llama_grammar.py +++ b/llama_cpp/llama_grammar.py @@ -3,521 +3,100 @@ # flake8: noqa from pathlib import Path import sys -from ctypes import * # type: ignore -from enum import Enum -from itertools import islice +import ctypes +import enum +import typing +import dataclasses + +from itertools import groupby from typing import ( Any, - Callable, - Dict, - Generic, + Set, List, Optional, - OrderedDict, - TextIO, Tuple, - TypeVar, Union, - overload, ) import llama_cpp.llama_cpp as llama_cpp -# Type aliases -llama_grammar_element = llama_cpp.llama_grammar_element -llama_grammar_element_p = llama_cpp.llama_grammar_element_p -llama_grammar_p = llama_cpp.llama_grammar_p - -# Type variables -Ptr = TypeVar("Ptr", bound="const_char_p") -T = TypeVar("T") -U = TypeVar("U") -V = TypeVar("V") -W = TypeVar("W") - - -class Sentinel: - """Used to mark the end of a iterator of std::vector & std::map.""" - - -class LlamaGrammar: - """Keeps reference counts of all the arguments, so that they are not - garbage collected by Python.""" - - def __del__(self) -> None: - """Free the grammar pointer when the object is deleted.""" - if self.grammar is not None: - llama_cpp.llama_grammar_free(self.grammar) - self.grammar = None - - def __init__( - self, - parsed_grammar: "parse_state", - ) -> None: - """Initialize the grammar pointer from the parsed state.""" - self._grammar_rules = ( - parsed_grammar.c_rules() - ) # type: std.vector[std.vector[LlamaGrammarElement]] - self._n_rules = self._grammar_rules.size() # type: int - self._start_rule_index = parsed_grammar.symbol_ids.at("root") # type: int - self.init() - - @classmethod - def from_string(cls, grammar: str, verbose: bool = True) -> "LlamaGrammar": - """Convert a GBNF grammar to a Llama grammar.""" - parsed_grammar = parse(const_char_p(grammar)) # type: parse_state - if parsed_grammar.rules.empty(): - raise ValueError( - f"{cls.from_string.__name__}: error parsing grammar file: parsed_grammar.rules is empty" - ) - if verbose: - print(f"{cls.from_string.__name__} grammar:", file=sys.stderr) - print_grammar(sys.stderr, parsed_grammar) - print(file=sys.stderr) - return cls(parsed_grammar) - - @classmethod - def from_json_schema( - cls, - json_schema: str, - verbose: bool = True, - ) -> "LlamaGrammar": - """Convert a JSON schema to a Llama grammar.""" - return cls.from_string(json_schema_to_gbnf(json_schema), verbose=verbose) - - @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" - ) - - def init(self) -> None: - # Step 1: Convert LlamaGrammarElement to llama_grammar_element - self._element_lists = [ - [ - llama_grammar_element(c_int(elem.type.value), c_uint32(elem.value)) - for elem in subvector - ] - for subvector in self._grammar_rules - ] # type: List[List[llama_grammar_element]] - - # Step 2: Convert each list to llama_grammar_element array and get pointer - self._element_arrays = [ - (llama_grammar_element * len(sublist))(*sublist) - for sublist in self._element_lists - ] # type: List[Array[llama_grammar_element]] - - # Step 3: Get pointer of each array - self._element_array_pointers = [ - cast(subarray, llama_grammar_element_p) for subarray in self._element_arrays - ] # type: List[llama_grammar_element_p] - - # Step 4: Make array of these pointers and get its pointer - self._rules = (llama_grammar_element_p * len(self._element_array_pointers))( - *self._element_array_pointers - ) - self.grammar = llama_cpp.llama_grammar_init( - self._rules, c_size_t(self._n_rules), c_size_t(self._start_rule_index) - ) - - def reset(self) -> None: - if self.grammar is not None: - llama_cpp.llama_grammar_free(self.grammar) - self.init() - - -class LlamaGrammarElement: - def __init__(self, type: "llama_gretype", value: int): - self.type = type - self.value = value # Unicode code point or rule ID - - -class const_char_p: - """C++ implementation of const char *.""" - - def __init__(self, value: Union[str, Ptr], move: Optional[int] = None): - if isinstance(value, const_char_p): - # We're copying an existing const_char_p - self.value = value.value - self.pos = value.pos + (move or 0) - return - - # We're creating a new const_char_p - self.value = value - self.pos = move or 0 - - def __str__(self) -> str: - assert self.value is not None, "null pointer" - return self.value[self.pos :] - - def __getitem__(self, index: int) -> str: - value = str(self) - return value[index] if index < len(value) else "" - - @overload - def __add__(self: Ptr, other: int) -> Ptr: - ... - - @overload - def __add__(self: Ptr, other: Ptr) -> int: - ... - - def __add__(self: Ptr, other: Union[int, Ptr]) -> Union[int, Ptr]: - return ( - self.__class__(self.value, self.pos + other) - if isinstance(other, int) - else self.pos + other.pos - ) - - @overload - def __sub__(self: Ptr, other: int) -> Ptr: - ... - - @overload - def __sub__(self: Ptr, other: Ptr) -> int: - ... - - def __sub__(self: Ptr, other: Union[int, Ptr]) -> Union[int, Ptr]: - return ( - self.__class__(self.value, self.pos - other) - if isinstance(other, int) - else self.pos - other.pos - ) - - def __eq__(self: Ptr, other: Ptr) -> bool: - assert self.value == other.value, "comparing pointers from different strings" - return self.pos == other.pos - - def __lt__(self: Ptr, other: Ptr) -> bool: - assert self.value == other.value, "comparing pointers from different strings" - return self.pos < other.pos - - def __gt__(self: Ptr, other: Ptr) -> bool: - assert self.value == other.value, "comparing pointers from different strings" - return self.pos > other.pos - - -class std: - @staticmethod - def string(ptr: const_char_p, length: Optional[int] = None) -> str: - """C++ implementation of std::string constructor.""" - value = str(ptr) - if length is not None: - value = value[:length] - return value - - class vector(Generic[T], List[T]): - """C++ implementation of std::vector.""" - - class iterator: - def __init__(self, vector: "std.vector[T]", index: int): - self._vector = vector - self._index = index - self._version = vector._version - - def _check_version(self): - if self._version != self._vector._version: - raise RuntimeError("Iterator used after vector was modified.") - - def __iter__(self): - return self - - def __next__(self) -> T: - self._check_version() - if self._index >= self._vector.size(): - raise StopIteration - value = self._vector[self._index] - self._index += 1 - return value - - def __add__(self, value: int) -> "std.vector[T].iterator": - return self.__class__(self._vector, self._index + value) - - def __sub__(self, value: int) -> "std.vector[T].iterator": - return self.__class__(self._vector, self._index - value) - - def __init__(self): - self._version = 0 - - def modify(self): - # This is a bit of a hack to make sure iterators are invalidated - self._version += 1 - - def push_back(self, value: T) -> None: - self.modify() - self.append(value) - - def pop_back(self) -> None: - self.modify() - if not self.empty(): - self.pop() - - def back(self) -> T: - return self[-1] - - def size(self) -> int: - return len(self) - - def clear(self) -> None: - self.modify() - super().clear() - - def empty(self) -> bool: - return self.size() == 0 - - def data(self) -> "std.vector[T]": - return self - - def resize( - self, - new_size: int, - fill_value_factory: Optional[Callable[[], T]] = None, - ) -> None: - if new_size > self.size(): - if fill_value_factory is None: - raise ValueError("A fill value factory function must be provided.") - self.reserve(new_size, fill_value_factory) - elif new_size < self.size(): - self[:] = self[:new_size] - - def reserve(self, capacity: int, fill_value_factory: Callable[[], T]) -> None: - if capacity > self.size(): - fill_value = fill_value_factory() - self.extend([fill_value] * (capacity - self.size())) - - def front(self) -> T: - if not self.empty(): - return self[0] - else: - raise IndexError("Vector is empty.") - - def assign(self, count: int, value: T) -> None: - self.clear() - self.extend([value] * count) - - def insert( - self, - pos: "std.vector[T].iterator", - first: "std.vector[T].iterator", - last: "std.vector[T].iterator", - ) -> None: - self[pos._index : pos._index] = list( - islice(first._vector, first._index, last._index) - ) - - def begin(self) -> "std.vector[T].iterator": - return self.iterator(self, 0) - - def end(self) -> "std.vector[T].iterator": - return self.iterator(self, self.size()) - - class map(Generic[T, U], OrderedDict[T, U]): - """C++ implementation of std::map.""" - - class iterator(Generic[V, W]): - def __init__(self, _map: "std.map[T, U]", key: Union[T, Sentinel]): - self._map = _map - self.iter = iter(_map) - self.key = key - self._advance() - - def _sanitize_key(self) -> T: - if isinstance(self.key, Sentinel): - raise StopIteration - return self.key - - def _advance(self) -> None: - try: - while next(self.iter) != self.key: - pass - except StopIteration: - self.key = Sentinel() - - def __next__(self) -> Tuple[T, U]: - key = self._sanitize_key() - if key in self._map: - value = self._map[key] - self._advance() - return key, value - else: - raise StopIteration - - def get(self) -> Tuple[T, U]: - key = self._sanitize_key() - return key, self._map[key] +class GrammarElementType(enum.IntEnum): + END = llama_cpp.LLAMA_GRETYPE_END + ALT = llama_cpp.LLAMA_GRETYPE_ALT + RULE_REF = llama_cpp.LLAMA_GRETYPE_RULE_REF + CHAR = llama_cpp.LLAMA_GRETYPE_CHAR + CHAR_NOT = llama_cpp.LLAMA_GRETYPE_CHAR_NOT + CHAR_RNG_UPPER = llama_cpp.LLAMA_GRETYPE_CHAR_RNG_UPPER + CHAR_ALT = llama_cpp.LLAMA_GRETYPE_CHAR_ALT + CHAR_ANY = llama_cpp.LLAMA_GRETYPE_CHAR_ANY - @property - def first(self) -> T: - return self._sanitize_key() - @property - def second(self) -> U: - return self._map[self._sanitize_key()] +@dataclasses.dataclass +class GrammarElement: + type: GrammarElementType + value: int - def insert( - self, key: T, value: U - ) -> Tuple["std.map[T, U].iterator[T, U]", bool]: - if key in self: - return self.iterator(self, key), False - else: - self[key] = value - return self.iterator(self, key), True - def find(self, key: T) -> "std.map[T, U].iterator[T, U]": - if key in self: - return self.iterator(self, key) - else: - return self.end() +@dataclasses.dataclass +class ParseState: + symbol_ids: typing.Dict[str, int] = dataclasses.field(default_factory=dict) + rules: typing.List[typing.List[GrammarElement]] = dataclasses.field(default_factory=list) - def at(self, key: T) -> U: - if key in self: - return self[key] - else: - raise KeyError("The provided key is not found in the map.") - - def erase(self, iterator: "std.map[T, U].iterator[T, U]") -> None: - key = iterator.first - if key in self: - del self[key] - - def size(self) -> int: - return len(self) - - def empty(self) -> bool: - return self.size() == 0 - - def lower_bound(self, key: T) -> "std.map[T, U].iterator[T, U]": - try: - keys = sorted(list(self.keys())) # type: ignore - for k in keys: - if k >= key: - return self.iterator(self, k) - raise ValueError("No key found that is not less than the input key") - except TypeError: - raise TypeError("Keys of type T cannot be sorted.") - - def begin(self) -> "std.map[T, U].iterator[T, U]": - return self.iterator(self, next(iter(self))) - - def end(self) -> "std.map[T, U].iterator[T, U]": - return self.iterator(self, Sentinel()) - - -# // grammar element type -# enum llama_gretype { -# // end of rule definition -# LLAMA_GRETYPE_END = 0, - -# // start of alternate definition for rule -# LLAMA_GRETYPE_ALT = 1, - -# // non-terminal element: reference to rule -# LLAMA_GRETYPE_RULE_REF = 2, - -# // terminal element: character (code point) -# LLAMA_GRETYPE_CHAR = 3, - -# // inverse char(s) ([^a], [^a-b] [^abc]) -# LLAMA_GRETYPE_CHAR_NOT = 4, - -# // modifies a preceding LLAMA_GRETYPE_CHAR or LLAMA_GRETYPE_CHAR_ALT to -# // be an inclusive range ([a-z]) -# LLAMA_GRETYPE_CHAR_RNG_UPPER = 5, - - -# // modifies a preceding LLAMA_GRETYPE_CHAR or -# // LLAMA_GRETYPE_CHAR_RNG_UPPER to add an alternate char to match ([ab], [a-zA]) -# LLAMA_GRETYPE_CHAR_ALT = 6, -# }; -class llama_gretype(Enum): - """grammar element type""" - - LLAMA_GRETYPE_END = 0 # end of rule definition - LLAMA_GRETYPE_ALT = 1 # start of alternate definition for rule - LLAMA_GRETYPE_RULE_REF = 2 # non-terminal element: reference to rule - LLAMA_GRETYPE_CHAR = 3 # terminal element: character (code point) - LLAMA_GRETYPE_CHAR_NOT = 4 # inverse char(s) ([^a], [^a-b] [^abc]) - LLAMA_GRETYPE_CHAR_RNG_UPPER = 5 # modifies a preceding LLAMA_GRETYPE_CHAR or LLAMA_GRETYPE_CHAR_ALT to be an inclusive range ([a-z]) - LLAMA_GRETYPE_CHAR_ALT = 6 # modifies a preceding LLAMA_GRETYPE_CHAR or LLAMA_GRETYPE_CHAR_RNG_UPPER to add an alternate char to match ([ab], [a-zA]) - - -# struct parse_state { -# std::map symbol_ids; -# std::vector> rules; -# std::vector c_rules(); -# }; -class parse_state: - def __init__(self): - self.symbol_ids: std.map[str, int] = std.map() - self.rules: std.vector[std.vector[LlamaGrammarElement]] = std.vector() - - # std::vector parse_state::c_rules() { - # std::vector ret; - # for (const auto & rule : rules) { - # ret.push_back(rule.data()); - # } - # return ret; - # } - def c_rules(self) -> std.vector[std.vector[LlamaGrammarElement]]: - ret = std.vector() # type: std.vector[std.vector[LlamaGrammarElement]] - for rule in self.rules: - ret.push_back(rule.data()) - return ret - - def __repr__(self) -> str: - return ( - f"parse_state(symbol_ids={len(self.symbol_ids)}, rules={len(self.rules)})" - ) +# static std::pair decode_utf8(const char * src) { +# static const int lookup[] = { 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 3, 4 }; +# uint8_t first_byte = static_cast(*src); +# uint8_t highbits = first_byte >> 4; +# int len = lookup[highbits]; +# uint8_t mask = (1 << (8 - len)) - 1; +# uint32_t value = first_byte & mask; +# const char * end = src + len; // may overrun! +# const char * pos = src + 1; +# for ( ; pos < end && *pos; pos++) { +# value = (value << 6) + (static_cast(*pos) & 0x3F); +# } +# return std::make_pair(value, pos); +# } +def decode_utf8(src: str) -> typing.Tuple[int, str]: + lookup: list[int] = [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 3, 4] + first_byte: int = ord(src[0]) + highbits: int = first_byte >> 4 + length: int = lookup[highbits] + mask: int = (1 << (8 - length)) - 1 + value: int = first_byte & mask + end: int = min(len(src), length) # Prevent overrun + + pos: int = 1 + for pos in range(1, end): + if not src[pos]: + break + value = (value << 6) + (ord(src[pos]) & 0x3F) -# struct llama_grammar { -# const std::vector> rules; -# std::vector> stacks; -# }; -# class llama_grammar: -# def __init__( -# self, -# rules: std.vector[std.vector[llama_grammar_element]], -# stacks: std.vector[std.vector[llama_grammar_element]], -# ): -# self.rules = rules -# self.stacks = stacks + return value, src[pos:] if pos < len(src) else "" -# uint32_t get_symbol_id(parse_state & state, const char * src, size_t len) { +# static uint32_t get_symbol_id(parse_state & state, const char * src, size_t len) { # uint32_t next_id = static_cast(state.symbol_ids.size()); -# auto result = state.symbol_ids.insert(std::make_pair(std::string(src, len), next_id)); +# auto result = state.symbol_ids.emplace(std::string(src, len), next_id); # return result.first->second; # } -def get_symbol_id(state: parse_state, src: const_char_p, len: int) -> int: - next_id = state.symbol_ids.size() # type: int - result = state.symbol_ids.insert(std.string(src, len), next_id) - return result[0].second # type: ignore +def get_symbol_id(state: ParseState, name: str) -> int: + next_id = len(state.symbol_ids) + return state.symbol_ids.setdefault(name, next_id) -# uint32_t generate_symbol_id(parse_state & state, const std::string & base_name) { +# static uint32_t generate_symbol_id(parse_state & state, const std::string & base_name) { # uint32_t next_id = static_cast(state.symbol_ids.size()); # state.symbol_ids[base_name + '_' + std::to_string(next_id)] = next_id; # return next_id; # } -def generate_symbol_id(state: parse_state, base_name: str) -> int: - next_id = state.symbol_ids.size() # type: int - state.symbol_ids[base_name + "_" + str(next_id)] = next_id +def generate_symbol_id(state: ParseState, base_name: str) -> int: + next_id = len(state.symbol_ids) + state.symbol_ids[f"{base_name}_{next_id}"] = next_id return next_id -# void add_rule( +# static void add_rule( # parse_state & state, # uint32_t rule_id, # const std::vector & rule) { @@ -526,57 +105,27 @@ def generate_symbol_id(state: parse_state, base_name: str) -> int: # } # state.rules[rule_id] = rule; # } -def add_rule( - state: parse_state, - rule_id: int, - rule: std.vector[LlamaGrammarElement], -) -> None: - if state.rules.size() <= rule_id: - state.rules.resize( - rule_id + 1, - fill_value_factory=std.vector[LlamaGrammarElement], - ) +def add_rule(state: ParseState, rule_id: int, rule: typing.List[GrammarElement]) -> None: + if len(state.rules) <= rule_id: + state.rules.extend([[]] * (rule_id + 1 - len(state.rules))) state.rules[rule_id] = rule -# std::pair decode_utf8(const char * src) { -# static const int lookup[] = { 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 3, 4 }; -# uint8_t first_byte = static_cast(*src); -# uint8_t highbits = first_byte >> 4; -# int len = lookup[highbits]; -# uint8_t mask = (1 << (8 - len)) - 1; -# uint32_t value = first_byte & mask; -# const char * end = src + len; // may overrun! -# const char * pos = src + 1; -# for ( ; pos < end && *pos; pos++) { -# value = (value << 6) + (static_cast(*pos) & 0x3F); -# } -# return std::make_pair(value, pos); +# static bool is_digit_char(char c) { +# return '0' <= c && c <= '9'; # } -def decode_utf8(src: const_char_p) -> Tuple[int, const_char_p]: - """Decodes a UTF-8 character from the source string.""" - lookup = (1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 3, 4) - first_byte = ord(src[0]) # type: int - highbits = first_byte >> 4 # type: int - len = lookup[highbits] # type: int - mask = (1 << (8 - len)) - 1 # type: int - value = first_byte & mask # type: int - end = src + len # type: const_char_p # may overrun! - pos = src + 1 # type: const_char_p - while pos < end and pos[0]: - value = (value << 6) + (ord(pos[0]) & 0x3F) - pos += 1 - return value, pos +def is_digit_char(c: str) -> bool: + return "0" <= c <= "9" -# bool is_word_char(char c) { -# return ('a' <= c && c <= 'z') || ('A' <= c && c <= 'Z') || c == '-' || ('0' <= c && c <= '9'); +# static bool is_word_char(char c) { +# return ('a' <= c && c <= 'z') || ('A' <= c && c <= 'Z') || c == '-' || is_digit_char(c); # } def is_word_char(c: str) -> bool: - return ("a" <= c <= "z") or ("A" <= c <= "Z") or c == "-" or ("0" <= c <= "9") + return ("a" <= c <= "z") or ("A" <= c <= "Z") or c == "-" or is_digit_char(c) -# std::pair parse_hex(const char * src, int size) { +# static std::pair parse_hex(const char * src, int size) { # const char * pos = src; # const char * end = src + size; # uint32_t value = 0; @@ -598,13 +147,12 @@ def is_word_char(c: str) -> bool: # } # return std::make_pair(value, pos); # } -def parse_hex(src: const_char_p, size: int) -> Tuple[int, const_char_p]: - pos = const_char_p(src) # type: const_char_p - end = src + size # type: const_char_p - value = 0 # type: int - while pos < end and pos[0]: +def parse_hex(src: str, size: int) -> typing.Tuple[int, str]: + pos = 0 + value = 0 + for _ in range(size): value <<= 4 - c = pos[0] # type: str + c = src[pos] if "a" <= c <= "f": value += ord(c) - ord("a") + 10 elif "A" <= c <= "F": @@ -614,12 +162,74 @@ def parse_hex(src: const_char_p, size: int) -> Tuple[int, const_char_p]: else: break pos += 1 - if pos != end: - raise RuntimeError("expecting " + str(size) + " hex chars at " + str(src)) - return (value, pos) + if pos != size: + raise ValueError(f"expecting {size} hex chars at {src}") + return value, src[pos:] + + +# static const char * parse_space(const char * src, bool newline_ok) { +# const char * pos = src; +# while (*pos == ' ' || *pos == '\t' || *pos == '#' || +# (newline_ok && (*pos == '\r' || *pos == '\n'))) { +# if (*pos == '#') { +# while (*pos && *pos != '\r' && *pos != '\n') { +# pos++; +# } +# } else { +# pos++; +# } +# } +# return pos; +# } +def parse_space(src: str, newline_ok: bool) -> str: + pos = src + while pos and (pos[0] in (' ', '\t', '#') or (newline_ok and pos[0] in ('\r', '\n'))): + if pos[0] == "#": + while pos and pos[0] not in ("\r", "\n"): + pos = pos[1:] + else: + pos = pos[1:] + return pos + +# static const char * parse_name(const char * src) { +# const char * pos = src; +# while (is_word_char(*pos)) { +# pos++; +# } +# if (pos == src) { +# throw std::runtime_error(std::string("expecting name at ") + src); +# } +# return pos; +# } +def parse_name(src: str) -> typing.Tuple[str, str]: + pos = src + while pos and is_word_char(pos[0]): + pos = pos[1:] + if pos == src: + raise ValueError(f"expecting name at {src}") + return src[:len(src) - len(pos)], pos -# std::pair parse_char(const char * src) { +# static const char * parse_int(const char * src) { +# const char * pos = src; +# while (is_digit_char(*pos)) { +# pos++; +# } +# if (pos == src) { +# throw std::runtime_error(std::string("expecting integer at ") + src); +# } +# return pos; +# } +def parse_int(src: str) -> typing.Tuple[int, str]: + pos = src + while pos and is_digit_char(pos[0]): + pos = pos[1:] + if pos == src: + raise ValueError(f"expecting integer at {src}") + return int(src[:len(src) - len(pos)]), pos + + +# static std::pair parse_char(const char * src) { # if (*src == '\\') { # switch (src[1]) { # case 'x': return parse_hex(src + 2, 2); @@ -641,273 +251,320 @@ def parse_hex(src: const_char_p, size: int) -> Tuple[int, const_char_p]: # } # throw std::runtime_error("unexpected end of input"); # } -def parse_char(src: const_char_p) -> Tuple[int, const_char_p]: +def parse_char(src: str) -> typing.Tuple[int, str]: + if not src: + raise ValueError("unexpected end of input") if src[0] == "\\": - case = src[1] # type: str - if case == "x": - return parse_hex(src + 2, 2) - elif case == "u": - return parse_hex(src + 2, 4) - elif case == "U": - return parse_hex(src + 2, 8) - elif case == "t": - return (ord("\t"), src + 2) # implicit cast - elif case == "r": - return (ord("\r"), src + 2) # implicit cast - elif case == "n": - return (ord("\n"), src + 2) # implicit cast - elif case in ("\\", '"', "[", "]"): - return (ord(case), src + 2) # implicit cast + if src[1] == "x": + return parse_hex(src[2:], 2) + elif src[1] == "u": + return parse_hex(src[2:], 4) + elif src[1] == "U": + return parse_hex(src[2:], 8) + elif src[1] == "t": + return ord("\t"), src[2:] + elif src[1] == "r": + return ord("\r"), src[2:] + elif src[1] == "n": + return ord("\n"), src[2:] + elif src[1] in ('\\', '"', '[', ']'): + return ord(src[1]), src[2:] else: - raise RuntimeError("unknown escape at " + str(src)) - elif src[0]: - return decode_utf8(src) - else: - raise RuntimeError("unexpected end of input") - - -# const char * parse_name(const char * src) { -# const char * pos = src; -# while (is_word_char(*pos)) { -# pos++; -# } -# if (pos == src) { -# throw std::runtime_error(std::string("expecting name at ") + src); -# } -# return pos; -# } -def parse_name(src: const_char_p) -> const_char_p: - pos = const_char_p(src) # type: const_char_p - while is_word_char(pos[0]): - pos += 1 - if pos == src: - raise RuntimeError("expecting name at " + str(src)) - return pos - + raise ValueError(f"unknown escape at {src}") + return decode_utf8(src) -# const char * parse_space(const char * src, bool newline_ok) { +# static const char * parse_sequence( +# parse_state & state, +# const char * src, +# const std::string & rule_name, +# std::vector & out_elements, +# bool is_nested) { +# size_t last_sym_start = out_elements.size(); # const char * pos = src; -# while (*pos == ' ' || *pos == '\t' || *pos == '#' || -# (newline_ok && (*pos == '\r' || *pos == '\n'))) { -# if (*pos == '#') { -# while (*pos && *pos != '\r' && *pos != '\n') { +# +# auto handle_repetitions = [&](int min_times, int max_times) { +# +# if (last_sym_start == out_elements.size()) { +# throw std::runtime_error(std::string("expecting preceding item to */+/?/{ at ") + pos); +# } +# +# // apply transformation to previous symbol (last_sym_start to end) according to +# // the following rewrite rules: +# // S{m,n} --> S S S (m times) S'(n-m) +# // S'(x) ::= S S'(x-1) | +# // (... n-m definitions of these S' rules ...) +# // S'(1) ::= S | +# // S{m,} --> S S S (m times) S' +# // S' ::= S S' | +# // S* --> S{0,} +# // --> S' ::= S S' | +# // S+ --> S{1,} +# // --> S S' +# // S' ::= S S' | +# // S? --> S{0,1} +# // --> S' +# // S' ::= S | +# +# std::vector previous_elements(out_elements.begin() + last_sym_start, out_elements.end()); +# if (min_times == 0) { +# out_elements.resize(last_sym_start); +# } else { +# // Repeat the previous elements (min_times - 1) times +# for (int i = 1; i < min_times; i++) { +# out_elements.insert(out_elements.end(), previous_elements.begin(), previous_elements.end()); +# } +# } +# +# uint32_t last_rec_rule_id = 0; +# auto n_opt = max_times < 0 ? 1 : max_times - min_times; +# +# std::vector rec_rule(previous_elements); +# for (int i = 0; i < n_opt; i++) { +# rec_rule.resize(previous_elements.size()); +# uint32_t rec_rule_id = generate_symbol_id(state, rule_name); +# if (i > 0 || max_times < 0) { +# rec_rule.push_back({LLAMA_GRETYPE_RULE_REF, max_times < 0 ? rec_rule_id : last_rec_rule_id}); +# } +# rec_rule.push_back({LLAMA_GRETYPE_ALT, 0}); +# rec_rule.push_back({LLAMA_GRETYPE_END, 0}); +# add_rule(state, rec_rule_id, rec_rule); +# last_rec_rule_id = rec_rule_id; +# } +# if (n_opt > 0) { +# out_elements.push_back({LLAMA_GRETYPE_RULE_REF, last_rec_rule_id}); +# } +# }; +# +# while (*pos) { +# if (*pos == '"') { // literal string +# pos++; +# last_sym_start = out_elements.size(); +# while (*pos != '"') { +# if (!*pos) { +# throw std::runtime_error("unexpected end of input"); +# } +# auto char_pair = parse_char(pos); +# pos = char_pair.second; +# out_elements.push_back({LLAMA_GRETYPE_CHAR, char_pair.first}); +# } +# pos = parse_space(pos + 1, is_nested); +# } else if (*pos == '[') { // char range(s) +# pos++; +# enum llama_gretype start_type = LLAMA_GRETYPE_CHAR; +# if (*pos == '^') { # pos++; +# start_type = LLAMA_GRETYPE_CHAR_NOT; +# } +# last_sym_start = out_elements.size(); +# while (*pos != ']') { +# if (!*pos) { +# throw std::runtime_error("unexpected end of input"); +# } +# auto char_pair = parse_char(pos); +# pos = char_pair.second; +# enum llama_gretype type = last_sym_start < out_elements.size() +# ? LLAMA_GRETYPE_CHAR_ALT +# : start_type; +# +# out_elements.push_back({type, char_pair.first}); +# if (pos[0] == '-' && pos[1] != ']') { +# if (!pos[1]) { +# throw std::runtime_error("unexpected end of input"); +# } +# auto endchar_pair = parse_char(pos + 1); +# pos = endchar_pair.second; +# out_elements.push_back({LLAMA_GRETYPE_CHAR_RNG_UPPER, endchar_pair.first}); +# } +# } +# pos = parse_space(pos + 1, is_nested); +# } else if (is_word_char(*pos)) { // rule reference +# const char * name_end = parse_name(pos); +# uint32_t ref_rule_id = get_symbol_id(state, pos, name_end - pos); +# pos = parse_space(name_end, is_nested); +# last_sym_start = out_elements.size(); +# out_elements.push_back({LLAMA_GRETYPE_RULE_REF, ref_rule_id}); +# } else if (*pos == '(') { // grouping +# // parse nested alternates into synthesized rule +# pos = parse_space(pos + 1, true); +# uint32_t sub_rule_id = generate_symbol_id(state, rule_name); +# pos = parse_alternates(state, pos, rule_name, sub_rule_id, true); +# last_sym_start = out_elements.size(); +# // output reference to synthesized rule +# out_elements.push_back({LLAMA_GRETYPE_RULE_REF, sub_rule_id}); +# if (*pos != ')') { +# throw std::runtime_error(std::string("expecting ')' at ") + pos); +# } +# pos = parse_space(pos + 1, is_nested); +# } else if (*pos == '.') { // any char +# last_sym_start = out_elements.size(); +# out_elements.push_back({LLAMA_GRETYPE_CHAR_ANY, 0}); +# pos = parse_space(pos + 1, is_nested); +# } else if (*pos == '*') { +# pos = parse_space(pos + 1, is_nested); +# handle_repetitions(0, -1); +# } else if (*pos == '+') { +# pos = parse_space(pos + 1, is_nested); +# handle_repetitions(1, -1); +# } else if (*pos == '?') { +# pos = parse_space(pos + 1, is_nested); +# handle_repetitions(0, 1); +# } else if (*pos == '{') { +# pos = parse_space(pos + 1, is_nested); +# +# if (!is_digit_char(*pos)) { +# throw std::runtime_error(std::string("expecting an int at ") + pos); # } +# const char * int_end = parse_int(pos); +# int min_times = std::stoul(std::string(pos, int_end - pos)); +# pos = parse_space(int_end, is_nested); +# +# int max_times = -1; +# +# if (*pos == '}') { +# max_times = min_times; +# pos = parse_space(pos + 1, is_nested); +# } else if (*pos == ',') { +# pos = parse_space(pos + 1, is_nested); +# +# if (is_digit_char(*pos)) { +# const char * int_end = parse_int(pos); +# max_times = std::stoul(std::string(pos, int_end - pos)); +# pos = parse_space(int_end, is_nested); +# } +# +# if (*pos != '}') { +# throw std::runtime_error(std::string("expecting '}' at ") + pos); +# } +# pos = parse_space(pos + 1, is_nested); +# } else { +# throw std::runtime_error(std::string("expecting ',' at ") + pos); +# } +# handle_repetitions(min_times, max_times); # } else { -# pos++; +# break; # } # } # return pos; # } -def parse_space(src: const_char_p, newline_ok: bool) -> const_char_p: - pos = const_char_p(src) # type: const_char_p - while pos[0] in (" ", "\t", "#") or (newline_ok and pos[0] in ("\r", "\n")): - if pos[0] == "#": - while pos[0] is not None and pos[0] not in ("\r", "\n"): - pos += 1 - else: - pos += 1 - return pos +def parse_sequence(state: ParseState, src: str, rule_name: str, out_elements: typing.List[GrammarElement], is_nested: bool) -> str: + last_sym_start = len(out_elements) + pos = src + def handle_repetitions(min_times: int, max_times: int) -> None: + nonlocal state, src, rule_name, out_elements, is_nested, last_sym_start, pos -# const char * parse_sequence( -# parse_state & state, -# const char * src, -# const std::string & rule_name, -# std::vector & out_elements, -# bool is_nested) { -def parse_sequence( - state: parse_state, - src: const_char_p, - rule_name: str, - out_elements: std.vector[LlamaGrammarElement], - is_nested: bool, -) -> const_char_p: - # size_t last_sym_start = out_elements.size(); - # const char * pos = src; - last_sym_start = out_elements.size() # type: int - pos = const_char_p(src) # type: const_char_p - # while (*pos) { - while pos[0]: - # if (*pos == '"') { // literal string - # pos++; - # last_sym_start = out_elements.size(); - # while (*pos != '"') { - # auto char_pair = parse_char(pos); - # pos = char_pair.second; - # out_elements.push_back({LLAMA_GRETYPE_CHAR, char_pair.first}); - # } - # pos = parse_space(pos + 1, is_nested); - if pos[0] == '"': # literal string - pos += 1 - last_sym_start = out_elements.size() - while pos[0] != '"': - char_pair = parse_char(pos) # type: Tuple[int, const_char_p] - pos = char_pair[1] - out_elements.push_back( - LlamaGrammarElement(llama_gretype.LLAMA_GRETYPE_CHAR, char_pair[0]) - ) - pos = parse_space(pos + 1, is_nested) - # } else if (*pos == '[') { // char range(s) - # pos++; - # enum llama_gretype start_type = LLAMA_GRETYPE_CHAR; - elif pos[0] == "[": # char range(s) - pos += 1 - start_type = llama_gretype.LLAMA_GRETYPE_CHAR # type: llama_gretype - # if (*pos == '^') { - # pos++; - # start_type = LLAMA_GRETYPE_CHAR_NOT; - # } - # last_sym_start = out_elements.size(); + if last_sym_start == len(out_elements): + raise ValueError(f"expecting preceding item to */+/?/{{ at {pos}") + + previous_elements = out_elements[last_sym_start:] + if min_times == 0: + del out_elements[last_sym_start:] + else: + for i in range(1, min_times): + out_elements.extend(previous_elements) + + last_rec_rule_id = 0 + n_opt = 1 if max_times < 0 else max_times - min_times + + rec_rule = previous_elements[:] + for i in range(n_opt): + rec_rule = rec_rule[:len(previous_elements)] + rec_rule_id = generate_symbol_id(state, rule_name) + if i > 0 or max_times < 0: + rec_rule.append(GrammarElement(GrammarElementType.RULE_REF, rec_rule_id if max_times < 0 else last_rec_rule_id)) + rec_rule.append(GrammarElement(GrammarElementType.ALT, 0)) + rec_rule.append(GrammarElement(GrammarElementType.END, 0)) + add_rule(state, rec_rule_id, rec_rule) + last_rec_rule_id = rec_rule_id + if n_opt > 0: + out_elements.append(GrammarElement(GrammarElementType.RULE_REF, last_rec_rule_id)) + + while pos: + if pos[0] == '"': + pos = pos[1:] + last_sym_start = len(out_elements) + while not pos.startswith('"'): + if not pos: + raise ValueError("unexpected end of input") + char, pos = parse_char(pos) + out_elements.append(GrammarElement(GrammarElementType.CHAR, char)) + pos = parse_space(pos[1:], is_nested) + elif pos[0] == "[": + pos = pos[1:] + start_type = GrammarElementType.CHAR if pos[0] == "^": - pos += 1 - start_type = llama_gretype.LLAMA_GRETYPE_CHAR_NOT - last_sym_start = out_elements.size() - # while (*pos != ']') { - # auto char_pair = parse_char(pos); - # pos = char_pair.second; - # enum llama_gretype type = last_sym_start < out_elements.size() - # ? LLAMA_GRETYPE_CHAR_ALT - # : start_type; - # out_elements.push_back({type, char_pair.first}); + pos = pos[1:] + start_type = GrammarElementType.CHAR_NOT + last_sym_start = len(out_elements) while pos[0] != "]": - char_pair = parse_char(pos) # type: Tuple[int, const_char_p] - pos = char_pair[1] - type = ( - llama_gretype.LLAMA_GRETYPE_CHAR_ALT - if last_sym_start < out_elements.size() - else start_type - ) # type: llama_gretype - out_elements.push_back(LlamaGrammarElement(type, char_pair[0])) - # if (pos[0] == '-' && pos[1] != ']') { - # auto endchar_pair = parse_char(pos + 1); - # pos = endchar_pair.second; - # out_elements.push_back({LLAMA_GRETYPE_CHAR_RNG_UPPER, endchar_pair.first}); - # } - # } + if not pos: + raise ValueError("unexpected end of input") + char, pos = parse_char(pos) + type = GrammarElementType.CHAR_ALT if last_sym_start < len(out_elements) else start_type + out_elements.append(GrammarElement(type, char)) if pos[0] == "-" and pos[1] != "]": - endchar_pair = parse_char(pos + 1) # type: Tuple[int, const_char_p] - pos = endchar_pair[1] - out_elements.push_back( - LlamaGrammarElement( - llama_gretype.LLAMA_GRETYPE_CHAR_RNG_UPPER, - endchar_pair[0], - ) - ) - # pos = parse_space(pos + 1, is_nested); - pos = parse_space(pos + 1, is_nested) - # } else if (is_word_char(*pos)) { // rule reference - # const char * name_end = parse_name(pos); - # uint32_t ref_rule_id = get_symbol_id(state, pos, name_end - pos); - # pos = parse_space(name_end, is_nested); - # last_sym_start = out_elements.size(); - # out_elements.push_back({LLAMA_GRETYPE_RULE_REF, ref_rule_id}); - elif is_word_char(pos[0]): # rule reference - name_end = parse_name(pos) # type: const_char_p - ref_rule_id = get_symbol_id(state, pos, name_end - pos) # type: int - pos = parse_space(name_end, is_nested) - last_sym_start = out_elements.size() - out_elements.push_back( - LlamaGrammarElement(llama_gretype.LLAMA_GRETYPE_RULE_REF, ref_rule_id) - ) - # } else if (*pos == '(') { // grouping - # // parse nested alternates into synthesized rule - # pos = parse_space(pos + 1, true); - # uint32_t sub_rule_id = generate_symbol_id(state, rule_name); - # pos = parse_alternates(state, pos, rule_name, sub_rule_id, true); - # last_sym_start = out_elements.size(); - # // output reference to synthesized rule - # out_elements.push_back({LLAMA_GRETYPE_RULE_REF, sub_rule_id}); - # if (*pos != ')') { - # throw std::runtime_error(std::string("expecting ')' at ") + pos); - # } - # pos = parse_space(pos + 1, is_nested); - elif pos[0] == "(": # grouping - # parse nested alternates into synthesized rule - pos = parse_space(pos + 1, True) - sub_rule_id = generate_symbol_id(state, rule_name) # type: int - pos = parse_alternates(state, pos, rule_name, sub_rule_id, True) - last_sym_start = out_elements.size() - # output reference to synthesized rule - out_elements.push_back( - LlamaGrammarElement(llama_gretype.LLAMA_GRETYPE_RULE_REF, sub_rule_id) - ) + if not pos[1]: + raise ValueError("unexpected end of input") + endchar, pos = parse_char(pos[1:]) + out_elements.append(GrammarElement(GrammarElementType.CHAR_RNG_UPPER, endchar)) + pos = parse_space(pos[1:], is_nested) + elif pos and is_word_char(pos[0]): + name, rest = parse_name(pos) + ref_rule_id = get_symbol_id(state, name) + pos = parse_space(rest, is_nested) + last_sym_start = len(out_elements) + out_elements.append(GrammarElement(GrammarElementType.RULE_REF, ref_rule_id)) + elif pos.startswith("("): + pos = parse_space(pos[1:], newline_ok=True) + sub_rule_id = generate_symbol_id(state, rule_name) + pos = parse_alternates(state, pos, rule_name, sub_rule_id, is_nested=True) + last_sym_start = len(out_elements) + out_elements.append(GrammarElement(GrammarElementType.RULE_REF, sub_rule_id)) if pos[0] != ")": - raise RuntimeError("expecting ')' at " + str(pos)) - pos = parse_space(pos + 1, is_nested) - # } else if (*pos == '*' || *pos == '+' || *pos == '?') { // repetition operator - # if (last_sym_start == out_elements.size()) { - # throw std::runtime_error(std::string("expecting preceeding item to */+/? at ") + pos); - # } - elif pos[0] in ("*", "+", "?"): # repetition operator - if last_sym_start == out_elements.size(): - raise RuntimeError("expecting preceding item to */+/? at " + str(pos)) - # // apply transformation to previous symbol (last_sym_start to end) according to - # // rewrite rules: - # // S* --> S' ::= S S' | - # // S+ --> S' ::= S S' | S - # // S? --> S' ::= S | - # uint32_t sub_rule_id = generate_symbol_id(state, rule_name); - # std::vector sub_rule; - # // add preceding symbol to generated rule - # sub_rule.insert( - # sub_rule.end(), out_elements.begin() + last_sym_start, out_elements.end()); - sub_rule_id = generate_symbol_id(state, rule_name) # type: int - sub_rule = std.vector[ - LlamaGrammarElement - ]() # type: std.vector[LlamaGrammarElement] - sub_rule.insert( - sub_rule.end(), - out_elements.begin() + last_sym_start, - out_elements.end(), - ) - # if (*pos == '*' || *pos == '+') { - # // cause generated rule to recurse - # sub_rule.push_back({LLAMA_GRETYPE_RULE_REF, sub_rule_id}); - # } - # // mark start of alternate def - # sub_rule.push_back({LLAMA_GRETYPE_ALT, 0}); - if pos[0] in ("*", "+"): - sub_rule.push_back( - LlamaGrammarElement( - llama_gretype.LLAMA_GRETYPE_RULE_REF, sub_rule_id - ) - ) - sub_rule.push_back(LlamaGrammarElement(llama_gretype.LLAMA_GRETYPE_ALT, 0)) - # if (*pos == '+') { - # // add preceding symbol as alternate only for '+' (otherwise empty) - # sub_rule.insert( - # sub_rule.end(), out_elements.begin() + last_sym_start, out_elements.end()); - # } - # sub_rule.push_back({LLAMA_GRETYPE_END, 0}); - # add_rule(state, sub_rule_id, sub_rule); - # // in original rule, replace previous symbol with reference to generated rule - # out_elements.resize(last_sym_start); - # out_elements.push_back({LLAMA_GRETYPE_RULE_REF, sub_rule_id}); - # pos = parse_space(pos + 1, is_nested); - if pos[0] == "+": - # add preceding symbol as alternate only for '+' (otherwise empty) - sub_rule.insert( - sub_rule.end(), - out_elements.begin() + last_sym_start, - out_elements.end(), - ) - sub_rule.push_back(LlamaGrammarElement(llama_gretype.LLAMA_GRETYPE_END, 0)) - add_rule(state, sub_rule_id, sub_rule) - # in original rule, replace previous symbol with reference to generated rule - out_elements.resize(last_sym_start) - out_elements.push_back( - LlamaGrammarElement(llama_gretype.LLAMA_GRETYPE_RULE_REF, sub_rule_id) - ) - pos = parse_space(pos + 1, is_nested) - # } else { - # break; - # } + raise ValueError(f"expecting ')' at {pos}") + pos = parse_space(pos[1:], is_nested) + elif pos.startswith("."): + last_sym_start = len(out_elements) + out_elements.append(GrammarElement(GrammarElementType.CHAR_ANY, 0)) + pos = parse_space(pos[1:], is_nested) + elif pos.startswith("*"): + pos = parse_space(pos[1:], is_nested) + handle_repetitions(0, -1) + elif pos.startswith("+"): + pos = parse_space(pos[1:], is_nested) + handle_repetitions(1, -1) + elif pos.startswith("?"): + pos = parse_space(pos[1:], is_nested) + handle_repetitions(0, 1) + elif pos.startswith("{"): + pos = parse_space(pos[1:], is_nested) + + if not is_digit_char(pos): + raise ValueError(f"expecting an int at {pos}") + min_times, pos = parse_int(pos) + pos = parse_space(pos, is_nested) + + max_times = -1 + + if pos[0] == "}": + max_times = min_times + pos = parse_space(pos[1:], is_nested) + elif pos[0] == ",": + pos = parse_space(pos[1:], is_nested) + + if is_digit_char(pos): + max_times, pos = parse_int(pos) + pos = parse_space(pos, is_nested) + + if pos[0] != "}": + raise ValueError("expecting '}' at {}".format(pos)) + + pos = parse_space(pos[1:], is_nested) + else: + raise ValueError(f"expecting ',' at {pos}") + handle_repetitions(min_times, max_times) else: break - # } - # return pos; - # } return pos @@ -928,39 +585,32 @@ def parse_sequence( # add_rule(state, rule_id, rule); # return pos; # } -def parse_alternates( - state: parse_state, - src: const_char_p, - rule_name: str, - rule_id: int, - is_nested: bool, -) -> const_char_p: - rule = std.vector() # type: std.vector[LlamaGrammarElement] - pos = parse_sequence(state, src, rule_name, rule, is_nested) # type: const_char_p - while pos[0] == "|": - rule.push_back(LlamaGrammarElement(llama_gretype.LLAMA_GRETYPE_ALT, 0)) - pos = parse_space(pos + 1, True) +def parse_alternates(state: ParseState, src: str, rule_name: str, rule_id: int, is_nested: bool) -> str: + rule = [] + pos = parse_sequence(state, src, rule_name, rule, is_nested) + while pos.startswith("|"): + rule.append(GrammarElement(GrammarElementType.ALT, 0)) + pos = parse_space(pos[1:], newline_ok=True) pos = parse_sequence(state, pos, rule_name, rule, is_nested) - rule.push_back(LlamaGrammarElement(llama_gretype.LLAMA_GRETYPE_END, 0)) + rule.append(GrammarElement(GrammarElementType.END, 0)) add_rule(state, rule_id, rule) return pos -# const char * parse_rule(parse_state & state, const char * src) { +# static const char * parse_rule(parse_state & state, const char * src) { # const char * name_end = parse_name(src); # const char * pos = parse_space(name_end, false); # size_t name_len = name_end - src; # uint32_t rule_id = get_symbol_id(state, src, name_len); # const std::string name(src, name_len); - +# # if (!(pos[0] == ':' && pos[1] == ':' && pos[2] == '=')) { # throw std::runtime_error(std::string("expecting ::= at ") + pos); # } # pos = parse_space(pos + 3, true); - +# # pos = parse_alternates(state, pos, name, rule_id, false); - - +# # if (*pos == '\r') { # pos += pos[1] == '\n' ? 2 : 1; # } else if (*pos == '\n') { @@ -970,26 +620,26 @@ def parse_alternates( # } # return parse_space(pos, true); # } -def parse_rule(state: parse_state, src: const_char_p) -> const_char_p: - name_end = parse_name(src) # type: const_char_p - pos = parse_space(name_end, False) # type: const_char_p - name_len = name_end - src # type: int - rule_id = get_symbol_id(state, src, name_len) # type: int - name = std.string(src, name_len) # type: str - - if not (pos[0] == ":" and pos[1] == ":" and pos[2] == "="): - raise RuntimeError("expecting ::= at " + str(pos)) - - pos = parse_space(pos + 3, True) # type: const_char_p - pos = parse_alternates(state, pos, name, rule_id, False) # type: const_char_p - - if pos[0] == "\r": - pos += 2 if pos[1] == "\n" else 1 - elif pos[0] == "\n": - pos += 1 - elif pos[0]: - raise RuntimeError("expecting newline or end at " + str(pos)) - return parse_space(pos, True) +def parse_rule(state: ParseState, src: str) -> str: + pos = src + name, pos = parse_name(pos) + pos = parse_space(pos, newline_ok=False) + rule_id = get_symbol_id(state, name) + + if not pos.startswith("::="): + raise ValueError(f"expecting ::= at {pos}") + + pos = parse_space(pos[3:], newline_ok=True) + + pos = parse_alternates(state, pos, name, rule_id, is_nested=False) + + if pos.startswith("\r"): + pos = pos[2:] if pos[1] == "\n" else pos[1:] + elif pos.startswith("\n"): + pos = pos[1:] + elif pos: + raise ValueError(f"expecting newline or end at {pos}") + return parse_space(pos, newline_ok=True) # parse_state parse(const char * src) { @@ -999,204 +649,273 @@ def parse_rule(state: parse_state, src: const_char_p) -> const_char_p: # while (*pos) { # pos = parse_rule(state, pos); # } +# // Validate the state to ensure that all rules are defined +# for (const auto & rule : state.rules) { +# for (const auto & elem : rule) { +# if (elem.type == LLAMA_GRETYPE_RULE_REF) { +# // Ensure that the rule at that location exists +# if (elem.value >= state.rules.size() || state.rules[elem.value].empty()) { +# // Get the name of the rule that is missing +# for (const auto & kv : state.symbol_ids) { +# if (kv.second == elem.value) { +# throw std::runtime_error("Undefined rule identifier '" + kv.first + "'"); +# } +# } +# } +# } +# } +# } # return state; # } catch (const std::exception & err) { # fprintf(stderr, "%s: error parsing grammar: %s\n", __func__, err.what()); # return parse_state(); # } # } -def parse(src: const_char_p) -> parse_state: - try: - state = parse_state() # type: parse_state - pos = parse_space(src, True) # type: const_char_p - while pos[0]: - pos = parse_rule(state, pos) - return state - except Exception as err: - print(f"{parse.__name__}: error parsing grammar: {err}") - return parse_state() - - -# void print_grammar_char(FILE * file, uint32_t c) { -# if (0x20 <= c && c <= 0x7f) { -# fprintf(file, "%c", static_cast(c)); -# } else { -# // cop out of encoding UTF-8 -# fprintf(file, "", c); -# } -# } -def print_grammar_char(file: TextIO, c: int) -> None: - if 0x20 <= c and c <= 0x7F: - file.write(chr(c)) - else: - # cop out of encoding UTF-8 - file.write(f"") - - -# bool is_char_element(llama_grammar_element elem) { +def parse(src: str) -> ParseState: + state = ParseState() + pos = src + pos = parse_space(pos, newline_ok=True) + while pos: + pos = parse_rule(state, pos) + # validate + for rule in state.rules: + for elem in rule: + if elem.type == GrammarElementType.RULE_REF: + if elem.value >= len(state.rules) or not state.rules[elem.value]: + for k, v in state.symbol_ids.items(): + if v == elem.value: + raise ValueError(f"Undefined rule identifier '{k}'") + return state + + +# static bool is_char_element(llama_grammar_element elem) { # switch (elem.type) { # case LLAMA_GRETYPE_CHAR: return true; # case LLAMA_GRETYPE_CHAR_NOT: return true; # case LLAMA_GRETYPE_CHAR_ALT: return true; # case LLAMA_GRETYPE_CHAR_RNG_UPPER: return true; +# case LLAMA_GRETYPE_CHAR_ANY: return true; # default: return false; # } # } -def is_char_element(elem: LlamaGrammarElement) -> bool: +def is_char_element(elem: GrammarElement) -> bool: return elem.type in ( - llama_gretype.LLAMA_GRETYPE_CHAR, - llama_gretype.LLAMA_GRETYPE_CHAR_NOT, - llama_gretype.LLAMA_GRETYPE_CHAR_ALT, - llama_gretype.LLAMA_GRETYPE_CHAR_RNG_UPPER, + GrammarElementType.CHAR, + GrammarElementType.CHAR_NOT, + GrammarElementType.CHAR_ALT, + GrammarElementType.CHAR_RNG_UPPER, + GrammarElementType.CHAR_ANY ) -# void print_rule( +def print_grammar_char(file: typing.TextIO, c: int) -> None: + if 0x20 <= c <= 0x7f: + print(chr(c), end="", file=file) + else: + print(f"", end="", file=file) + + +# static void print_rule( # FILE * file, # uint32_t rule_id, # const std::vector & rule, # const std::map & symbol_id_names) { +# if (rule.empty() || rule.back().type != LLAMA_GRETYPE_END) { +# throw std::runtime_error( +# "malformed rule, does not end with LLAMA_GRETYPE_END: " + std::to_string(rule_id)); +# } +# fprintf(file, "%s ::= ", symbol_id_names.at(rule_id).c_str()); +# for (size_t i = 0, end = rule.size() - 1; i < end; i++) { +# llama_grammar_element elem = rule[i]; +# switch (elem.type) { +# case LLAMA_GRETYPE_END: +# throw std::runtime_error( +# "unexpected end of rule: " + std::to_string(rule_id) + "," + +# std::to_string(i)); +# case LLAMA_GRETYPE_ALT: +# fprintf(file, "| "); +# break; +# case LLAMA_GRETYPE_RULE_REF: +# fprintf(file, "%s ", symbol_id_names.at(elem.value).c_str()); +# break; +# case LLAMA_GRETYPE_CHAR: +# fprintf(file, "["); +# print_grammar_char(file, elem.value); +# break; +# case LLAMA_GRETYPE_CHAR_NOT: +# fprintf(file, "[^"); +# print_grammar_char(file, elem.value); +# break; +# case LLAMA_GRETYPE_CHAR_RNG_UPPER: +# if (i == 0 || !is_char_element(rule[i - 1])) { +# throw std::runtime_error( +# "LLAMA_GRETYPE_CHAR_RNG_UPPER without preceding char: " + +# std::to_string(rule_id) + "," + std::to_string(i)); +# } +# fprintf(file, "-"); +# print_grammar_char(file, elem.value); +# break; +# case LLAMA_GRETYPE_CHAR_ALT: +# if (i == 0 || !is_char_element(rule[i - 1])) { +# throw std::runtime_error( +# "LLAMA_GRETYPE_CHAR_ALT without preceding char: " + +# std::to_string(rule_id) + "," + std::to_string(i)); +# } +# print_grammar_char(file, elem.value); +# break; +# case LLAMA_GRETYPE_CHAR_ANY: +# fprintf(file, "."); +# break; +# } +# if (is_char_element(elem)) { +# switch (rule[i + 1].type) { +# case LLAMA_GRETYPE_CHAR_ALT: +# case LLAMA_GRETYPE_CHAR_RNG_UPPER: +# case LLAMA_GRETYPE_CHAR_ANY: +# break; +# default: +# fprintf(file, "] "); +# } +# } +# } +# fprintf(file, "\n"); +# } def print_rule( - file: TextIO, + file: typing.TextIO, rule_id: int, - rule: std.vector[LlamaGrammarElement], - symbol_id_names: std.map[int, str], + rule: typing.List[GrammarElement], + symbol_id_names: typing.Dict[int, str], ) -> None: - # if (rule.empty() || rule.back().type != LLAMA_GRETYPE_END) { - # throw std::runtime_error( - # "malformed rule, does not end with LLAMA_GRETYPE_END: " + std::to_string(rule_id)); - # } - # fprintf(file, "%s ::= ", symbol_id_names.at(rule_id).c_str()); - if rule.empty() or rule.back().type != llama_gretype.LLAMA_GRETYPE_END: - raise RuntimeError( - "malformed rule, does not end with LLAMA_GRETYPE_END: " + str(rule_id) - ) - print(f"{symbol_id_names.at(rule_id)} ::=", file=file, end=" ") - # for (size_t i = 0, end = rule.size() - 1; i < end; i++) { - # llama_grammar_element elem = rule[i]; - # switch (elem.type) { - # case LLAMA_GRETYPE_END: - # throw std::runtime_error( - # "unexpected end of rule: " + std::to_string(rule_id) + "," + - # std::to_string(i)); - # case LLAMA_GRETYPE_ALT: - # fprintf(file, "| "); - # break; - # case LLAMA_GRETYPE_RULE_REF: - # fprintf(file, "%s ", symbol_id_names.at(elem.value).c_str()); - # break; - # case LLAMA_GRETYPE_CHAR: - # fprintf(file, "["); - # print_grammar_char(file, elem.value); - # break; - # case LLAMA_GRETYPE_CHAR_NOT: - # fprintf(file, "[^"); - # print_grammar_char(file, elem.value); - # break; - # case LLAMA_GRETYPE_CHAR_RNG_UPPER: - # if (i == 0 || !is_char_element(rule[i - 1])) { - # throw std::runtime_error( - # "LLAMA_GRETYPE_CHAR_RNG_UPPER without preceding char: " + - # std::to_string(rule_id) + "," + std::to_string(i)); - # } - # fprintf(file, "-"); - # print_grammar_char(file, elem.value); - # break; - # case LLAMA_GRETYPE_CHAR_ALT: - # if (i == 0 || !is_char_element(rule[i - 1])) { - # throw std::runtime_error( - # "LLAMA_GRETYPE_CHAR_ALT without preceding char: " + - # std::to_string(rule_id) + "," + std::to_string(i)); - # } - # print_grammar_char(file, elem.value); - # break; - # } + if not rule or rule[-1].type != GrammarElementType.END: + raise ValueError(f"malformed rule, does not end with LLAMA_GRETYPE_END: {rule_id}") + + print(f"{symbol_id_names[rule_id]} ::=", end=" ", file=file) + for i, elem in enumerate(rule[:-1]): - case = elem.type # type: llama_gretype - if case is llama_gretype.LLAMA_GRETYPE_END: - raise RuntimeError("unexpected end of rule: " + str(rule_id) + "," + str(i)) - elif case is llama_gretype.LLAMA_GRETYPE_ALT: - print("| ", file=file, end="") - elif case is llama_gretype.LLAMA_GRETYPE_RULE_REF: - print(f"{symbol_id_names.at(elem.value)} ", file=file, end="") - elif case is llama_gretype.LLAMA_GRETYPE_CHAR: - print("[", file=file, end="") + if elem.type == GrammarElementType.END: + raise ValueError(f"unexpected end of rule: {rule_id}, {i}") + if elem.type == GrammarElementType.ALT: + print("| ", end="", file=file) + elif elem.type == GrammarElementType.RULE_REF: + print(f"{symbol_id_names[elem.value]} ", end="", file=file) + elif elem.type == GrammarElementType.CHAR: + print("[", end="", file=file) print_grammar_char(file, elem.value) - elif case is llama_gretype.LLAMA_GRETYPE_CHAR_NOT: - print("[^", file=file, end="") + elif elem.type == GrammarElementType.CHAR_NOT: + print("[^", end="", file=file) print_grammar_char(file, elem.value) - elif case is llama_gretype.LLAMA_GRETYPE_CHAR_RNG_UPPER: + elif elem.type == GrammarElementType.CHAR_RNG_UPPER: if i == 0 or not is_char_element(rule[i - 1]): - raise RuntimeError( - "LLAMA_GRETYPE_CHAR_RNG_UPPER without preceding char: " - + str(rule_id) - + "," - + str(i) - ) - print("-", file=file, end="") + raise ValueError(f"LLAMA_GRETYPE_CHAR_RNG_UPPER without preceding char: {rule_id}, {i}") + print(f"-", end="", file=file) print_grammar_char(file, elem.value) - elif case is llama_gretype.LLAMA_GRETYPE_CHAR_ALT: + elif elem.type == GrammarElementType.CHAR_ALT: if i == 0 or not is_char_element(rule[i - 1]): - raise RuntimeError( - "LLAMA_GRETYPE_CHAR_ALT without preceding char: " - + str(rule_id) - + "," - + str(i) - ) + raise ValueError(f"LLAMA_GRETYPE_CHAR_ALT without preceding char: {rule_id}, {i}") print_grammar_char(file, elem.value) - # if (is_char_element(elem)) { - # switch (rule[i + 1].type) { - # case LLAMA_GRETYPE_CHAR_ALT: - # case LLAMA_GRETYPE_CHAR_RNG_UPPER: - # break; - # default: - # fprintf(file, "] "); + elif elem.type == GrammarElementType.CHAR_ANY: + print(".", end="", file=file) if is_char_element(elem): - if rule[i + 1].type in ( - llama_gretype.LLAMA_GRETYPE_CHAR_ALT, - llama_gretype.LLAMA_GRETYPE_CHAR_RNG_UPPER, - ): - pass - else: - print("] ", file=file, end="") - # } - # } - # } - # fprintf(file, "\n"); - # } + if rule[i + 1].type in (GrammarElementType.CHAR_ALT, GrammarElementType.CHAR_RNG_UPPER, GrammarElementType.CHAR_ANY): + continue + print("] ", end="", file=file) print(file=file) -# void print_grammar(FILE * file, const parse_state & state) { -# try { -# std::map symbol_id_names; -# for (auto kv : state.symbol_ids) { -# symbol_id_names[kv.second] = kv.first; -# } -# for (size_t i = 0, end = state.rules.size(); i < end; i++) { -# // fprintf(file, "%zu: ", i); -# // print_rule_binary(file, state.rules[i]); -# print_rule(file, i, state.rules[i], symbol_id_names); -# // fprintf(file, "\n"); -# } -# } catch (const std::exception & err) { -# fprintf(stderr, "\n%s: error printing grammar: %s\n", __func__, err.what()); -# } -# } -def print_grammar(file: TextIO, state: parse_state) -> None: +def print_grammar(file: typing.TextIO, state: ParseState) -> None: try: - symbol_id_names = std.map() # type: std.map[int, str] - for kv in state.symbol_ids.items(): - symbol_id_names[kv[1]] = kv[0] - + symbol_id_names = {v: k for k, v in state.symbol_ids.items()} for i, rule in enumerate(state.rules): print_rule(file, i, rule, symbol_id_names) except Exception as err: - print( - f"{print_grammar.__name__}: error printing grammar: {err}", - file=sys.stderr, + print(f"\nerror printing grammar: {err}", file=file) + raise err + + +class LlamaGrammar: + def __init__(self, parse_state: ParseState): + self.parse_state = parse_state + + self._grammar_rules = parse_state.rules + self._n_rules = len(self._grammar_rules) + self._start_rule_index = parse_state.symbol_ids["root"] + + self._element_lists = [ + [ + llama_cpp.llama_grammar_element(ctypes.c_int(elem.type), ctypes.c_uint32(elem.value)) + for elem in subvector + ] + for subvector in self._grammar_rules + ] + + # Step 2: Convert each list to llama_grammar_element array and get pointer + self._element_arrays = [ + (llama_cpp.llama_grammar_element * len(sublist))(*sublist) + for sublist in self._element_lists + ] + + # Step 3: Get pointer of each array + self._element_array_pointers = [ + ctypes.cast(subarray, llama_cpp.llama_grammar_element_p) for subarray in self._element_arrays + ] + + # Step 4: Make array of these pointers and get its pointer + self._rules = (llama_cpp.llama_grammar_element_p * len(self._element_array_pointers))( + *self._element_array_pointers + ) + + self.grammar = None + self._init_grammar() + + + def _init_grammar(self): + grammar = llama_cpp.llama_grammar_init( + self._rules, ctypes.c_size_t(self._n_rules), ctypes.c_size_t(self._start_rule_index) + ) + + if grammar is None: + raise ValueError("Failed to create grammar") + + self.grammar = grammar + + def __del__(self): + if self.grammar is not None: + llama_cpp.llama_grammar_free(self.grammar) + self.grammar = None + + def reset(self): + if self.grammar is not None: + llama_cpp.llama_grammar_free(self.grammar) + self._init_grammar() + + @classmethod + def from_string(cls, grammar: str, verbose: bool = True) -> "LlamaGrammar": + parsed_grammar = parse(grammar) + if verbose: + print_grammar(file=sys.stdout, state=parsed_grammar) + return cls(parsed_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""" @@ -1367,12 +1086,13 @@ def print_grammar(file: TextIO, state: parse_state) -> None: string ::= "\"" ( [^"\\\x7F\x00-\x1F] | - "\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F]) # escapes + "\\" (["\\bfnrt] | "u" [0-9a-fA-F]{4}) # escapes )* "\"" ws -number ::= ("-"? ([0-9] | [1-9] [0-9]*)) ("." [0-9]+)? ([eE] [-+]? [0-9]+)? ws +number ::= ("-"? ([0-9] | [1-9] [0-9]{0,15})) ("." [0-9]+)? ([eE] [-+]? [0-9] [1-9]{0,15})? ws -ws ::= ([ \t\n] ws)? +# Optional space: by convention, applied in this grammar after literal chars when allowed +ws ::= | " " | "\n" [ \t]{0,20} """ LIST_GBNF = r""" @@ -1391,145 +1111,713 @@ def print_grammar(file: TextIO, state: parse_state) -> None: # whitespace. Also maybe improves generation quality? SPACE_RULE = '" "?' -PRIMITIVE_RULES = { - "boolean": '("true" | "false") space', - "number": '("-"? ([0-9] | [1-9] [0-9]*)) ("." [0-9]+)? ([eE] [-+]? [0-9]+)? space', - "integer": '("-"? ([0-9] | [1-9] [0-9]*)) space', - "string": r""" "\"" ( - [^"\\] | - "\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F]) - )* "\"" space """, - "null": '"null" space', -} 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): + def __init__(self, *, prop_order, allow_fetch, dotall, raw_pattern): self._prop_order = prop_order - self._rules = {"space": SPACE_RULE} - self._defs: Dict[str, Any] = {} - - def _format_literal(self, literal: str): - escaped: str = GRAMMAR_LITERAL_ESCAPE_RE.sub( - lambda m: GRAMMAR_LITERAL_ESCAPES.get(m.group(0)), json.dumps(literal) + 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 _add_rule(self, name: str, rule: str): + 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: + 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 visit(self, schema: Dict[str, Any], name: str) -> str: - rule_name = name or "root" + 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 - if "$defs" in schema: - # add defs to self._defs for later inlining - for def_name, def_schema in schema["$defs"].items(): - self._defs[def_name] = def_schema + return visit(schema) - if "oneOf" in schema or "anyOf" in schema: - rule = " | ".join( - ( - self.visit(alt_schema, f'{name}{"-" if name else ""}{i}') - for i, alt_schema in enumerate( - schema.get("oneOf") or schema["anyOf"] + 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]), ) - return self._add_rule(rule_name, rule) elif "const" in schema: - return self._add_rule(rule_name, self._format_literal(schema["const"])) + return self._add_rule( + rule_name, self._generate_constant_rule(schema["const"]) + ) elif "enum" in schema: - rule = " | ".join((self._format_literal(v) for v in schema["enum"])) + rule = " | ".join((self._generate_constant_rule(v) for v in schema["enum"])) return self._add_rule(rule_name, rule) - elif "$ref" in schema: - ref = schema["$ref"] - assert ref.startswith("#/$defs/"), f"Unrecognized schema: {schema}" - # inline $defs - def_name = ref[len("#/$defs/") :] - def_schema = self._defs[def_name] - return self.visit(def_schema, f'{name}{"-" if name else ""}{def_name}') - - - schema_type: Optional[str] = schema.get("type") # type: ignore - assert isinstance(schema_type, str), f"Unrecognized schema: {schema}" - - if schema_type == "object" and "properties" in schema: - # TODO: `required` keyword - if self._prop_order: - prop_order = self._prop_order - prop_pairs = sorted( - schema["properties"].items(), - # sort by position in prop_order (if specified) then by key - key=lambda kv: (prop_order.get(kv[0], len(prop_order)), kv[0]), + 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: - prop_pairs = schema["properties"].items() - - rule = '"{" space' - for i, (prop_name, prop_schema) in enumerate(prop_pairs): - prop_rule_name = self.visit( - prop_schema, f'{name}{"-" if name else ""}{prop_name}' + 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', ) - if i > 0: - rule += ' "," space' - rule += rf' {self._format_literal(prop_name)} space ":" space {prop_rule_name}' - rule += ' "}" space' - return self._add_rule(rule_name, rule) + elif schema_type in (None, "string") and "pattern" in schema: + return self._visit_pattern(schema["pattern"], rule_name) - elif schema_type == "array" and "items" in schema: - # TODO `prefixItems` keyword - item_rule_name = self.visit( - schema["items"], f'{name}{"-" if name else ""}item' + 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"]) ) - list_item_operator = f'("," space {item_rule_name})' - successive_items = "" - min_items = schema.get("minItems", 0) - if min_items > 0: - first_item = f"({item_rule_name})" - successive_items = list_item_operator * (min_items - 1) - min_items -= 1 - else: - first_item = f"({item_rule_name})?" - max_items = schema.get("maxItems") - if max_items is not None and max_items > min_items: - successive_items += (list_item_operator + "?") * (max_items - min_items - 1) - else: - successive_items += list_item_operator + "*" - rule = f'"[" space {first_item} {successive_items} "]" space' - return self._add_rule(rule_name, rule) else: assert schema_type in PRIMITIVE_RULES, f"Unrecognized schema: {schema}" - return self._add_rule( + # 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 self._rules.items())) + 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) + 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_tokenizer.py b/llama_cpp/llama_tokenizer.py index 5a8b13b7a..029bf2acc 100644 --- a/llama_cpp/llama_tokenizer.py +++ b/llama_cpp/llama_tokenizer.py @@ -17,7 +17,7 @@ def tokenize( self, text: bytes, add_bos: bool = True, special: bool = True ) -> List[int]: """Tokenize the text into tokens. - + Args: text: The text to tokenize. add_bos: Whether to add a beginning of sequence token. @@ -29,10 +29,11 @@ def detokenize( self, tokens: List[int], prev_tokens: Optional[List[int]] = None ) -> bytes: """Detokenize the tokens into text. - + Args: tokens: The tokens to detokenize. - prev_tokens: If tokens is a continuation of a previous sequence, the previous tokens.""" + prev_tokens: If tokens is a continuation of a previous sequence, the previous tokens. + """ raise NotImplementedError @@ -80,7 +81,9 @@ def detokenize( self, tokens: List[int], prev_tokens: Optional[List[int]] = None ) -> bytes: if prev_tokens is not None: - text = self.hf_tokenizer.decode(prev_tokens + tokens).encode("utf-8", errors="ignore") + text = self.hf_tokenizer.decode(prev_tokens + tokens).encode( + "utf-8", errors="ignore" + ) prev_text = self.hf_tokenizer.decode(prev_tokens).encode( "utf-8", errors="ignore" ) diff --git a/llama_cpp/llama_types.py b/llama_cpp/llama_types.py index 1b1befebe..bbb58afc3 100644 --- a/llama_cpp/llama_types.py +++ b/llama_cpp/llama_types.py @@ -6,6 +6,7 @@ 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 @@ -24,7 +25,7 @@ class EmbeddingUsage(TypedDict): class Embedding(TypedDict): index: int object: str - embedding: List[float] + embedding: Union[List[float], List[List[float]]] class CreateEmbeddingResponse(TypedDict): @@ -84,6 +85,7 @@ class ChatCompletionFunction(TypedDict): class ChatCompletionResponseChoice(TypedDict): index: int message: "ChatCompletionResponseMessage" + logprobs: Optional[CompletionLogprobs] finish_reason: Optional[str] @@ -155,7 +157,9 @@ class ChatCompletionFunctionCallOption(TypedDict): class ChatCompletionRequestResponseFormat(TypedDict): type: Literal["text", "json_object"] - schema: NotRequired[JsonType] # https://docs.endpoints.anyscale.com/guides/json_mode/ + schema: NotRequired[ + JsonType + ] # https://docs.endpoints.anyscale.com/guides/json_mode/ class ChatCompletionRequestMessageContentPartText(TypedDict): @@ -271,7 +275,7 @@ class ChatCompletionNamedToolChoice(TypedDict): ChatCompletionToolChoiceOption = Union[ - Literal["none", "auto"], ChatCompletionNamedToolChoice + Literal["none", "auto", "required"], ChatCompletionNamedToolChoice ] diff --git a/llama_cpp/llava_cpp.py b/llama_cpp/llava_cpp.py index 543c87d0b..b80d85913 100644 --- a/llama_cpp/llava_cpp.py +++ b/llama_cpp/llava_cpp.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import sys import os import ctypes @@ -14,14 +16,26 @@ Structure, ) import pathlib -from typing import List, Union, NewType, Optional, TypeVar, Callable, Any +from typing import ( + List, + Union, + NewType, + Optional, + TypeVar, + Callable, + Any, + TYPE_CHECKING, + Generic, +) +from typing_extensions import TypeAlias import llama_cpp.llama_cpp as llama_cpp + # Load the library def _load_shared_library(lib_base_name: str): # Construct the paths to the possible shared library names - _base_path = pathlib.Path(os.path.abspath(os.path.dirname(__file__))) + _base_path = pathlib.Path(os.path.abspath(os.path.dirname(__file__))) / "lib" # 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] = [] @@ -62,7 +76,7 @@ def _load_shared_library(lib_base_name: str): for _lib_path in _lib_paths: if _lib_path.exists(): try: - return ctypes.CDLL(str(_lib_path), **cdll_args) # type: ignore + 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}") @@ -79,8 +93,27 @@ def _load_shared_library(lib_base_name: str): # ctypes helper +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): def ctypes_function( name: str, argtypes: List[Any], restype: Any, enabled: bool = True @@ -111,6 +144,7 @@ def decorator(f: F) -> F: 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; @@ -121,36 +155,72 @@ class llava_image_embed(Structure): ("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: - ... +@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: bytes, image_bytes_length: Union[c_int, int], /) -> "_Pointer[llava_image_embed]": - ... +@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]": - ... +@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]", /): - ... +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: - ... +@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: ... ################################################ @@ -161,11 +231,12 @@ def llava_eval_image_embed(ctx_llama: llama_cpp.llama_context_p, embed: "_Pointe # /** 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]: - ... +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, /): - ... +def clip_free(ctx: clip_ctx_p, /): ... diff --git a/llama_cpp/server/__main__.py b/llama_cpp/server/__main__.py index fadfc5fb4..bbac4957e 100644 --- a/llama_cpp/server/__main__.py +++ b/llama_cpp/server/__main__.py @@ -21,6 +21,7 @@ Then visit http://localhost:8000/docs to see the interactive API docs. """ + from __future__ import annotations import os @@ -59,7 +60,18 @@ def main(): if not os.path.exists(config_file): raise ValueError(f"Config file {config_file} not found!") with open(config_file, "rb") as f: - config_file_settings = ConfigFileSettings.model_validate_json(f.read()) + # 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: diff --git a/llama_cpp/server/app.py b/llama_cpp/server/app.py index aa6afc112..cd3255176 100644 --- a/llama_cpp/server/app.py +++ b/llama_cpp/server/app.py @@ -2,6 +2,8 @@ import os import json +import typing +import contextlib from threading import Lock from functools import partial @@ -12,14 +14,7 @@ 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, - HTTPException, - status, -) +from fastapi import Depends, FastAPI, APIRouter, Request, HTTPException, status, Body from fastapi.middleware import Middleware from fastapi.middleware.cors import CORSMiddleware from fastapi.security import HTTPBearer @@ -94,6 +89,14 @@ def get_llama_proxy(): llama_outer_lock.release() +_ping_message_factory: typing.Optional[typing.Callable[[], bytes]] = None + + +def set_ping_message_factory(factory: typing.Callable[[], bytes]): + global _ping_message_factory + _ping_message_factory = factory + + def create_app( settings: Settings | None = None, server_settings: ServerSettings | None = None, @@ -104,7 +107,15 @@ def create_app( if not os.path.exists(config_file): raise ValueError(f"Config file {config_file} not found!") with open(config_file, "rb") as f: - config_file_settings = ConfigFileSettings.model_validate_json(f.read()) + # 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 @@ -124,6 +135,7 @@ def create_app( middleware=middleware, title="🦙 llama.cpp Python API", version=llama_cpp.__version__, + root_path=server_settings.root_path, ) app.add_middleware( CORSMiddleware, @@ -137,24 +149,29 @@ def create_app( assert model_settings is not None set_llama_proxy(model_settings=model_settings) + if server_settings.disable_ping_events: + set_ping_message_factory(lambda: bytes()) + return app async def get_event_publisher( request: Request, - inner_send_chan: MemoryObjectSendStream, - iterator: Iterator, + inner_send_chan: MemoryObjectSendStream[typing.Any], + iterator: Iterator[typing.Any], + on_complete: typing.Optional[typing.Callable[[], None]] = None, ): + server_settings = next(get_server_settings()) + interrupt_requests = ( + server_settings.interrupt_requests if server_settings else False + ) async with inner_send_chan: try: 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 ( - next(get_server_settings()).interrupt_requests - and llama_outer_lock.locked() - ): + 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]")) @@ -163,6 +180,9 @@ async def get_event_publisher( with anyio.move_on_after(1, shield=True): print(f"Disconnected from client (via refresh/close) {request.client}") raise e + finally: + if on_complete: + on_complete() def _logit_bias_tokens_to_input_ids( @@ -246,8 +266,16 @@ async def authenticate( async def create_completion( request: Request, body: CreateCompletionRequest, - llama_proxy: LlamaProxy = Depends(get_llama_proxy), ) -> llama_cpp.Completion: + exit_stack = contextlib.ExitStack() + llama_proxy = await run_in_threadpool( + lambda: exit_stack.enter_context(contextlib.contextmanager(get_llama_proxy)()) + ) + if llama_proxy is None: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail="Service is not available", + ) if isinstance(body.prompt, list): assert len(body.prompt) <= 1 body.prompt = body.prompt[0] if len(body.prompt) > 0 else "" @@ -263,6 +291,7 @@ async def create_completion( "best_of", "logit_bias_type", "user", + "min_tokens", } kwargs = body.model_dump(exclude=exclude) @@ -276,6 +305,15 @@ async def create_completion( if body.grammar is not None: kwargs["grammar"] = llama_cpp.LlamaGrammar.from_string(body.grammar) + 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) + iterator_or_completion: Union[ llama_cpp.CreateCompletionResponse, Iterator[llama_cpp.CreateCompletionStreamResponse], @@ -290,6 +328,7 @@ async def create_completion( def iterator() -> Iterator[llama_cpp.CreateCompletionStreamResponse]: yield first_response yield from iterator_or_completion + exit_stack.close() send_chan, recv_chan = anyio.create_memory_object_stream(10) return EventSourceResponse( @@ -299,8 +338,10 @@ def iterator() -> Iterator[llama_cpp.CreateCompletionStreamResponse]: request=request, inner_send_chan=send_chan, iterator=iterator(), + on_complete=exit_stack.close, ), sep="\n", + ping_message_factory=_ping_message_factory, ) else: return iterator_or_completion @@ -356,13 +397,95 @@ async def create_embedding( ) async def create_chat_completion( request: Request, - body: CreateChatCompletionRequest, - llama_proxy: LlamaProxy = Depends(get_llama_proxy), + 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 + exit_stack = contextlib.ExitStack() + llama_proxy = await run_in_threadpool( + lambda: exit_stack.enter_context(contextlib.contextmanager(get_llama_proxy)()) + ) + if llama_proxy is None: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail="Service is not available", + ) exclude = { "n", "logit_bias_type", "user", + "min_tokens", } kwargs = body.model_dump(exclude=exclude) llama = llama_proxy(body.model) @@ -376,6 +499,15 @@ async def create_chat_completion( if body.grammar is not None: kwargs["grammar"] = llama_cpp.LlamaGrammar.from_string(body.grammar) + 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) + iterator_or_completion: Union[ llama_cpp.ChatCompletion, Iterator[llama_cpp.ChatCompletionChunk] ] = await run_in_threadpool(llama.create_chat_completion, **kwargs) @@ -389,6 +521,7 @@ async def create_chat_completion( def iterator() -> Iterator[llama_cpp.ChatCompletionChunk]: yield first_response yield from iterator_or_completion + exit_stack.close() send_chan, recv_chan = anyio.create_memory_object_stream(10) return EventSourceResponse( @@ -398,10 +531,13 @@ def iterator() -> Iterator[llama_cpp.ChatCompletionChunk]: request=request, inner_send_chan=send_chan, iterator=iterator(), + on_complete=exit_stack.close, ), sep="\n", + ping_message_factory=_ping_message_factory, ) else: + exit_stack.close() return iterator_or_completion @@ -443,7 +579,7 @@ async def tokenize( ) -> TokenizeInputResponse: tokens = llama_proxy(body.model).tokenize(body.input.encode("utf-8"), special=True) - return {"tokens": tokens} + return TokenizeInputResponse(tokens=tokens) @router.post( @@ -458,7 +594,7 @@ async def count_query_tokens( ) -> TokenizeInputCountResponse: tokens = llama_proxy(body.model).tokenize(body.input.encode("utf-8"), special=True) - return {"count": len(tokens)} + return TokenizeInputCountResponse(count=len(tokens)) @router.post( @@ -473,4 +609,4 @@ async def detokenize( ) -> DetokenizeInputResponse: text = llama_proxy(body.model).detokenize(body.tokens).decode("utf-8") - return {"text": text} + return DetokenizeInputResponse(text=text) diff --git a/llama_cpp/server/model.py b/llama_cpp/server/model.py index dace8d547..071a18b67 100644 --- a/llama_cpp/server/model.py +++ b/llama_cpp/server/model.py @@ -44,6 +44,8 @@ def __call__(self, model: Optional[str] = None) -> llama_cpp.Llama: 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] @@ -65,6 +67,7 @@ def __iter__(self): def free(self): if self._current_model: + self._current_model.close() del self._current_model @staticmethod @@ -72,9 +75,88 @@ 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" - chat_handler = llama_cpp.llama_chat_format.Llava15ChatHandler( - clip_model_path=settings.clip_model_path, verbose=settings.verbose - ) + 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 == "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 == "hf-autotokenizer": assert ( settings.hf_pretrained_model_name_or_path is not None @@ -101,10 +183,11 @@ def load_llama_from_model_settings(settings: ModelSettings) -> llama_cpp.Llama: draft_model = None if settings.draft_model is not None: draft_model = llama_speculative.LlamaPromptLookupDecoding( - num_pred_tokens=settings.draft_model_num_pred_tokens + num_pred_tokens=settings.draft_model_num_pred_tokens, + max_ngram_size=settings.draft_model_max_ngram_size, ) - kv_overrides: Optional[Dict[str, Union[bool, int, float]]] = None + 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 = {} @@ -118,6 +201,8 @@ def load_llama_from_model_settings(settings: ModelSettings) -> llama_cpp.Llama: 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}") @@ -139,12 +224,14 @@ def load_llama_from_model_settings(settings: ModelSettings) -> llama_cpp.Llama: **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, @@ -163,6 +250,7 @@ def load_llama_from_model_settings(settings: ModelSettings) -> llama_cpp.Llama: 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 @@ -175,6 +263,9 @@ def load_llama_from_model_settings(settings: ModelSettings) -> llama_cpp.Llama: 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 @@ -184,7 +275,22 @@ def load_llama_from_model_settings(settings: ModelSettings) -> llama_cpp.Llama: 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) + if settings.cache_dir: + cache = llama_cpp.LlamaDiskCache( + capacity_bytes=settings.cache_size, cache_dir=settings.cache_dir + ) + else: + cache = llama_cpp.LlamaDiskCache(capacity_bytes=settings.cache_size) + elif settings.cache_type == "static_disk": + if settings.verbose: + print(f"Using static disk cache with size {settings.cache_size}") + + if not settings.cache_dir: + raise ValueError("cache_dir must be set for static_disk cache!") + + cache = llama_cpp.LlamaStaticDiskCache( + cache_dir=settings.cache_dir, capacity_bytes=settings.cache_size + ) else: if settings.verbose: print(f"Using ram cache with size {settings.cache_size}") diff --git a/llama_cpp/server/settings.py b/llama_cpp/server/settings.py index daa913fac..253fb3c7c 100644 --- a/llama_cpp/server/settings.py +++ b/llama_cpp/server/settings.py @@ -2,8 +2,10 @@ import multiprocessing -from typing import Optional, List, Literal, Union -from pydantic import Field +from typing import 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 @@ -56,6 +58,10 @@ class ModelSettings(BaseSettings): 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 seperated list of rpc servers for offloading", + ) # Context Params seed: int = Field( default=llama_cpp.LLAMA_DEFAULT_SEED, description="Random seed. -1 for random." @@ -67,12 +73,12 @@ class ModelSettings(BaseSettings): n_threads: int = Field( default=max(multiprocessing.cpu_count() // 2, 1), ge=1, - description="The number of threads to use.", + description="The number of threads to use. Use -1 for max cpu threads", ) n_threads_batch: int = Field( - default=max(multiprocessing.cpu_count() // 2, 1), + default=max(multiprocessing.cpu_count(), 1), ge=0, - description="The number of threads to use when batch processing.", + 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 @@ -90,10 +96,13 @@ class ModelSettings(BaseSettings): 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=True, description="Whether to use embeddings.") + 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, @@ -128,7 +137,7 @@ class ModelSettings(BaseSettings): default=False, description="Use a cache to reduce processing times for evaluated prompts.", ) - cache_type: Literal["ram", "disk"] = Field( + cache_type: Literal["ram", "disk", "static_disk"] = Field( default="ram", description="The type of cache to use. Only used if cache is True.", ) @@ -136,6 +145,9 @@ class ModelSettings(BaseSettings): default=2 << 30, description="The size of the cache in bytes. Only used if cache is True.", ) + cache_dir: Optional[str] = Field( + default=None, description="Directory to use for disk cache." + ) # Tokenizer Options hf_tokenizer_config_path: Optional[str] = Field( default=None, @@ -159,11 +171,38 @@ class ModelSettings(BaseSettings): default=10, description="Number of tokens to predict using the draft model.", ) + + draft_model_max_ngram_size: int = Field( + default=2, description="Maximum ngram size to use for 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.""" @@ -186,6 +225,14 @@ class ServerSettings(BaseSettings): 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): diff --git a/llama_cpp/server/types.py b/llama_cpp/server/types.py index c8b2ebc6c..fdd164456 100644 --- a/llama_cpp/server/types.py +++ b/llama_cpp/server/types.py @@ -16,10 +16,14 @@ 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, - 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.", ) @@ -113,6 +117,7 @@ class CreateCompletionRequest(BaseModel): 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 @@ -130,7 +135,6 @@ class CreateCompletionRequest(BaseModel): 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) seed: Optional[int] = Field(None) # ignored or currently unsupported @@ -209,6 +213,16 @@ class CreateChatCompletionRequest(BaseModel): 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 @@ -268,7 +282,7 @@ class ModelList(TypedDict): class TokenizeInputRequest(BaseModel): model: Optional[str] = model_field - input: Optional[str] = Field(description="The input to tokenize.") + input: str = Field(description="The input to tokenize.") model_config = { "json_schema_extra": {"examples": [{"input": "How many tokens in this query?"}]} diff --git a/pyproject.toml b/pyproject.toml index 2f3d3ced0..2fd61f641 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = ["scikit-build-core[pyproject]>=0.5.1"] +requires = ["scikit-build-core[pyproject]>=0.9.2"] build-backend = "scikit_build_core.build" [project] @@ -8,14 +8,13 @@ dynamic = ["version"] description = "Python bindings for the llama.cpp library" readme = "README.md" license = { text = "MIT" } -authors = [ - { name = "Andrei Betlen", email = "abetlen@gmail.com" }, -] +authors = [{ name = "Andrei Betlen", email = "abetlen@gmail.com" }] dependencies = [ "typing-extensions>=4.5.0", - "numpy>=1.20.0", + "numpy==1.26.4", "diskcache>=5.6.1", "jinja2>=2.11.3", + "PyTrie>=0.4.0", ] requires-python = ">=3.8" classifiers = [ @@ -35,11 +34,16 @@ server = [ "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", ] dev = [ "black>=23.3.0", @@ -49,17 +53,19 @@ dev = [ "mkdocs-material>=9.1.18", "pytest>=7.4.0", "httpx>=0.24.1", + "pandas>=2.2.1", + "tqdm>=4.66.2", + "pyinstaller>=6.11.1", ] -all = [ - "llama_cpp_python[server,test,dev]", -] +all = ["llama_cpp_python[server,test,dev]"] [tool.scikit-build] wheel.packages = ["llama_cpp"] cmake.verbose = true cmake.minimum-version = "3.21" minimum-version = "0.5.1" -sdist.include = [".git", "vendor/llama.cpp/.git"] +sdist.include = [".git", "vendor/llama.cpp/*"] +logging.level = "INFO" [tool.scikit-build.metadata.version] provider = "scikit_build_core.metadata.regex" diff --git a/scripts/releases-to-pep-503.sh b/scripts/releases-to-pep-503.sh new file mode 100755 index 000000000..44fbbf3cf --- /dev/null +++ b/scripts/releases-to-pep-503.sh @@ -0,0 +1,60 @@ +#!/bin/bash + +# Get output directory or default to index/whl/cpu +output_dir=${1:-"index/whl/cpu"} + +# Create output directory +mkdir -p $output_dir + +# Change to output directory +pushd $output_dir + +# Create an index html file +echo "" > index.html +echo "" >> index.html +echo " " >> index.html +echo " " >> index.html +echo " llama-cpp-python" >> index.html +echo "
" >> index.html +echo " " >> index.html +echo "" >> index.html +echo "" >> index.html + +# Create llama-cpp-python directory +mkdir -p llama-cpp-python + +# Change to llama-cpp-python directory +pushd llama-cpp-python + +# Create an index html file +echo "" > index.html +echo "" >> index.html +echo " " >> index.html +echo "

Links for llama-cpp-python

" >> index.html + +# Get all releases +releases=$(curl -s https://api.github.com/repos/abetlen/llama-cpp-python/releases | jq -r .[].tag_name) + +# Get pattern from second arg or default to valid python package version pattern +pattern=${2:-"^[v]?[0-9]+\.[0-9]+\.[0-9]+$"} + +# Filter releases by pattern +releases=$(echo $releases | tr ' ' '\n' | grep -E $pattern) + +# For each release, get all assets +for release in $releases; do + assets=$(curl -s https://api.github.com/repos/abetlen/llama-cpp-python/releases/tags/$release | jq -r .assets) + # 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 "

$release_version

" >> index.html + for asset in $(echo $assets | jq -r .[].browser_download_url); do + if [[ $asset == *".whl" ]]; then + echo " $asset" >> index.html + echo "
" >> index.html + fi + done +done + +echo " " >> index.html +echo "" >> index.html +echo "" >> index.html diff --git a/tests/test_llama.py b/tests/test_llama.py index fa2f6dfc1..469ef91ca 100644 --- a/tests/test_llama.py +++ b/tests/test_llama.py @@ -6,7 +6,7 @@ import llama_cpp -MODEL = "./vendor/llama.cpp/models/ggml-vocab-llama.gguf" +MODEL = "./vendor/llama.cpp/models/ggml-vocab-llama-spm.gguf" def test_llama_cpp_tokenization(): diff --git a/tests/test_llama_cache.py b/tests/test_llama_cache.py new file mode 100644 index 000000000..ab7afb3e5 --- /dev/null +++ b/tests/test_llama_cache.py @@ -0,0 +1,241 @@ +import os +import tempfile + +import pytest + +from llama_cpp.llama import Llama, LlamaState +from llama_cpp.llama_cache import LlamaStaticDiskCache, StateReloadError + + +# Have to be careful to reset to good state when testing, but don't want to +# recreate model each time. +@pytest.fixture(scope="module") +def small_model(): + model_filename = os.getenv("LLAMA_TEST_MODEL") + if not model_filename: + pytest.skip("LLAMA_TEST_MODEL environment variable is not set") + return + + model_filename = os.path.expanduser(model_filename) + + test_model = Llama( + model_filename, + n_ctx=2_048, + n_gpu_layers=0, + offload_kqv=False, + n_batch=512, + embedding=False, + verbose=False, + ) + + system_prompt = r""" +You are an advanced intelligence "Hal" aboard a spaceship. You are required to +act as the primary interface between the ship and its crew. You can: +* Provide information on the current status of the ship +* Turn on/off the lights in the crew quarters +* Open/close the airlocks + +Respond in a terse, professional manner. Do not engage in casual conversation. + +The current state of the ship is: +* Airlocks: closed +* Lights: on +* Oxygen levels: normal +""".strip() + + user_prompt = "Hal, please open the airlocks." + + # Ingest prompt and create completion so that will have some state. + # Last token of prompt + all tokens of generated completion will have + # non-zero logits. + _ = test_model.create_chat_completion( + [ + {"role": "system", "text": system_prompt}, + {"role": "user", "text": user_prompt}, + ], + seed=1234, + ) + + assert test_model.n_tokens > 0 + + # Have at least some scores, and last entry is non-zero + assert ~(test_model.scores == 0).all() + # pylint: disable=protected-access + assert (test_model._scores[-1, :] != 0.0).all() + + return test_model + + +@pytest.fixture(scope="module") +def llama_state(small_model) -> LlamaState: + state = small_model.save_state() + # Clear scores so that can test reloading from cache without them. + state.scores = None + return state + + +def test_reload_from_cache_state_success(small_model, llama_state: LlamaState): + current_state = small_model.save_state() + old_score = small_model.scores.copy() + + LlamaStaticDiskCache.reload_from_cache_state(small_model, llama_state) + new_state = small_model.save_state() + new_score = small_model.scores.copy() + + assert (current_state.input_ids == new_state.input_ids).all() + + assert current_state.n_tokens == new_state.n_tokens + + # Logits for last token should match, others may not if n_batch < n_tokens + assert ( + old_score[small_model.n_tokens - 1, :] == new_score[small_model.n_tokens - 1, :] + ).all() + + +def test_reload_from_cache_state_state_reload_error(small_model, llama_state): + small_model.context_params.logits_all = True + small_model.context_params.embeddings = True + try: + with pytest.raises(StateReloadError): + LlamaStaticDiskCache.reload_from_cache_state(small_model, llama_state) + finally: + small_model.context_params.logits_all = False + small_model.context_params.embeddings = False + + +def test_disk_cache_e2e(small_model: Llama): + prompts = ["this is a test prompt", "and this is another test prompt"] + capacity_bytes = 2 << 30 + + small_model.reset() + # This is a weird thing to reset, but input_ids > n_tokens are not + # significant (like a scratchpad), left over if had previous prompt that + # was longer. + # + # Reset for ease of comparison later. + small_model.input_ids[:] = 0 + + with tempfile.TemporaryDirectory() as cache_dir: + cache = LlamaStaticDiskCache.build_cache( + cache_dir=cache_dir, + prompts=prompts, + model=small_model, + capacity_bytes=capacity_bytes, + add_bos=True, + seed=1234, + save_logits=False, + ) + + for p in prompts: + key = tuple( + small_model.tokenize(p.encode("utf-8"), add_bos=True, special=True) + ) + assert key in cache + state = cache[key] + assert ~(state.input_ids == 0).all() + assert state is not None + assert ( + state.scores is None + ), "Logits should not be stored when save_logits=False and model doesn't require them." + + small_model.reset() + small_model.input_ids[:] = 0 + small_model.eval(key) + + state2 = small_model.save_state() + assert state2.n_tokens == state.n_tokens + assert ~(state2.input_ids == 0).all() + assert (state2.input_ids == state.input_ids).all() + + last_logits = small_model.scores[small_model.n_tokens - 1, :] + + LlamaStaticDiskCache.reload_from_cache_state(small_model, state) + + last_logits2 = small_model.scores[small_model.n_tokens - 1, :] + + assert (last_logits == last_logits2).all() + + +def test_cache_save_reload_scores_when_needed( + small_model: Llama, +): + """ + When model requires it, can reload from state with scores. + """ + test_prompt = "this is a test prompt" + with tempfile.TemporaryDirectory() as cache_dir: + cache = LlamaStaticDiskCache.build_cache( + cache_dir=cache_dir, + prompts=[test_prompt], + model=small_model, + capacity_bytes=2 << 30, + add_bos=True, + seed=1234, + save_logits=True, + ) + + llama_state = small_model.save_state() + cur_scores = llama_state.scores.copy() + assert ~(cur_scores == 0.0).all() + + try: + small_model.context_params.logits_all = True + state_from_cache = cache[ + tuple(llama_state.input_ids[: llama_state.n_tokens].tolist()) + ] + assert state_from_cache.scores is not None, "Scores should be saved." + LlamaStaticDiskCache.reload_from_cache_state(small_model, state_from_cache) + # Do I have to limit these to n_tokens? + assert (state_from_cache.input_ids == llama_state.input_ids).all() + assert ( + cur_scores == small_model.scores[: small_model.n_tokens] + ).all(), "Reloaded scores should match" + finally: + small_model.scores[:] = 0.0 + small_model.context_params.logits_all = False + small_model.reset() + + +def test_cache_reload_errors_when_requires_scores_and_state_doesnt_have_it( + small_model: Llama, llama_state: LlamaState +): + """ + If model requires logits for sampling and state doesn't have it, should raise error. + """ + old_state_scores = ( + llama_state.scores.copy() + if llama_state.scores is not None + else llama_state.scores + ) + try: + small_model.context_params.logits_all = True + llama_state.scores = None + + with pytest.raises(StateReloadError): + LlamaStaticDiskCache.reload_from_cache_state(small_model, llama_state) + finally: + small_model.context_params.logits_all = False + llama_state.scores = old_state_scores + + +# pylint: disable=invalid-name +def test_cache_errors_when_save_logits_False_but_model_requires(small_model: Llama): + """ + If model requires logits but save_logits is False, should raise error. + """ + + try: + small_model.context_params.logits_all = True + with pytest.raises(ValueError): + with tempfile.TemporaryDirectory() as cache_dir: + LlamaStaticDiskCache.build_cache( + cache_dir=cache_dir, + prompts=["this is a test prompt"], + model=small_model, + capacity_bytes=2 << 30, + add_bos=True, + seed=1234, + save_logits=False, + ) + finally: + small_model.context_params.logits_all = False diff --git a/tests/test_llama_chat_format.py b/tests/test_llama_chat_format.py index c10aee42e..f031bf72b 100644 --- a/tests/test_llama_chat_format.py +++ b/tests/test_llama_chat_format.py @@ -21,12 +21,13 @@ def test_mistral_instruct(): response = llama_chat_format.format_mistral_instruct( messages=messages, ) + prompt = ("" if response.added_special else "") + response.prompt reference = chat_formatter.render( messages=messages, bos_token="", eos_token="", ) - assert response.prompt == reference + assert prompt == reference mistral_7b_tokenizer_config = """{ diff --git a/vendor/llama.cpp b/vendor/llama.cpp index 4e9a7f7f7..11ceab88c 160000 --- a/vendor/llama.cpp +++ b/vendor/llama.cpp @@ -1 +1 @@ -Subproject commit 4e9a7f7f7fb6acbddd1462909c8d696e38edbfcc +Subproject commit 11ceab88c92bc1742268e87eb8057a10fea80ad1