initial commit

.devcontainer/devcontainer.json

@@ -0,0 +1,32 @@
+// For format details, see https://aka.ms/devcontainer.json. For config options, see the
+// README at: https://github.com/devcontainers/templates/tree/main/src/python
+{
+    "name": "linalg-zero",
+    // Or use a Dockerfile or Docker Compose file. More info: https://containers.dev/guide/dockerfile
+    "image": "mcr.microsoft.com/devcontainers/python:1-3.11-bullseye",
+    "runArgs": [
+        // avoid UID/GID remapping under rootless Podman
+        // This flag works with Podman but not with Docker:
+        // "--userns=keep-id",
+        "-v", "${localEnv:HOME}/.ssh:/home/vscode/.ssh:ro",
+        "--user", "1000:1000"
+    ],
+    "features": {},
+
+    // Use 'postCreateCommand' to run commands after the container is created.
+    "postCreateCommand": "./.devcontainer/postCreateCommand.sh",
+
+    // Configure tool-specific properties.
+    "customizations": {
+        "vscode": {
+            "extensions": ["ms-python.python", "editorconfig.editorconfig"],
+            "settings": {
+                "python.testing.pytestArgs": ["tests"],
+                "python.testing.unittestEnabled": false,
+                "python.testing.pytestEnabled": true,
+                "python.defaultInterpreterPath": "/workspaces/linalg-zero/.venv/bin/python",
+                "python.testing.pytestPath": "/workspaces/linalg-zero/.venv/bin/pytest"
+            }
+        }
+    }
+}

.devcontainer/postCreateCommand.sh

@@ -0,0 +1,10 @@
+#! /usr/bin/env bash
+
+# Install uv
+curl -LsSf https://astral.sh/uv/install.sh | sh
+
+# Install Dependencies
+uv sync
+
+# Install pre-commit hooks
+uv run pre-commit install --install-hooks

.gemini/config.yaml

@@ -0,0 +1,10 @@
+have_fun: true
+code_review:
+  disable: false
+  comment_severity_threshold: MEDIUM
+  max_review_comments: -1
+  pull_request_opened:
+    help: false
+    summary: true
+    code_review: true
+ignore_patterns: []

.gemini/styleguide.md

@@ -0,0 +1,144 @@
+# Project LinAlgZero Python Style Guide
+
+# Introduction
+This style guide outlines the coding conventions for Python code.
+It's based on PEP 8, but with some modifications to address specific needs and
+preferences within our organization.
+
+# Key Principles
+* **Readability:** Code should be easy to understand for all team members.
+* **Maintainability:** Code should be easy to modify and extend.
+* **Consistency:** Adhering to a consistent style across all projects improves
+  collaboration and reduces errors.
+* **Performance:** While readability is paramount, code should be efficient.
+
+# Deviations from PEP 8
+
+## Line Length
+* **Maximum line length:** 119 characters (instead of PEP 8's 79).
+    * Modern screens allow for wider lines, improving code readability in many cases.
+    * Many common patterns in our codebase, like long strings or URLs, often exceed 79 characters.
+
+## Indentation
+* **Use 4 spaces per indentation level.** (PEP 8 recommendation)
+
+## Imports
+* **Group imports:**
+    * Standard library imports
+    * Related third party imports
+    * Local application/library specific imports
+* **Absolute imports:** Always use absolute imports for clarity.
+* **Import order within groups:**  Sort alphabetically.
+
+## Naming Conventions
+
+* **Variables:** Use lowercase with underscores (snake_case): `user_name`, `total_count`
+* **Constants:**  Use uppercase with underscores: `MAX_VALUE`, `DATABASE_NAME`
+* **Functions:** Use lowercase with underscores (snake_case): `calculate_total()`, `process_data()`
+* **Classes:** Use CapWords (CamelCase): `UserManager`, `PaymentProcessor`
+* **Modules:** Use lowercase with underscores (snake_case): `user_utils`, `payment_gateway`
+
+## Docstrings
+* **Use triple double quotes (`"""Docstring goes here."""`) for all docstrings.**
+* **First line:** Concise summary of the object's purpose.
+* **For complex functions/classes:** Include detailed descriptions of parameters, return values,
+  attributes, and exceptions.
+* **Use Google style docstrings:** This helps with automated documentation generation.
+    ```python
+    def my_function(param1, param2):
+        """Single-line summary.
+
+        More detailed description, if necessary.
+
+        Args:
+            param1 (int): The first parameter.
+            param2 (str): The second parameter.
+
+        Returns:
+            bool: The return value. True for success, False otherwise.
+
+        Raises:
+            ValueError: If `param2` is invalid.
+        """
+        # function body here
+    ```
+
+## Type Hints
+* **Use type hints:**  Type hints improve code readability and help catch errors early.
+* **Follow PEP 484:**  Use the standard type hinting syntax.
+
+## Comments
+* **Write clear and concise comments:** Explain the "why" behind the code, not just the "what".
+* **Comment sparingly:** Well-written code should be self-documenting where possible.
+* **Use complete sentences:** Start comments with a capital letter and use proper punctuation.
+
+## Logging
+* **Use a standard logging framework:**  Project LinAlgZero uses the built-in `logging` module.
+* **Log at appropriate levels:** DEBUG, INFO, WARNING, ERROR, CRITICAL
+* **Provide context:** Include relevant information in log messages to aid debugging.
+
+## Error Handling
+* **Use specific exceptions:** Avoid using broad exceptions like `Exception`.
+* **Handle exceptions gracefully:** Provide informative error messages and avoid crashing the program.
+* **Use `try...except` blocks:**  Isolate code that might raise exceptions.
+
+# Tooling
+* **Code formatter:**  Ruff (ruff-format) - Enforces consistent formatting automatically.
+* **Linter:**  Ruff - Identifies potential issues and style violations.
+
+# Example
+```python
+"""Module for user authentication."""
+
+import hashlib
+import logging
+import os
+
+from companyx.db import user_database
+
+LOGGER = logging.getLogger(__name__)
+
+def hash_password(password: str) -> str:
+  """Hashes a password using SHA-256.
+
+  Args:
+      password (str): The password to hash.
+
+  Returns:
+      str: The hashed password.
+  """
+  salt = os.urandom(16)
+  salted_password = salt + password.encode('utf-8')
+  hashed_password = hashlib.sha256(salted_password).hexdigest()
+  return f"{salt.hex()}:{hashed_password}"
+
+def authenticate_user(username: str, password: str) -> bool:
+  """Authenticates a user against the database.
+
+  Args:
+      username (str): The user's username.
+      password (str): The user's password.
+
+  Returns:
+      bool: True if the user is authenticated, False otherwise.
+  """
+  try:
+      user = user_database.get_user(username)
+      if user is None:
+          LOGGER.warning("Authentication failed: User not found - %s", username)
+          return False
+
+      stored_hash = user.password_hash
+      salt, hashed_password = stored_hash.split(':')
+      salted_password = bytes.fromhex(salt) + password.encode('utf-8')
+      calculated_hash = hashlib.sha256(salted_password).hexdigest()
+
+      if calculated_hash == hashed_password:
+          LOGGER.info("User authenticated successfully - %s", username)
+          return True
+      else:
+          LOGGER.warning("Authentication failed: Incorrect password - %s", username)
+          return False
+  except Exception as e:
+      LOGGER.error("An error occurred during authentication: %s", e)
+      return False

.github/actions/setup-python-env/action.yml

@@ -0,0 +1,41 @@
+name: "Setup Python Environment"
+description: "Set up Python environment for the given Python version"
+
+inputs:
+  python-version:
+    description: "Python version to use"
+    required: true
+    default: "3.12"
+  uv-version:
+    description: "uv version to use"
+    required: true
+    default: "0.6.14"
+  dependency-groups:
+    description: "Space-separated dependency groups to install (e.g. 'dev test')"
+    required: false
+    default: "dev"
+
+runs:
+  using: "composite"
+  steps:
+    - uses: actions/setup-python@v5
+      with:
+        python-version: ${{ inputs.python-version }}
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v2
+      with:
+        version: ${{ inputs.uv-version }}
+        enable-cache: 'true'
+        cache-suffix: ${{ matrix.python-version }}
+
+    - name: Install Python dependencies
+      run: |
+        args=()
+        if [ -n "${{ inputs.dependency-groups }}" ]; then
+          for group in ${{ inputs.dependency-groups }}; do
+            args+=(--group "$group")
+          done
+        fi
+        uv sync --frozen "${args[@]}"
+      shell: bash

.github/dependabot.yml

@@ -0,0 +1,11 @@
+# To get started with Dependabot version updates, you'll need to specify which
+# package ecosystems to update and where the package manifests are located.
+# Please see the documentation for all configuration options:
+# https://docs.github.com/code-security/dependabot/dependabot-version-updates/configuration-options-for-the-dependabot.yml-file
+
+version: 2
+updates:
+  - package-ecosystem: "uv" # See documentation for possible values
+    directory: "/" # Location of package manifests
+    schedule:
+      interval: "weekly"

.github/workflows/codeql.yml

@@ -0,0 +1,100 @@
+# For most projects, this workflow file will not need changing; you simply need
+# to commit it to your repository.
+#
+# You may wish to alter this file to override the set of languages analyzed,
+# or to provide custom queries or build logic.
+#
+# ******** NOTE ********
+# We have attempted to detect the languages in your repository. Please check
+# the `language` matrix defined below to confirm you have the correct set of
+# supported CodeQL languages.
+#
+name: "CodeQL Advanced"
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+  schedule:
+    - cron: '42 21 * * 2'
+
+jobs:
+  analyze:
+    name: Analyze (${{ matrix.language }})
+    # Runner size impacts CodeQL analysis time. To learn more, please see:
+    #   - https://gh.io/recommended-hardware-resources-for-running-codeql
+    #   - https://gh.io/supported-runners-and-hardware-resources
+    #   - https://gh.io/using-larger-runners (GitHub.com only)
+    # Consider using larger runners or machines with greater resources for possible analysis time improvements.
+    runs-on: ${{ (matrix.language == 'swift' && 'macos-latest') || 'ubuntu-latest' }}
+    permissions:
+      # required for all workflows
+      security-events: write
+
+      # required to fetch internal or private CodeQL packs
+      packages: read
+
+      # only required for workflows in private repositories
+      actions: read
+      contents: read
+
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+        - language: actions
+          build-mode: none
+        - language: python
+          build-mode: none
+        # CodeQL supports the following values keywords for 'language': 'actions', 'c-cpp', 'csharp', 'go', 'java-kotlin', 'javascript-typescript', 'python', 'ruby', 'swift'
+        # Use `c-cpp` to analyze code written in C, C++ or both
+        # Use 'java-kotlin' to analyze code written in Java, Kotlin or both
+        # Use 'javascript-typescript' to analyze code written in JavaScript, TypeScript or both
+        # To learn more about changing the languages that are analyzed or customizing the build mode for your analysis,
+        # see https://docs.github.com/en/code-security/code-scanning/creating-an-advanced-setup-for-code-scanning/customizing-your-advanced-setup-for-code-scanning.
+        # If you are analyzing a compiled language, you can modify the 'build-mode' for that language to customize how
+        # your codebase is analyzed, see https://docs.github.com/en/code-security/code-scanning/creating-an-advanced-setup-for-code-scanning/codeql-code-scanning-for-compiled-languages
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@v4
+
+    # Add any setup steps before running the `github/codeql-action/init` action.
+    # This includes steps like installing compilers or runtimes (`actions/setup-node`
+    # or others). This is typically only required for manual builds.
+    # - name: Setup runtime (example)
+    #   uses: actions/setup-example@v1
+
+    # Initializes the CodeQL tools for scanning.
+    - name: Initialize CodeQL
+      uses: github/codeql-action/init@v3
+      with:
+        languages: ${{ matrix.language }}
+        build-mode: ${{ matrix.build-mode }}
+        # If you wish to specify custom queries, you can do so here or in a config file.
+        # By default, queries listed here will override any specified in a config file.
+        # Prefix the list here with "+" to use these queries and those in the config file.
+
+        # For more details on CodeQL's query packs, refer to: https://docs.github.com/en/code-security/code-scanning/automatically-scanning-your-code-for-vulnerabilities-and-errors/configuring-code-scanning#using-queries-in-ql-packs
+        # queries: security-extended,security-and-quality
+
+    # If the analyze step fails for one of the languages you are analyzing with
+    # "We were unable to automatically build your code", modify the matrix above
+    # to set the build mode to "manual" for that language. Then modify this step
+    # to build your code.
+    # ℹ️ Command-line programs to run using the OS shell.
+    # 📚 See https://docs.github.com/en/actions/using-workflows/workflow-syntax-for-github-actions#jobsjob_idstepsrun
+    - if: matrix.build-mode == 'manual'
+      shell: bash
+      run: |
+        echo 'If you are using a "manual" build mode for one or more of the' \
+          'languages you are analyzing, replace this with the commands to build' \
+          'your code, for example:'
+        echo '  make bootstrap'
+        echo '  make release'
+        exit 1
+
+    - name: Perform CodeQL Analysis
+      uses: github/codeql-action/analyze@v3
+      with:
+        category: "/language:${{matrix.language}}"

.github/workflows/conventional-pr.yml

@@ -0,0 +1,17 @@
+name: conventional-pr
+on:
+  pull_request:
+    types:
+      - opened
+      - edited
+      - synchronize
+    branches:
+      - main
+jobs:
+  lint-pr:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - uses: CondeNast/conventional-pull-request-action@v0.2.0
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/main.yml

@@ -0,0 +1,70 @@
+name: Main
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    types: [opened, synchronize, reopened, ready_for_review]
+
+jobs:
+  quality:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out
+        uses: actions/checkout@v4
+
+      - uses: actions/cache@v4
+        with:
+          path: ~/.cache/pre-commit
+          key: pre-commit-${{ hashFiles('.pre-commit-config.yaml') }}
+
+      - name: Set up the environment
+        uses: ./.github/actions/setup-python-env
+
+      - name: Run checks
+        run: make check
+
+  tests-and-type-check:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.11", "3.12"]
+      fail-fast: false
+    defaults:
+      run:
+        shell: bash
+    steps:
+      - name: Check out
+        uses: actions/checkout@v4
+
+      - name: Set up the environment
+        uses: ./.github/actions/setup-python-env
+        with:
+          python-version: ${{ matrix.python-version }}
+          dependency-groups: "dev test"
+
+      - name: Run tests
+        run: uv run python -m pytest tests --cov --cov-config=pyproject.toml --cov-report=xml
+
+      - name: Check typing
+        run: uv run mypy
+
+
+      - name: Upload coverage reports to Codecov with GitHub Action on Python 3.11
+        uses: codecov/codecov-action@v4
+        if: ${{ matrix.python-version == '3.11' }}
+        env:
+          CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
+
+  check-docs:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out
+        uses: actions/checkout@v4
+
+      - name: Set up the environment
+        uses: ./.github/actions/setup-python-env
+
+      - name: Check if documentation can be built
+        run: uv run mkdocs build -s

.github/workflows/on-release-main.yml

@@ -0,0 +1,102 @@
+name: release-main
+
+on:
+  workflow_dispatch:
+
+jobs:
+  semantic-release:
+    runs-on: ubuntu-latest
+    concurrency: release
+    permissions:
+      id-token: write
+      contents: write
+    outputs:
+      released: ${{ steps.release.outputs.released }}
+      tag: ${{ steps.release.outputs.tag }}
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ github.ref_name }}
+          fetch-depth: 0
+          persist-credentials: false
+
+      - name: Setup | Force release branch to be at workflow sha
+        run: |
+          git reset --hard ${{ github.sha }}
+
+      - name: Action | Python Semantic Release
+        id: release
+        uses: python-semantic-release/python-semantic-release@v9.21.1
+        with:
+          github_token: ${{ secrets.GITHUB_TOKEN }}
+
+      # The following is done through the on-release-main.yml action.
+      # Leaving this here for reference.
+      # - name: Publish | Upload package to PyPI
+      #   uses: pypa/gh-action-pypi-publish@release/v1
+      #   # NOTE: DO NOT wrap the conditional in ${{ }} as it will always evaluate to true.
+      #   # See https://github.com/actions/runner/issues/1173
+      #   if: steps.release.outputs.released == 'true'
+
+      - name: Publish | Upload to GitHub Release Assets
+        uses: python-semantic-release/publish-action@v9.21.1
+        if: steps.release.outputs.released == 'true'
+        with:
+          github_token: ${{ secrets.GITHUB_TOKEN }}
+          tag: ${{ steps.release.outputs.tag }}
+
+
+  set-version:
+    needs: [semantic-release]
+    if: needs.semantic-release.outputs.released == 'true'
+    runs-on: ubuntu-24.04
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Update project version
+        run: |
+          sed -i "s/^version = \".*\"/version = \"$RELEASE_VERSION\"/" pyproject.toml
+        env:
+          RELEASE_VERSION: ${{ needs.semantic-release.outputs.tag }}
+
+      - name: Upload updated pyproject.toml
+        uses: actions/upload-artifact@v4
+        with:
+          name: pyproject-toml
+          path: pyproject.toml
+
+  publish:
+    runs-on: ubuntu-latest
+    needs: [set-version]
+    steps:
+      - name: Check out
+        uses: actions/checkout@v4
+
+      - name: Set up the environment
+        uses: ./.github/actions/setup-python-env
+
+      - name: Download updated pyproject.toml
+        uses: actions/download-artifact@v4
+        with:
+          name: pyproject-toml
+
+      - name: Build package
+        run: uv build
+
+      - name: Publish package
+        run: uv publish
+        env:
+          UV_PUBLISH_TOKEN: ${{ secrets.PYPI_TOKEN }}
+
+  deploy-docs:
+    needs: publish
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out
+        uses: actions/checkout@v4
+
+      - name: Set up the environment
+        uses: ./.github/actions/setup-python-env
+
+      - name: Deploy documentation
+        run: uv run mkdocs gh-deploy --force

.github/workflows/validate-codecov-config.yml

@@ -0,0 +1,15 @@
+name: validate-codecov-config
+
+on:
+  pull_request:
+    paths: [codecov.yaml]
+  push:
+    branches: [main]
+
+jobs:
+  validate-codecov-config:
+    runs-on: ubuntu-22.04
+    steps:
+      - uses: actions/checkout@v4
+      - name: Validate codecov configuration
+        run: curl -sSL --fail-with-body --data-binary @codecov.yaml https://codecov.io/validate

.gitignore

@@ -0,0 +1,155 @@
+docs/source
+
+# From https://raw.githubusercontent.com/github/gitignore/main/Python.gitignore
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+venv/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# Vscode config files
+.vscode/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  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/
+
+logs
+models
+notebooks
+*DE.md
+wandb/
+results/
+*_debug.py
+*_debug.sh
+distilabel_cache/
+hf-cache/
+vllm-cache/
+uv-cache/
+unsloth_compiled_cache/
+.*o
+outputs
+500
+distillation-cache
+_unsloth_temporary_saved_buffers

.pre-commit-config.yaml

@@ -0,0 +1,26 @@
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: "v5.0.0"
+    hooks:
+      - id: check-case-conflict
+      - id: check-merge-conflict
+      - id: check-toml
+      - id: check-yaml
+      - id: check-json
+        exclude: ^.devcontainer/devcontainer.json
+      - id: pretty-format-json
+        exclude: ^.devcontainer/devcontainer.json
+        args: [--autofix, --no-sort-keys]
+      - id: end-of-file-fixer
+        exclude: ^linalg_zero/distillation/vllm/.*\.jinja$
+      - id: trailing-whitespace
+        exclude: ^linalg_zero/distillation/vllm/.*\.jinja$
+
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: "v0.11.5"
+    hooks:
+      - id: ruff
+        args: [--exit-non-zero-on-fix]
+        exclude: ^linalg_zero/distillation/vllm/.*\.jinja$
+      - id: ruff-format
+        exclude: ^linalg_zero/distillation/vllm/.*\.jinja$

CHANGELOG.md

@@ -0,0 +1,389 @@
+# CHANGELOG
+
+
+## v1.0.0 (2026-02-17)
+
+### Bug Fixes
+
+- **generate**: Remove ambiguities in problem generation ([`5339d04`](https://github.com/atomwalk12/linalg-zero/commit/5339d04bd7f1a8253e744a209f47179f67a47394))
+- **distillation**: Remove vllm chat templates ([`57be920`](https://github.com/atomwalk12/linalg-zero/commit/57be920fdb2ca9524593b67b735dfee8716affc0))
+- **grpo**: Adjust launch script parameters ([`a40be0e`](https://github.com/atomwalk12/linalg-zero/commit/a40be0ebcf73657a4d0c695ba933d379a31eef2b))
+- **grpo**: Log statistics ([`b983abd`](https://github.com/atomwalk12/linalg-zero/commit/b983abd1760b26b62c2e101ac76e1e3290359ec9))
+- **sft**: Adjust optional accuracy callback ([`16b7d32`](https://github.com/atomwalk12/linalg-zero/commit/16b7d322f5a50fd830c1ea5acd51e3ca20702454))
+
+### Documentation
+
+- **config**: Add sample config files ([`62290aa`](https://github.com/atomwalk12/linalg-zero/commit/62290aa84d011ef52078272d241204518a78cf16))
+- **report**: Finalize project ([`ffbfcb2`](https://github.com/atomwalk12/linalg-zero/commit/ffbfcb24a0cc54bcd888b4b7c19148e75ffce4f7))
+
+### Features
+
+- **distillation**: Add data generation configs ([`6aab283`](https://github.com/atomwalk12/linalg-zero/commit/6aab283a343c6273bfdc274e7397dff76bb3fa5c))
+- **distillation**: Add validation script ([`8344dba`](https://github.com/atomwalk12/linalg-zero/commit/8344dbad775768a7a74ffd83dec56587514e1c49))
+- **distillation**: Improve multi-turn generation with progress tracking ([`f58eadf`](https://github.com/atomwalk12/linalg-zero/commit/f58eadf0083f135f2fee2405b43eca1c2705d84e))
+- **grpo**: Add environment abstraction layer with base classes ([`edc9615`](https://github.com/atomwalk12/linalg-zero/commit/edc9615163385df75ad18aa15f918e534c45f7fa))
+- **grpo**: Add linear algebra environment with reward computation ([`4e46e7c`](https://github.com/atomwalk12/linalg-zero/commit/4e46e7c9e0040ee34e5fca060f652272f0c6ed1f))
+- **grpo**: Add linear algebra tools for matrix operations ([`d2f08b1`](https://github.com/atomwalk12/linalg-zero/commit/d2f08b1231c2c0f493bbea1f3f6718cb6829c0fa))
+- **grpo**: Add reward model and RL training utilities ([`aba4bae`](https://github.com/atomwalk12/linalg-zero/commit/aba4baefdbbe9c21c01ed82d2a062f788062f9d6))
+- **grpo**: Add task selection and training infrastructure ([`d82f8c6`](https://github.com/atomwalk12/linalg-zero/commit/d82f8c69506056c227825741ae7098eabbcb2067))
+- **grpo**: Add tool calling agent implementation ([`b95ec50`](https://github.com/atomwalk12/linalg-zero/commit/b95ec50ce7956f897a1bde16481633cb2395c77b))
+- **grpo**: Add training and evaluation scripts ([`38b99fd`](https://github.com/atomwalk12/linalg-zero/commit/38b99fd20ecf08d0fa7ce185b6094b1d782d90d9))
+- **grpo**: Add utility modules for training infrastructure ([`7362699`](https://github.com/atomwalk12/linalg-zero/commit/73626995bdcfd365b40fcd6813eca5ee99f2b8c0))
+- **grpo**: Add yaml config files ([`bcb5b68`](https://github.com/atomwalk12/linalg-zero/commit/bcb5b68ef026110a8fcee14da23e5bf4736c8251))
+- **sft**: Add model evaluation and dataset preparation scripts ([`5e8ee4c`](https://github.com/atomwalk12/linalg-zero/commit/5e8ee4c9d8fdf90e825a8a907216dadf78e8103b))
+- **sft**: Add yaml training configs ([`65ff234`](https://github.com/atomwalk12/linalg-zero/commit/65ff2348dc592d790bef130ecf593ce157ebdaf3))
+- **sft**: Adjust tool evaluation callback ([`c7f2c9d`](https://github.com/atomwalk12/linalg-zero/commit/c7f2c9daf2b8506e3687bd92d1d90d6599c99aa7))
+- **sft**: Improve tool calling accuracy callback with weave logging ([`823b263`](https://github.com/atomwalk12/linalg-zero/commit/823b263a447470a30993c809264dbb20494d1331))
+- **sft**: Refactor diagnostics and callbacks and improve evaluation logging ([`10ae836`](https://github.com/atomwalk12/linalg-zero/commit/10ae836eb0c7af9b7756bc0690c4461f2eaf3c8d))
+- **system_prompts**: Add SFT system prompt and improve tool usage guidelines ([`061d21a`](https://github.com/atomwalk12/linalg-zero/commit/061d21aa6de9dc5e0f5ff3ac97a8f71be7a83c2e))
+
+### Refactoring
+
+- Reorganize training entry points ([`5a65164`](https://github.com/atomwalk12/linalg-zero/commit/5a65164ceb460701ba946c745ddf5e8a1631b73e))
+- **grpo**: Improve XML parser API and validation logic ([`fa80cba`](https://github.com/atomwalk12/linalg-zero/commit/fa80cba680bd6af515bc66a66b69a255b4e3f3d1))
+
+
+## v0.3.0 (2025-09-29)
+
+### Bug Fixes
+
+- Normalize grpo dataset schema
+  ([`8590141`](https://github.com/atomwalk12/linalg-zero/commit/8590141149580635669ffcb40b3ad89c09849772))
+
+- **context**: Modify entropy budget validation precision
+  ([`b27d237`](https://github.com/atomwalk12/linalg-zero/commit/b27d2376c8399c491c52ded7bfbf0b838824d306))
+
+- **distillation**: Align parsed_messages size with the conversations size
+  ([`6a0a6bd`](https://github.com/atomwalk12/linalg-zero/commit/6a0a6bd173a48de5036700379c75da4ea51eb6a6))
+
+- this ensures simpler reasoning since the two arrays contain similar data - fix small problems
+  around malformed messages, which are not silently skipped
+
+- **distillation**: Solve accumulation problem with tool-use statistics
+  ([`523f2ec`](https://github.com/atomwalk12/linalg-zero/commit/523f2ecf0156b805bcb1b2d9aa14852de00e5acd))
+
+- **entropy**: Allow data generation using varying precision levels
+  ([`08bbddd`](https://github.com/atomwalk12/linalg-zero/commit/08bbddd6bf699bfa0327381cc24da303afea9dd9))
+
+- **generation**: Add difficulty rating based number of tool calls performed to reach a solution
+  ([`aa2bace`](https://github.com/atomwalk12/linalg-zero/commit/aa2bacea495661d2cbf23117bd751f079e1e3069))
+
+- **generation**: Adjust matrices format and introduce better organization for the difficulty levels
+  ([`00bb08f`](https://github.com/atomwalk12/linalg-zero/commit/00bb08f44a4ac4a6bda4559a7ada28bc73084811))
+
+- **generation**: Fix circular dependency by moving the library types to a separate file
+  ([`33ee32a`](https://github.com/atomwalk12/linalg-zero/commit/33ee32a3e2c26faa7da170b3adba812c1f6661ce))
+
+- **generator**: Use Dirichlet distribution for entropy allocation
+  ([`0a16a52`](https://github.com/atomwalk12/linalg-zero/commit/0a16a52590aa038e3d2780e6378be275bd5e8b1b))
+
+### Documentation
+
+- Improve generated statistics on tool execution
+  ([`7c62398`](https://github.com/atomwalk12/linalg-zero/commit/7c62398a065a2c746dd95ed5b2bd21314b2bab29))
+
+### Features
+
+- Enhance distillation pipeline with new model configurations and diagnostics support
+  ([`0559239`](https://github.com/atomwalk12/linalg-zero/commit/05592398c9be27b9c364643ddc5b88cb19f4cdfe))
+
+- Update distillation configurations
+  ([`b43e149`](https://github.com/atomwalk12/linalg-zero/commit/b43e149a411f3f4587d76fbdd67d135cac91896c))
+
+- **base_generator**: Add ability to redirect arbitrarily matrix component across composition
+  contexts
+  ([`1c57a61`](https://github.com/atomwalk12/linalg-zero/commit/1c57a61ca13327ec2bde9be7539cb5e331452f9f))
+
+- **distillation**: Add base script to run the distillation pipeline
+  ([`2189bb9`](https://github.com/atomwalk12/linalg-zero/commit/2189bb9303c0a2bb84331a81e200d9dae19724bd))
+
+- simplifies the previous approach to use only 1 task for multi-turn conversations
+
+- **generation**: Add entropy allocation using dirichlet distribution
+  ([`03d4bcb`](https://github.com/atomwalk12/linalg-zero/commit/03d4bcb88a8ad7f3c2cb9b6e89da3f30e71e816d))
+
+- **generation**: Add matrix vector multiplication generator
+  ([`a1e0ff1`](https://github.com/atomwalk12/linalg-zero/commit/a1e0ff1c2f44a48aef44d53ec68ffa27f3494788))
+
+- **generation**: Add result verification step involving sympy-primitive conversion
+  ([`07f8d51`](https://github.com/atomwalk12/linalg-zero/commit/07f8d517e886e46821fc6c8fa09a4bba00ab6e78))
+
+- **generation**: Generate a fixed number of examples per class
+  ([`4eb251a`](https://github.com/atomwalk12/linalg-zero/commit/4eb251ac933e708520a46e22ae5f6e9dd47cb113))
+
+- **generator**: Add composition infrastructure to allow mixture of problems
+  ([`d08820b`](https://github.com/atomwalk12/linalg-zero/commit/d08820bfa258f85cb5d533c9ccac5bb26940bcce))
+
+- **generator**: Add constrained matrix generation for composed components
+  ([`76faaf1`](https://github.com/atomwalk12/linalg-zero/commit/76faaf10984ff94b911252147c77c9a41a633512))
+
+- **generator**: Add difficulty data-class for managing difficulty parameters
+  ([`33eae33`](https://github.com/atomwalk12/linalg-zero/commit/33eae3355afb9255f5252eb9c661fecec612130d))
+
+- **generator**: Add frobenius norm generator atomic operation
+  ([`c3c53b7`](https://github.com/atomwalk12/linalg-zero/commit/c3c53b7e6f036f99c8ae142dd0aad25c7113c1b4))
+
+- refactor to differentiate between independent and dependent components - add tests for composing
+  linear-system -> matrix multiplication -> frobenius norm
+
+- **generator**: Add frobenius norm solver
+  ([`7e346b3`](https://github.com/atomwalk12/linalg-zero/commit/7e346b36aad034a14d0a1ac71f6a2f25115a783d))
+
+- **generator**: Add generator for solving linear systems
+  ([`37faf45`](https://github.com/atomwalk12/linalg-zero/commit/37faf45a86b1c0068d19a1391174c1618c7ca77b))
+
+- **generator**: Add integer/rational number generators using on controllable complexity
+  ([`731100f`](https://github.com/atomwalk12/linalg-zero/commit/731100f609c6591a1fe2297ff0c64294195035d2))
+
+- **generator**: Add matrix cofactor generator and related components
+  ([`347e8b3`](https://github.com/atomwalk12/linalg-zero/commit/347e8b3388ae2ed5858a3b78f1fadb2ec15a44ba))
+
+- **generator**: Add matrix inverse generator and related components
+  ([`0154fef`](https://github.com/atomwalk12/linalg-zero/commit/0154fef50dfaa6a31dea0eeb78abd51144eca211))
+
+- **generator**: Add matrix trace generator and related functionality
+  ([`6cee5cd`](https://github.com/atomwalk12/linalg-zero/commit/6cee5cdc31b82b0606f1b3655e3bd592831e2514))
+
+- **generator**: Add matrix transpose generator
+  ([`32eae5c`](https://github.com/atomwalk12/linalg-zero/commit/32eae5cacb8b7472da3a3db7c3e5d890883b7df6))
+
+- **generator**: Add new problem solver for finding matrix determinant
+  ([`c700f05`](https://github.com/atomwalk12/linalg-zero/commit/c700f0537f49b75e59793fbe26aa861d1a36598f))
+
+- **generator**: Add rank generator
+  ([`6fdedf4`](https://github.com/atomwalk12/linalg-zero/commit/6fdedf41ec9437bc3500595fe21c36ca5dc603de))
+
+- **generator**: Add SymPy problem generation base class with context management and template
+  handling
+  ([`c94e860`](https://github.com/atomwalk12/linalg-zero/commit/c94e860ad06ff8298a6fcc1233a3de04521187d3))
+
+- **generator**: Add template engine to help with diverse problem generation
+  ([`d52b4eb`](https://github.com/atomwalk12/linalg-zero/commit/d52b4ebe9844c8ad808feef50fb0cfa64b020060))
+
+- **generator**: Add validation and wrappers for generating dependent components
+  ([`d557dfe`](https://github.com/atomwalk12/linalg-zero/commit/d557dfefbeb2c421d3e9f13bf11612972680622e))
+
+- **generator**: Allow dependencies between components to be referenced across problems
+  ([`ebf0d4a`](https://github.com/atomwalk12/linalg-zero/commit/ebf0d4adc3209be869a7bc88bf51aa2206ee8c7d))
+
+- **generator**: Improve dataset generation with optimized registry and added constraints to integer
+  generation
+  ([`6352e35`](https://github.com/atomwalk12/linalg-zero/commit/6352e35b4550eb65778cdbaa82c8f6f3536d704f))
+
+- **generator**: Improve factory registration API
+  ([`5a13a11`](https://github.com/atomwalk12/linalg-zero/commit/5a13a1144405f9add7b0a9e271871d53448b89e1))
+
+- **generator**: Separate matrix-matrix multiplication from matrix-vector multiplication generators
+  ([`b07e5bb`](https://github.com/atomwalk12/linalg-zero/commit/b07e5bb23a8d709ff46f6bae88380747ec2f7c59))
+
+- **generator**: Simplify component logic by transferring the responsibility to the base class for
+  question/answer generation and verification
+  ([`8f880f5`](https://github.com/atomwalk12/linalg-zero/commit/8f880f5d81232fe2361fac3669d1155adff0d754))
+
+- **grpo**: Add GRPO configuration files and reward functions
+  ([`04557c4`](https://github.com/atomwalk12/linalg-zero/commit/04557c46114c56fb4ad4b7cd91cbf1a3403a763e))
+
+- Added GRPO training scripts and configuration files. - Introduced new linear algebra functions and
+  tools for GRPO. - Implemented dataset preparation and processing for GRPO training. - Added
+  support for multi-turn interactions in training. - Created a comprehensive debug dataset for
+  testing and validation.
+
+- **multi-turn**: Implement multi-turn generation with tool use and verification
+  ([`85f9314`](https://github.com/atomwalk12/linalg-zero/commit/85f93149fc0c45fd00ec9f2d806460b051f68979))
+
+- **templates**: Adjust the templates to ensure a consistent format
+  ([`7cde039`](https://github.com/atomwalk12/linalg-zero/commit/7cde039a9dacaaf417c5e57443ef50290f52db16))
+
+- **verification**: Add checks to ensure generated data is valid
+  ([`af961f8`](https://github.com/atomwalk12/linalg-zero/commit/af961f83090714757b7fe53f889835201c39d02f))
+
+### Refactoring
+
+- Add enum for available problem types
+  ([`ca9bc90`](https://github.com/atomwalk12/linalg-zero/commit/ca9bc903188b9b6ba33d718548b5f24300e58eaa))
+
+- Auto entropy allocation in base class based on difficulty
+  ([`6f56296`](https://github.com/atomwalk12/linalg-zero/commit/6f56296e69e39224017446808ffb56d71a981d4b))
+
+- Centralise using a mixin class the template generation method to ensure consistent generation
+  across all components
+  ([`0b6b3c7`](https://github.com/atomwalk12/linalg-zero/commit/0b6b3c727674fbf6ab319c0196f26a99684789ba))
+
+- Customise question layout
+  ([`d527db0`](https://github.com/atomwalk12/linalg-zero/commit/d527db0befed2da79a2cde57d0d894a05ebb13af))
+
+- Minor tweaks
+  ([`d72ec79`](https://github.com/atomwalk12/linalg-zero/commit/d72ec79e23bf9049262ac05dd1207423f8a9aa35))
+
+- Simplify tool accuracy callback
+  ([`8305221`](https://github.com/atomwalk12/linalg-zero/commit/8305221c2759d9b9df45c2e8c31baf2355b24e15))
+
+- **composition**: Remove unnecessary code
+  ([`50f0258`](https://github.com/atomwalk12/linalg-zero/commit/50f02586d080b4edde21014d09d345d49ef5b959))
+
+- **generator**: Add GenerationConstraints for matrix generation constraints
+  ([`a5f763f`](https://github.com/atomwalk12/linalg-zero/commit/a5f763f0f077345edb400b496f0e148ccd7b5942))
+
+- **generator**: Encapsulate generation parameters in the config
+  ([`b89cabb`](https://github.com/atomwalk12/linalg-zero/commit/b89cabbc62c04fa44041cd5d5c56c1703410af8f))
+
+- **generator**: Improve type safety and update function signatures and imports
+  ([`16cd9ea`](https://github.com/atomwalk12/linalg-zero/commit/16cd9eade56b3f44476af1ea3da0c820c71e6ba3))
+
+- **generator**: Improve type safety and update function signatures and imports
+  ([`39f080e`](https://github.com/atomwalk12/linalg-zero/commit/39f080e4809ed96ed3b72f7a668d2893bdd1d1ef))
+
+- **generator**: Sample entropy in a centralized place within the context
+  ([`d920fcf`](https://github.com/atomwalk12/linalg-zero/commit/d920fcfb1c81241c14eb3118e28df411fa87671a))
+
+- **generator**: Unify problem types and improve entropy management
+  ([`fc309c2`](https://github.com/atomwalk12/linalg-zero/commit/fc309c247ac07d0ed05de3227a275c408db6bfea))
+
+- **generator**: Update matrix operations to use centralized library calls
+  ([`bdd537b`](https://github.com/atomwalk12/linalg-zero/commit/bdd537be0b628b7b8f7f5da5528f0aaed00fa7ab))
+
+
+## v0.2.0 (2025-08-03)
+
+### Bug Fixes
+
+- Add logging utilities for files and stdout
+  ([`d2692e6`](https://github.com/atomwalk12/linalg-zero/commit/d2692e68bbedabb5d110582737b3d995e26d64a2))
+
+- Be granular about exceptions being thrown
+  ([`4c88f9f`](https://github.com/atomwalk12/linalg-zero/commit/4c88f9fc8a14effb11b49ce90fb323122e0c4916))
+
+Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
+
+- Ensure an exception is raised if the subprocess returns a non-zero exit code
+  ([`4695ab9`](https://github.com/atomwalk12/linalg-zero/commit/4695ab9c9afa4b56ec4167d99777b6b230d01ac7))
+
+Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
+
+- Finetune the number of Llama-cpp GPU layers offloaded to the GPU
+  ([`a37fc13`](https://github.com/atomwalk12/linalg-zero/commit/a37fc130ab9b04be0ecccfe1278c6e1fe68d3350))
+
+- **dependencies**: Pin llama-cpp-python version to 0.3.13), update distilabel dependencies and add
+  lock file
+  ([`ef18d53`](https://github.com/atomwalk12/linalg-zero/commit/ef18d53a9fda569c6148ab8130557ce689d7c7ba))
+
+- **distillation**: Add script to push debugging dataset to huggingface
+  ([`6d6c91a`](https://github.com/atomwalk12/linalg-zero/commit/6d6c91ac275aaaee1d145bbb2d99296da823c617))
+
+- **inference**: Add hf_pretrained_model_name_or_path and remove redundant installation in launch
+  script
+  ([`75be580`](https://github.com/atomwalk12/linalg-zero/commit/75be58003bf3fd321deb185d117c7fc43ba11a57))
+
+### Build System
+
+- Remove python 3.9 support
+  ([`e7421e6`](https://github.com/atomwalk12/linalg-zero/commit/e7421e611029b15f21ac9dcef2c4e14d97e99f9b))
+
+### Features
+
+- Add configuration parameters, tasks for running distributed training and pin dependencies
+  ([`180fa5f`](https://github.com/atomwalk12/linalg-zero/commit/180fa5f4ccabf0f45d5c02371b80d0774a9cace5))
+
+- Add dataset generator utility
+  ([`279b0bd`](https://github.com/atomwalk12/linalg-zero/commit/279b0bdcf867b586c3b201afe425da1e93f2bfea))
+
+- Add workflow to generate a new dataset
+  ([`e0ba35d`](https://github.com/atomwalk12/linalg-zero/commit/e0ba35d7792309627dcca34293d11cb758ac3def))
+
+- Implement question generation factories for arithmetic and linear algebra
+  ([`f20fb2c`](https://github.com/atomwalk12/linalg-zero/commit/f20fb2c97d825d059eb6612d1f7cd4504c41f74c))
+
+- make use of a global registry to keep track of the various problem definitions
+
+- **distillation**: Add centralised control for launching the inference server
+  ([`23a4078`](https://github.com/atomwalk12/linalg-zero/commit/23a4078dc59d1706989850f32d67786db32191ec))
+
+- **distillation**: Add filter to easily track and discard incorrect results
+  ([`d991534`](https://github.com/atomwalk12/linalg-zero/commit/d991534a1f6eee0c691520f7d0a53bbdd0fbb341))
+
+- **distillation**: Add generation pipeline
+  ([`d18317f`](https://github.com/atomwalk12/linalg-zero/commit/d18317f32e6e03dbb04e124560c58c731af57ba9))
+
+- **distillation**: Add local script for llama.cpp inference server
+  ([`bf1c3e1`](https://github.com/atomwalk12/linalg-zero/commit/bf1c3e18ff1ff5cf39b47b8df3319898b2f41c3a))
+
+- **distillation**: Add planning and tool selection
+  ([`6b9e3e2`](https://github.com/atomwalk12/linalg-zero/commit/6b9e3e2e952267ce9cd5796bcb0d448a48f98a3e))
+
+- **distillation**: Add result synthesiser
+  ([`017a47c`](https://github.com/atomwalk12/linalg-zero/commit/017a47ca3204fb0a27c21532ba1e9fd480dae7a6))
+
+- **distillation**: Add the argilla components for simpler result inspection
+  ([`20b7649`](https://github.com/atomwalk12/linalg-zero/commit/20b7649f88621cd5d5fea11624ce4e5fde8441de))
+
+- **distillation**: Add the planner component
+  ([`ed5175a`](https://github.com/atomwalk12/linalg-zero/commit/ed5175aca518ed6d593ae863e00f6e84a02f6f4e))
+
+- **distillation**: Add verl dependency for GRPO
+  ([`9d7ccb6`](https://github.com/atomwalk12/linalg-zero/commit/9d7ccb65c8488760e566cda5aabb345ef41a78f6))
+
+- **distillation**: Code execution component
+  ([`24072e7`](https://github.com/atomwalk12/linalg-zero/commit/24072e75338875335d1efd226df6b72339f0c3c4))
+
+- **distillation**: Customise the chat generation pipeline to preserve input/output results
+  ([`9973a10`](https://github.com/atomwalk12/linalg-zero/commit/9973a1009b1400a3e0cc0e0562b1464ffbd291b8))
+
+- **distillation**: Demonstrate the tool selection component and update planner to use
+  ChatGeneration
+  ([`8fe4e0a`](https://github.com/atomwalk12/linalg-zero/commit/8fe4e0a840dfa665b80a89b2a62e3f07fb5806d5))
+
+- **distillation**: Implement function calling pipeline using Llama-cpp
+  ([`f8776ce`](https://github.com/atomwalk12/linalg-zero/commit/f8776cedf82002f1c08a18f9252a929d70991dd1))
+
+- **distillation**: Improve launch script to download models from the hf-hub and tune configuration
+  parameters
+  ([`6fc9503`](https://github.com/atomwalk12/linalg-zero/commit/6fc950382725f00268b4b9b671c90abed1f3ab02))
+
+- **distillation**: Integrate all related components to generate new data (completed pipeline)
+  ([`c8532a4`](https://github.com/atomwalk12/linalg-zero/commit/c8532a4d5a6da3bfb7ce82e6de10f81f033ab826))
+
+- **distillation**: Integrate math-verify for formal evaluation of the output
+  ([`a2c3757`](https://github.com/atomwalk12/linalg-zero/commit/a2c3757e216d462ed9915a179e517382c68639bd))
+
+- **sft**: Add additional callbacks (i.e. evaluation, push revision to hub, early stopping)
+  ([`90e7f35`](https://github.com/atomwalk12/linalg-zero/commit/90e7f3532e0ca502f0edda571e5b33889cb41b89))
+
+- **sft**: Complete evaluation with the ability to resume training, log results via wandb, create
+  model cards and save to the huggingface hub
+  ([`a7d86a3`](https://github.com/atomwalk12/linalg-zero/commit/a7d86a3e5c59ab527e9d3ebc5779f4d35065b2e7))
+
+### Testing
+
+- Add dataset configuration file
+  ([`5372677`](https://github.com/atomwalk12/linalg-zero/commit/5372677debfcb91b17c0bbc5af665bd19872c3ba))
+
+- Check default registry configuration
+  ([`47f190a`](https://github.com/atomwalk12/linalg-zero/commit/47f190a94e045650b3c1cffee3d4888d26b16164))
+
+
+## v0.1.0 (2025-07-08)
+
+### Bug Fixes
+
+- Add container fix for Docker
+  ([`bf3884b`](https://github.com/atomwalk12/linalg-zero/commit/bf3884b10d03f6dfa253944c2497297bc91d32d2))
+
+- Improve release pipeline to include semantic releases
+  ([`888738c`](https://github.com/atomwalk12/linalg-zero/commit/888738c4b914c65e001025ca3ba4473c1712f235))
+
+### Features
+
+- Add additional tasks
+  ([`c036f31`](https://github.com/atomwalk12/linalg-zero/commit/c036f31ba63af43b205e7a73464935f98f602773))
+
+- Add codeql support
+  ([`9b337c6`](https://github.com/atomwalk12/linalg-zero/commit/9b337c6eb7892943e2e6e1d004ec8611f868e0bd))
+
+- Add gemini style guide
+  ([`0f65ed1`](https://github.com/atomwalk12/linalg-zero/commit/0f65ed1847f72f900cf3d37e0766e1d06269c451))
+
+- Set up semantic release
+  ([`a97ecd3`](https://github.com/atomwalk12/linalg-zero/commit/a97ecd319d69c348b1a3a760c00ad6678a7b7339))
+
+- Update mkdocs theme
+  ([`8dd0a40`](https://github.com/atomwalk12/linalg-zero/commit/8dd0a40dd627c64ce868d557b8c889c7c6a24d0d))

CONTRIBUTING.md

@@ -0,0 +1,126 @@
+# Contributing to `linalg-zero`
+
+Contributions are welcome, and they are greatly appreciated!
+Every little bit helps, and credit will always be given.
+
+You can contribute in many ways:
+
+# Types of Contributions
+
+## Report Bugs
+
+Report bugs at https://github.com/atomwalk12/linalg-zero/issues
+
+If you are reporting a bug, please include:
+
+- Your operating system name and version.
+- Any details about your local setup that might be helpful in troubleshooting.
+- Detailed steps to reproduce the bug.
+
+## Fix Bugs
+
+Look through the GitHub issues for bugs.
+Anything tagged with "bug" and "help wanted" is open to whoever wants to implement a fix for it.
+
+## Implement Features
+
+Look through the GitHub issues for features.
+Anything tagged with "enhancement" and "help wanted" is open to whoever wants to implement it.
+
+## Write Documentation
+
+linalg-zero could always use more documentation, whether as part of the official docs, in docstrings, or even on the web in blog posts, articles, and such.
+
+## Submit Feedback
+
+The best way to send feedback is to file an issue at https://github.com/atomwalk12/linalg-zero/issues.
+
+If you are proposing a new feature:
+
+- Explain in detail how it would work.
+- Keep the scope as narrow as possible, to make it easier to implement.
+- Remember that this is a volunteer-driven project, and that contributions
+  are welcome :)
+
+# Get Started!
+
+Ready to contribute? Here's how to set up `linalg-zero` for local development.
+Please note this documentation assumes you already have `uv` and `Git` installed and ready to go.
+
+1. Fork the `linalg-zero` repo on GitHub.
+
+2. Clone your fork locally:
+
+```bash
+cd <directory_in_which_repo_should_be_created>
+git clone git@github.com:YOUR_NAME/linalg-zero.git
+```
+
+3. Now we need to install the environment. Navigate into the directory
+
+```bash
+cd linalg-zero
+```
+
+Then, install and activate the environment with:
+
+```bash
+uv sync --group dev --group test
+```
+
+4. Install pre-commit to run linters/formatters at commit time:
+
+```bash
+uv run pre-commit install
+```
+
+5. Create a branch for local development:
+
+```bash
+git checkout -b name-of-your-bugfix-or-feature
+```
+
+Now you can make your changes locally.
+
+6. Don't forget to add test cases for your added functionality to the `tests` directory.
+
+7. When you're done making changes, check that your changes pass the formatting tests.
+
+```bash
+make check
+```
+
+Now, validate that all unit tests are passing:
+
+```bash
+make test
+```
+
+9. Before raising a pull request you should also run tox.
+   This will run the tests across different versions of Python:
+
+```bash
+tox
+```
+
+This requires you to have multiple versions of python installed.
+This step is also triggered in the CI/CD pipeline, so you could also choose to skip this step locally.
+
+10. Commit your changes and push your branch to GitHub:
+
+```bash
+git add .
+git commit -m "Your detailed description of your changes."
+git push origin name-of-your-bugfix-or-feature
+```
+
+11. Submit a pull request through the GitHub website.
+
+# Pull Request Guidelines
+
+Before you submit a pull request, check that it meets these guidelines:
+
+1. The pull request should include tests.
+
+2. If the pull request adds functionality, the docs should be updated.
+   Put your new functionality into a function with a docstring, and add the feature to the list in `README.md`.

Dockerfile

@@ -0,0 +1,21 @@
+# Install uv
+FROM python:3.12-slim
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /bin/uv
+
+# Change the working directory to the `app` directory
+WORKDIR /app
+
+# Copy the lockfile and `pyproject.toml` into the image
+COPY uv.lock /app/uv.lock
+COPY pyproject.toml /app/pyproject.toml
+
+# Install dependencies
+RUN uv sync --frozen --no-install-project
+
+# Copy the project into the image
+COPY . /app
+
+# Sync the project
+RUN uv sync --frozen
+
+CMD [ "python", "linalg_zero/generate.py" ]

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Razvan Florian Vasile
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

Makefile

@@ -0,0 +1,166 @@
+## NOTE: the llama-cpp server is used to startup the inference server using `make distillation-server`
+# Fixing the llama-cpp server version to 0.3.13 as the upstream repository gets updated frequently
+# leading to incompatibility issues. If bumping the version don't forget to update pyproject.toml.
+.PHONY: install-data-gen
+install-data-gen: ## Install the virtual environment and install the pre-commit hooks.
+	@echo "🚀 Creating virtual environment using uv"
+	@uv sync --locked --group data-gen
+	@uv pip install setuptools flash-attn==2.7.3 --no-build-isolation
+	@uv run pre-commit install
+
+.PHONY: install-sft
+install-sft: ## Install the virtual environment and install the pre-commit hooks.
+	@echo "🚀 Creating virtual environment using uv"
+	@uv sync --locked --group sft
+	@uv pip install --upgrade setuptools ninja
+	@uv pip install --upgrade --torch-backend cu128 "xformers==0.0.33.post1"
+	@TORCH_CUDA_ARCH_LIST=8.9 MAX_JOBS=3 uv pip install --upgrade --no-build-isolation "flash-attn==2.8.3"
+	@uv run pre-commit install
+
+.PHONY: install-grpo
+install-grpo: ## Install the virtual environment and install the pre-commit hooks.
+	@echo "🚀 Creating virtual environment using uv"
+	@uv sync --locked --group grpo
+	@uv run pre-commit install
+
+.PHONY: setup-dev
+setup-dev: ## Setup the development environment
+	@echo "🚀 Setting up development environment"
+	@uv run linalg_zero/distillation/scripts/push_debug_dataset.py --dataset-name atomwalk12/linalg-debug --private
+
+.PHONY: check
+check: ## Run code quality tools.
+	@echo "🚀 Checking lock file consistency with 'pyproject.toml'"
+	@uv lock --locked
+	@echo "🚀 Linting code: Running pre-commit"
+ifeq ($(CI),true)
+	@echo "🔍 CI detected: Running ruff in check mode"
+	@uv run ruff check .
+	@uv run ruff format --check .
+	@SKIP=ruff,ruff-format uv run pre-commit run -a
+else
+	@uv run pre-commit run -a
+endif
+	@echo "🚀 Static type checking: Running mypy"
+	@uv run mypy
+	@echo "🚀 Checking for obsolete dependencies: Running deptry"
+	@uv run deptry .
+
+.PHONY: test
+test: ## Test the code with pytest
+	@echo "🚀 Testing code: Running pytest"
+	@uv run python -m pytest --cov --cov-config=pyproject.toml --cov-report=xml
+
+.PHONY: coverage-site
+coverage-site: ## Generate coverage report in HTML format
+	@echo "🚀 Generating coverage report in HTML format"
+	@uv run coverage html
+
+.PHONY: build
+build: clean-build ## Build wheel file
+	@echo "🚀 Creating wheel file"
+	@uvx --from build pyproject-build --installer uv
+
+.PHONY: clean-build
+clean-build: ## Clean build artifacts
+	@echo "🚀 Removing build artifacts"
+	@uv run python -c "import shutil; import os; shutil.rmtree('dist') if os.path.exists('dist') else None"
+
+.PHONY: publish
+publish: ## Publish a release to PyPI.
+	@echo "🚀 Publishing."
+	@uvx twine upload --repository-url https://upload.pypi.org/legacy/ dist/*
+
+.PHONY: build-and-publish
+build-and-publish: build publish ## Build and publish.
+
+.PHONY: docs-test
+docs-test: ## Test if documentation can be built without warnings or errors
+	@echo "🚀 Testing documentation build"
+	@uv run mkdocs build -s
+
+.PHONY: docs
+docs: ## Build and serve the documentation
+	@echo "🚀 Building and serving documentation"
+	@uv run mkdocs serve
+
+.PHONY: semantic-release
+semantic-release: ## Test semantic release
+	@echo "🚀 Testing semantic release"
+	@uv run semantic-release -vv --noop version --print
+
+.PHONY: gh-deploy
+gh-deploy: ## Deploy the documentation to GitHub Pages
+	@echo "🚀 Deploying documentation to GitHub Pages"
+	@uv run mkdocs gh-deploy --force
+
+LLAMACPP_CONFIG=linalg_zero/config/distillation/llamacpp_qwen3_32b_instruct.yaml
+VLLM_CONFIG=linalg_zero/config/distillation/vllm_qwen3_32b.yaml
+
+.PHONY: distillation-llamacpp
+distillation-llamacpp: ## Start the llama.cpp server
+	@echo "🚀 Starting llama.cpp server"
+	@INFERENCE_BACKEND=llamacpp uv run python linalg_zero/distillation/launch_server.py --config $(LLAMACPP_CONFIG)
+
+.PHONY: distillation-vllm
+distillation-vllm: ## Start the vLLM server
+	@echo "🚀 Starting vLLM server"
+	@. ./env.sh && uv run python linalg_zero/distillation/launch_server.py --config ${VLLM_CONFIG}
+
+.PHONY: distillation
+distillation: ## Run the distillation pipeline using the vllm config
+	@echo "🚀 Running distillation pipeline"
+	@. ./env.sh && uv run python linalg_zero/distillation.py --config ${VLLM_CONFIG}
+
+.PHONY: distillation-vllm-local
+distillation-vllm-local: ## Start the vLLM server
+	@echo "🚀 Starting vLLM server"
+	@export USING_VLLM=true INFERENCE_BACKEND=vllm && uv run python linalg_zero/distillation/launch_server.py --config linalg_zero/config/distillation/vllm_qwen3_4b_think.yaml
+
+
+.PHONY: distillation-local
+distillation-debug: ## Start the vLLM server
+	@echo "🚀 Starting vLLM server"
+	@export USING_VLLM=true && uv run python linalg_zero/distillation.py --config linalg_zero/config/distillation/vllm_qwen3_4b_think.yaml
+
+
+# SFT Training Commands
+SFT_CONFIG=linalg_zero/config/sft/sft_debug_config.yaml
+# SFT_CONFIG=linalg_zero/config/sft/sft_config.yaml
+ACCELERATE_CONFIG=linalg_zero/config/sft/accelerate/zero3.yaml
+
+.PHONY: sft-debug
+sft-debug: ## Run SFT training on single GPU
+	@echo "🚀 Running SFT training on single GPU"
+	@uv run python linalg_zero/sft.py --config $(SFT_CONFIG)
+
+
+.PHONY: sft-distributed
+sft-distributed: ## Run SFT training with distributed setup using DeepSpeed ZeroStage 3
+	@echo "🚀 Running distributed SFT training with DeepSpeed"
+	@uv run accelerate launch --config_file=$(ACCELERATE_CONFIG) linalg_zero/sft.py --config $(SFT_CONFIG)
+
+.PHONY: prepare-grpo-dataset
+prepare-grpo-dataset: ## Prepare the GRPO dataset
+	@echo "🚀 Creating GRPO dataset"
+	@uv run linalg_zero/grpo/process_dataset.py
+
+.PHONY: generate-optimised-config
+generate-optimised-config: ## Generate the optimised config
+	@echo "🚀 Generating optimised config"
+	@uv run linalg_zero/generator/analysis/analyse.py
+
+.PHONY: run-training
+run-training: ## Run the training pipeline
+	@echo "🚀 Running training pipeline"
+	@$(MAKE) setup-dev
+	@$(MAKE) prepare-grpo-dataset
+	@echo "🚀 Training pipeline completed"
+
+
+.PHONY: help
+help:
+	@uv run python -c "import re; \
+	[[print(f'\033[36m{m[0]:<20}\033[0m {m[1]}') for m in re.findall(r'^([a-zA-Z_-]+):.*?## (.*)$$', open(makefile).read(), re.M)] for makefile in ('$(MAKEFILE_LIST)').strip().split()]"
+
+.DEFAULT_GOAL := help

README.md

@@ -1,12 +1,155 @@
 ---
-title: Linalg Zero
-emoji: 🐨
-colorFrom: pink
-colorTo: blue
+title: Linalg-Zero
+emoji: 🧠
+colorFrom: indigo
+colorTo: purple
 sdk: gradio
-sdk_version: 6.9.0
-app_file: app.py
+sdk_version: "6.9.0"
+python_version: "3.12.12"
+app_file: linalg_zero/demo/app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+[![Release](https://img.shields.io/github/v/release/atomwalk12/linalg-zero)](https://img.shields.io/github/v/release/atomwalk12/linalg-zero)
+[![Build status](https://img.shields.io/github/actions/workflow/status/atomwalk12/linalg-zero/main.yml?branch=main)](https://github.com/atomwalk12/linalg-zero/actions/workflows/main.yml?query=branch%3Amain)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
+
+# Linalg-Zero
+
+## Overview
+
+<details>
+  <summary>Table of Contents</summary>
+  <ol>
+    <li><a href="#overview">Overview</a></li>
+    <li><a href="#main-phases">Main Phases</a></li>
+    <li><a href="#installation">Installation</a></li>
+    <li><a href="#quickstart">Quickstart</a></li>
+    <li><a href="#results">Results</a></li>
+    <li><a href="#artifacts">Artifacts</a></li>
+    <li><a href="#reproducibility">Reproducibility</a></li>
+    <li><a href="#acknowledgements">Acknowledgements</a></li>
+  </ol>
+</details>
+
+This repository offers tools for generating a linear algebra problem dataset and training an open-source base model (i.e. Qwen2.5-3B), aiming to explore planning and tool use using SFT and RL, distinct from Deepseek-R1's primary emphasis on reasoning.
+
+The project is simple by design and mostly consists of:
+
+- `linalg_zero/`: contains the scripts to train models as well as generate synthetic data:
+    - `generate.py`: generates the linear algebra dataset and splits.
+    - `distillation.py`: runs the distillation pipeline to create multi-turn tool-use data.
+    - `sft_train.py`: performs a simple SFT of a model on a dataset.
+    - `grpo_train.py`: trains a model with GRPO on a given dataset.
+- `Makefile`: contains easy-to-run commands for the dataset and training workflows using previous scripts.
+
+## Main Phases
+
+We use the DeepSeek-R1 [tech report](https://github.com/deepseek-ai/DeepSeek-R1) as a loose guide, but the project phases are:
+
+* Step 1: generate a linear algebra dataset with controlled difficulty and tool-call metadata.
+* Step 2: distill multi-turn tool-use data from a teacher model.
+* Step 3: SFT the base model on the dataset to teach the tool-calling format.
+* Step 4: GRPO fine-tune on the tool-use tasks, using a curriculum.
+
+
+## Installation
+
+We use `uv` as the dependency management tool.
+First, to install `uv`, follow the [UV Installation Guide](https://docs.astral.sh/uv/getting-started/installation/).
+
+To run the experiments install the dependencies using:
+
+* For generation/distillation: `make install-data-gen`
+* For SFT: `make install-sft`
+* For RL: `make install-grpo`
+
+Next, log into your Hugging Face and Weights and Biases accounts as follows:
+
+```shell
+huggingface-cli login
+wandb login
+```
+
+## Quickstart
+
+After installing dependencies above, run the commands below. For modifications, see the config files.
+
+```shell
+# Phase 1: Generate dataset
+uv run python linalg_zero/generate.py --dataset_name atomwalk12/linalgzero --push_dataset
+
+# Phase 2: Distillation (setup once)
+cp linalg_zero/config/distillation/env.example.sh env.sh
+# Edit env.sh to set HF_TOKEN and ARGILLA_API_KEY.
+source env.sh
+
+# Terminal A
+uv run python linalg_zero/distillation/launch_server.py --config linalg_zero/config/distillation/vllm_qwen3_32b.yaml
+
+# Terminal B (new terminal; source env.sh again)
+source env.sh
+uv run python linalg_zero/distillation.py --config linalg_zero/config/distillation/vllm_qwen3_32b.yaml
+
+# Phase 3: SFT
+uv run python linalg_zero/sft_train.py --config linalg_zero/config/sft/qwen2.5-3B/lora.yaml
+
+# Phase 4: GRPO
+uv run python linalg_zero/grpo_train.py --config-name runpod.yaml
+```
+
+Training requires the dataset to follow the strict OpenAI tool-calling format (see [this link](https://huggingface.co/docs/trl/en/dataset_formats#tool-calling)). We provide scripts to prepare and validate the data accordingly:
+
+- `linalg_zero/`
+    - `sft/scripts/prepare_dataset.py`: prepares the SFT dataset.
+    - `grpo/scripts/prepare_dataset.py`: prepares and validates the GRPO dataset.
+
+## Results
+
+We provide a recipe to encourage planning and tool-use capabilities in the [Qwen2.5-3B](https://huggingface.co/Qwen/Qwen2.5-3B) model, starting from a pre-trained (not instruction-tuned) base model.
+
+This yields models like [Linalg-Zero-SFT](https://huggingface.co/atomwalk12/LinalgZero-SFT) and [Linalg-Zero-GRPO](https://huggingface.co/atomwalk12/LinalgZero-GRPO), with the following downstream performance on the test set:
+
+
+| Metric             | LinAlgZero-SFT | LinAlgZero-GRPO |
+|--------------------|----------------|-----------------|
+| Optimal Trajectory | 89.87%         | 90.26%          |
+| Correctness        | 91.86%         | 92.63%          |
+| Format Validity    | 96.15%         | 96.66%          |
+| Tool Success       | 100.00%        | 100.00%         |
+
+### Artifacts
+
+| Artifact | Link |
+|---|---|
+| SFT checkpoint | [atomwalk12/LinalgZero-SFT](https://huggingface.co/atomwalk12/LinalgZero-SFT) |
+| GRPO checkpoint | [atomwalk12/LinAlgZero-GRPO](https://huggingface.co/atomwalk12/LinAlgZero-GRPO) |
+| Base dataset | [atomwalk12/linalgzero](https://huggingface.co/datasets/atomwalk12/linalgzero) |
+| Distilled dataset (clean) | [atomwalk12/linalgzero-distilled-clean](https://huggingface.co/datasets/atomwalk12/linalgzero-distilled-clean) |
+| SFT dataset | [atomwalk12/linalgzero-sft](https://huggingface.co/datasets/atomwalk12/linalgzero-sft) |
+| GRPO dataset | [atomwalk12/linalgzero-grpo](https://huggingface.co/datasets/atomwalk12/linalgzero-grpo) |
+
+## Reproducibility
+- **Distillation:** H100 80GB on [Runpod](https://www.runpod.io/) with [Qwen/Qwen3-32B-FP8](https://huggingface.co/Qwen/Qwen3-32B-FP8); 15 hours at $2.39/hr (~$25).
+- **SFT:** Local 24GB RTX 4090 with [Qwen/Qwen2.5-3B](https://huggingface.co/Qwen/Qwen2.5-3B).
+- **GRPO:** RTX 6000 Ada on [Runpod](https://www.runpod.io/), improving on the SFT baseline; 57 hours at $0.77/hr (~$50).
+- **Total:** ~$75 using a mix of cloud GPUs and local training.
+
+## Acknowledgements
+- We base our distillation pipeline on [distilabel](https://github.com/argilla-io/distilabel).
+- We base the RL experiment on [ART](https://deepwiki.com/OpenPipe/ART).
+- We use Qwen2.5 series base model [Qwen2.5](https://github.com/QwenLM/Qwen2.5).
+
+## Citation
+
+If you find this project is useful in your own work, please consider citing as follows:
+
+```bibtex
+@misc{linalg-zero,
+    title = {Linalg-Zero: Distilling Neurosymbolic Reasoning for Linear Algebra in Small Language Models},
+    url = {https://github.com/atomwalk12/linalg-zero},
+    author = {{Razvan F. Vasile}},
+    month = {March},
+    year = {2026}
+}
+```

app.py

@@ -1,7 +0,0 @@
-import gradio as gr
-
-def greet(name):
-    return "Hello " + name + "!!"
-
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()

codecov.yaml

@@ -0,0 +1,9 @@
+coverage:
+  range: 70..100
+  round: down
+  precision: 1
+  status:
+    project:
+      default:
+        target: 90%
+        threshold: 0.5%

env/.env.distillation.example

@@ -0,0 +1,14 @@
+#!/bin/bash
+
+# Distillation
+export HF_HOME=/workspace/linalg-zero/hf-cache
+export HUGGINGFACE_HUB_CACHE=/workspace/linalg-zero/hf-cache
+export VLLM_WORKDIR=/workspace/linalg-zero/vllm-cache
+export UV_CACHE_DIR=/workspace/linalg-zero/uv-cache
+export XDG_CACHE_HOME=/workspace/linalg-zero/.cache
+export PIP_CACHE_DIR=/workspace/linalg-zero/pip-cache
+export USING_VLLM=true
+export INFERENCE_BACKEND=vllm
+export ARGILLA_API_URL=https://atomwalk12-linalgzero-distilled.hf.space
+export ARGILLA_API_KEY=<...>
+export HF_TOKEN=<...>

env/.env.grpo.example

@@ -0,0 +1,9 @@
+OPENPIPE_API_KEY=<insert-key>
+WANDB_API_KEY=<insert-key>
+VLLM_API_KEY="my-secret-key"
+ACCELERATE_MIXED_PRECISION=bf16
+VLLM_NO_USAGE_STATS=1
+
+HF_HUB_NAMESPACE=<...>
+HF_TOKEN=<...>
+HF_REPO_PRIVATE=1

env/.env.grpo.runpod.example

@@ -0,0 +1,16 @@
+export HF_HOME=/workspace/linalg-zero/hf-cache
+export HUGGINGFACE_HUB_CACHE=/workspace/linalg-zero/hf-cache
+export VLLM_WORKDIR=/workspace/linalg-zero/vllm-cache
+export UV_CACHE_DIR=/workspace/linalg-zero/uv-cache
+export XDG_CACHE_HOME=/workspace/linalg-zero/.cache
+export PIP_CACHE_DIR=/workspace/linalg-zero/pip-cache
+
+export OPENPIPE_API_KEY=<...>
+export WANDB_API_KEY=<...>
+export VLLM_API_KEY="my-secret-key"
+export ACCELERATE_MIXED_PRECISION=bf16
+export VLLM_NO_USAGE_STATS=1
+
+export HF_HUB_NAMESPACE=<...>
+export HF_TOKEN=<...>
+export HF_REPO_PRIVATE=1

linalg_zero/config/cleaning_config.yaml

@@ -0,0 +1,42 @@
+# Initial dirty: 49 entries
+
+##############
+# Tool calls #
+##############
+120:
+  to_remove: [4, 5]
+  remove_reason:
+    "4,5": "exact msg/tool call with 2,3"
+  # Note that when counting these fields, they start from 0.
+  initial_msg_count: 9
+  expected_final_msg_count: 7
+
+
+419:
+  to_remove: [8, 9]
+  remove_reason:
+    "8,9": "exact msg/tool call with 6,7"
+  initial_msg_count: 11
+  expected_final_msg_count: 9
+
+
+1203:
+  to_remove: [4,5]
+  remove_reason:
+    "4,5": "exact msg/tool call with 2,3"
+  initial_msg_count: 7
+  expected_final_msg_count: 5
+
+################
+# Final answer #
+################
+1052:
+  to_remove: [4,5]
+  remove_reason:
+    "4,5": "exact msg/tool call with 2,3"
+  to_replace:
+    6:
+      "content":
+        "think": "<think>This tool call completes the task. Now, I can provide the final answer.</think>"
+  initial_msg_count: 7
+  expected_final_msg_count: 5

linalg_zero/config/data.py

@@ -0,0 +1,341 @@
+from dataclasses import dataclass, field
+
+from trl.scripts.utils import ScriptArguments as ScriptArgs
+from trl.trainer.model_config import ModelConfig
+
+
+@dataclass
+class ScriptArguments(ScriptArgs):
+    """
+    Extended version of ScriptArguments with support for dataset mixtures.
+    """
+
+    dataset_name: str | None = field(
+        default=None, metadata={"help": "Training dataset name. Contains chain-of-thought solutions."}
+    )
+    eval_dataset_config: str | None = field(default=None, metadata={"help": "Evaluation dataset config."})
+    take_n: int | None = field(default=None, metadata={"help": "Number of examples to take from the dataset."})
+
+
+@dataclass
+class SFTModelConfig(ModelConfig):
+    enforce_eager: bool | None = field(default=None, metadata={"help": "Whether to enforce eager execution."})
+
+
+@dataclass
+class SFTRunConfig:
+    add_special_tokens: bool = field(
+        metadata={"help": "Whether to add special tokens to the model."},
+    )
+
+    early_stopping_patience: int = field(
+        default=3, metadata={"help": "The number of epochs to wait before early stopping."}
+    )
+
+    early_stopping_threshold: float = field(
+        default=0.0, metadata={"help": "Minimum improvement required to reset patience counter."}
+    )
+
+    benchmarks: list[str] = field(
+        default_factory=lambda: [],
+        metadata={"help": "The benchmarks to run after training."},
+    )
+    callbacks: list[str] = field(
+        default_factory=lambda: [],
+        metadata={"help": "The callbacks to run during training."},
+    )
+    chat_template: str | None = field(default=None, metadata={"help": "The chat template to use."})
+    system_prompt: str | None = field(
+        default=None,
+        metadata={"help": "The optional system prompt to use for benchmarking."},
+    )
+    hub_model_revision: str | None = field(
+        default="main",
+        metadata={"help": "The Hub model branch to push the model to."},
+    )
+    overwrite_hub_revision: bool = field(default=False, metadata={"help": "Whether to overwrite the Hub revision."})
+    push_to_hub_revision: bool = field(default=False, metadata={"help": "Whether to push to a Hub revision/branch."})
+    wandb_entity: str | None = field(
+        default=None,
+        metadata={"help": ("The entity to store runs under.")},
+    )
+    wandb_project: str | None = field(
+        default=None,
+        metadata={"help": ("The project to store runs under.")},
+    )
+    wandb_run_group: str | None = field(
+        default=None,
+        metadata={"help": ("The group to store runs under.")},
+    )
+    wandb_run_id: str | None = field(default=None, metadata={"help": {"The wandb run id."}})
+
+    eval_max_new_tokens: int | None = field(
+        default=None,
+        metadata={"help": "Max new tokens for evaluation callbacks (does not affect training)."},
+    )
+
+    max_seq_length: int | None = field(
+        default=None,
+        metadata={"help": "Max sequence length for evaluation callbacks (does not affect training)."},
+    )
+
+    gpu_memory_utilization: float | None = field(
+        default=0.95, metadata={"help": "Fraction of GPU memory to be used by vLLM (0-1)"}
+    )
+
+    # Evaluation sampling
+    max_eval_samples: int | None = field(
+        default=None,
+        metadata={"help": "Maximum number of eval samples for periodic evaluations. None/-1 for full dataset."},
+    )
+    final_eval_max_samples: int | None = field(
+        default=None,
+        metadata={"help": "Maximum number of eval samples for periodic evaluations. None/-1 for full dataset."},
+    )
+
+
+@dataclass
+class DatasetGenerationConfig:
+    """
+    Data class that stores the dataset generation parameters.
+
+    Args:
+        dataset_name (str): The name of the dataset to generate.
+    """
+
+    dataset_name: str | None = field(
+        metadata={"help": "Should be the name used to store the dataset on the Hugging Face Hub."},
+    )
+
+
+@dataclass
+class LlamaCppServerConfig:
+    """
+    Data class that stores LlamaCPP server parameters with llama_cpp_ prefix.
+    """
+
+    def __post_init__(self) -> None:
+        pass
+
+    # Server parameters
+    host: str = field(
+        metadata={"help": "Host address to bind to"},
+    )
+    port: int = field(
+        metadata={"help": "Port to listen on"},
+    )
+    n_ctx: int = field(
+        metadata={"help": "Context size"},
+    )
+    split_mode: int = field(
+        metadata={"help": "Split mode (0=none, 1=layer, 2=row)"},
+    )
+
+    # Model parameters
+    n_gpu_layers: int = field(
+        metadata={"help": "Number of GPU layers to offload"},
+    )
+
+    model: str = field(
+        metadata={"help": "Model URL to download (GGUF format)"},
+    )
+
+    hf_pretrained_model_name_or_path: str | None = field(
+        default=None,
+        metadata={"help": "Huggingface repository ID to ensure that the correct tokenizer is used."},
+    )
+
+    hf_model_repo_id: str | None = field(
+        default=None,
+        metadata={"help": "Path to the repository where the model is stored."},
+    )
+
+
+@dataclass
+class VllmServerConfig:
+    """
+    Data class that stores vLLM server parameters with vllm_ prefix.
+    """
+
+    # Model parameters
+    model: str = field(
+        metadata={"help": "Model name (HuggingFace format)"},
+    )
+
+    # Server parameters
+    host: str = field(
+        metadata={"help": "Host address to bind to"},
+    )
+    port: int = field(
+        metadata={"help": "Port to listen on"},
+    )
+    enable_auto_tool_choice: bool = field(
+        metadata={"help": "Enable automatic tool choice"},
+    )
+    tool_call_parser: str = field(
+        metadata={"help": "Tool call parser to use"},
+    )
+    chat_template: str | None = field(
+        default=None,
+        metadata={"help": "Chat template to use"},
+    )
+    quantization: str | None = field(
+        default=None,
+        metadata={"help": "Quantization to use"},
+    )
+    api_key: str = field(
+        default="not-used",
+        metadata={"help": "API key for authentication (use 'not-used' for local development)"},
+    )
+
+    # Memory / performance tuning parameters
+    dtype: str | None = field(
+        default=None,
+        metadata={"help": "Computation dtype for model weights and activations (e.g., float16)"},
+    )
+    kv_cache_dtype: str | None = field(
+        default=None,
+        metadata={"help": "KV cache dtype (auto, fp8, fp8_e4m3, fp8_e5m2)"},
+    )
+    max_model_len: int | None = field(
+        default=None,
+        metadata={"help": "Maximum model context length (tokens)"},
+    )
+    max_num_seqs: int | None = field(
+        default=None,
+        metadata={"help": "Maximum number of concurrent sequences"},
+    )
+    gpu_memory_utilization: float | None = field(
+        default=None,
+        metadata={"help": "Fraction of GPU memory to be used by vLLM (0-1)"},
+    )
+    enforce_eager: bool | None = field(
+        default=None,
+        metadata={"help": "Disable CUDA graphs to reduce memory usage"},
+    )
+    swap_space: int | None = field(
+        default=None,
+        metadata={"help": "CPU swap space in GB per GPU for paging KV cache"},
+    )
+    max_num_batched_tokens: int | None = field(
+        default=None,
+        metadata={"help": "Limit number of tokens processed per batch (prefill)"},
+    )
+    tensor_parallel_size: int | None = field(
+        default=None,
+        metadata={"help": "Tensor parallelism degree"},
+    )
+    enable_chunked_prefill: bool | None = field(
+        default=None,
+        metadata={"help": "Enable chunked prefill to reduce peak prefill memory"},
+    )
+
+    # Model parameters
+    reasoning_parser: str | None = field(
+        default=None,
+        metadata={"help": "Reasoning parser to use"},
+    )
+
+
+@dataclass
+class DistillationConfig:
+    """
+    Data class that stores the distillation pipeline parameters.
+    """
+
+    # Dataset parameters
+    dataset_name: str | None = field(
+        metadata={"help": "HuggingFace dataset to load"},
+    )
+
+    # Prompt parameters
+    prompt_column: str = field(
+        metadata={"help": "Column name for prompt data"},
+    )
+    prompt_template: str = field(
+        metadata={"help": "Template string for formatting prompts"},
+    )
+
+    # Generation parameters (non-defaults first)
+    model_type: str | None = field(metadata={"help": "Model type for generation"})
+    enable_reasoning: bool = field(metadata={"help": "Whether to enable thinking"})
+    max_new_tokens: int = field(
+        metadata={"help": "Maximum number of new tokens to generate"},
+    )
+    num_generations: int = field(
+        metadata={"help": "Number of generations per problem"},
+    )
+
+    # Processing parameters
+    input_batch_size: int = field(
+        metadata={"help": "Batch size for input processing"},
+    )
+    use_cache: bool = field(
+        metadata={"help": "Whether to use cache for the pipeline. This can enable error recovery."},
+    )
+
+    timeout: int = field(
+        metadata={"help": "Request timeout in seconds"},
+    )
+    retries: int = field(
+        metadata={"help": "Number of retries for failed requests"},
+    )
+
+    # Output parameters
+    hf_output_dataset: str | None = field(
+        metadata={"help": "HuggingFace repo to push results to"},
+    )
+    argilla_output_dataset: str | None = field(
+        metadata={"help": "Argilla dataset to push results to. This is used for manual annotation."},
+    )
+    private: bool = field(
+        metadata={"help": "Whether to make the output dataset private when pushing to HF Hub"},
+    )
+
+    # Generation parameters
+    n_turns: int = field(
+        metadata={"help": "Number of turns to generate"},
+    )
+    min_successful_completions: int = field(
+        default=-1,
+        metadata={"help": "Minimum number of successful completions to generate"},
+    )
+    strip_think_prefix: bool = field(
+        default=True,
+        metadata={"help": "Whether to strip the think prefix from the conversation. This is needed for Qwen3 models."},
+    )
+    # Optional stopping sequences (must come after non-default fields)
+    stop: list[str] | None = field(
+        default=None,
+        metadata={"help": "Stop sequences for generation (each string is a stop token)"},
+    )
+
+    debug_mode: bool = field(
+        default=False,
+        metadata={"help": "Whether to do evaluation"},
+    )
+
+    take_n: int | None = field(
+        default=None,
+        metadata={"help": "Number of examples to take from the dataset."},
+    )
+
+    structured_output: bool = field(
+        default=False,
+        metadata={"help": "Whether to use structured output"},
+    )
+
+    deterministic: bool = field(
+        default=True,
+        metadata={"help": "Make generation deterministic (temperature=0, top_p=1)"},
+    )
+
+    client_replicas: int | None = field(
+        default=None,
+        metadata={"help": "Number of client replicas for parallel processing"},
+    )
+
+    dataset_config: str | None = field(
+        default=None,
+        metadata={"help": "Dataset config to use"},
+    )

linalg_zero/config/dataset/default_debug.yml

@@ -0,0 +1 @@
+dataset_name: "linalg_zero"

linalg_zero/config/distillation/env.example.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+
+export HF_HOME=/workspace/linalg-zero/hf-cache
+export HUGGINGFACE_HUB_CACHE=/workspace/linalg-zero/hf-cache
+export VLLM_WORKDIR=/workspace/linalg-zero/vllm-cache
+export UV_CACHE_DIR=/workspace/linalg-zero/uv-cache
+export XDG_CACHE_HOME=/workspace/linalg-zero/.cache
+export PIP_CACHE_DIR=/workspace/linalg-zero/pip-cache
+export USING_VLLM=true
+export INFERENCE_BACKEND=vllm
+export ARGILLA_API_URL=https://atomwalk12-linalgzero-distilled.hf.space
+export ARGILLA_API_KEY=<token>
+export HF_TOKEN=<token>

linalg_zero/config/distillation/llamacpp_qwen3_30b_A3B_instruct.yaml

@@ -0,0 +1,49 @@
+######################
+# DistillationConfig #
+######################
+
+# Dataset parameters
+dataset_name: "atomwalk12/linalgzero"
+debug_mode: True
+take_n: 20
+n_turns: 6
+
+# Prompt parameters
+prompt_column: "prompt"
+prompt_template: "{{ instruction }}"
+
+# Generation parameters
+model_type: "default"
+deterministic: true
+max_new_tokens: 8192
+num_generations: 1
+
+# Processing parameters
+input_batch_size: 5
+use_cache: true
+client_replicas: 1
+timeout: 600
+retries: 2
+
+# Output parameters
+hf_output_dataset: "atomwalk12/linalgzero-distilled-debug"
+argilla_output_dataset: "atomwalk12/linalgzero-distilled-debug"
+private: false
+
+########################
+# LlamaCppServerConfig #
+########################
+
+# Server parameters
+host: "localhost"
+port: 8000
+
+# Model
+model: "Qwen3-30B-A3B-Instruct-2507-UD-Q5_K_XL.gguf"
+hf_model_repo_id: "unsloth/Qwen3-30B-A3B-Instruct-2507-GGUF"
+hf_pretrained_model_name_or_path: "Qwen/Qwen3-32B"
+enable_reasoning: true
+
+n_gpu_layers: 45
+n_ctx: 8192
+split_mode: 2

linalg_zero/config/distillation/llamacpp_qwen3_30b_A3B_think.yaml

@@ -0,0 +1,49 @@
+######################
+# DistillationConfig #
+######################
+
+# Dataset parameters
+dataset_name: "atomwalk12/linalgzero"
+debug_mode: True
+take_n: 20
+n_turns: 6
+
+# Prompt parameters
+prompt_column: "prompt"
+prompt_template: "{{ instruction }}"
+
+# Generation parameters
+model_type: "default"
+deterministic: true
+max_new_tokens: 8192
+num_generations: 1
+
+# Processing parameters
+input_batch_size: 5
+use_cache: true
+client_replicas: 1
+timeout: 600
+retries: 2
+
+# Output parameters
+hf_output_dataset: "atomwalk12/linalgzero-distilled-debug"
+argilla_output_dataset: "atomwalk12/linalgzero-distilled-debug"
+private: false
+
+########################
+# LlamaCppServerConfig #
+########################
+
+# Server parameters
+host: "localhost"
+port: 8000
+
+# Model
+model: "Qwen3-30B-A3B-Thinking-2507-UD-Q5_K_XL.gguf"
+hf_model_repo_id: "unsloth/Qwen3-30B-A3B-Thinking-2507-GGUF"
+hf_pretrained_model_name_or_path: "Qwen/Qwen3-32B"
+enable_reasoning: true
+
+n_gpu_layers: 45
+n_ctx: 8192
+split_mode: 2

linalg_zero/config/distillation/llamacpp_qwen3_32b_instruct.yaml

@@ -0,0 +1,49 @@
+######################
+# DistillationConfig #
+######################
+
+# Dataset parameters
+dataset_name: "atomwalk12/linalgzero"
+debug_mode: False
+take_n: 20
+n_turns: 6
+
+# Prompt parameters
+prompt_column: "prompt"
+prompt_template: "{{ instruction }}"
+
+# Generation parameters
+model_type: "default"
+deterministic: true
+max_new_tokens: 8192
+num_generations: 1
+
+# Processing parameters
+input_batch_size: 7
+use_cache: true
+client_replicas: 1
+timeout: 600
+retries: 2
+
+# Output parameters
+hf_output_dataset: "atomwalk12/linalgzero-distilled-debug"
+argilla_output_dataset: "atomwalk12/linalgzero-distilled-debug"
+private: false
+
+########################
+# LlamaCppServerConfig #
+########################
+
+# Server parameters
+host: "localhost"
+port: 8000
+
+# Model
+model: "Qwen3-32B-Q4_K_M.gguf"
+hf_model_repo_id: "unsloth/Qwen3-32B-GGUF"
+hf_pretrained_model_name_or_path: "Qwen/Qwen3-32B"
+enable_reasoning: true
+
+n_gpu_layers: 60
+n_ctx: 8192
+split_mode: 2

linalg_zero/config/distillation/qwen3_think_vllm_debug.yaml

@@ -0,0 +1,65 @@
+# Distillation Configuration for vLLM inference
+
+# Dataset parameters
+dataset_name: "atomwalk12/linalgzero"
+# dataset_name: "atomwalk12/linalg-distilled-debug"
+debug_mode: True
+dataset_config: null
+take_n: 20
+
+# Prompt parameters
+prompt_column: "prompt"
+prompt_template: "{{ instruction }}"
+
+# Generation parameters
+model_type: "default"
+deterministic: true
+max_new_tokens: 8192 #  32768
+num_generations: 1
+
+# Processing parameters
+input_batch_size: 3
+use_cache: true
+client_replicas: 1
+timeout: 600
+retries: 2
+
+# Output parameters
+hf_output_dataset: "atomwalk12/linalgzero-distilled-debug"
+argilla_output_dataset: "atomwalk12/linalgzero-distilled-debug"
+private: false
+
+# Server parameters
+# To inspect the various model/chat template combinations:
+# https://github.com/vllm-project/vllm/blob/main/docs/features/tool_calling.md
+
+model: "Eslzzyl/Qwen3-4B-Thinking-2507-AWQ"
+tool_call_parser: "hermes"
+reasoning_parser: "qwen3" # "deepseek_r1"
+
+#model: "lurker18/Llama_3.1_8B_Instruct_AWQ_4bit"
+#tool_call_parser: "llama3_json"
+#chat_template: "linalg_zero/distillation/vllm/tool_chat_template_llama3.1_json.jinja"
+
+#model: "solidrust/Mistral-7B-Instruct-v0.3-AWQ"
+#tool_call_parser: "mistral"
+#chat_template: "linalg_zero/distillation/vllm/tool_chat_template_mistral_parallel.jinja"
+
+enable_auto_tool_choice: true
+quantization: "awq"
+host: "localhost"
+port: 8000
+
+# Memory / performance tuning for 24GB GPU
+dtype: "float16"
+kv_cache_dtype: "fp8"
+max_num_seqs: 1
+gpu_memory_utilization: 0.85
+enforce_eager: true
+swap_space: 4
+max_num_batched_tokens: 2048
+max_model_len: 32768 # 131072
+enable_chunked_prefill: true
+
+# Distillation parameters
+n_turns: 6

linalg_zero/config/distillation/vllm_qwen3_32b.yaml

@@ -0,0 +1,46 @@
+# Dataset parameters
+dataset_name: "atomwalk12/linalgzero"
+debug_mode: True
+dataset_config: null
+n_turns: 8
+min_successful_completions: -1
+
+# Prompt parameters
+prompt_column: "prompt"
+prompt_template: "{{ instruction }}"
+
+# Generation parameters
+model_type: "default"
+deterministic: false
+max_new_tokens: 4096
+num_generations: 1
+enable_reasoning: true
+
+# Processing parameters
+input_batch_size: 24
+use_cache: true
+timeout: 600
+retries: 5
+
+# Output parameters
+hf_output_dataset: "atomwalk12/linalgzero-distilled"
+argilla_output_dataset: "atomwalk12/linalgzero-distilled"
+private: false
+
+# Server parameters
+model: "Qwen/Qwen3-32B-FP8"
+tool_call_parser: "hermes"
+reasoning_parser: "qwen3"
+
+enable_auto_tool_choice: true
+host: "localhost"
+port: 8000
+
+dtype: "bfloat16"
+kv_cache_dtype: "auto"
+max_num_seqs: 24
+max_model_len: 24576
+gpu_memory_utilization: 0.93
+enable_chunked_prefill: true
+swap_space: 8
+max_num_batched_tokens: 1024

linalg_zero/config/distillation/vllm_qwen3_32b_debug.yaml

@@ -0,0 +1,46 @@
+# Dataset parameters
+dataset_name: "atomwalk12/linalgzero"
+debug_mode: False
+dataset_config: null
+n_turns: 8
+take_n: 500
+
+# Prompt parameters
+prompt_column: "prompt"
+prompt_template: "{{ instruction }}"
+
+# Generation parameters
+model_type: "default"
+deterministic: false
+max_new_tokens: 4096
+num_generations: 1
+enable_reasoning: true
+
+# Processing parameters
+input_batch_size: 24
+use_cache: true
+timeout: 600
+retries: 5
+
+# Output parameters
+hf_output_dataset: "atomwalk12/linalgzero-distilled-500"
+argilla_output_dataset: "atomwalk12/linalgzero-distilled-500"
+private: false
+
+# Server parameters
+model: "Qwen/Qwen3-32B-FP8"
+tool_call_parser: "hermes"
+reasoning_parser: "qwen3"
+
+enable_auto_tool_choice: true
+host: "localhost"
+port: 8000
+
+dtype: "bfloat16"
+kv_cache_dtype: "auto"
+max_num_seqs: 24
+max_model_len: 24576
+gpu_memory_utilization: 0.93
+enable_chunked_prefill: true
+swap_space: 8
+max_num_batched_tokens: 1024

linalg_zero/config/distillation/vllm_qwen3_4b_think.yaml

@@ -0,0 +1,52 @@
+# Dataset parameters
+dataset_name: "atomwalk12/linalgzero"
+debug_mode: true
+dataset_config: null
+n_turns: 6
+take_n: 7
+min_successful_completions: 8200
+
+# Prompt parameters
+prompt_column: "prompt"
+prompt_template: "{{ instruction }}"
+
+# Generation parameters
+model_type: "default"
+deterministic: false
+max_new_tokens: 2048
+num_generations: 1
+enable_reasoning: true
+strip_think_prefix: false
+
+# Processing parameters
+input_batch_size: 3
+use_cache: true
+client_replicas: 1
+timeout: 600
+retries: 2
+
+# Output parameters
+hf_output_dataset: "atomwalk12/linalgzero-distilled-debug"
+argilla_output_dataset: "atomwalk12/linalgzero-distilled-debug"
+private: false
+
+# Server parameters
+model: "Eslzzyl/Qwen3-4B-Thinking-2507-AWQ"
+tool_call_parser: "hermes"
+reasoning_parser: "qwen3"
+
+enable_auto_tool_choice: true
+quantization: "awq"
+host: "localhost"
+port: 8000
+
+# Memory / performance tuning for 24GB GPU
+dtype: "float16"
+kv_cache_dtype: "auto"
+max_num_seqs: 1
+gpu_memory_utilization: 0.75
+enforce_eager: true
+swap_space: 4
+max_num_batched_tokens: 2048
+max_model_len: 8096
+enable_chunked_prefill: true

linalg_zero/config/grpo/Qwen/Qwen2.5-3B/eval/linalgzero-grpo-110.yaml

@@ -0,0 +1,62 @@
+init:
+  max_lora_rank: 32
+  dtype: bfloat16
+  load_in_4bit: false
+  load_in_8bit: false
+  max_seq_length: 6144
+
+training:
+  trajectories_per_group: 4
+  groups_per_step: 8
+  learning_rate: 2e-6
+  beta: 0.03
+  eval_steps: 100
+  num_epochs: 2
+  train_mode: sync_rl
+  chat_template: "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- else %}\n        {{- 'You are a helpful assistant.' }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- else %}\n        {{- '<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n"
+
+run:
+  project_id: 001-linalgzero-sft-110
+  project: linalgzero-eval
+  base_model: "results/LinalgZero-SFT-110-checkpoint-300/checkpoint-300"
+  model: "results/LinalgZero-SFT-110-checkpoint-300/checkpoint-300"
+  dataset_path: "atomwalk12/linalgzero-grpo"
+  env: linear_algebra
+  model_provider: hosted_vllm
+  user_model_provider: openai
+  user_strategy: mathematician
+  user_model: gpt-4o
+  agent_strategy: tool-calling-rl
+  temperature: 0.7
+  top_p: 0.8
+  repetition_penalty: 1.0
+  task_split: test
+  log_dir: rl_results
+  skip_eval: False
+  in_process: False
+  max_assistant_turns: 5
+  max_completion_tokens: 800
+  skip_special_tokens: False
+  stop: [
+    "<tool_response>",
+    "</tool_response>",
+    "User:",
+  ]
+  seed: 42
+
+trainer:
+  max_prompt_length: 5000
+  max_completion_length: 800
+  report_to: ["wandb"]
+  seed: 42
+  data_seed: 42
+
+engine:
+  gpu_memory_utilization: 0.6
+  max_lora_rank: 32
+  enable_sleep_mode: true
+  max_model_len: 6144
+  dtype: 'bfloat16'
+  quantization: null
+  load_format: 'auto'
+  seed: 42

linalg_zero/config/grpo/Qwen/Qwen2.5-3B/eval/linalgzero-grpo.yaml

@@ -0,0 +1,62 @@
+init:
+  max_lora_rank: 32
+  dtype: bfloat16
+  load_in_4bit: false
+  load_in_8bit: false
+  max_seq_length: 10240
+
+training:
+  trajectories_per_group: 4
+  groups_per_step: 8
+  learning_rate: 2e-6
+  beta: 0.03
+  eval_steps: 100
+  num_epochs: 2
+  train_mode: sync_rl
+  chat_template: "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- else %}\n        {{- 'You are a helpful assistant.' }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- else %}\n        {{- '<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n"
+
+run:
+  project_id: LinAlgZero-GRPO-001
+  project: linalgzero-eval
+  base_model: "atomwalk12/LinAlgZero-GRPO"
+  model: "atomwalk12/LinAlgZero-GRPO"
+  dataset_path: "atomwalk12/linalgzero-grpo"
+  env: linear_algebra
+  model_provider: hosted_vllm
+  user_model_provider: openai
+  user_strategy: mathematician
+  user_model: gpt-4o
+  agent_strategy: tool-calling-rl
+  temperature: 0.0
+  top_p: null
+  repetition_penalty: 1.0
+  task_split: test
+  log_dir: rl_results
+  skip_eval: False
+  in_process: False
+  max_assistant_turns: 5
+  max_completion_tokens: 1536
+  skip_special_tokens: False
+  stop: [
+    "<tool_response>",
+    "</tool_response>",
+    "User:",
+  ]
+  seed: 42
+
+trainer:
+  max_prompt_length: 8192
+  max_completion_length: 1536
+  report_to: ["wandb"]
+  seed: 42
+  data_seed: 42
+
+engine:
+  gpu_memory_utilization: 0.6
+  max_lora_rank: 32
+  enable_sleep_mode: true
+  max_model_len: 10240
+  dtype: 'bfloat16'
+  quantization: null
+  load_format: 'auto'
+  seed: 42

linalg_zero/config/grpo/Qwen/Qwen2.5-3B/eval/linalgzero-sft-110.yaml

@@ -0,0 +1,62 @@
+init:
+  max_lora_rank: 32
+  dtype: bfloat16
+  load_in_4bit: false
+  load_in_8bit: false
+  max_seq_length: 6144
+
+training:
+  trajectories_per_group: 4
+  groups_per_step: 8
+  learning_rate: 2e-6
+  beta: 0.03
+  eval_steps: 100
+  num_epochs: 2
+  train_mode: sync_rl
+  chat_template: "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- else %}\n        {{- 'You are a helpful assistant.' }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- else %}\n        {{- '<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n"
+
+run:
+  project_id: 001-linalgzero-sft-110
+  project: linalgzero-eval
+  base_model: "results/LinalgZero-SFT-110-checkpoint-300/checkpoint-300"
+  model: "results/LinalgZero-SFT-110-checkpoint-300/checkpoint-300"
+  dataset_path: "atomwalk12/linalgzero-grpo"
+  env: linear_algebra
+  model_provider: hosted_vllm
+  user_model_provider: openai
+  user_strategy: mathematician
+  user_model: gpt-4o
+  agent_strategy: tool-calling-rl
+  temperature: 0.7
+  top_p: 0.8
+  repetition_penalty: 1.0
+  task_split: test
+  log_dir: rl_results
+  skip_eval: False
+  in_process: False
+  max_assistant_turns: 5
+  max_completion_tokens: 800
+  skip_special_tokens: False
+  stop: [
+    "<tool_response>",
+    "</tool_response>",
+    "User:",
+  ]
+  seed: 42
+
+trainer:
+  max_prompt_length: 5000
+  max_completion_length: 800
+  report_to: ["wandb"]
+  seed: 42
+  data_seed: 42
+
+engine:
+  gpu_memory_utilization: 0.6
+  max_lora_rank: 32
+  enable_sleep_mode: true
+  max_model_len: 6144
+  dtype: 'bfloat16'
+  quantization: null
+  load_format: 'auto'
+  seed: 42

linalg_zero/config/grpo/Qwen/Qwen2.5-3B/eval/linalgzero-sft.yaml

@@ -0,0 +1,62 @@
+init:
+  max_lora_rank: 32
+  dtype: bfloat16
+  load_in_4bit: false
+  load_in_8bit: false
+  max_seq_length: 6144
+
+training:
+  trajectories_per_group: 4
+  groups_per_step: 8
+  learning_rate: 2e-6
+  beta: 0.03
+  eval_steps: 100
+  num_epochs: 2
+  train_mode: sync_rl
+  chat_template: "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- else %}\n        {{- 'You are a helpful assistant.' }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- else %}\n        {{- '<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n"
+
+run:
+  project_id: 004-linalgzero-sft
+  project: linalgzero-eval
+  base_model: "atomwalk12/LinAlgZero-SFT"
+  model: "atomwalk12/LinAlgZero-SFT"
+  dataset_path: "atomwalk12/linalgzero-grpo"
+  env: linear_algebra
+  model_provider: hosted_vllm
+  user_model_provider: openai
+  user_strategy: mathematician
+  user_model: gpt-4o
+  agent_strategy: tool-calling-rl
+  temperature: 0.0
+  top_p: null
+  repetition_penalty: 1.0
+  task_split: test
+  log_dir: rl_results
+  skip_eval: False
+  in_process: False
+  max_assistant_turns: 5
+  max_completion_tokens: 800
+  skip_special_tokens: False
+  stop: [
+    "<tool_response>",
+    "</tool_response>",
+    "User:",
+  ]
+  seed: 42
+
+trainer:
+  max_prompt_length: 5000
+  max_completion_length: 800
+  report_to: ["wandb"]
+  seed: 42
+  data_seed: 42
+
+engine:
+  gpu_memory_utilization: 0.6
+  max_lora_rank: 32
+  enable_sleep_mode: true
+  max_model_len: 6144
+  dtype: 'bfloat16'
+  quantization: null
+  load_format: 'auto'
+  seed: 42

linalg_zero/config/grpo/Qwen/Qwen2.5-3B/local.yaml

@@ -0,0 +1,95 @@
+defaults:
+  - /training@training: training_schema
+  - /run@run: run_schema
+  - _self_
+
+
+init:
+  max_lora_rank: 32
+  dtype: bfloat16
+  load_in_4bit: false
+  load_in_8bit: false
+  max_seq_length: 10240
+
+training:
+  eval_retries: 3
+  trajectories_per_group: 8
+  groups_per_step: 16
+  learning_rate: 1e-6
+  beta: 0.0
+  eval_steps: 100
+  num_epochs: 3
+  train_mode: sync_rl
+  importance_sampling_level: sequence
+  scale_rewards: True
+  epsilon: 3e-4
+  epsilon_high: 4e-4
+  chat_template: "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- else %}\n        {{- 'You are a helpful assistant.' }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- else %}\n        {{- '<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n"
+
+run:
+  project_id: test
+  project: linalgzero-temp
+  base_model: "atomwalk12/LinalgZero-SFT"
+  model: "atomwalk12/LinalgZero-SFT"
+  dataset_path: "atomwalk12/linalgzero-grpo"
+  env: linear_algebra
+  curriculum:
+    enabled: true
+    metric: tool_calls
+    sampling: mixture
+    initial_max_tool_calls: 1
+    final_max_tool_calls: 3
+    mixture_sigma: 0.5
+    mixture_min_prob_easiest: 0.1
+    fraction_at_start: 1.0
+    fraction_at_end: 1.0
+    min_total_tasks: 64
+  model_provider: hosted_vllm
+  user_strategy: mathematician
+  agent_strategy: tool-calling-rl
+  temperature: 1.0
+  top_p: 1.0
+  repetition_penalty: 1.0
+  task_split: test
+  log_dir: rl_results
+  skip_eval: False
+  in_process: False
+  max_assistant_turns: 5
+  max_completion_tokens: 1536
+  skip_special_tokens: False
+  stop: [
+    "<tool_response>",
+    "</tool_response>",
+    "User:",
+  ]
+  seed: 42
+
+peft:
+  r: 32
+  lora_alpha: 64
+  use_gradient_checkpointing: "unsloth"
+  target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
+  random_state: 42
+
+trainer:
+  learning_rate: 1e-6
+  per_device_train_batch_size: 8
+  gradient_accumulation_steps: 1   # if != 1, deadlock occurs
+  num_generations: 2
+  max_prompt_length: 8192
+  max_completion_length: 1536
+  report_to: ["wandb"]
+  save_total_limit: 10
+  seed: 42
+  data_seed: 42
+  max_grad_norm: 1.0
+
+engine:
+  gpu_memory_utilization: 0.33
+  max_lora_rank: 32
+  enable_sleep_mode: true
+  max_model_len: 10240
+  dtype: 'bfloat16'
+  quantization: null
+  load_format: 'auto'
+  seed: 42

linalg_zero/config/grpo/Qwen/Qwen2.5-3B/runpod.yaml

@@ -0,0 +1,95 @@
+defaults:
+  - /training@training: training_schema
+  - /run@run: run_schema
+  - _self_
+
+
+init:
+  max_lora_rank: 32
+  dtype: bfloat16
+  load_in_4bit: false
+  load_in_8bit: false
+  max_seq_length: 10240
+
+training:
+  eval_retries: 3
+  trajectories_per_group: 8
+  groups_per_step: 16
+  learning_rate: 1e-6
+  beta: 0.0
+  eval_steps: 100
+  num_epochs: 3
+  train_mode: sync_rl
+  importance_sampling_level: sequence
+  scale_rewards: True
+  epsilon: 3e-4
+  epsilon_high: 4e-4
+  chat_template: "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- else %}\n        {{- 'You are a helpful assistant.' }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- else %}\n        {{- '<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n"
+
+run:
+  project_id: linalgzero-3b-grpo-runpod
+  project: linalgzero-grpo
+  base_model: "atomwalk12/LinalgZero-SFT"
+  model: "atomwalk12/LinalgZero-SFT"
+  dataset_path: "atomwalk12/linalgzero-grpo"
+  env: linear_algebra
+  curriculum:
+    enabled: true
+    metric: tool_calls
+    sampling: mixture
+    initial_max_tool_calls: 1
+    final_max_tool_calls: 3
+    mixture_sigma: 0.5
+    mixture_min_prob_easiest: 0.1
+    fraction_at_start: 1.0
+    fraction_at_end: 1.0
+    min_total_tasks: 64
+  model_provider: hosted_vllm
+  user_strategy: mathematician
+  agent_strategy: tool-calling-rl
+  temperature: 1.0
+  top_p: 1.0
+  repetition_penalty: 1.0
+  task_split: test
+  log_dir: rl_results
+  skip_eval: False
+  in_process: False
+  max_assistant_turns: 5
+  max_completion_tokens: 1536
+  skip_special_tokens: False
+  stop: [
+    "<tool_response>",
+    "</tool_response>",
+    "User:",
+  ]
+  seed: 42
+
+peft:
+  r: 32
+  lora_alpha: 64
+  use_gradient_checkpointing: "unsloth"
+  target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
+  random_state: 42
+
+trainer:
+  learning_rate: 1e-6
+  per_device_train_batch_size: 8
+  gradient_accumulation_steps: 1   # if != 1, deadlock occurs
+  num_generations: 2
+  max_prompt_length: 8192
+  max_completion_length: 1536
+  report_to: ["wandb"]
+  save_total_limit: 10
+  seed: 42
+  data_seed: 42
+  max_grad_norm: 1.0
+
+engine:
+  gpu_memory_utilization: 0.9
+  max_lora_rank: 32
+  enable_sleep_mode: true
+  max_model_len: 10240
+  dtype: 'bfloat16'
+  quantization: null
+  load_format: 'auto'
+  seed: 42

linalg_zero/config/sft/accelerate/zero2.yaml

@@ -0,0 +1,21 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+deepspeed_config:
+  deepspeed_multinode_launcher: standard
+  offload_optimizer_device: none
+  offload_param_device: none
+  zero3_init_flag: false
+  zero_stage: 2
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 8
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

linalg_zero/config/sft/accelerate/zero3.yaml

@@ -0,0 +1,22 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+deepspeed_config:
+  deepspeed_multinode_launcher: standard
+  offload_optimizer_device: none
+  offload_param_device: none
+  zero3_init_flag: true
+  zero3_save_16bit_model: true
+  zero_stage: 3
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 1 # NOTE: The number of GPUs on one node
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

linalg_zero/config/sft/qwen2.5-3B/instruct.yaml

@@ -0,0 +1,78 @@
+# SFT config for a light touch-up of the Instruct model,
+# targeting only the input/output layers (embed_tokens and lm_head).
+
+# Model arguments
+model_name_or_path: Qwen/Qwen2.5-3B-Instruct
+model_revision: main
+torch_dtype: bfloat16
+attn_implementation: flash_attention_2
+
+# Quantization / PEFT
+load_in_4bit: false
+load_in_8bit: false            # Use 16bit training
+bnb_4bit_quant_type: nf4
+use_peft: false                # IMPORTANT: do not create new LoRA adapters
+
+# Data training arguments
+dataset_name: atomwalk12/linalgzero-sft
+dataset_config: default
+eval_dataset_config: default
+dataset_num_proc: 12
+dataset_text_field: text
+eos_token: <|im_end|>
+add_special_tokens: true
+
+# SFT trainer config
+bf16: true
+
+## Evaluation parameters
+do_eval: true                  # final evaluation at the end of training
+eval_on_start: true            # evaluation at the start of training
+eval_strategy: 'steps'
+eval_steps: 100
+max_eval_samples: 520
+final_eval_max_samples: 520
+
+metric_for_best_model: eval_format_accuracy
+greater_is_better: true
+load_best_model_at_end: true
+save_total_limit: 5
+eval_max_new_tokens: 800
+
+callbacks:
+  - tool_calling_accuracy
+
+## Training parameters
+gradient_accumulation_steps: 16
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+hub_model_id: atomwalk12/LinalgZero-SFT-Instruct
+hub_strategy: every_save
+learning_rate: 2e-5 # Alternatives: 5e-5/4e-5/3e-5/2e-5 and train for 2 epochs
+log_level: info
+logging_steps: 5
+logging_strategy: steps
+lr_scheduler_type: cosine_with_min_lr
+lr_scheduler_kwargs:
+  min_lr_rate: 0.1
+packing: false
+max_grad_norm: 0.2
+max_seq_length: 8192
+max_length: 8192
+max_steps: -1
+num_train_epochs: 1
+output_dir: results/LinalgZero-SFT-Instruct
+overwrite_output_dir: true
+per_device_eval_batch_size: 1
+per_device_train_batch_size: 1
+push_to_hub: true
+report_to:
+- wandb
+save_strategy: steps
+save_steps: 100
+seed: 42
+use_liger_kernel: false
+warmup_ratio: 0.03
+chat_template: "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- else %}\n        {{- 'You are a helpful assistant.' }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- else %}\n        {{- '<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n"
+optim: adamw_torch_fused

linalg_zero/config/sft/qwen2.5-3B/lora.yaml

@@ -0,0 +1,97 @@
+# NOTE: ADJUST gradient_accumulation_steps, per_device_eval_batch_size, per_device_train_batch_size
+
+# Model arguments
+model_name_or_path: Qwen/Qwen2.5-3B
+model_revision: main
+torch_dtype: bfloat16
+attn_implementation: flash_attention_2
+
+# Model arguments
+load_in_4bit: false
+load_in_8bit: false            # Use 16bit training
+bnb_4bit_quant_type: nf4
+use_peft: true
+lora_r: 32                     # NOTE: Alternatively 16 and 64 if underfitting
+lora_alpha: 64
+lora_dropout: 0.1
+lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
+lora_task_type: "CAUSAL_LM"
+# vllm parameters
+# gpu_memory_utilization: 0.95
+# enforce_eager: true
+# fast_inference: true
+
+# Data training arguments
+dataset_name: atomwalk12/linalgzero-sft
+dataset_config: default
+eval_dataset_config: default
+dataset_num_proc: 12
+dataset_text_field: text
+eos_token: <|im_end|>
+add_special_tokens: true
+
+# SFT trainer config
+bf16: true
+
+## Evaluation parameters
+do_eval: true                  # does a final evaluation step at the end of training
+eval_on_start: true           # does an evaluation step at the start of training
+eval_strategy: 'steps'         # can be "no","epoch","steps"
+eval_steps: 100                # NOTE: Check also save_steps
+max_eval_samples: 520          # use all evaluation samples (720 format checks -- 40*6*2+40*4*3+40*3*4 (520 samples))
+final_eval_max_samples: 520
+
+# NOTE: if we use a customized metric for best model selection, ensure to adjust:
+# - metric_for_best_model: eval_format_accuracy or eval_loss
+# - greater_is_better: true or false (if eval_loss is set)
+# - load_best_model_at_end: true
+metric_for_best_model: eval_format_accuracy
+greater_is_better: true
+load_best_model_at_end: true
+save_total_limit: 5
+eval_max_new_tokens: 800
+# NOTE: for early stopping:
+#  - early_stopping_patience (3)
+#  - early_stopping_threshold (0.01)
+#  - callback (early_stopping)
+
+callbacks:
+  - tool_calling_accuracy
+#  - push_to_hub_revision
+
+
+## Training parameters
+# NOTE: adjust batch size, lr
+gradient_accumulation_steps: 16   # Recommended batch size: 16
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+hub_model_id: atomwalk12/LinalgZero-SFT-LoRA
+hub_strategy: every_save
+learning_rate: 2e-5 # Alternatives: 5e-5/4e-5/3e-5/2e-5 and train for 2 epochs
+log_level: info
+logging_steps: 5
+logging_strategy: steps
+lr_scheduler_type: cosine_with_min_lr
+lr_scheduler_kwargs:
+  min_lr_rate: 0.1
+packing: false
+max_grad_norm: 0.2
+max_seq_length: 8192
+max_length: 8192
+max_steps: -1
+num_train_epochs: 1               # NOTE: Alternatively 2 if underfitting
+output_dir: results/LinalgZero-SFT-LoRA
+overwrite_output_dir: true
+per_device_eval_batch_size: 1 # Overridden to 1 by unsloth
+per_device_train_batch_size: 1
+push_to_hub: true
+report_to:
+- wandb
+save_strategy: steps
+save_steps: 100
+seed: 42
+use_liger_kernel: false     # Supported only for Llama, Mistral, Mixtral, and Gemma models
+warmup_ratio: 0.03
+chat_template: "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- else %}\n        {{- 'You are a helpful assistant.' }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- else %}\n        {{- '<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n"
+optim: adamw_torch_fused

linalg_zero/config/sft/qwen2.5-3B/merged.yaml

@@ -0,0 +1,78 @@
+# SFT config for a light touch-up of the merged model,
+# targeting only the input/output layers (embed_tokens and lm_head).
+
+# Model arguments
+model_name_or_path: results/LinalgZero-SFT-merged
+model_revision: main
+torch_dtype: bfloat16
+attn_implementation: flash_attention_2
+
+# Quantization / PEFT
+load_in_4bit: false
+load_in_8bit: false            # Use 16bit training
+bnb_4bit_quant_type: nf4
+use_peft: false                # IMPORTANT: do not create new LoRA adapters
+
+# Data training arguments
+dataset_name: atomwalk12/linalgzero-sft
+dataset_config: default
+eval_dataset_config: default
+dataset_num_proc: 12
+dataset_text_field: text
+eos_token: <|im_end|>
+add_special_tokens: true
+
+# SFT trainer config
+bf16: true
+
+## Evaluation parameters
+do_eval: true                  # final evaluation at the end of training
+eval_on_start: true            # evaluation at the start of training
+eval_strategy: 'steps'
+eval_steps: 100
+max_eval_samples: 520
+final_eval_max_samples: 520
+
+metric_for_best_model: eval_format_accuracy
+greater_is_better: true
+load_best_model_at_end: true
+save_total_limit: 5
+eval_max_new_tokens: 800
+
+callbacks:
+  - tool_calling_accuracy
+
+## Training parameters
+gradient_accumulation_steps: 16
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+hub_model_id: atomwalk12/LinalgZero-SFT
+hub_strategy: every_save
+learning_rate: 2e-5 # Alternatives: 5e-5/4e-5/3e-5/2e-5 and train for 2 epochs
+log_level: info
+logging_steps: 5
+logging_strategy: steps
+lr_scheduler_type: cosine_with_min_lr
+lr_scheduler_kwargs:
+  min_lr_rate: 0.1
+packing: false
+max_grad_norm: 0.2
+max_seq_length: 8192
+max_length: 8192
+max_steps: -1
+num_train_epochs: 1
+output_dir: results/LinalgZero-SFT
+overwrite_output_dir: true
+per_device_eval_batch_size: 1
+per_device_train_batch_size: 1
+push_to_hub: true
+report_to:
+- wandb
+save_strategy: steps
+save_steps: 100
+seed: 42
+use_liger_kernel: false
+warmup_ratio: 0.03
+chat_template: "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- else %}\n        {{- 'You are a helpful assistant.' }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- else %}\n        {{- '<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n"
+optim: adamw_torch_fused

linalg_zero/config/sft/qwen2.5-3B/nst.yaml

@@ -0,0 +1,99 @@
+# NOTE: ADJUST gradient_accumulation_steps, per_device_eval_batch_size, per_device_train_batch_size
+
+# Model arguments
+model_name_or_path: Qwen/Qwen2.5-3B
+model_revision: main
+torch_dtype: bfloat16
+attn_implementation: flash_attention_2
+wandb_project: linalgzero-sft
+wandb_run_id: run001-sft-nst
+
+# Model arguments
+load_in_4bit: false
+load_in_8bit: false            # Use 16bit training
+bnb_4bit_quant_type: nf4
+use_peft: true
+lora_r: 32                     # NOTE: Alternatively 16 and 64 if underfitting
+lora_alpha: 64
+lora_dropout: 0.1
+lora_target_modules: ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
+lora_task_type: "CAUSAL_LM"
+# vllm parameters
+# gpu_memory_utilization: 0.95
+# enforce_eager: true
+# fast_inference: true
+
+# Data training arguments
+dataset_name: atomwalk12/linalgzero-sft-110
+dataset_config: default
+eval_dataset_config: default
+dataset_num_proc: 12
+dataset_text_field: text
+eos_token: <|im_end|>
+
+# SFT trainer config
+bf16: true
+
+## Evaluation parameters
+do_eval: true                  # does a final evaluation step at the end of training
+eval_on_start: false           # does an evaluation step at the start of training
+eval_strategy: 'steps'         # can be "no","epoch","steps"
+eval_steps: 110                # NOTE: Check also save_steps
+max_eval_samples: 260          # use all evaluation samples (720 format checks -- 40*6*2+40*4*3+40*3*4 (520 samples))
+final_eval_max_samples: 260
+add_special_tokens: False
+
+# NOTE: if we use a customized metric for best model selection, ensure to adjust:
+# - metric_for_best_model: eval_format_accuracy or eval_loss
+# - greater_is_better: true or false (if eval_loss is set)
+# - load_best_model_at_end: true
+metric_for_best_model: eval_format_accuracy
+greater_is_better: true
+load_best_model_at_end: true
+save_total_limit: 30
+eval_max_new_tokens: 800
+# NOTE: for early stopping:
+#  - early_stopping_patience (3)
+#  - early_stopping_threshold (0.01)
+#  - callback (early_stopping)
+
+callbacks:
+  - tool_calling_accuracy
+#  - push_to_hub_revision
+
+
+## Training parameters
+# NOTE: adjust batch size, lr
+gradient_accumulation_steps: 16   # Recommended batch size: 16
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+hub_model_id: atomwalk12/LinalgZero-SFT-LoRA-NST
+hub_strategy: every_save
+learning_rate: 2e-5 # Alternatives: 5e-5/4e-5/3e-5/2e-5 and train for 2 epochs
+log_level: info
+logging_steps: 5
+logging_strategy: steps
+lr_scheduler_type: cosine_with_min_lr
+lr_scheduler_kwargs:
+  min_lr_rate: 0.1
+packing: false
+max_grad_norm: 0.2
+max_seq_length: 8192
+max_length: 8192
+max_steps: -1
+num_train_epochs: 1               # NOTE: Alternatively 2 if underfitting
+output_dir: results/LinalgZero-SFT-LoRA-NST
+overwrite_output_dir: true
+per_device_eval_batch_size: 1 # Overridden to 1 by unsloth
+per_device_train_batch_size: 1
+push_to_hub: true
+report_to:
+- wandb
+save_strategy: steps
+save_steps: 110
+seed: 42
+use_liger_kernel: false     # Supported only for Llama, Mistral, Mixtral, and Gemma models
+warmup_ratio: 0.03
+chat_template: "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- else %}\n        {{- 'You are a helpful assistant.' }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n    {%- else %}\n        {{- '<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {{- '<|im_start|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + message.content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n"
+optim: adamw_torch_fused

linalg_zero/demo/app.py

@@ -0,0 +1,968 @@
+import argparse
+import gc
+import json
+import os
+import threading
+from typing import Any
+
+import gradio as gr
+import spaces  # Import spaces early to enable ZeroGPU support
+import torch
+from datasets import load_dataset
+from huggingface_hub import snapshot_download
+from linalg_zero.distillation.components.models import DefaultConfig
+from linalg_zero.distillation.data import FunctionInvocationInfo, ThoughtSchema
+from linalg_zero.grpo.verifiers.xml_parser import XMLParser
+from linalg_zero.shared.lib import get_lib, get_tools
+from linalg_zero.shared.system_prompts import (ANSWER_CLOSE, ANSWER_OPEN,
+                                               THINK_CLOSE, THINK_OPEN,
+                                               TOOL_CALL_CLOSE, TOOL_CALL_OPEN,
+                                               TOOL_RESPONSE_CLOSE,
+                                               TOOL_RESPONSE_OPEN,
+                                               get_math_system_prompt)
+from torch.utils._pytree import tree_map
+from transformers import (AutoTokenizer, StoppingCriteria,
+                          StoppingCriteriaList, TextIteratorStreamer, pipeline)
+
+# Global event to signal cancellation from the UI thread to the generation thread
+cancel_event = threading.Event()
+
+access_token = os.environ["HF_TOKEN"]
+DEMO_DIR = os.path.dirname(__file__)
+ASSISTANT_AVATAR_PATH = os.path.join(DEMO_DIR, "assets", "linalgzero-avatar.svg")
+
+# Optional: Disable GPU visibility if you wish to force CPU usage
+# os.environ["CUDA_VISIBLE_DEVICES"] = " "
+
+# ------------------------------
+# Allowed model definitions
+# ------------------------------
+MODELS = {
+    "Qwen3-1.7B": {
+        "repo_id": "Qwen/Qwen3-1.7B",
+        "description": "Dense causal language model with 1.7B parameters.",
+        "params_b": 1.7,
+    },
+    "atomwalk12/LinAlgZero-GRPO": {
+        "repo_id": "atomwalk12/LinAlgZero-GRPO-merged",
+        "description": "LinAlgZero GRPO fine-tuned model.",
+        "params_b": 3.0,
+    },
+    "atomwalk12/LinalgZero-SFT": {
+        "repo_id": "atomwalk12/LinalgZero-SFT",
+        "description": "LinAlgZero SFT fine-tuned model.",
+        "params_b": 3.0,
+    },
+}
+# Global cache for pipelines to avoid re-loading.
+PIPELINES = {}
+MODEL_SNAPSHOT_PATHS: dict[str, str] = {}
+SYSTEM_PROMPT = get_math_system_prompt(include_examples=False)
+TOOL_LIBRARY = get_lib()
+TOOL_SCHEMAS = get_tools()
+PARSER = XMLParser()
+MODEL_MESSAGE_CONFIG = DefaultConfig()
+TRACE_TITLES = {"💭 Thought", "🛠 Tool Call", "📦 Tool Response"}
+KEEP_SPECIAL_TOKENS = {
+    ANSWER_OPEN,
+    ANSWER_CLOSE,
+    THINK_OPEN,
+    THINK_CLOSE,
+    TOOL_CALL_OPEN,
+    TOOL_CALL_CLOSE,
+    TOOL_RESPONSE_OPEN,
+    TOOL_RESPONSE_CLOSE,
+}
+SIMPLE_EXAMPLE_QUERIES = [
+    "What is the rank of matrix A = [[2, 3], [2, -4]]?",
+    "Step 1: find the transpose of A = [[1, 2], [3, 4]]. Step 2: find the trace of the result.",
+    "Find the determinant of [[3, 1], [2, 5]].",
+    "Find the Frobenius norm of [[3, -2], [-1, 5]].",
+    "Find the cofactor matrix of [[5, 2], [1, 3]].",
+]
+
+
+def _parse_launch_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(add_help=False)
+    parser.add_argument(
+        "--share",
+        action="store_true",
+        help="Enable a public Gradio share link.",
+    )
+    args, _ = parser.parse_known_args()
+    return args
+
+
+def _extract_example_query(example: dict[str, Any]) -> str | None:
+    query = example.get("query")
+    if isinstance(query, str) and query.strip():
+        return query.strip()
+
+    messages = example.get("messages")
+    if isinstance(messages, str):
+        try:
+            messages = json.loads(messages)
+        except Exception:
+            return None
+
+    if isinstance(messages, list):
+        for message in reversed(messages):
+            if not isinstance(message, dict):
+                continue
+            if message.get("role") != "user":
+                continue
+            content = message.get("content")
+            if isinstance(content, str) and content.strip():
+                return content.strip()
+    return None
+
+
+def _truncate_example_label(query: str, max_chars: int = 96) -> str:
+    one_line_query = " ".join(query.split())
+    if len(one_line_query) <= max_chars:
+        return one_line_query
+    return one_line_query[: max_chars - 3] + "..."
+
+
+def _load_dataset_example_queries(limit: int = 10) -> tuple[list[str], str | None, str | None]:
+    dataset_candidates = {
+        "atomwalk12/linalgzero-grpo": ["test", "validation", "train"],
+        "atomwalk12/linalgzero": ["validation", "test", "train"],
+        "atomwalk12/linalgzero-sft": ["test", "validation", "train"],
+    }
+
+    for repo, splits in dataset_candidates.items():
+        for split in splits:
+            try:
+                dataset = load_dataset(repo, split=split, streaming=True)
+            except Exception:
+                continue
+
+            queries: list[str] = []
+            seen_queries: set[str] = set()
+            for example in dataset:
+                query = _extract_example_query(example)
+                if not query or query in seen_queries:
+                    continue
+                seen_queries.add(query)
+                queries.append(query)
+                if len(queries) >= limit:
+                    return queries, repo, split
+
+            if queries:
+                return queries, repo, split
+
+    return [], None, None
+
+
+DATASET_EXAMPLE_QUERIES, DATASET_EXAMPLE_REPO, DATASET_EXAMPLE_SPLIT = _load_dataset_example_queries()
+DATASET_EXAMPLE_CHOICES = [(_truncate_example_label(query), query) for query in DATASET_EXAMPLE_QUERIES]
+SIMPLE_EXAMPLE_ROWS = [[query] for query in SIMPLE_EXAMPLE_QUERIES]
+if DATASET_EXAMPLE_REPO is not None and DATASET_EXAMPLE_SPLIT is not None:
+    DATASET_EXAMPLE_INFO = f"Loaded from {DATASET_EXAMPLE_REPO} ({DATASET_EXAMPLE_SPLIT})"
+else:
+    DATASET_EXAMPLE_INFO = "Dataset examples unavailable"
+
+# Match the current GRPO/SFT evaluation configs:
+# - temperature = 0.0
+# - top_p = null (functionally ignored in deterministic mode)
+# - repetition_penalty = 1.0
+# In the actual GRPO eval agent, non-train splits force deterministic decoding
+# (do_sample=False, temperature=0.0, top_p=None), so top_k is effectively unused.
+EVAL_TEMPERATURE = 0.0
+UI_TOP_K_DEFAULT = 40
+UI_TOP_P_DEFAULT = 1.0
+EVAL_REPETITION_PENALTY = 1.0
+DEFAULT_MAX_TOKENS = 1024
+ZERO_GPU_DURATION_SECONDS = 120
+
+def prepare_model_artifacts(model_name: str, *, local_files_only: bool = False) -> str:
+    cached_path = MODEL_SNAPSHOT_PATHS.get(model_name)
+    if cached_path is not None and os.path.exists(cached_path):
+        return cached_path
+
+    repo = MODELS[model_name]["repo_id"]
+    model_path = snapshot_download(
+        repo_id=repo,
+        token=access_token,
+        local_files_only=local_files_only,
+    )
+    MODEL_SNAPSHOT_PATHS[model_name] = model_path
+    return model_path
+
+
+def load_pipeline(model_name):
+    """
+    Load and cache a transformers pipeline for text generation.
+    Tries bfloat16, falls back to float16 or float32 if unsupported.
+    """
+    global PIPELINES
+    if model_name in PIPELINES:
+        return PIPELINES[model_name]
+    model_path = prepare_model_artifacts(model_name, local_files_only=True)
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_path,
+        token=access_token,
+        local_files_only=True,
+    )
+    for dtype in (torch.bfloat16, torch.float16, torch.float32):
+        try:
+            pipe = pipeline(
+                task="text-generation",
+                model=model_path,
+                tokenizer=tokenizer,
+                trust_remote_code=True,
+                dtype=dtype, # Use `dtype` instead of deprecated `torch_dtype`
+                device_map="auto",
+                use_cache=True,      # Enable past-key-value caching
+                token=access_token,
+                local_files_only=True,
+            )
+            PIPELINES[model_name] = pipe
+            return pipe
+        except Exception:
+            continue
+    # Final fallback
+    pipe = pipeline(
+        task="text-generation",
+        model=model_path,
+        tokenizer=tokenizer,
+        trust_remote_code=True,
+        device_map="auto",
+        use_cache=True,
+        local_files_only=True,
+    )
+    PIPELINES[model_name] = pipe
+    return pipe
+
+def _extract_visible_title(message: dict[str, Any]) -> str | None:
+    metadata = message.get("metadata")
+    if isinstance(metadata, dict):
+        title = metadata.get("title")
+        if isinstance(title, str):
+            return title
+    return None
+
+
+def _build_llm_messages(chat_history: list[dict[str, Any]] | None, user_msg: str) -> list[dict[str, Any]]:
+    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
+    for msg in chat_history or []:
+        role = msg.get("role")
+        content = msg.get("content")
+        if role not in {"user", "assistant"} or not isinstance(content, str):
+            continue
+        if role == "assistant" and _extract_visible_title(msg) in TRACE_TITLES:
+            continue
+        messages.append({"role": role, "content": content})
+    messages.append({"role": "user", "content": user_msg})
+    return messages
+
+
+def _render_prompt(messages: list[dict[str, Any]], tokenizer: AutoTokenizer) -> str:
+    if hasattr(tokenizer, "chat_template") and tokenizer.chat_template:
+        kwargs: dict[str, Any] = {
+            "tokenize": False,
+            "add_generation_prompt": True,
+            "tools": TOOL_SCHEMAS,
+        }
+        try:
+            return tokenizer.apply_chat_template(messages, enable_thinking=True, **kwargs)
+        except TypeError:
+            return tokenizer.apply_chat_template(messages, **kwargs)
+
+    prompt = messages[0]["content"].strip() + "\n"
+    for msg in messages[1:]:
+        if msg["role"] == "user":
+            prompt += "User: " + str(msg["content"]).strip() + "\n"
+        elif msg["role"] == "assistant":
+            prompt += "Assistant: " + str(msg.get("content", "")).strip() + "\n"
+            for tool_call in msg.get("tool_calls", []):
+                function_info = tool_call.get("function", {})
+                prompt += (
+                    f"{TOOL_CALL_OPEN}"
+                    + json.dumps({
+                        "name": function_info.get("name"),
+                        "arguments": function_info.get("arguments"),
+                    })
+                    + f"{TOOL_CALL_CLOSE}\n"
+                )
+        elif msg["role"] == "tool":
+            prompt += f"User: {TOOL_RESPONSE_OPEN}{msg['content']}{TOOL_RESPONSE_CLOSE}\n"
+    if not prompt.strip().endswith("Assistant:"):
+        prompt += "Assistant: "
+    return prompt
+
+
+class StopOnSequences(StoppingCriteria):
+    def __init__(self, stop_token_sequences: list[list[int]]) -> None:
+        super().__init__()
+        self.stop_token_sequences = [seq for seq in stop_token_sequences if seq]
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs: Any) -> bool:
+        if cancel_event.is_set():
+            return True
+
+        current_tokens = input_ids[0].tolist()
+        for seq in self.stop_token_sequences:
+            if len(current_tokens) >= len(seq) and current_tokens[-len(seq):] == seq:
+                return True
+        return False
+
+
+def _get_stop_token_sequences(tokenizer: AutoTokenizer) -> list[list[int]]:
+    stop_strings = [TOOL_CALL_CLOSE, ANSWER_CLOSE, TOOL_RESPONSE_OPEN, TOOL_RESPONSE_CLOSE]
+    return [tokenizer.encode(stop_text, add_special_tokens=False) for stop_text in stop_strings]
+
+
+def _decode_generated_tokens(tokenizer: AutoTokenizer, generated_tokens: torch.Tensor) -> str:
+    raw_output = tokenizer.decode(generated_tokens, skip_special_tokens=False)
+    cleaned_output = raw_output
+    for special_token in tokenizer.all_special_tokens:
+        if special_token not in KEEP_SPECIAL_TOKENS:
+            cleaned_output = cleaned_output.replace(special_token, "")
+    return cleaned_output
+
+
+def _clean_decoded_text(tokenizer: AutoTokenizer, text: str) -> str:
+    cleaned_text = text
+    for special_token in tokenizer.all_special_tokens:
+        if special_token not in KEEP_SPECIAL_TOKENS:
+            cleaned_text = cleaned_text.replace(special_token, "")
+    return cleaned_text
+
+
+def _extract_partial_tag_content(message: str, tag: str) -> str | None:
+    open_tag = f"<{tag}>"
+    close_tag = f"</{tag}>"
+    start_idx = message.rfind(open_tag)
+    if start_idx == -1:
+        return None
+    content = message[start_idx + len(open_tag):]
+    end_idx = content.find(close_tag)
+    if end_idx != -1:
+        content = content[:end_idx]
+    return content.strip()
+
+
+def _extract_live_sections(message: str) -> dict[str, str | None]:
+    return {
+        "thought": _extract_partial_tag_content(message, "think"),
+        "tool_call": _extract_partial_tag_content(message, "tool_call"),
+        "answer": _extract_partial_tag_content(message, "answer"),
+    }
+
+
+def _set_live_message(
+    history: list[dict[str, Any]],
+    live_indices: dict[str, int],
+    key: str,
+    content: str | None,
+    *,
+    title: str | None = None,
+    code_block_lang: str | None = None,
+) -> None:
+    if not content:
+        return
+
+    rendered_content = content
+    if code_block_lang is not None:
+        rendered_content = f"```{code_block_lang}\n{content}\n```"
+
+    idx = live_indices.get(key)
+    if idx is None:
+        message: dict[str, Any] = {"role": "assistant", "content": rendered_content}
+        if title is not None:
+            message["metadata"] = {"title": title}
+        history.append(message)
+        live_indices[key] = len(history) - 1
+        return
+
+    history[idx]["content"] = rendered_content
+
+
+def _build_debug_info() -> str:
+    tool_lines = "\n".join(f"- `{tool_name}`" for tool_name in sorted(TOOL_LIBRARY.keys()))
+    return f"### Available Tools\n{tool_lines}\n\n### System Prompt\n```text\n{SYSTEM_PROMPT}\n```"
+
+
+def _stream_assistant_turn(
+    pipe: Any,
+    messages: list[dict[str, Any]],
+    max_tokens: int,
+    temperature: float,
+    top_k: int,
+    top_p: float,
+    repeat_penalty: float,
+):
+    prompt = _render_prompt(messages, pipe.tokenizer)
+    inputs = pipe.tokenizer(
+        prompt,
+        return_tensors="pt",
+        truncation=True,
+        padding=bool(getattr(pipe.tokenizer, "pad_token_id", None)),
+    )
+    device = getattr(pipe, "device", None) or pipe.model.device
+    inputs = tree_map(lambda x: x.to(device) if hasattr(x, "to") else x, inputs)
+
+    streamer = TextIteratorStreamer(
+        pipe.tokenizer,
+        skip_prompt=True,
+        skip_special_tokens=False,
+    )
+
+    generation_kwargs: dict[str, Any] = {
+        "input_ids": inputs["input_ids"],
+        "attention_mask": inputs["attention_mask"],
+        "max_new_tokens": int(max_tokens),
+        "repetition_penalty": float(repeat_penalty),
+        "use_cache": True,
+        "pad_token_id": getattr(pipe.tokenizer, "pad_token_id", pipe.tokenizer.eos_token_id),
+        "stopping_criteria": StoppingCriteriaList([StopOnSequences(_get_stop_token_sequences(pipe.tokenizer))]),
+        "streamer": streamer,
+    }
+    eos_token_id = getattr(pipe.tokenizer, "eos_token_id", None)
+    if eos_token_id is not None:
+        generation_kwargs["eos_token_id"] = eos_token_id
+    if temperature > 0:
+        generation_kwargs.update({
+            "do_sample": True,
+            "temperature": float(temperature),
+            "top_k": int(top_k),
+            "top_p": float(top_p),
+        })
+    else:
+        generation_kwargs["do_sample"] = False
+
+    error_holder: dict[str, Exception] = {}
+
+    def _run_generation() -> None:
+        try:
+            with torch.inference_mode():
+                pipe.model.generate(**generation_kwargs)
+        except Exception as exc:
+            error_holder["error"] = exc
+        finally:
+            streamer.end()
+
+    generation_thread = threading.Thread(target=_run_generation, daemon=True)
+    generation_thread.start()
+
+    raw_output = ""
+    yield raw_output, prompt
+    for chunk in streamer:
+        raw_output = _clean_decoded_text(pipe.tokenizer, raw_output + chunk)
+        yield raw_output, prompt
+
+    generation_thread.join()
+    if "error" in error_holder:
+        raise error_holder["error"]
+
+
+def _parse_assistant_output(raw_output: str, context: list[dict[str, Any]]) -> tuple[ThoughtSchema | None, dict[str, Any]]:
+    analysis = PARSER.analyze_message_in_context(
+        context,
+        message=raw_output,
+        tool_names=list(TOOL_LIBRARY.keys()),
+    )
+    if not bool(analysis["is_valid_think_then_tool_or_answer"]):
+        return None, analysis
+    if analysis["has_answer"] and not bool(analysis["answer_policy_valid"]):
+        return None, analysis
+
+    tool_call: FunctionInvocationInfo | None = None
+    tool_info = analysis["tool"]
+    if analysis["has_tool_call"]:
+        if not bool(tool_info["json_valid"]):
+            return None, analysis
+        tool_call = FunctionInvocationInfo(
+            name=str(tool_info["name"]),
+            arguments=dict(tool_info["arguments"]),
+        )
+
+    return ThoughtSchema(
+        thought=analysis["thought"] or "",
+        tool_call=tool_call,
+        final_answer=analysis["answer"],
+        completed=analysis["answer"] is not None,
+    ), analysis
+
+
+def _execute_tool_call(message: ThoughtSchema) -> dict[str, str]:
+    if message.tool_call is None:
+        raise ValueError("Tool call is required")
+
+    name = message.tool_call.name
+    arguments = message.tool_call.arguments
+    try:
+        if name not in TOOL_LIBRARY:
+            return {
+                "function_name": name,
+                "execution_result": f"ERROR: Function '{name}' not found in library",
+            }
+        result = TOOL_LIBRARY[name](**arguments)
+        return {"function_name": name, "execution_result": str(result)}
+    except Exception as exc:
+        return {
+            "function_name": name,
+            "execution_result": f"ERROR: {type(exc).__name__}: {exc}",
+        }
+
+
+# Keep a single GPU allocation for the whole multi-turn solve, with a fixed
+# budget that is simpler and more predictable than a heuristic estimate.
+@spaces.GPU(duration=ZERO_GPU_DURATION_SECONDS)
+def chat_response(
+    user_msg,
+    chat_history,
+    show_tool_trace,
+    enable_streaming,
+    max_tool_turns,
+    model_name,
+    max_tokens,
+    temperature,
+    top_k,
+    top_p,
+    repeat_penalty,
+):
+    """
+    Generates responses by iteratively calling linear algebra tools until a final answer is produced.
+    """
+    cancel_event.clear()
+
+    history = list(chat_history or [])
+    history.append({"role": "user", "content": user_msg})
+    debug_info = _build_debug_info()
+
+    try:
+        pipe = load_pipeline(model_name)
+        context = _build_llm_messages(chat_history, user_msg)
+
+        yield history, debug_info
+
+        for step_idx in range(int(max_tool_turns)):
+            raw_output = ""
+            live_indices: dict[str, int] = {}
+
+            turn_stream = _stream_assistant_turn(
+                pipe=pipe,
+                messages=context,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                top_k=top_k,
+                top_p=top_p,
+                repeat_penalty=repeat_penalty,
+            )
+
+            if enable_streaming:
+                for partial_output, _current_prompt in turn_stream:
+                    raw_output = partial_output
+
+                    live_sections = _extract_live_sections(raw_output)
+                    if show_tool_trace:
+                        _set_live_message(
+                            history,
+                            live_indices,
+                            "thought",
+                            live_sections["thought"],
+                            title="💭 Thought",
+                        )
+                        _set_live_message(
+                            history,
+                            live_indices,
+                            "tool_call",
+                            live_sections["tool_call"],
+                            title="🛠 Tool Call",
+                            code_block_lang="json",
+                        )
+                    _set_live_message(
+                        history,
+                        live_indices,
+                        "answer",
+                        live_sections["answer"],
+                    )
+
+                    yield history, debug_info
+            else:
+                for partial_output, _current_prompt in turn_stream:
+                    raw_output = partial_output
+
+            if cancel_event.is_set():
+                history.append({"role": "assistant", "content": "[Generation Canceled]"})
+                yield history, debug_info
+                break
+
+            parsed, analysis = _parse_assistant_output(raw_output, context)
+            if parsed is None:
+                history.append({
+                    "role": "assistant",
+                    "content": "I couldn't produce a valid tool call or final answer for this problem.",
+                })
+                yield history, debug_info
+                break
+
+            assistant_message = MODEL_MESSAGE_CONFIG.format_assistant_message(parsed)
+            if assistant_message is None:
+                history.append({"role": "assistant", "content": "The model returned an empty action."})
+                yield history, debug_info
+                break
+
+            context.append(assistant_message)
+
+            if show_tool_trace and parsed.thought:
+                _set_live_message(
+                    history,
+                    live_indices,
+                    "thought",
+                    parsed.thought,
+                    title="💭 Thought",
+                )
+
+            if parsed.tool_call is not None:
+                if show_tool_trace:
+                    _set_live_message(
+                        history,
+                        live_indices,
+                        "tool_call",
+                        json.dumps(
+                            {
+                                "name": parsed.tool_call.name,
+                                "arguments": parsed.tool_call.arguments,
+                            },
+                            indent=2,
+                        ),
+                        title="🛠 Tool Call",
+                        code_block_lang="json",
+                    )
+
+                tool_result = _execute_tool_call(parsed)
+                context.append(MODEL_MESSAGE_CONFIG.create_tool_message(context, tool_result))
+
+                if show_tool_trace:
+                    history.append({
+                        "role": "assistant",
+                        "content": tool_result["execution_result"],
+                        "metadata": {"title": "📦 Tool Response"},
+                    })
+
+                yield history, debug_info
+                continue
+
+            if parsed.final_answer is not None:
+                answer_idx = live_indices.get("answer")
+                if answer_idx is not None:
+                    history[answer_idx]["content"] = parsed.final_answer
+                else:
+                    history.append({"role": "assistant", "content": parsed.final_answer})
+                yield history, debug_info
+                break
+        else:
+            history.append({
+                "role": "assistant",
+                "content": "I hit the tool-turn limit before finishing the problem.",
+            })
+            yield history, debug_info
+    except GeneratorExit:
+        print("Chat response cancelled.")
+        return
+    except Exception as e:
+        history.append({"role": "assistant", "content": f"Error: {e}"})
+        yield history, debug_info
+    finally:
+        gc.collect()
+
+
+demo_theme = gr.themes.Soft(
+    primary_hue="indigo",
+    secondary_hue="purple",
+    neutral_hue="slate",
+    radius_size="lg",
+    font=[gr.themes.GoogleFont("Inter"), "Arial", "sans-serif"]
+)
+
+demo_css = """
+    .chatbot { border-radius: 12px; box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1); }
+    button.primary { font-weight: 600; }
+    .gradio-accordion { margin-bottom: 12px; }
+"""
+
+with gr.Blocks(title="LLM Inference with ZeroGPU") as demo:
+    # Header
+    gr.Markdown("""
+    # 🧠 LinAlgZero Demo
+    ### Multi-turn linear algebra solving with tool calling
+    """)
+
+    with gr.Row():
+        # Left Panel - Configuration
+        with gr.Column(scale=3):
+            # Core Settings (Always Visible)
+            with gr.Group():
+                gr.Markdown("### ⚙️ Core Settings")
+                model_dd = gr.Dropdown(
+                    label="🤖 Model",
+                    choices=list(MODELS.keys()),
+                    value="atomwalk12/LinAlgZero-GRPO",
+                    info="Select the language model to use"
+                )
+                trace_chk = gr.Checkbox(
+                    label="🔍 Show Tool Trace",
+                    value=True,
+                    info="Show thoughts, tool calls, and tool responses in the chat"
+                )
+                stream_chk = gr.Checkbox(
+                    label="⚡ Enable Live Streaming",
+                    value=True,
+                    info="Stream partial model output while each tool-calling turn is being generated"
+                )
+                max_tool_turns = gr.Slider(
+                    1, 12, value=6, step=1,
+                    label="Max Tool Turns",
+                    info="Maximum number of tool-calling turns before stopping"
+                )
+
+            example_dropdown = gr.Dropdown(
+                label="📚 Dataset Example",
+                choices=DATASET_EXAMPLE_CHOICES,
+                value=None,
+                info=DATASET_EXAMPLE_INFO,
+                allow_custom_value=False,
+            )
+
+            # Advanced Settings (Collapsible)
+            with gr.Accordion("🎛️ Advanced Generation Parameters", open=True):
+                max_tok = gr.Slider(
+                    64, 16384, value=DEFAULT_MAX_TOKENS, step=32,
+                    label="Max Tokens",
+                    info="Maximum length of generated response"
+                )
+                temp = gr.Slider(
+                    0.0, 2.0, value=EVAL_TEMPERATURE, step=0.1,
+                    label="Temperature",
+                    info="Higher = more creative, Lower = more focused"
+                )
+                with gr.Row():
+                    k = gr.Slider(
+                        1, 100, value=UI_TOP_K_DEFAULT, step=1,
+                        label="Top-K",
+                        info="Number of top tokens to consider",
+                        interactive=False,
+                    )
+                    p = gr.Slider(
+                        0.1, 1.0, value=UI_TOP_P_DEFAULT, step=0.05,
+                        label="Top-P",
+                        info="Nucleus sampling threshold",
+                        interactive=False,
+                    )
+                rp = gr.Slider(
+                    1.0, 2.0, value=EVAL_REPETITION_PENALTY, step=0.1,
+                    label="Repetition Penalty",
+                    info="Penalize repeated tokens"
+                )
+
+            # Actions
+            with gr.Column():
+                reset_params_btn = gr.Button("↺ Restore Default Params", variant="secondary")
+                clr = gr.Button("🗑️ Clear Chat", variant="secondary", scale=1)
+
+        # Right Panel - Chat Interface
+        with gr.Column(scale=7):
+            chat = gr.Chatbot(
+                height=1200,
+                label="💬 Conversation",
+                buttons=["copy"],
+                avatar_images=(None, ASSISTANT_AVATAR_PATH),
+                layout="bubble"
+            )
+
+            # Input Area
+            with gr.Row():
+                txt = gr.Textbox(
+                    placeholder="💭 Type your message here... (Press Enter to send)",
+                    scale=9,
+                    container=False,
+                    show_label=False,
+                    lines=1,
+                    max_lines=5
+                )
+                with gr.Column(scale=1, min_width=120):
+                    submit_btn = gr.Button("📤 Send", variant="primary", size="lg", interactive=False)
+                    cancel_btn = gr.Button("⏹️ Stop", variant="stop", visible=False, size="lg")
+
+            # Example Prompts
+            gr.Examples(
+                examples=SIMPLE_EXAMPLE_ROWS,
+                inputs=txt,
+                label="💡 Example Prompts"
+            )
+
+            # Debug/Status Info (Collapsible)
+            with gr.Accordion("🔍 Debug Info", open=False):
+                dbg = gr.Markdown(value=_build_debug_info())
+
+    # --- Event Listeners ---
+
+    # Group all inputs for cleaner event handling
+    chat_inputs = [txt, chat, trace_chk, stream_chk, max_tool_turns, model_dd, max_tok, temp, k, p, rp]
+    # Group all UI components that can be updated.
+    ui_components = [chat, dbg, txt, submit_btn, cancel_btn]
+
+    def submit_and_manage_ui(user_msg, chat_history, *args):
+        """
+        Orchestrator function that manages UI state and calls the backend chat function.
+        It uses a try...finally block to ensure the UI is always reset.
+        """
+        if not user_msg.strip():
+            # If the message is empty, do nothing.
+            # We yield an empty dict to avoid any state changes.
+            yield {}
+            return
+
+        # 1. Update UI to "generating" state.
+        #    Crucially, we do NOT update the `chat` component here, as the backend
+        #    will provide the correctly formatted history in the first response chunk.
+        yield {
+            txt: gr.update(value="", interactive=False),
+            submit_btn: gr.update(interactive=False),
+            cancel_btn: gr.update(visible=True),
+        }
+
+        cancelled = False
+        try:
+            model_name = args[3]
+            prepare_model_artifacts(model_name)
+
+            # 2. Call the backend and stream updates
+            backend_args = [user_msg, chat_history] + list(args)
+            for response_chunk in chat_response(*backend_args):
+                yield {
+                    chat: response_chunk[0],
+                    dbg: response_chunk[1],
+                }
+        except GeneratorExit:
+            # Mark as cancelled and re-raise to prevent "generator ignored GeneratorExit"
+            cancelled = True
+            print("Generation cancelled by user.")
+            raise
+        except Exception as e:
+            print(f"An error occurred during generation: {e}")
+            # If an error happens, add it to the chat history to inform the user.
+            error_history = (chat_history or []) + [
+                {'role': 'user', 'content': user_msg},
+                {'role': 'assistant', 'content': f"**An error occurred:** {str(e)}"}
+            ]
+            yield {chat: error_history}
+        finally:
+            # Only reset UI if not cancelled (to avoid "generator ignored GeneratorExit")
+            if not cancelled:
+                print("Resetting UI state.")
+                yield {
+                    txt: gr.update(interactive=True),
+                    submit_btn: gr.update(interactive=False),
+                    cancel_btn: gr.update(visible=False),
+                }
+
+    def set_cancel_flag():
+        """Called by the cancel button, sets the global event."""
+        cancel_event.set()
+        print("Cancellation signal sent.")
+
+    def reset_ui_after_cancel():
+        """Reset UI components after cancellation."""
+        cancel_event.clear()  # Clear the flag for next generation
+        print("UI reset after cancellation.")
+        return {
+            txt: gr.update(interactive=True),
+            submit_btn: gr.update(interactive=False),
+            cancel_btn: gr.update(visible=False),
+        }
+
+    def clear_chat_state():
+        """Clear the chat and restore default UI state."""
+        cancel_event.clear()
+        return {
+            chat: [],
+            dbg: _build_debug_info(),
+            txt: "",
+            submit_btn: gr.update(interactive=False),
+            cancel_btn: gr.update(visible=False),
+        }
+
+    def set_example_text(example_query: str | None) -> str:
+        return example_query or ""
+
+    def apply_example_selection(evt: gr.SelectData):
+        if evt is None or not evt.selected:
+            return gr.skip(), gr.skip()
+        example_text = set_example_text(evt.value)
+        return gr.update(value=example_text), toggle_submit_button(example_text)
+
+    def toggle_submit_button(text: str | None):
+        return gr.update(interactive=bool(text and text.strip()))
+
+    def toggle_sampling_controls(temperature: float):
+        sampling_enabled = float(temperature) > 0.0
+        return [
+            gr.update(interactive=sampling_enabled),
+            gr.update(interactive=sampling_enabled),
+        ]
+
+    def restore_generation_defaults():
+        return (
+            gr.update(value=DEFAULT_MAX_TOKENS),
+            gr.update(value=EVAL_TEMPERATURE),
+            gr.update(value=UI_TOP_K_DEFAULT, interactive=False),
+            gr.update(value=UI_TOP_P_DEFAULT, interactive=False),
+            gr.update(value=EVAL_REPETITION_PENALTY),
+        )
+
+    # Event for submitting text via Enter key or Submit button
+    submit_event = txt.submit(
+        fn=submit_and_manage_ui,
+        inputs=chat_inputs,
+        outputs=ui_components,
+    )
+    submit_click_event = submit_btn.click(
+        fn=submit_and_manage_ui,
+        inputs=chat_inputs,
+        outputs=ui_components,
+    )
+
+    # Event for the "Cancel" button.
+    # It sets the cancel flag, cancels the submit event, then resets the UI.
+    cancel_btn.click(
+        fn=set_cancel_flag,
+        cancels=[submit_event, submit_click_event]
+    ).then(
+        fn=reset_ui_after_cancel,
+        outputs=ui_components
+    )
+
+    txt.input(fn=toggle_submit_button, inputs=txt, outputs=submit_btn)
+    txt.change(fn=toggle_submit_button, inputs=txt, outputs=submit_btn)
+    example_dropdown.select(
+        fn=apply_example_selection,
+        outputs=[txt, submit_btn],
+    )
+    temp.change(fn=toggle_sampling_controls, inputs=temp, outputs=[k, p])
+    reset_params_btn.click(
+        fn=restore_generation_defaults,
+        outputs=[max_tok, temp, k, p, rp],
+    )
+
+    # Clear chat action
+    clr.click(
+        fn=set_cancel_flag,
+        cancels=[submit_event, submit_click_event]
+    ).then(
+        fn=clear_chat_state,
+        outputs=ui_components
+    )
+
+if __name__ == "__main__":
+    launch_args = _parse_launch_args()
+    demo.launch(
+        theme=demo_theme,
+        css=demo_css,
+        share=launch_args.share,
+        ssr_mode=False,
+    )