├── .DS_Store
├── .github
    ├── FUNDING.yml
    ├── ISSUE_TEMPLATE
    │   ├── bug_report.md
    │   └── feature_request.md
    ├── PULL_REQUEST_TEMPLATE.yml
    ├── dependabot.yml
    └── workflows
    │   ├── code_quality_control.yml
    │   ├── cos_integration.yml
    │   ├── docs.yml
    │   ├── docs_test.yml
    │   ├── label.yml
    │   ├── lints.yml
    │   ├── pr_request_checks.yml
    │   ├── pull-request-links.yml
    │   ├── pylint.yml
    │   ├── python-publish.yml
    │   ├── quality.yml
    │   ├── ruff.yml
    │   ├── run_test.yml
    │   ├── stale.yml
    │   ├── test.yml
    │   ├── testing.yml
    │   ├── unit-test.yml
    │   └── welcome.yml
├── .gitignore
├── .pre-commit-config.yaml
├── .readthedocs.yml
├── LICENSE
├── Makefile
├── README.md
├── agorabanner.png
├── example.py
├── pali.png
├── pali3
    ├── __init__.py
    ├── main.py
    ├── tokenizer.py
    ├── train_vit_siglip.py
    └── ul2.py
├── pyproject.toml
└── requirements.txt


/.DS_Store:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/kyegomez/PALI3/355fd477e04efcc2dc7abd3676569f18b050de63/.DS_Store


--------------------------------------------------------------------------------
/.github/FUNDING.yml:
--------------------------------------------------------------------------------
 1 | # These are supported funding model platforms
 2 | 
 3 | github: [kyegomez]
 4 | patreon: # Replace with a single Patreon username
 5 | open_collective: # Replace with a single Open Collective username
 6 | ko_fi: # Replace with a single Ko-fi username
 7 | tidelift: # Replace with a single Tidelift platform-name/package-name e.g., npm/babel
 8 | community_bridge: # Replace with a single Community Bridge project-name e.g., cloud-foundry
 9 | liberapay: # Replace with a single Liberapay username
10 | issuehunt: # Replace with a single IssueHunt username
11 | otechie: # Replace with a single Otechie username
12 | lfx_crowdfunding: # Replace with a single LFX Crowdfunding project-name e.g., cloud-foundry
13 | custom: #Nothing
14 | 


--------------------------------------------------------------------------------
/.github/ISSUE_TEMPLATE/bug_report.md:
--------------------------------------------------------------------------------
 1 | ---
 2 | name: Bug report
 3 | about: Create a detailed report on the bug and it's root cause. Conduct root cause error analysis
 4 | title: "[BUG] "
 5 | labels: bug
 6 | assignees: kyegomez
 7 | 
 8 | ---
 9 | 
10 | **Describe the bug**
11 | A clear and concise description of what the bug is and what the main root cause error is. Test very thoroughly before submitting.
12 | 
13 | **To Reproduce**
14 | Steps to reproduce the behavior:
15 | 1. Go to '...'
16 | 2. Click on '....'
17 | 3. Scroll down to '....'
18 | 4. See error
19 | 
20 | **Expected behavior**
21 | A clear and concise description of what you expected to happen.
22 | 
23 | **Screenshots**
24 | If applicable, add screenshots to help explain your problem.
25 | 
26 | **Additional context**
27 | Add any other context about the problem here.
28 | 


--------------------------------------------------------------------------------
/.github/ISSUE_TEMPLATE/feature_request.md:
--------------------------------------------------------------------------------
 1 | ---
 2 | name: Feature request
 3 | about: Suggest an idea for this project
 4 | title: ''
 5 | labels: ''
 6 | assignees: 'kyegomez'
 7 | 
 8 | ---
 9 | 
10 | **Is your feature request related to a problem? Please describe.**
11 | A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
12 | 
13 | **Describe the solution you'd like**
14 | A clear and concise description of what you want to happen.
15 | 
16 | **Describe alternatives you've considered**
17 | A clear and concise description of any alternative solutions or features you've considered.
18 | 
19 | **Additional context**
20 | Add any other context or screenshots about the feature request here.
21 | 


--------------------------------------------------------------------------------
/.github/PULL_REQUEST_TEMPLATE.yml:
--------------------------------------------------------------------------------
 1 | <!-- Thank you for contributing to Zeta!
 2 | 
 3 | Replace this comment with:
 4 |   - Description: a description of the change, 
 5 |   - Issue: the issue # it fixes (if applicable),
 6 |   - Dependencies: any dependencies required for this change,
 7 |   - Tag maintainer: for a quicker response, tag the relevant maintainer (see below),
 8 |   - Twitter handle: we announce bigger features on Twitter. If your PR gets announced and you'd like a mention, we'll gladly shout you out!
 9 | 
10 | If you're adding a new integration, please include:
11 |   1. a test for the integration, preferably unit tests that do not rely on network access,
12 |   2. an example notebook showing its use.
13 | 
14 | Maintainer responsibilities:
15 |   - nn / Misc / if you don't know who to tag: kye@apac.ai
16 |   - tokenizers: kye@apac.ai
17 |   - training / Prompts: kye@apac.ai
18 |   - models: kye@apac.ai
19 | 
20 | If no one reviews your PR within a few days, feel free to kye@apac.ai
21 | 
22 | See contribution guidelines for more information on how to write/run tests, lint, etc: https://github.com/kyegomez/zeta


--------------------------------------------------------------------------------
/.github/dependabot.yml:
--------------------------------------------------------------------------------
 1 | # https://docs.github.com/en/code-security/supply-chain-security/keeping-your-dependencies-updated-automatically/configuration-options-for-dependency-updates
 2 | 
 3 | version: 2
 4 | updates:
 5 |   - package-ecosystem: "github-actions"
 6 |     directory: "/"
 7 |     schedule:
 8 |       interval: "weekly"
 9 | 
10 |   - package-ecosystem: "pip"
11 |     directory: "/"
12 |     schedule:
13 |       interval: "weekly"
14 | 
15 | 


--------------------------------------------------------------------------------
/.github/workflows/code_quality_control.yml:
--------------------------------------------------------------------------------
 1 | name: Linting and Formatting
 2 | 
 3 | on:
 4 |   push:
 5 |     branches:
 6 |       - main
 7 | 
 8 | jobs:
 9 |   lint_and_format:
10 |     runs-on: ubuntu-latest
11 | 
12 |     steps:
13 |       - name: Checkout code
14 |         uses: actions/checkout@v4
15 | 
16 |       - name: Set up Python
17 |         uses: actions/setup-python@v5
18 |         with:
19 |           python-version: 3.x
20 | 
21 |       - name: Install dependencies
22 |         run: pip install -r requirements.txt
23 | 
24 |       - name: Find Python files
25 |         run: find swarms_torch -name "*.py" -type f -exec autopep8 --in-place --aggressive --aggressive {} +
26 | 
27 |       - name: Push changes
28 |         uses: ad-m/github-push-action@master
29 |         with:
30 |           github_token: ${{ secrets.GITHUB_TOKEN }}


--------------------------------------------------------------------------------
/.github/workflows/cos_integration.yml:
--------------------------------------------------------------------------------
 1 | name: Continuous Integration
 2 | 
 3 | on:
 4 |   push:
 5 |     branches:
 6 |       - main
 7 | 
 8 | jobs:
 9 |   test:
10 |     runs-on: ubuntu-latest
11 |     steps:
12 |       - name: Checkout code
13 |         uses: actions/checkout@v4
14 | 
15 |       - name: Set up Python
16 |         uses: actions/setup-python@v5
17 |         with:
18 |           python-version: 3.x
19 | 
20 |       - name: Install dependencies
21 |         run: pip install -r requirements.txt
22 | 
23 |       - name: Run unit tests
24 |         run: pytest tests/unit
25 | 
26 |       - name: Run integration tests
27 |         run: pytest tests/integration
28 | 
29 |       - name: Run code coverage
30 |         run: pytest --cov=swarms tests/
31 | 
32 |       - name: Run linters
33 |         run: pylint swarms
34 | 
35 |       - name: Build documentation
36 |         run: make docs
37 | 
38 |       - name: Validate documentation
39 |         run: sphinx-build -b linkcheck docs build/docs
40 | 
41 |       - name: Run performance tests
42 |         run: pytest tests/performance


--------------------------------------------------------------------------------
/.github/workflows/docs.yml:
--------------------------------------------------------------------------------
 1 | name: Docs WorkFlow
 2 | 
 3 | on:
 4 |   push:
 5 |     branches:
 6 |       - master
 7 |       - main
 8 |       - develop
 9 | jobs:
10 |   deploy:
11 |     runs-on: ubuntu-latest
12 |     steps:
13 |       - uses: actions/checkout@v4
14 |       - uses: actions/setup-python@v5
15 |         with:
16 |           python-version: 3.x
17 |       - run: pip install mkdocs-material
18 |       - run: pip install "mkdocstrings[python]"
19 |       - run: mkdocs gh-deploy --force


--------------------------------------------------------------------------------
/.github/workflows/docs_test.yml:
--------------------------------------------------------------------------------
 1 | name: Documentation Tests
 2 | 
 3 | on:
 4 |   push:
 5 |     branches:
 6 |       - master
 7 | 
 8 | jobs:
 9 |   test:
10 |     runs-on: ubuntu-latest
11 | 
12 |     steps:
13 |       - name: Checkout code
14 |         uses: actions/checkout@v4
15 | 
16 |       - name: Set up Python
17 |         uses: actions/setup-python@v5
18 |         with:
19 |           python-version: 3.x
20 | 
21 |       - name: Install dependencies
22 |         run: pip install -r requirements.txt
23 | 
24 |       - name: Build documentation
25 |         run: make docs
26 | 
27 |       - name: Validate documentation
28 |         run: sphinx-build -b linkcheck docs build/docs


--------------------------------------------------------------------------------
/.github/workflows/label.yml:
--------------------------------------------------------------------------------
 1 | # This workflow will triage pull requests and apply a label based on the
 2 | # paths that are modified in the pull request.
 3 | #
 4 | # To use this workflow, you will need to set up a .github/labeler.yml
 5 | # file with configuration.  For more information, see:
 6 | # https://github.com/actions/labeler
 7 | 
 8 | name: Labeler
 9 | on: [pull_request_target]
10 | 
11 | jobs:
12 |   label:
13 | 
14 |     runs-on: ubuntu-latest
15 |     permissions:
16 |       contents: read
17 |       pull-requests: write
18 | 
19 |     steps:
20 |     - uses: actions/labeler@v5
21 |       with:
22 |         repo-token: "${{ secrets.GITHUB_TOKEN }}"
23 | 


--------------------------------------------------------------------------------
/.github/workflows/lints.yml:
--------------------------------------------------------------------------------
 1 | name: Linting
 2 | 
 3 | on:
 4 |   push:
 5 |     branches:
 6 |       - master
 7 | 
 8 | jobs:
 9 |   lint:
10 |     runs-on: ubuntu-latest
11 | 
12 |     steps:
13 |       - name: Checkout code
14 |         uses: actions/checkout@v4
15 | 
16 |       - name: Set up Python
17 |         uses: actions/setup-python@v5
18 |         with:
19 |           python-version: 3.x
20 | 
21 |       - name: Install dependencies
22 |         run: pip install -r requirements.txt
23 | 
24 |       - name: Run linters
25 |         run: pylint swarms_torch


--------------------------------------------------------------------------------
/.github/workflows/pr_request_checks.yml:
--------------------------------------------------------------------------------
 1 | name: Pull Request Checks
 2 | 
 3 | on:
 4 |   pull_request:
 5 |     branches:
 6 |       - master
 7 | 
 8 | jobs:
 9 |   test:
10 |     runs-on: ubuntu-latest
11 | 
12 |     steps:
13 |       - name: Checkout code
14 |         uses: actions/checkout@v4
15 | 
16 |       - name: Set up Python
17 |         uses: actions/setup-python@v5
18 |         with:
19 |           python-version: 3.x
20 | 
21 |       - name: Install dependencies
22 |         run: pip install -r requirements.txt
23 | 
24 |       - name: Run tests and checks
25 |         run: |
26 |           pytest tests/
27 |           pylint swarms_torch


--------------------------------------------------------------------------------
/.github/workflows/pull-request-links.yml:
--------------------------------------------------------------------------------
 1 | name: readthedocs/actions
 2 | on:
 3 |   pull_request_target:
 4 |     types:
 5 |       - opened
 6 |     paths:
 7 |       - "docs/**"
 8 | 
 9 | permissions:
10 |   pull-requests: write
11 | 
12 | jobs:
13 |   pull-request-links:
14 |     runs-on: ubuntu-latest
15 |     steps:
16 |       - uses: readthedocs/actions/preview@v1
17 |         with:
18 |           project-slug: swarms_torch


--------------------------------------------------------------------------------
/.github/workflows/pylint.yml:
--------------------------------------------------------------------------------
 1 | name: Pylint
 2 | 
 3 | on: [push]
 4 | 
 5 | jobs:
 6 |   build:
 7 |     runs-on: ubuntu-latest
 8 |     strategy:
 9 |       matrix:
10 |         python-version: ["3.8", "3.9", "3.10"]
11 |     steps:
12 |     - uses: actions/checkout@v4
13 |     - name: Set up Python ${{ matrix.python-version }}
14 |       uses: actions/setup-python@v5
15 |       with:
16 |         python-version: ${{ matrix.python-version }}
17 |     - name: Install dependencies
18 |       run: |
19 |         python -m pip install --upgrade pip
20 |         pip install pylint
21 |     - name: Analysing the code with pylint
22 |       run: |
23 |         pylint $(git ls-files '*.py')
24 | 


--------------------------------------------------------------------------------
/.github/workflows/python-publish.yml:
--------------------------------------------------------------------------------
 1 | 
 2 | name: Upload Python Package
 3 | 
 4 | on:
 5 |   release:
 6 |     types: [published]
 7 | 
 8 | permissions:
 9 |   contents: read
10 | 
11 | jobs:
12 |   deploy:
13 | 
14 |     runs-on: ubuntu-latest
15 | 
16 |     steps:
17 |     - uses: actions/checkout@v4
18 |     - name: Set up Python
19 |       uses: actions/setup-python@v5
20 |       with:
21 |         python-version: '3.x'
22 |     - name: Install dependencies
23 |       run: |
24 |         python -m pip install --upgrade pip
25 |         pip install build
26 |     - name: Build package
27 |       run: python -m build
28 |     - name: Publish package
29 |       uses: pypa/gh-action-pypi-publish@e53eb8b103ffcb59469888563dc324e3c8ba6f06
30 |       with:
31 |         user: __token__
32 |         password: ${{ secrets.PYPI_API_TOKEN }}


--------------------------------------------------------------------------------
/.github/workflows/quality.yml:
--------------------------------------------------------------------------------
 1 | name: Quality
 2 | 
 3 | on:
 4 |   push:
 5 |     branches: [ "main" ]
 6 |   pull_request:
 7 |     branches: [ "main" ]
 8 | 
 9 | jobs:
10 |   lint:
11 |     runs-on: ubuntu-latest
12 |     strategy:
13 |       fail-fast: false
14 |     steps:
15 |       - name: Checkout actions
16 |         uses: actions/checkout@v4
17 |         with:
18 |           fetch-depth: 0
19 |       - name: Init environment 
20 |         uses: ./.github/actions/init-environment 
21 |       - name: Run linter
22 |         run: |
23 |           pylint `git diff --name-only --diff-filter=d origin/main HEAD | grep -E '\.py$' | tr '\n' ' '`


--------------------------------------------------------------------------------
/.github/workflows/ruff.yml:
--------------------------------------------------------------------------------
1 | name: Ruff
2 | on: [ push, pull_request ]
3 | jobs:
4 |   ruff:
5 |     runs-on: ubuntu-latest
6 |     steps:
7 |       - uses: actions/checkout@v4
8 |       - uses: chartboost/ruff-action@v1
9 | 


--------------------------------------------------------------------------------
/.github/workflows/run_test.yml:
--------------------------------------------------------------------------------
 1 | name: Python application test
 2 | 
 3 | on: [push]
 4 | 
 5 | jobs:
 6 |   build:
 7 | 
 8 |     runs-on: ubuntu-latest
 9 | 
10 |     steps:
11 |     - uses: actions/checkout@v4
12 |     - name: Set up Python 3.8
13 |       uses: actions/setup-python@v5
14 |       with:
15 |         python-version: 3.8
16 |     - name: Install dependencies
17 |       run: |
18 |         python -m pip install --upgrade pip
19 |         pip install pytest
20 |         if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
21 |     - name: Run tests with pytest
22 |       run: |
23 |         pytest tests/
24 | 


--------------------------------------------------------------------------------
/.github/workflows/stale.yml:
--------------------------------------------------------------------------------
 1 | # This workflow warns and then closes issues and PRs that have had no activity for a specified amount of time.
 2 | #
 3 | # You can adjust the behavior by modifying this file.
 4 | # For more information, see:
 5 | # https://github.com/actions/stale
 6 | name: Mark stale issues and pull requests
 7 | 
 8 | on:
 9 |   schedule:
10 |   - cron: '26 12 * * *'
11 | 
12 | jobs:
13 |   stale:
14 | 
15 |     runs-on: ubuntu-latest
16 |     permissions:
17 |       issues: write
18 |       pull-requests: write
19 | 
20 |     steps:
21 |     - uses: actions/stale@v9
22 |       with:
23 |         repo-token: ${{ secrets.GITHUB_TOKEN }}
24 |         stale-issue-message: 'Stale issue message'
25 |         stale-pr-message: 'Stale pull request message'
26 |         stale-issue-label: 'no-issue-activity'
27 |         stale-pr-label: 'no-pr-activity'


--------------------------------------------------------------------------------
/.github/workflows/test.yml:
--------------------------------------------------------------------------------
 1 | name: test
 2 | 
 3 | on:
 4 |   push:
 5 |     branches: [master]
 6 |   pull_request:
 7 |   workflow_dispatch:
 8 | 
 9 | env:
10 |   POETRY_VERSION: "1.4.2"
11 | 
12 | jobs:
13 |   build:
14 |     runs-on: ubuntu-latest
15 |     strategy:
16 |       matrix:
17 |         python-version:
18 |           - "3.8"
19 |           - "3.9"
20 |           - "3.10"
21 |           - "3.11"
22 |         test_type:
23 |           - "core"
24 |           - "extended"
25 |     name: Python ${{ matrix.python-version }} ${{ matrix.test_type }}
26 |     steps:
27 |       - uses: actions/checkout@v4
28 |       - name: Set up Python ${{ matrix.python-version }}
29 |         uses: "./.github/actions/poetry_setup"
30 |         with:
31 |           python-version: ${{ matrix.python-version }}
32 |           poetry-version: "1.4.2"
33 |           cache-key: ${{ matrix.test_type }}
34 |           install-command: |
35 |               if [ "${{ matrix.test_type }}" == "core" ]; then
36 |                 echo "Running core tests, installing dependencies with poetry..."
37 |                 poetry install
38 |               else
39 |                 echo "Running extended tests, installing dependencies with poetry..."
40 |                 poetry install -E extended_testing
41 |               fi
42 |       - name: Run ${{matrix.test_type}} tests
43 |         run: |
44 |           if [ "${{ matrix.test_type }}" == "core" ]; then
45 |             make test
46 |           else
47 |             make extended_tests
48 |           fi
49 |         shell: bash


--------------------------------------------------------------------------------
/.github/workflows/testing.yml:
--------------------------------------------------------------------------------
 1 | name: Unit Tests
 2 | 
 3 | on:
 4 |   push:
 5 |     branches:
 6 |       - master
 7 | 
 8 | jobs:
 9 |   test:
10 |     runs-on: ubuntu-latest
11 | 
12 |     steps:
13 |       - name: Checkout code
14 |         uses: actions/checkout@v4
15 | 
16 |       - name: Set up Python
17 |         uses: actions/setup-python@v5
18 |         with:
19 |           python-version: 3.x
20 | 
21 |       - name: Install dependencies
22 |         run: pip install -r requirements.txt
23 | 
24 |       - name: Run unit tests
25 |         run: pytest tests/


--------------------------------------------------------------------------------
/.github/workflows/unit-test.yml:
--------------------------------------------------------------------------------
 1 | name: build
 2 | 
 3 | on:
 4 |   push:
 5 |     branches: [ main ]
 6 |   pull_request:
 7 |     branches: [ main ]
 8 | 
 9 | jobs:
10 | 
11 |   build:
12 | 
13 |     runs-on: ubuntu-latest
14 | 
15 |     steps:
16 |     - uses: actions/checkout@v4
17 | 
18 |     - name: Setup Python
19 |       uses: actions/setup-python@v5
20 |       with:
21 |         python-version: '3.10'
22 | 
23 |     - name: Install dependencies
24 |       run: pip install -r requirements.txt
25 | 
26 |     - name: Run Python unit tests
27 |       run: python3 -m unittest tests/
28 | 
29 |     - name: Verify that the Docker image for the action builds
30 |       run: docker build . --file Dockerfile
31 | 
32 |     - name: Verify integration test results
33 |       run: python3 -m unittest tests/
34 | 


--------------------------------------------------------------------------------
/.github/workflows/welcome.yml:
--------------------------------------------------------------------------------
 1 | name: Welcome WorkFlow
 2 | 
 3 | on:
 4 |   issues:
 5 |     types: [opened]
 6 |   pull_request_target:
 7 |     types: [opened]
 8 | 
 9 | jobs:
10 |   build:
11 |     name: 👋 Welcome
12 |     runs-on: ubuntu-latest
13 |     steps:
14 |       - uses: actions/first-interaction@v1.3.0
15 |         with:
16 |           repo-token: ${{ secrets.GITHUB_TOKEN }}
17 |           issue-message: "Hello there, thank you for opening an Issue ! 🙏🏻 The team was notified and they will get back to you asap."
18 |           pr-message:  "Hello there, thank you for opening an PR ! 🙏🏻 The team was notified and they will get back to you asap."


--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
  1 | # Byte-compiled / optimized / DLL files
  2 | __pycache__/
  3 | *.py[cod]
  4 | *$py.class
  5 | 
  6 | # C extensions
  7 | *.so
  8 | 
  9 | # Distribution / packaging
 10 | .Python
 11 | build/
 12 | develop-eggs/
 13 | dist/
 14 | downloads/
 15 | eggs/
 16 | .eggs/
 17 | lib/
 18 | lib64/
 19 | parts/
 20 | .ruff_cache/
 21 | sdist/
 22 | var/
 23 | wheels/
 24 | share/python-wheels/
 25 | *.egg-info/
 26 | .installed.cfg
 27 | *.egg
 28 | MANIFEST
 29 | 
 30 | # PyInstaller
 31 | #  Usually these files are written by a python script from a template
 32 | #  before PyInstaller builds the exe, so as to inject date/other infos into it.
 33 | *.manifest
 34 | *.spec
 35 | 
 36 | # Installer logs
 37 | pip-log.txt
 38 | pip-delete-this-directory.txt
 39 | 
 40 | # Unit test / coverage reports
 41 | htmlcov/
 42 | .tox/
 43 | .nox/
 44 | .coverage
 45 | .coverage.*
 46 | .cache
 47 | nosetests.xml
 48 | coverage.xml
 49 | *.cover
 50 | *.py,cover
 51 | .hypothesis/
 52 | .pytest_cache/
 53 | cover/
 54 | 
 55 | # Translations
 56 | *.mo
 57 | *.pot
 58 | 
 59 | # Django stuff:
 60 | *.log
 61 | local_settings.py
 62 | db.sqlite3
 63 | db.sqlite3-journal
 64 | 
 65 | # Flask stuff:
 66 | instance/
 67 | .webassets-cache
 68 | 
 69 | # Scrapy stuff:
 70 | .scrapy
 71 | 
 72 | # Sphinx documentation
 73 | docs/_build/
 74 | 
 75 | # PyBuilder
 76 | .pybuilder/
 77 | target/
 78 | 
 79 | # Jupyter Notebook
 80 | .ipynb_checkpoints
 81 | 
 82 | # IPython
 83 | profile_default/
 84 | ipython_config.py
 85 | 
 86 | # pyenv
 87 | #   For a library or package, you might want to ignore these files since the code is
 88 | #   intended to run in multiple environments; otherwise, check them in:
 89 | # .python-version
 90 | 
 91 | # pipenv
 92 | #   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
 93 | #   However, in case of collaboration, if having platform-specific dependencies or dependencies
 94 | #   having no cross-platform support, pipenv may install dependencies that don't work, or not
 95 | #   install all needed dependencies.
 96 | #Pipfile.lock
 97 | 
 98 | # poetry
 99 | #   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
100 | #   This is especially recommended for binary packages to ensure reproducibility, and is more
101 | #   commonly ignored for libraries.
102 | #   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
103 | #poetry.lock
104 | 
105 | # pdm
106 | #   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
107 | #pdm.lock
108 | #   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
109 | #   in version control.
110 | #   https://pdm.fming.dev/#use-with-ide
111 | .pdm.toml
112 | 
113 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
114 | __pypackages__/
115 | 
116 | # Celery stuff
117 | celerybeat-schedule
118 | celerybeat.pid
119 | 
120 | # SageMath parsed files
121 | *.sage.py
122 | 
123 | # Environments
124 | .env
125 | .venv
126 | env/
127 | venv/
128 | ENV/
129 | env.bak/
130 | venv.bak/
131 | 
132 | # Spyder project settings
133 | .spyderproject
134 | .spyproject
135 | 
136 | # Rope project settings
137 | .ropeproject
138 | 
139 | # mkdocs documentation
140 | /site
141 | 
142 | # mypy
143 | .mypy_cache/
144 | .dmypy.json
145 | dmypy.json
146 | 
147 | # Pyre type checker
148 | .pyre/
149 | 
150 | # pytype static type analyzer
151 | .pytype/
152 | 
153 | # Cython debug symbols
154 | cython_debug/
155 | 
156 | # PyCharm
157 | #  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
158 | #  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
159 | #  and can be added to the global gitignore or merged into this file.  For a more nuclear
160 | #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
161 | #.idea/
162 | 


--------------------------------------------------------------------------------
/.pre-commit-config.yaml:
--------------------------------------------------------------------------------
 1 | repos:
 2 |   - repo: https://github.com/ambv/black
 3 |     rev: 22.3.0
 4 |     hooks:
 5 |     - id: black
 6 |   - repo: https://github.com/charliermarsh/ruff-pre-commit
 7 |     rev: 'v0.0.255'
 8 |     hooks:
 9 |       - id: ruff
10 |         args: [--fix]
11 |   - repo: https://github.com/nbQA-dev/nbQA
12 |     rev: 1.6.3
13 |     hooks:
14 |     - id: nbqa-black
15 |       additional_dependencies: [ipython==8.12, black]
16 |     - id: nbqa-ruff 
17 |       args: ["--ignore=I001"]
18 |       additional_dependencies: [ipython==8.12, ruff]


--------------------------------------------------------------------------------
/.readthedocs.yml:
--------------------------------------------------------------------------------
 1 | version: 2
 2 | 
 3 | build:
 4 |   os: ubuntu-22.04
 5 |   tools:
 6 |     python: "3.11"
 7 | 
 8 | mkdocs:
 9 |   configuration: mkdocs.yml
10 | 
11 | python:
12 |    install:
13 |    - requirements: requirements.txt


--------------------------------------------------------------------------------
/LICENSE:
--------------------------------------------------------------------------------
 1 | MIT License
 2 | 
 3 | Copyright (c) 2023 Eternal Reclaimer
 4 | 
 5 | Permission is hereby granted, free of charge, to any person obtaining a copy
 6 | of this software and associated documentation files (the "Software"), to deal
 7 | in the Software without restriction, including without limitation the rights
 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 9 | copies of the Software, and to permit persons to whom the Software is
10 | furnished to do so, subject to the following conditions:
11 | 
12 | The above copyright notice and this permission notice shall be included in all
13 | copies or substantial portions of the Software.
14 | 
15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21 | SOFTWARE.
22 | 


--------------------------------------------------------------------------------
/Makefile:
--------------------------------------------------------------------------------
 1 | .PHONY: style check_code_quality
 2 | 
 3 | export PYTHONPATH = .
 4 | check_dirs := src
 5 | 
 6 | style:
 7 | 	black  $(check_dirs)
 8 | 	isort --profile black $(check_dirs)
 9 | 
10 | check_code_quality:
11 | 	black --check $(check_dirs)
12 | 	isort --check-only --profile black $(check_dirs)
13 | 	# stop the build if there are Python syntax errors or undefined names
14 | 	flake8 $(check_dirs) --count --select=E9,F63,F7,F82 --show-source --statistics
15 | 	# exit-zero treats all errors as warnings. E203 for black, E501 for docstring, W503 for line breaks before logical operators 
16 | 	flake8 $(check_dirs) --count --max-line-length=88 --exit-zero  --ignore=D --extend-ignore=E203,E501,W503  --statistics
17 | 	
18 | publish:
19 | 	python setup.py sdist bdist_wheel
20 | 	twine upload -r testpypi dist/* -u ${PYPI_USERNAME} -p ${PYPI_TEST_PASSWORD} --verbose 
21 | 	twine check dist/*
22 | 	twine upload dist/* -u ${PYPI_USERNAME} -p ${PYPI_PASSWORD} --verbose 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
  1 | [![Multi-Modality](agorabanner.png)](https://discord.gg/qUtxnK2NMf)
  2 | 
  3 | # Pali3
  4 | ![pali](pali.png)
  5 | 
  6 | "Figure 1: Overview of the PaLI-3 (5B) model: images are encoded into visual tokens individually
  7 | by the contrastively pretrained 2B SigLIP vision model. Along with a query, these visual tokens
  8 | are passed to an 3B encoder-decoder UL2 Transformer which produces the desired answer."
  9 | 
 10 | 
 11 | Vit trained with siglip loss -> embeddings -> ul2 -> text tokens
 12 | 
 13 | text -> tokenizer -> embeddings -> ul2 -> text tokens
 14 | 
 15 | [ARXVIV PAPER LINK](https://arxiv.org/pdf/2310.09199v1.pdf)
 16 | 
 17 | --------
 18 | 
 19 | ## Installation
 20 | 
 21 | `pip install pali3`
 22 | 
 23 | -------
 24 | 
 25 | ## Usage:
 26 | 
 27 | ```python
 28 | import torch
 29 | from pali3.main import Pali3
 30 | 
 31 | model = Pali3()
 32 | 
 33 | img = torch.randn(1, 3, 256, 256)
 34 | prompt = torch.randint(0, 256, (1, 1024))
 35 | mask = torch.ones(1, 1024).bool()
 36 | output_text = torch.randint(0, 256, (1, 1024))
 37 | 
 38 | result = model.process(img, prompt, output_text, mask)
 39 | print(result)
 40 | 
 41 | 
 42 | ```
 43 | 
 44 | 
 45 | -------
 46 | 
 47 | ## Architecture
 48 | 
 49 | Here is the ASCII representation of the model architecture and the stages of training:
 50 | 
 51 | ```
 52 | Model Architecture:
 53 | 
 54 | Image Input
 55 |     |
 56 |     V
 57 | Contrastive Vision Encoder (ViT-G/14)
 58 |     |
 59 |     V
 60 | Transformer Encoder
 61 |     |
 62 |     V
 63 | Transformer Decoder
 64 |     |
 65 |     V
 66 | Text Output
 67 | 
 68 | Stages of Training:
 69 | 
 70 | Stage 0: Unimodal pretraining
 71 |     |
 72 |     V
 73 | Stage 1: Multimodal training
 74 |     |
 75 |     V
 76 | Stage 2: Resolution increase
 77 |     |
 78 |     V
 79 | Task specialization (transfer)
 80 | 
 81 | ```
 82 | 
 83 | 
 84 | # Model Training Phases
 85 | The model architecture consists of a contrastive vision encoder (ViT-G/14) that encodes the image into tokens. These tokens are passed to a transformer encoder and then to a transformer decoder that generates a text output.
 86 | 
 87 | The training procedure consists of multiple stages:
 88 | 
 89 | -   Stage 0: Unimodal pretraining. The image encoder is pretrained contrastively on image-text pairs from the web, following the SigLIP training protocol. The text encoder-decoder is a 3B UL2 model trained following the mixture of denoisers procedure.
 90 | 
 91 | -   Stage 1: Multimodal training. The image encoder is combined with the text encoder-decoder and trained on a multimodal task and data mixture, keeping the image encoder frozen and using its native resolution.
 92 | 
 93 | -   Stage 2: Resolution increase. The resolution of the model is increased by fine-tuning the whole model with a short curriculum of increasing resolutions.
 94 | 
 95 | -   Task specialization (transfer). Finally, for each individual task, the model is fine-tuned with frozen ViT image encoder on the task's training data.
 96 | 
 97 | Please note that this is a high-level representation and the actual implementation might involve more details and complexities.
 98 | 
 99 | 
100 | 
101 | ------
102 | 
103 | # Vit Architecture
104 | Here are the ASCII diagrams for the ViT (Vision Transformer)
105 | 
106 | ```
107 | ViT (Vision Transformer):
108 | 
109 | Image Input
110 |     |
111 |     V
112 | Patch Extraction
113 |     |
114 |     V
115 | Linear Embedding
116 |     |
117 |     V
118 | Positional Encoding
119 |     |
120 |     V
121 | Transformer Encoder Blocks (Multiple Layers)
122 |     |
123 |     V
124 | Classification Head (Optional)
125 |     |
126 |     V
127 | Output (Image Embeddings)
128 | 
129 | ```
130 | 
131 | The ViT starts with patch extraction from the input image. These patches are then linearly embedded and positional encodings are added. The resulting sequence of patch embeddings is passed through multiple layers of transformer encoders. Optionally, a classification head can be added at the end to get class probabilities for image classification tasks. The output of the ViT is the image embeddings.
132 | 
133 | -------
134 | 
135 | # UL2 Encoder/Decoder Transformer
136 | ```
137 | Encoder-Decoder Architecture:
138 | 
139 | Input (Image + Text Tokens)
140 |     |
141 |     V
142 | Transformer Encoder
143 |     |
144 |     V
145 | Encoder Output (Context for Decoder)
146 |     |
147 |     V
148 | Transformer Decoder
149 |     |
150 |     V
151 | Output (Generated Text)
152 | 
153 | ```
154 | 
155 | The encoder-decoder architecture starts with the input, which is a combination of image and text tokens in this case. The input is passed through a transformer encoder, which generates a context for the decoder. The transformer decoder then uses this context to generate the output text.
156 | 
157 | 
158 | # Dataset Strategy
159 | Here is a table summarizing the key datasets mentioned in the paper along with their metadata and source links:
160 | 
161 | - Made with claude so links could be fake
162 | 
163 | | Dataset | Type | Size | Tasks | Source |
164 | |-|-|-|-|-|
165 | | ImageNet-22k | Image Classification | 14M images, 21,841 classes | Pretraining | https://github.com/google-research-datasets/ImageNet-21k-P |
166 | | MS COCO | Image Captioning, VQA | 330K images, 80 object categories | Evaluation | https://cocodataset.org | 
167 | | Flickr30k | Image Captioning | 31K images | Evaluation | https://www.kaggle.com/dataset/flickr30k |
168 | | VQAv2 | Visual QA | 204K images, 1.1M questions | Evaluation | https://visualqa.org/download.html |  
169 | | GQA | Visual QA | 22M graph-based questions | Evaluation | https://cs.stanford.edu/people/dorarad/gqa/download.html |
170 | | RefCOCO/RefCOCO+ | Referring Expression | 19,994/19,992 images | Evaluation | https://github.com/lichengunc/refer |
171 | | TextCaps | Image Captioning | 31,014 images | Evaluation | https://textvqa.org/textcaps |
172 | | TextVQA | Visual QA | 28,408 images | Evaluation | https://textvqa.org/index.html |
173 | | STVQA | Visual QA | 249,991 QA pairs | Evaluation | https://tvqa.cs.unc.edu/ |
174 | | OCR-VQA | Visual QA | 45,336 images | Evaluation | https://ocrvqa.cloudcv.org/ |
175 | | DocVQA | Visual QA | 5,000 document images | Evaluation | https://github.com/doc-vqa/docvqa |
176 | | InfographiVQA | Visual QA | 10,047 infographic images | Evaluation | https://github.com/doc-vqa/InfoVQA |
177 | | WebLI | Image-Text Pairs | 72M image-text pairs in 100+ languages | Pretraining | https://laion.ai/blogs/webli/ |
178 | | JFT-300M | Image Classification | 303M images, 18,291 classes | Pretraining | https://github.com/google-research-datasets/jft300m |
179 | | CrossModal-3600 | Image-Text Retrieval | 31K images, 3600 lang-image pairs | Evaluation | https://laion.ai/crossmodal-3600/ |
180 | 
181 | -----
182 | 
183 | # License
184 | MIT
185 | 
186 | # Todo
187 | 
188 | - [x] Implement sig_lip vit model with training recipe
189 | - [x] Implement the text tokenizer, maybe use token monster 
190 | - [x] Implement the UL2 Transformer Encoder and Decoder
191 | - [ ] Implement the pooling layer after vit then linear
192 | - [ ] Implement the prepending the visual token embeddings to the text embeddings
193 | - [ ] Implement training scripts for the full pali3 model
194 | 
195 | 
196 | # Citation
197 | 
198 | ```bibtex
199 | @misc{2310.09199,
200 | Author = {Xi Chen and Xiao Wang and Lucas Beyer and Alexander Kolesnikov and Jialin Wu and Paul Voigtlaender and Basil Mustafa and Sebastian Goodman and Ibrahim Alabdulmohsin and Piotr Padlewski and Daniel Salz and Xi Xiong and Daniel Vlasic and Filip Pavetic and Keran Rong and Tianli Yu and Daniel Keysers and Xiaohua Zhai and Radu Soricut},
201 | Title = {PaLI-3 Vision Language Models: Smaller, Faster, Stronger},
202 | Year = {2023},
203 | Eprint = {arXiv:2310.09199},
204 | }
205 | ```
206 | 
207 | 
208 | 
209 | 


--------------------------------------------------------------------------------
/agorabanner.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/kyegomez/PALI3/355fd477e04efcc2dc7abd3676569f18b050de63/agorabanner.png


--------------------------------------------------------------------------------
/example.py:
--------------------------------------------------------------------------------
 1 | import torch
 2 | from pali3.main import Pali3
 3 | 
 4 | 
 5 | model = Pali3()
 6 | 
 7 | img = torch.randn(1, 3, 256, 256)
 8 | prompt = torch.randint(0, 256, (1, 1024))
 9 | mask = torch.ones(1, 1024).bool()
10 | output_text = torch.randint(0, 256, (1, 1024))
11 | 
12 | result = model.process(img, prompt, output_text, mask)
13 | print(result)
14 | 


--------------------------------------------------------------------------------
/pali.png:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/kyegomez/PALI3/355fd477e04efcc2dc7abd3676569f18b050de63/pali.png


--------------------------------------------------------------------------------
/pali3/__init__.py:
--------------------------------------------------------------------------------
1 | from pali3.tokenizer import tokenize, detokenize
2 | from pali3.ul2 import UL2, ViTransformerWrapper, Encoder, Decoder, TransformerWrapper
3 | from pali3.main import VitModel, Pali3
4 | 


--------------------------------------------------------------------------------
/pali3/main.py:
--------------------------------------------------------------------------------
  1 | import torch
  2 | from torch import nn
  3 | from pali3.ul2 import UL2, ViTransformerWrapper, Encoder
  4 | 
  5 | 
  6 | class PrependTokens(nn.Module):
  7 |     """
  8 | 
  9 |     # Initialize models
 10 |     vit_model = ViTModel()
 11 |     text_embedding = TextEmbedding("bert-base-uncased")
 12 | 
 13 |     # Initialize PrependVisualTokens
 14 |     prepend_visual_tokens = PrependVisualTokens(vit_model, text_embedding)
 15 | 
 16 |     # Process image and text
 17 |     img = torch.randn(1, 3, 256, 256)  # dummy image
 18 |     text = "This is a sample text"
 19 |     combined_tokens = prepend_visual_tokens.process(img, text)
 20 | 
 21 |     """
 22 | 
 23 |     def __init__(
 24 |         self,
 25 |         vit,
 26 |         text_embedding,
 27 |     ):
 28 |         super().__init__()
 29 |         self.vit = vit
 30 |         self.text_embedding = text_embedding
 31 | 
 32 |     def forward(self, x):
 33 |         visual_tokens = self.vit.process(x)
 34 |         text_tokens = self.text_embedding.process(x)
 35 |         combined_tokens = torch.cat((visual_tokens, text_tokens), dim=1)
 36 |         return combined_tokens
 37 | 
 38 | 
 39 | class VitModel:
 40 |     """
 41 |     VitModel is a wrapper around the ViT model from the PyTorch Image Models library.
 42 | 
 43 |     Args:
 44 |         image_size (int, optional): Size of the image. Defaults to 256.
 45 |         patch_size (int, optional): Size of the patch. Defaults to 32.
 46 |         dim (int, optional): Dimension of the model. Defaults to 512.
 47 |         depth (int, optional): Depth of the model. Defaults to 6.
 48 |         heads (int, optional): Number of heads in the model. Defaults to 8.
 49 | 
 50 |     Raises:
 51 |         ValueError: If the input image is None.
 52 |         ValueError: If the input image shape is not [*, 3, image_size, image_size].
 53 | 
 54 |     Examples:
 55 |     x = torch.randn(1, 3, 256, 256)
 56 |     model = VitModel()
 57 |     out = model.process(x)
 58 |     print(out)
 59 | 
 60 | 
 61 | 
 62 | 
 63 |     """
 64 | 
 65 |     def __init__(
 66 |         self, image_size=256, patch_size=32, dim=512, depth=6, heads=8, *args, **kwargs
 67 |     ):
 68 |         self.image_size = image_size
 69 |         self.patch_size = patch_size
 70 |         self.dim = dim
 71 | 
 72 |         self.depth = depth
 73 |         self.heads = heads
 74 | 
 75 |         self.vit = ViTransformerWrapper(
 76 |             image_size=image_size,
 77 |             patch_size=patch_size,
 78 |             attn_layers=Encoder(dim=dim, depth=depth, heads=heads),
 79 |         )
 80 |         # adaptive avg pool
 81 |         # self.pool = nn.MaxPool1d(kernel_size=2, stride=2)
 82 |         self.linear_projection = nn.Linear(dim, dim)
 83 | 
 84 |     def process(self, img):
 85 |         if img is None:
 86 |             raise ValueError("Input image cannot be None")
 87 |         if img.shape[1:] != (3, self.image_size, self.image_size):
 88 |             raise ValueError(
 89 |                 "Input image must have the shape [*, 3, {}, {}]".format(
 90 |                     self.image_size, self.image_size
 91 |                 )
 92 |             )
 93 |         vit_output = self.vit(img, return_embeddings=True)
 94 |         # missing pool add later
 95 |         # pooled_output = self.pool(vit_output)
 96 |         projected_output = self.linear_projection(vit_output)
 97 |         return projected_output
 98 | 
 99 | 
100 | class Pali3:
101 |     """
102 |     Pali3 is a vit model with a transformer encoder and decoder.
103 |     It is a wrapper around the UL2 model from the PyTorch Image Models library.
104 | 
105 | 
106 |     Args:
107 |         model_name (str, optional): Name of the model. Defaults to None.
108 |         image_size (int, optional): Size of the image. Defaults to 256.
109 |         patch_size (int, optional): Size of the patch. Defaults to 32.
110 |         dim (int, optional): Dimension of the model. Defaults to 512.
111 |         depth (int, optional): Depth of the model. Defaults to 6.
112 |         heads (int, optional): Number of heads in the model. Defaults to 8.
113 |         enc_num_tokens (int, optional): Number of tokens in the encoder. Defaults to 256.
114 |         enc_max_seq_len (int, optional): Maximum sequence length in the encoder. Defaults to 1024.
115 |         dec_num_tokens (int, optional): Number of tokens in the decoder. Defaults to 256.
116 |         dec_max_seq_len (int, optional): Maximum sequence length in the decoder. Defaults to 1024.
117 |         enc_depth (int, optional): Depth of the encoder. Defaults to 6.
118 |         enc_heads (int, optional): Number of heads in the encoder. Defaults to 8.
119 |         dec_depth (int, optional): Depth of the decoder. Defaults to 6.
120 |         dec_heads (int, optional): Number of heads in the decoder. Defaults to 8.
121 |         seq_len (int, optional): Length of the sequence. Defaults to 1024.
122 | 
123 |     Raises:
124 |         ValueError: If the model name is None.
125 |         ValueError: If the tokenizer is None.
126 | 
127 |     Examples:
128 | 
129 |     model = Pali3()
130 |     img = torch.randn(1, 3, 256, 256)  # dummy image
131 |     prompt = "This is a sample text"
132 |     output = "This is a sample text"
133 |     mask = None
134 |     result = model.process(img, prompt, output, mask)
135 | 
136 |     """
137 | 
138 |     def __init__(
139 |         self,
140 |         model_name=None,
141 |         image_size=256,
142 |         patch_size=32,
143 |         dim=512,
144 |         depth=6,
145 |         heads=8,
146 |         enc_num_tokens=256,
147 |         enc_max_seq_len=1024,
148 |         dec_num_tokens=256,
149 |         dec_max_seq_len=1024,
150 |         enc_depth=6,
151 |         enc_heads=8,
152 |         dec_depth=6,
153 |         dec_heads=8,
154 |         seq_len=1024,
155 |     ):
156 |         self.model_name = model_name
157 |         self.vit_model = VitModel(
158 |             image_size=image_size,
159 |             patch_size=patch_size,
160 |             dim=dim,
161 |             depth=depth,
162 |             heads=heads,
163 |         )
164 | 
165 |         self.pali_model = UL2(
166 |             dim=dim,
167 |             enc_num_tokens=enc_num_tokens,
168 |             enc_depth=enc_depth,
169 |             enc_heads=enc_heads,
170 |             enc_max_seq_len=enc_max_seq_len,
171 |             dec_num_tokens=dec_num_tokens,
172 |             dec_depth=dec_depth,
173 |             dec_heads=dec_heads,
174 |             dec_max_seq_len=dec_max_seq_len,
175 |         )
176 | 
177 |     def process(self, img, prompt, output, mask):
178 |         img_embeds = self.vit_model.process(img)
179 | 
180 |         result = self.pali_model(
181 |             prompt, output, mask=mask, src_prepend_embeds=img_embeds
182 |         )
183 |         return result
184 | 
185 |     def generate(self, text, mask=None, attn_mask=None, model_name=None):
186 |         if model_name:
187 |             self.model_name = model_name
188 | 
189 |         if not self.model_name:
190 |             raise ValueError(
191 |                 "model_name must be specidfied either in the class constructor or in the generate method"
192 |             )
193 | 
194 |         seq_out_start = torch.zeros(1, 1).long()
195 |         result = self.pali_model.generate(
196 |             text, seq_out_start, self.seq_len, mask, attn_mask
197 |         )
198 |         result_text = self.tokenizer.decode(result[0], skip_special_tokens=True)
199 |         return result_text
200 | 


--------------------------------------------------------------------------------
/pali3/tokenizer.py:
--------------------------------------------------------------------------------
 1 | # multi modal tokenizer
 2 | from zeta.tokenizers import TokenMonster
 3 | 
 4 | 
 5 | def tokenize(text):
 6 |     tokenizer = TokenMonster("englishcode-32000-consistent-v1")
 7 |     tokens = tokenizer.tokenize(text)
 8 |     return tokens
 9 | 
10 | 
11 | def detokenize(tokens):
12 |     tokenizer = TokenMonster("englishcode-32000-consistent-v1")
13 |     text = tokenizer.detokenize(tokens)
14 |     return text
15 | 
16 | 
17 | # tokenize = tokenize("Hello my name is Kye what's your name")
18 | # print(tokenize)
19 | 


--------------------------------------------------------------------------------
/pali3/train_vit_siglip.py:
--------------------------------------------------------------------------------
 1 | import torch
 2 | from torch.utils.data import DataLoader
 3 | from torchvision import datasets, transforms
 4 | from torch.optim.lr_scheduler import StepLR
 5 | from pali3.main import VitModel
 6 | from zeta.nn.modules import SigLipLoss
 7 | 
 8 | 
 9 | def pretrain(
10 |     model,
11 |     train_loader,
12 |     val_loader,
13 |     optimizer,
14 |     loss_fn,
15 |     device,
16 |     scheduler,
17 |     num_epochs,
18 |     model_path,
19 | ):
20 |     best_val_loss = float("inf")
21 |     for epoch in range(num_epochs):
22 |         model.train()
23 |         for images, labels in train_loader:
24 |             images = images.to(device)
25 |             labels = labels.to(device)
26 |             optimizer.zero_grad()
27 |             embeddings = model.process(images)
28 |             loss = loss_fn(embeddings, labels)
29 |             loss.backward()
30 |             optimizer.step()
31 |         scheduler.step()
32 |         val_loss = validate(model, val_loader, loss_fn, device)
33 |         print(f"Epoch: {epoch+1}, Train Loss: {loss.item()}, Val Loss: {val_loss}")
34 |         if val_loss < best_val_loss:
35 |             best_val_loss = val_loss
36 |             torch.save(model.state_dict(), model_path)
37 | 
38 | 
39 | def validate(model, dataloader, loss_fn, device):
40 |     model.eval()
41 |     total_loss = 0
42 |     with torch.no_grad():
43 |         for images, labels in dataloader:
44 |             images = images.to(device)
45 |             labels = labels.to(device)
46 |             embeddings = model.process(images)
47 |             loss = loss_fn(embeddings, labels)
48 |             total_loss += loss.item()
49 |     return total_loss / len(dataloader)
50 | 
51 | 
52 | # Dataset and DataLoader
53 | transform = transforms.Compose(
54 |     [transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
55 | )
56 | train_dataset = datasets.CIFAR10(
57 |     root="./data", train=True, download=True, transform=transform
58 | )
59 | val_dataset = datasets.CIFAR10(
60 |     root="./data", train=False, download=True, transform=transform
61 | )
62 | train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=2)
63 | val_loader = DataLoader(val_dataset, batch_size=64, shuffle=False, num_workers=2)
64 | 
65 | # Model, Optimizer, Scheduler and Loss
66 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
67 | model = VitModel(image_size=32, patch_size=4, dim=512, depth=6, heads=8).to(device)
68 | optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
69 | scheduler = StepLR(optimizer, step_size=10, gamma=0.1)
70 | loss_fn = SigLipLoss()
71 | 
72 | # Pretrain the model
73 | model_path = "./model.pth"
74 | num_epochs = 100
75 | pretrain(
76 |     model,
77 |     train_loader,
78 |     val_loader,
79 |     optimizer,
80 |     loss_fn,
81 |     device,
82 |     scheduler,
83 |     num_epochs,
84 |     model_path,
85 | )
86 | 


--------------------------------------------------------------------------------
/pali3/ul2.py:
--------------------------------------------------------------------------------
   1 | import math
   2 | from random import random
   3 | 
   4 | import torch
   5 | from torch import nn, einsum, Tensor
   6 | import torch.nn.functional as F
   7 | 
   8 | from functools import partial, wraps
   9 | from inspect import isfunction
  10 | from dataclasses import dataclass
  11 | from typing import List, Callable, Optional
  12 | 
  13 | from einops import rearrange, repeat, reduce, pack, unpack
  14 | 
  15 | from zeta.nn.attention.attend import Attend, Intermediates
  16 | from zeta.structs.auto_regressive_wrapper import AutoregressiveWrapper
  17 | 
  18 | 
  19 | # constants
  20 | 
  21 | DEFAULT_DIM_HEAD = 64
  22 | 
  23 | 
  24 | @dataclass
  25 | class LayerIntermediates:
  26 |     hiddens: Optional[List[Tensor]] = None
  27 |     attn_intermediates: Optional[List[Intermediates]] = None
  28 |     layer_hiddens: Optional[List[Tensor]] = None
  29 |     attn_z_loss: Optional[Tensor] = None
  30 |     mems: Optional[Tensor] = None
  31 | 
  32 | 
  33 | # helpers
  34 | 
  35 | 
  36 | def exists(val):
  37 |     return val is not None
  38 | 
  39 | 
  40 | def default(val, d):
  41 |     if exists(val):
  42 |         return val
  43 |     return d() if isfunction(d) else d
  44 | 
  45 | 
  46 | def cast_tuple(val, depth):
  47 |     return val if isinstance(val, tuple) else (val,) * depth
  48 | 
  49 | 
  50 | def divisible_by(num, den):
  51 |     return (num % den) == 0
  52 | 
  53 | 
  54 | def maybe(fn):
  55 |     @wraps(fn)
  56 |     def inner(x, *args, **kwargs):
  57 |         if not exists(x):
  58 |             return x
  59 |         return fn(x, *args, **kwargs)
  60 | 
  61 |     return inner
  62 | 
  63 | 
  64 | class always:
  65 |     def __init__(self, val):
  66 |         self.val = val
  67 | 
  68 |     def __call__(self, *args, **kwargs):
  69 |         return self.val
  70 | 
  71 | 
  72 | class not_equals:
  73 |     def __init__(self, val):
  74 |         self.val = val
  75 | 
  76 |     def __call__(self, x, *args, **kwargs):
  77 |         return x != self.val
  78 | 
  79 | 
  80 | class equals:
  81 |     def __init__(self, val):
  82 |         self.val = val
  83 | 
  84 |     def __call__(self, x, *args, **kwargs):
  85 |         return x == self.val
  86 | 
  87 | 
  88 | def Sequential(*modules):
  89 |     return nn.Sequential(*filter(exists, modules))
  90 | 
  91 | 
  92 | # tensor helpers
  93 | 
  94 | 
  95 | def max_neg_value(tensor):
  96 |     return -torch.finfo(tensor.dtype).max
  97 | 
  98 | 
  99 | def l2norm(t, groups=1):
 100 |     t = rearrange(t, "... (g d) -> ... g d", g=groups)
 101 |     t = F.normalize(t, p=2, dim=-1)
 102 |     return rearrange(t, "... g d -> ... (g d)")
 103 | 
 104 | 
 105 | def pad_at_dim(t, pad, dim=-1, value=0.0):
 106 |     dims_from_right = (-dim - 1) if dim < 0 else (t.ndim - dim - 1)
 107 |     zeros = (0, 0) * dims_from_right
 108 |     return F.pad(t, (*zeros, *pad), value=value)
 109 | 
 110 | 
 111 | def or_reduce(masks):
 112 |     head, *body = masks
 113 |     for rest in body:
 114 |         head = head | rest
 115 |     return head
 116 | 
 117 | 
 118 | # auxiliary loss helpers
 119 | 
 120 | 
 121 | def calc_z_loss(pre_softmax_attns: List[Tensor], mask=None, weight=1.0):
 122 |     # the same loss applied to the mixture of experts router logits in https://arxiv.org/abs/2202.08906
 123 |     # in the paper, in a tiny footnote, they mention using it on attention logits with stabilizing effects
 124 |     # also used in PaLM as one of the measures
 125 | 
 126 |     lse = 0.0
 127 | 
 128 |     for attn in pre_softmax_attns:
 129 |         lse = lse + attn.logsumexp(dim=-1)
 130 | 
 131 |     loss = torch.square(lse)
 132 |     loss = reduce(loss, "b h n -> b n", "sum")
 133 | 
 134 |     if not exists(mask):
 135 |         return loss.mean() * weight
 136 | 
 137 |     loss = loss[mask].sum() / mask.sum().clamp(min=1e-5)
 138 |     return loss * weight
 139 | 
 140 | 
 141 | # init helpers
 142 | 
 143 | 
 144 | def init_zero_(layer):
 145 |     nn.init.constant_(layer.weight, 0.0)
 146 |     if exists(layer.bias):
 147 |         nn.init.constant_(layer.bias, 0.0)
 148 | 
 149 | 
 150 | # keyword argument helpers
 151 | 
 152 | 
 153 | def pick_and_pop(keys, d):
 154 |     values = list(map(lambda key: d.pop(key), keys))
 155 |     return dict(zip(keys, values))
 156 | 
 157 | 
 158 | def group_dict_by_key(cond, d):
 159 |     return_val = [dict(), dict()]
 160 |     for key in d.keys():
 161 |         match = bool(cond(key))
 162 |         ind = int(not match)
 163 |         return_val[ind][key] = d[key]
 164 |     return (*return_val,)
 165 | 
 166 | 
 167 | def string_begins_with(prefix, str):
 168 |     return str.startswith(prefix)
 169 | 
 170 | 
 171 | def group_by_key_prefix(prefix, d):
 172 |     return group_dict_by_key(partial(string_begins_with, prefix), d)
 173 | 
 174 | 
 175 | def groupby_prefix_and_trim(prefix, d):
 176 |     kwargs_with_prefix, kwargs = group_dict_by_key(
 177 |         partial(string_begins_with, prefix), d
 178 |     )
 179 |     kwargs_without_prefix = dict(
 180 |         map(lambda x: (x[0][len(prefix) :], x[1]), tuple(kwargs_with_prefix.items()))
 181 |     )
 182 |     return kwargs_without_prefix, kwargs
 183 | 
 184 | 
 185 | # structured dropout, more effective than traditional attention dropouts
 186 | 
 187 | 
 188 | def dropout_seq(seq, mask, dropout):
 189 |     b, n, *_, device = *seq.shape, seq.device
 190 |     logits = torch.randn(b, n, device=device)
 191 | 
 192 |     if exists(mask):
 193 |         mask_value = max_neg_value(logits)
 194 |         logits = logits.masked_fill(~mask, mask_value)
 195 | 
 196 |     keep_prob = 1.0 - dropout
 197 |     num_keep = max(1, int(keep_prob * n))
 198 |     keep_indices = logits.topk(num_keep, dim=1).indices
 199 | 
 200 |     batch_indices = torch.arange(b, device=device)
 201 |     batch_indices = rearrange(batch_indices, "b -> b 1")
 202 | 
 203 |     seq = seq[batch_indices, keep_indices]
 204 | 
 205 |     if exists(mask):
 206 |         seq_counts = mask.sum(dim=-1)
 207 |         seq_keep_counts = torch.ceil(seq_counts * keep_prob).int()
 208 |         keep_mask = torch.arange(num_keep, device=device) < rearrange(
 209 |             seq_keep_counts, "b -> b 1"
 210 |         )
 211 | 
 212 |         mask = mask[batch_indices, keep_indices] & keep_mask
 213 | 
 214 |     return seq, mask
 215 | 
 216 | 
 217 | # activations
 218 | 
 219 | 
 220 | class ReluSquared(nn.Module):
 221 |     def forward(self, x):
 222 |         return F.relu(x) ** 2
 223 | 
 224 | 
 225 | # embedding
 226 | 
 227 | 
 228 | class TokenEmbedding(nn.Module):
 229 |     def __init__(self, dim, num_tokens, l2norm_embed=False):
 230 |         super().__init__()
 231 |         self.l2norm_embed = l2norm_embed
 232 |         self.emb = nn.Embedding(num_tokens, dim)
 233 | 
 234 |     def forward(self, x):
 235 |         token_emb = self.emb(x)
 236 |         return l2norm(token_emb) if self.l2norm_embed else token_emb
 237 | 
 238 | 
 239 | # positional embeddings
 240 | 
 241 | 
 242 | class AbsolutePositionalEmbedding(nn.Module):
 243 |     def __init__(self, dim, max_seq_len, l2norm_embed=False):
 244 |         super().__init__()
 245 |         self.scale = dim**-0.5 if not l2norm_embed else 1.0
 246 |         self.max_seq_len = max_seq_len
 247 |         self.l2norm_embed = l2norm_embed
 248 |         self.emb = nn.Embedding(max_seq_len, dim)
 249 | 
 250 |     def forward(self, x, pos=None, seq_start_pos=None):
 251 |         seq_len, device = x.shape[1], x.device
 252 |         assert (
 253 |             seq_len <= self.max_seq_len
 254 |         ), f"you are passing in a sequence length of {seq_len} but your absolute positional embedding has a max sequence length of {self.max_seq_len}"
 255 | 
 256 |         if not exists(pos):
 257 |             pos = torch.arange(seq_len, device=device)
 258 | 
 259 |         if exists(seq_start_pos):
 260 |             pos = (pos - seq_start_pos[..., None]).clamp(min=0)
 261 | 
 262 |         pos_emb = self.emb(pos)
 263 |         pos_emb = pos_emb * self.scale
 264 |         return l2norm(pos_emb) if self.l2norm_embed else pos_emb
 265 | 
 266 | 
 267 | class ScaledSinusoidalEmbedding(nn.Module):
 268 |     def __init__(self, dim, theta=10000):
 269 |         super().__init__()
 270 |         assert divisible_by(dim, 2)
 271 |         self.scale = nn.Parameter(torch.ones(1) * dim**-0.5)
 272 | 
 273 |         half_dim = dim // 2
 274 |         freq_seq = torch.arange(half_dim).float() / half_dim
 275 |         inv_freq = theta**-freq_seq
 276 |         self.register_buffer("inv_freq", inv_freq, persistent=False)
 277 | 
 278 |     def forward(self, x, pos=None, seq_start_pos=None):
 279 |         seq_len, device = x.shape[1], x.device
 280 | 
 281 |         if not exists(pos):
 282 |             pos = torch.arange(seq_len, device=device)
 283 | 
 284 |         if exists(seq_start_pos):
 285 |             pos = pos - seq_start_pos[..., None]
 286 | 
 287 |         emb = einsum("i, j -> i j", pos, self.inv_freq)
 288 |         emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
 289 |         return emb * self.scale
 290 | 
 291 | 
 292 | class RelativePositionBias(nn.Module):
 293 |     def __init__(self, scale, causal=False, num_buckets=32, max_distance=128, heads=8):
 294 |         super().__init__()
 295 |         self.scale = scale
 296 |         self.causal = causal
 297 |         self.num_buckets = num_buckets
 298 |         self.max_distance = max_distance
 299 |         self.relative_attention_bias = nn.Embedding(num_buckets, heads)
 300 | 
 301 |     @staticmethod
 302 |     def _relative_position_bucket(
 303 |         relative_position, causal=True, num_buckets=32, max_distance=128
 304 |     ):
 305 |         ret = 0
 306 |         n = -relative_position
 307 |         if not causal:
 308 |             num_buckets //= 2
 309 |             ret += (n < 0).long() * num_buckets
 310 |             n = torch.abs(n)
 311 |         else:
 312 |             n = torch.max(n, torch.zeros_like(n))
 313 | 
 314 |         max_exact = num_buckets // 2
 315 |         is_small = n < max_exact
 316 | 
 317 |         val_if_large = (
 318 |             max_exact
 319 |             + (
 320 |                 torch.log(n.float() / max_exact)
 321 |                 / math.log(max_distance / max_exact)
 322 |                 * (num_buckets - max_exact)
 323 |             ).long()
 324 |         )
 325 |         val_if_large = torch.min(
 326 |             val_if_large, torch.full_like(val_if_large, num_buckets - 1)
 327 |         )
 328 | 
 329 |         ret += torch.where(is_small, n, val_if_large)
 330 |         return ret
 331 | 
 332 |     @property
 333 |     def device(self):
 334 |         return next(self.parameters()).device
 335 | 
 336 |     def forward(self, i, j):
 337 |         device = self.device
 338 |         q_pos = torch.arange(j - i, j, dtype=torch.long, device=device)
 339 |         k_pos = torch.arange(j, dtype=torch.long, device=device)
 340 |         rel_pos = k_pos[None, :] - q_pos[:, None]
 341 |         rp_bucket = self._relative_position_bucket(
 342 |             rel_pos,
 343 |             causal=self.causal,
 344 |             num_buckets=self.num_buckets,
 345 |             max_distance=self.max_distance,
 346 |         )
 347 |         values = self.relative_attention_bias(rp_bucket)
 348 |         bias = rearrange(values, "i j h -> h i j")
 349 |         return bias * self.scale
 350 | 
 351 | 
 352 | class DynamicPositionBias(nn.Module):
 353 |     def __init__(self, dim, *, heads, depth, log_distance=False, norm=False):
 354 |         super().__init__()
 355 |         assert (
 356 |             depth >= 1
 357 |         ), "depth for dynamic position bias MLP must be greater or equal to 1"
 358 |         self.log_distance = log_distance
 359 | 
 360 |         self.mlp = nn.ModuleList([])
 361 | 
 362 |         self.mlp.append(
 363 |             Sequential(
 364 |                 nn.Linear(1, dim), nn.LayerNorm(dim) if norm else None, nn.SiLU()
 365 |             )
 366 |         )
 367 | 
 368 |         for _ in range(depth - 1):
 369 |             self.mlp.append(
 370 |                 Sequential(
 371 |                     nn.Linear(dim, dim), nn.LayerNorm(dim) if norm else None, nn.SiLU()
 372 |                 )
 373 |             )
 374 | 
 375 |         self.mlp.append(nn.Linear(dim, heads))
 376 | 
 377 |     @property
 378 |     def device(self):
 379 |         return next(self.parameters()).device
 380 | 
 381 |     def forward(self, i, j):
 382 |         assert i == j
 383 |         n, device = j, self.device
 384 | 
 385 |         # get the (n x n) matrix of distances
 386 |         seq_arange = torch.arange(n, device=device)
 387 |         context_arange = torch.arange(n, device=device)
 388 |         indices = rearrange(seq_arange, "i -> i 1") - rearrange(
 389 |             context_arange, "j -> 1 j"
 390 |         )
 391 |         indices += n - 1
 392 | 
 393 |         # input to continuous positions MLP
 394 |         pos = torch.arange(-n + 1, n, device=device).float()
 395 |         pos = rearrange(pos, "... -> ... 1")
 396 | 
 397 |         if self.log_distance:
 398 |             pos = torch.sign(pos) * torch.log(
 399 |                 pos.abs() + 1
 400 |             )  # log of distance is sign(rel_pos) * log(abs(rel_pos) + 1)
 401 | 
 402 |         for layer in self.mlp:
 403 |             pos = layer(pos)
 404 | 
 405 |         # get position biases
 406 |         bias = pos[indices]
 407 |         bias = rearrange(bias, "i j h -> h i j")
 408 |         return bias
 409 | 
 410 | 
 411 | class AlibiPositionalBias(nn.Module):
 412 |     def __init__(self, heads, total_heads, **kwargs):
 413 |         super().__init__()
 414 |         self.heads = heads
 415 |         self.total_heads = total_heads
 416 | 
 417 |         slopes = Tensor(self._get_slopes(heads))
 418 |         slopes = rearrange(slopes, "h -> h 1 1")
 419 |         self.register_buffer("slopes", slopes, persistent=False)
 420 |         self.register_buffer("bias", None, persistent=False)
 421 | 
 422 |     def get_bias(self, i, j, device):
 423 |         i_arange = torch.arange(j - i, j, device=device)
 424 |         j_arange = torch.arange(j, device=device)
 425 |         bias = -torch.abs(
 426 |             rearrange(j_arange, "j -> 1 1 j") - rearrange(i_arange, "i -> 1 i 1")
 427 |         )
 428 |         return bias
 429 | 
 430 |     @staticmethod
 431 |     def _get_slopes(heads):
 432 |         def get_slopes_power_of_2(n):
 433 |             start = 2 ** (-(2 ** -(math.log2(n) - 3)))
 434 |             ratio = start
 435 |             return [start * ratio**i for i in range(n)]
 436 | 
 437 |         if math.log2(heads).is_integer():
 438 |             return get_slopes_power_of_2(heads)
 439 | 
 440 |         closest_power_of_2 = 2 ** math.floor(math.log2(heads))
 441 |         return (
 442 |             get_slopes_power_of_2(closest_power_of_2)
 443 |             + get_slopes_power_of_2(2 * closest_power_of_2)[0::2][
 444 |                 : heads - closest_power_of_2
 445 |             ]
 446 |         )
 447 | 
 448 |     @property
 449 |     def device(self):
 450 |         return next(self.buffers()).device
 451 | 
 452 |     def forward(self, i, j):
 453 |         h, device = self.total_heads, self.device
 454 | 
 455 |         if exists(self.bias) and self.bias.shape[-1] >= j and self.bias.shape[-2] >= i:
 456 |             return self.bias[..., -i:, -j:]
 457 | 
 458 |         bias = self.get_bias(i, j, device)
 459 |         bias = bias * self.slopes
 460 | 
 461 |         num_heads_unalibied = h - bias.shape[0]
 462 |         bias = pad_at_dim(bias, (0, num_heads_unalibied), dim=0)
 463 |         self.register_buffer("bias", bias, persistent=False)
 464 | 
 465 |         return self.bias
 466 | 
 467 | 
 468 | class RotaryEmbedding(nn.Module):
 469 |     def __init__(
 470 |         self,
 471 |         dim,
 472 |         use_xpos=False,
 473 |         scale_base=512,
 474 |         interpolation_factor=1.0,
 475 |         base=10000,
 476 |         base_rescale_factor=1.0,
 477 |     ):
 478 |         super().__init__()
 479 |         # proposed by reddit user bloc97, to rescale rotary embeddings to longer sequence length without fine-tuning
 480 |         # has some connection to NTK literature
 481 |         # https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/
 482 |         base *= base_rescale_factor ** (dim / (dim - 2))
 483 | 
 484 |         inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
 485 |         self.register_buffer("inv_freq", inv_freq)
 486 | 
 487 |         assert interpolation_factor >= 1.0
 488 |         self.interpolation_factor = interpolation_factor
 489 | 
 490 |         if not use_xpos:
 491 |             self.register_buffer("scale", None)
 492 |             return
 493 | 
 494 |         scale = (torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim)
 495 | 
 496 |         self.scale_base = scale_base
 497 |         self.register_buffer("scale", scale)
 498 | 
 499 |     def forward(self, seq_len):
 500 |         device = self.inv_freq.device
 501 |         t = torch.arange(seq_len, device=device).type_as(self.inv_freq)
 502 | 
 503 |         t = t / self.interpolation_factor
 504 | 
 505 |         freqs = torch.einsum("i , j -> i j", t, self.inv_freq)
 506 |         freqs = torch.cat((freqs, freqs), dim=-1)
 507 | 
 508 |         if not exists(self.scale):
 509 |             return freqs, 1.0
 510 | 
 511 |         power = (
 512 |             torch.arange(seq_len, device=device) - (seq_len // 2)
 513 |         ) / self.scale_base
 514 |         scale = self.scale ** rearrange(power, "n -> n 1")
 515 |         scale = torch.cat((scale, scale), dim=-1)
 516 | 
 517 |         return freqs, scale
 518 | 
 519 | 
 520 | def rotate_half(x):
 521 |     x = rearrange(x, "... (j d) -> ... j d", j=2)
 522 |     x1, x2 = x.unbind(dim=-2)
 523 |     return torch.cat((-x2, x1), dim=-1)
 524 | 
 525 | 
 526 | def apply_rotary_pos_emb(t, freqs, scale=1):
 527 |     rot_dim, seq_len = freqs.shape[-1], t.shape[-2]
 528 |     freqs = freqs[-seq_len:, :]
 529 | 
 530 |     if t.ndim == 4 and freqs.ndim == 3:
 531 |         freqs = rearrange(freqs, "b n d -> b 1 n d")
 532 | 
 533 |     # partial rotary embeddings, Wang et al. GPT-J
 534 |     t, t_unrotated = t[..., :rot_dim], t[..., rot_dim:]
 535 |     t = (t * freqs.cos() * scale) + (rotate_half(t) * freqs.sin() * scale)
 536 |     return torch.cat((t, t_unrotated), dim=-1)
 537 | 
 538 | 
 539 | # norms
 540 | 
 541 | 
 542 | class Scale(nn.Module):
 543 |     def __init__(self, value, fn):
 544 |         super().__init__()
 545 |         self.value = value
 546 |         self.fn = fn
 547 | 
 548 |     def forward(self, x, **kwargs):
 549 |         out = self.fn(x, **kwargs)
 550 | 
 551 |         def scale_fn(t):
 552 |             return t * self.value
 553 | 
 554 |         if not isinstance(out, tuple):
 555 |             return scale_fn(out)
 556 | 
 557 |         return (scale_fn(out[0]), *out[1:])
 558 | 
 559 | 
 560 | class ScaleNorm(nn.Module):
 561 |     def __init__(self, dim, eps=1e-5):
 562 |         super().__init__()
 563 |         self.eps = eps
 564 |         self.g = nn.Parameter(torch.ones(1) * (dim**-0.5))
 565 | 
 566 |     def forward(self, x):
 567 |         norm = torch.norm(x, dim=-1, keepdim=True)
 568 |         return x / norm.clamp(min=self.eps) * self.g
 569 | 
 570 | 
 571 | class RMSNorm(nn.Module):
 572 |     def __init__(self, dim):
 573 |         super().__init__()
 574 |         self.scale = dim**0.5
 575 |         self.g = nn.Parameter(torch.ones(dim))
 576 | 
 577 |     def forward(self, x):
 578 |         return F.normalize(x, dim=-1) * self.scale * self.g
 579 | 
 580 | 
 581 | class SimpleRMSNorm(nn.Module):
 582 |     def __init__(self, dim):
 583 |         super().__init__()
 584 |         self.scale = dim**0.5
 585 | 
 586 |     def forward(self, x):
 587 |         return F.normalize(x, dim=-1) * self.scale
 588 | 
 589 | 
 590 | # residual and residual gates
 591 | 
 592 | 
 593 | class Residual(nn.Module):
 594 |     def __init__(self, dim, scale_residual=False, scale_residual_constant=1.0):
 595 |         super().__init__()
 596 |         self.residual_scale = nn.Parameter(torch.ones(dim)) if scale_residual else None
 597 |         self.scale_residual_constant = scale_residual_constant
 598 | 
 599 |     def forward(self, x, residual):
 600 |         if exists(self.residual_scale):
 601 |             residual = residual * self.residual_scale
 602 | 
 603 |         if self.scale_residual_constant != 1:
 604 |             residual = residual * self.scale_residual_constant
 605 | 
 606 |         return x + residual
 607 | 
 608 | 
 609 | class GRUGating(nn.Module):
 610 |     def __init__(self, dim, scale_residual=False, **kwargs):
 611 |         super().__init__()
 612 |         self.gru = nn.GRUCell(dim, dim)
 613 |         self.residual_scale = nn.Parameter(torch.ones(dim)) if scale_residual else None
 614 | 
 615 |     def forward(self, x, residual):
 616 |         if exists(self.residual_scale):
 617 |             residual = residual * self.residual_scale
 618 | 
 619 |         gated_output = self.gru(
 620 |             rearrange(x, "b n d -> (b n) d"), rearrange(residual, "b n d -> (b n) d")
 621 |         )
 622 | 
 623 |         return gated_output.reshape_as(x)
 624 | 
 625 | 
 626 | # token shifting
 627 | 
 628 | 
 629 | def shift(t, amount, mask=None):
 630 |     if amount == 0:
 631 |         return t
 632 |     else:
 633 |         amount = min(amount, t.shape[1])
 634 | 
 635 |     if exists(mask):
 636 |         t = t.masked_fill(~mask[..., None], 0.0)
 637 | 
 638 |     return pad_at_dim(t, (amount, -amount), dim=-2, value=0.0)
 639 | 
 640 | 
 641 | class ShiftTokens(nn.Module):
 642 |     def __init__(self, shifts, fn):
 643 |         super().__init__()
 644 |         self.fn = fn
 645 |         self.shifts = tuple(shifts)
 646 | 
 647 |     def forward(self, x, **kwargs):
 648 |         mask = kwargs.get("mask", None)
 649 |         shifts = self.shifts
 650 |         segments = len(shifts)
 651 |         
 652 |         feats_per_shift = x.shape[-1] // segments
 653 |         splitted = x.split(feats_per_shift, dim=-1)
 654 |         segments_to_shift, rest = splitted[:segments], splitted[segments:]
 655 |         segments_to_shift = list(
 656 |             map(lambda args: shift(*args, mask=mask), zip(segments_to_shift, shifts))
 657 |         )
 658 |         x = torch.cat((*segments_to_shift, *rest), dim=-1)
 659 |         return self.fn(x, **kwargs)
 660 | 
 661 | 
 662 | # feedforward
 663 | 
 664 | 
 665 | class GLU(nn.Module):
 666 |     def __init__(self, dim_in, dim_out, activation: Callable, mult_bias=False):
 667 |         super().__init__()
 668 |         self.act = activation
 669 |         self.proj = nn.Linear(dim_in, dim_out * 2)
 670 |         self.mult_bias = nn.Parameter(torch.ones(dim_out)) if mult_bias else 1.0
 671 | 
 672 |     def forward(self, x):
 673 |         x, gate = self.proj(x).chunk(2, dim=-1)
 674 |         return x * self.act(gate) * self.mult_bias
 675 | 
 676 | 
 677 | class FeedForward(nn.Module):
 678 |     def __init__(
 679 |         self,
 680 |         dim,
 681 |         dim_out=None,
 682 |         mult=4,
 683 |         glu=False,
 684 |         glu_mult_bias=False,
 685 |         swish=False,
 686 |         relu_squared=False,
 687 |         post_act_ln=False,
 688 |         dropout=0.0,
 689 |         no_bias=False,
 690 |         zero_init_output=False,
 691 |     ):
 692 |         super().__init__()
 693 |         inner_dim = int(dim * mult)
 694 |         dim_out = default(dim_out, dim)
 695 | 
 696 |         if relu_squared:
 697 |             activation = ReluSquared()
 698 |         elif swish:
 699 |             activation = nn.SiLU()
 700 |         else:
 701 |             activation = nn.GELU()
 702 | 
 703 |         if glu:
 704 |             project_in = GLU(dim, inner_dim, activation, mult_bias=glu_mult_bias)
 705 |         else:
 706 |             project_in = nn.Sequential(
 707 |                 nn.Linear(dim, inner_dim, bias=not no_bias), activation
 708 |             )
 709 | 
 710 |         self.ff = Sequential(
 711 |             project_in,
 712 |             nn.LayerNorm(inner_dim) if post_act_ln else None,
 713 |             nn.Dropout(dropout),
 714 |             nn.Linear(inner_dim, dim_out, bias=not no_bias),
 715 |         )
 716 | 
 717 |         # init last linear layer to 0
 718 |         if zero_init_output:
 719 |             init_zero_(self.ff[-1])
 720 | 
 721 |     def forward(self, x):
 722 |         return self.ff(x)
 723 | 
 724 | 
 725 | # attention. it is all we need
 726 | 
 727 | 
 728 | class Attention(nn.Module):
 729 |     def __init__(
 730 |         self,
 731 |         dim,
 732 |         dim_head=DEFAULT_DIM_HEAD,
 733 |         heads=8,
 734 |         causal=False,
 735 |         flash=False,
 736 |         talking_heads=False,
 737 |         head_scale=False,
 738 |         sparse_topk=None,
 739 |         num_mem_kv=0,
 740 |         dropout=0.0,
 741 |         on_attn=False,
 742 |         gate_value_heads=False,
 743 |         gate_values=False,
 744 |         zero_init_output=False,
 745 |         max_attend_past=None,
 746 |         qk_norm=False,
 747 |         qk_norm_groups=1,
 748 |         qk_norm_scale=10,
 749 |         qk_norm_dim_scale=False,
 750 |         one_kv_head=False,
 751 |         kv_heads=None,
 752 |         shared_kv=False,
 753 |         value_dim_head=None,
 754 |         tensor_product=False,  # https://arxiv.org/abs/2208.06061
 755 |         add_zero_kv=False,  # same as add_zero_attn in pytorch
 756 |         rotary_embed_values=False,
 757 |         onnxable=False,
 758 |     ):
 759 |         super().__init__()
 760 |         self.scale = dim_head**-0.5
 761 | 
 762 |         self.heads = heads
 763 |         self.causal = causal
 764 |         self.max_attend_past = max_attend_past
 765 | 
 766 |         assert not (
 767 |             exists(kv_heads) and one_kv_head
 768 |         ), "either attn_one_kv_head is set to True (in which case kv_heads is set to 1), or attn_kv_heads is set, but not both"
 769 | 
 770 |         value_dim_head = default(value_dim_head, dim_head)
 771 |         kv_heads = default(kv_heads, heads)
 772 | 
 773 |         kv_heads = 1 if one_kv_head else kv_heads
 774 |         assert divisible_by(heads, kv_heads)
 775 | 
 776 |         self.kv_heads = kv_heads
 777 | 
 778 |         q_dim = dim_head * heads
 779 |         k_dim = dim_head * kv_heads
 780 |         v_dim = value_dim_head * kv_heads
 781 |         out_dim = value_dim_head * heads
 782 | 
 783 |         self.to_q = nn.Linear(dim, q_dim, bias=False)
 784 |         self.to_k = nn.Linear(dim, k_dim, bias=False)
 785 | 
 786 |         # shared key / values, for further memory savings during inference
 787 |         assert not (
 788 |             shared_kv and value_dim_head != dim_head
 789 |         ), "key and value head dimensions must be equal for shared key / values"
 790 |         self.to_v = nn.Linear(dim, v_dim, bias=False) if not shared_kv else None
 791 | 
 792 |         # relations projection from tp-attention
 793 |         self.to_r = nn.Linear(dim, v_dim, bias=False) if tensor_product else None
 794 | 
 795 |         # add GLU gating for aggregated values, from alphafold2
 796 |         self.to_v_gate = None
 797 |         if gate_values:
 798 |             self.to_v_gate = nn.Linear(dim, out_dim)
 799 |             nn.init.constant_(self.to_v_gate.weight, 0)
 800 |             nn.init.constant_(self.to_v_gate.bias, 10)
 801 | 
 802 |         # add per head gating of the output values, from 'Attend to nothing' paper
 803 |         self.to_v_head_gate = None
 804 |         if gate_value_heads:
 805 |             self.to_v_head_gate = nn.Linear(dim, heads)
 806 |             nn.init.constant_(self.to_v_head_gate.weight, 0)
 807 |             nn.init.constant_(self.to_v_head_gate.bias, 10)
 808 | 
 809 |         # cosine sim attention
 810 |         self.qk_norm = qk_norm
 811 |         self.qk_norm_groups = qk_norm_groups
 812 |         self.qk_norm_scale = qk_norm_scale
 813 | 
 814 |         # whether to use the rmsnorm (equivalent to cosine sim attention when scale is equal to 1) - https://arxiv.org/abs/2302.05442
 815 |         self.qk_norm_dim_scale = qk_norm_dim_scale
 816 | 
 817 |         self.qk_norm_q_scale = self.qk_norm_k_scale = 1
 818 |         if qk_norm and qk_norm_dim_scale:
 819 |             self.qk_norm_q_scale = nn.Parameter(torch.ones(heads, 1, dim_head))
 820 |             self.qk_norm_k_scale = nn.Parameter(torch.ones(heads, 1, dim_head))
 821 | 
 822 |         assert (not qk_norm) or divisible_by(
 823 |             dim_head, qk_norm_groups
 824 |         ), "dimension per attention head must be divisible by the qk norm groups"
 825 |         assert not (
 826 |             qk_norm and (dim_head // qk_norm_groups) <= 2
 827 |         ), "the group dimension may be too small (2 was too small in my tests, but 4 still works, surprisingly)"
 828 | 
 829 |         # attend class - includes core attention algorithm + talking heads
 830 | 
 831 |         self.attend = Attend(
 832 |             heads=heads,
 833 |             causal=causal,
 834 |             talking_heads=talking_heads,
 835 |             dropout=dropout,
 836 |             sparse_topk=sparse_topk,
 837 |             qk_norm=qk_norm,
 838 |             scale=qk_norm_scale if qk_norm else self.scale,
 839 |             add_zero_kv=add_zero_kv,
 840 |             flash=flash,
 841 |             onnxable=onnxable,
 842 |         )
 843 | 
 844 |         # head scaling
 845 |         self.head_scale = head_scale
 846 |         if head_scale:
 847 |             self.head_scale_params = nn.Parameter(torch.ones(1, heads, 1, 1))
 848 | 
 849 |         # explicit topk sparse attention
 850 |         self.sparse_topk = sparse_topk
 851 | 
 852 |         # add memory key / values
 853 |         self.num_mem_kv = num_mem_kv
 854 |         if num_mem_kv > 0:
 855 |             self.mem_k = nn.Parameter(torch.randn(heads, num_mem_kv, dim_head))
 856 |             self.mem_v = nn.Parameter(torch.randn(heads, num_mem_kv, dim_head))
 857 | 
 858 |         # attention on attention
 859 |         self.attn_on_attn = on_attn
 860 |         self.to_out = (
 861 |             nn.Sequential(nn.Linear(out_dim, dim * 2, bias=False), nn.GLU())
 862 |             if on_attn
 863 |             else nn.Linear(out_dim, dim, bias=False)
 864 |         )
 865 | 
 866 |         # whether to rotate positions into values, for absolute positions in addition to relative
 867 |         self.rotary_embed_values = rotary_embed_values
 868 | 
 869 |         # init output projection 0
 870 |         if zero_init_output:
 871 |             init_zero_(self.to_out)
 872 | 
 873 |     def forward(
 874 |         self,
 875 |         x,
 876 |         context=None,
 877 |         mask=None,
 878 |         context_mask=None,
 879 |         attn_mask=None,
 880 |         rel_pos=None,
 881 |         rotary_pos_emb=None,
 882 |         prev_attn=None,
 883 |         mem=None,
 884 |         return_intermediates=False,
 885 |         cache: Optional[Intermediates] = None,
 886 |     ):
 887 |         b, n, _, h, kv_h, head_scale, device, has_context = (
 888 |             *x.shape,
 889 |             self.heads,
 890 |             self.kv_heads,
 891 |             self.head_scale,
 892 |             x.device,
 893 |             exists(context),
 894 |         )
 895 |         kv_input = default(context, x)
 896 | 
 897 |         q_input = x
 898 |         k_input = kv_input
 899 |         v_input = kv_input
 900 |         r_input = x
 901 | 
 902 |         if exists(mem):
 903 |             k_input, mem_packed_shape = pack([mem, k_input], "b * d")
 904 |             v_input, _ = pack([mem, v_input], "b * d")
 905 | 
 906 |         q = self.to_q(q_input)
 907 |         k = self.to_k(k_input)
 908 |         v = self.to_v(v_input) if exists(self.to_v) else k
 909 |         r = self.to_r(r_input) if exists(self.to_r) else None
 910 | 
 911 |         q = rearrange(q, "b n (h d) -> b h n d", h=h)
 912 | 
 913 |         k, v, r = map(
 914 |             lambda t: maybe(rearrange)(t, "b n (h d) -> b h n d", h=kv_h), (k, v, r)
 915 |         )
 916 | 
 917 |         if exists(cache) and not has_context:
 918 |             ck, cv = cache.cached_kv
 919 | 
 920 |             if exists(mem):
 921 |                 mk, k = unpack(k, mem_packed_shape, "b h * d")
 922 |                 mv, v = unpack(v, mem_packed_shape, "b h * d")
 923 | 
 924 |             k = torch.cat((ck, k), dim=-2)
 925 |             v = torch.cat((cv, v), dim=-2)
 926 | 
 927 |             if exists(mem):
 928 |                 k = torch.cat((mk, k), dim=-2)
 929 |                 v = torch.cat((mv, v), dim=-2)
 930 | 
 931 |         if return_intermediates:
 932 |             mem_len = mem.shape[-2] if exists(mem) else 0
 933 |             cached_kv = (k[..., mem_len:, :], v[..., mem_len:, :])
 934 | 
 935 |         if self.qk_norm:
 936 |             qk_l2norm = partial(l2norm, groups=self.qk_norm_groups)
 937 |             q, k = map(qk_l2norm, (q, k))
 938 | 
 939 |             q = q * self.qk_norm_q_scale
 940 |             k = k * self.qk_norm_k_scale
 941 | 
 942 |         if exists(rotary_pos_emb) and not has_context:
 943 |             freqs, xpos_scale = rotary_pos_emb
 944 |             q_xpos_scale, k_xpos_scale = (
 945 |                 (xpos_scale, xpos_scale**-1.0) if exists(xpos_scale) else (1.0, 1.0)
 946 |             )
 947 | 
 948 |             q = apply_rotary_pos_emb(q, freqs, q_xpos_scale)
 949 |             k = apply_rotary_pos_emb(k, freqs, k_xpos_scale)
 950 | 
 951 |             if self.rotary_embed_values:
 952 |                 v = apply_rotary_pos_emb(v, freqs, k_xpos_scale)
 953 | 
 954 |         input_mask = context_mask
 955 | 
 956 |         if not exists(input_mask) and not has_context:
 957 |             input_mask = mask
 958 | 
 959 |         if self.num_mem_kv > 0:
 960 |             mem_k, mem_v = map(
 961 |                 lambda t: repeat(t, "h n d -> b h n d", b=b), (self.mem_k, self.mem_v)
 962 |             )
 963 | 
 964 |             if self.qk_norm:
 965 |                 mem_k = l2norm(mem_k)
 966 |                 mem_k = mem_k * self.qk_norm_k_scale
 967 | 
 968 |             k = torch.cat((mem_k, k), dim=-2)
 969 |             v = torch.cat((mem_v, v), dim=-2)
 970 | 
 971 |             if exists(input_mask):
 972 |                 input_mask = pad_at_dim(
 973 |                     input_mask, (self.num_mem_kv, 0), dim=-1, value=True
 974 |                 )
 975 | 
 976 |         i, j = map(lambda t: t.shape[-2], (q, k))
 977 | 
 978 |         # determine masking
 979 | 
 980 |         max_neg_value(q)
 981 |         masks = []
 982 |         final_attn_mask = None
 983 | 
 984 |         if exists(input_mask):
 985 |             input_mask = rearrange(input_mask, "b j -> b 1 1 j")
 986 |             masks.append(~input_mask)
 987 | 
 988 |         if exists(attn_mask):
 989 |             assert (
 990 |                 2 <= attn_mask.ndim <= 4
 991 |             ), "attention mask must have greater than 2 dimensions but less than or equal to 4"
 992 |             if attn_mask.ndim == 2:
 993 |                 attn_mask = rearrange(attn_mask, "i j -> 1 1 i j")
 994 |             elif attn_mask.ndim == 3:
 995 |                 attn_mask = rearrange(attn_mask, "h i j -> 1 h i j")
 996 |             masks.append(~attn_mask)
 997 | 
 998 |         if exists(self.max_attend_past):
 999 |             range_q = torch.arange(j - i, j, device=device)
1000 |             range_k = torch.arange(j, device=device)
1001 |             dist = rearrange(range_q, "i -> 1 1 i 1") - rearrange(
1002 |                 range_k, "j -> 1 1 1 j"
1003 |             )
1004 |             max_attend_past_mask = dist > self.max_attend_past
1005 |             masks.append(max_attend_past_mask)
1006 | 
1007 |         if len(masks) > 0:
1008 |             final_attn_mask = ~or_reduce(masks)
1009 | 
1010 |         # prepare relative positional bias, if needed
1011 | 
1012 |         attn_bias = None
1013 |         if exists(rel_pos):
1014 |             attn_bias = rel_pos(i, j)
1015 | 
1016 |         # attention is all we need
1017 | 
1018 |         out, intermediates = self.attend(
1019 |             q, k, v, mask=final_attn_mask, attn_bias=attn_bias, prev_attn=prev_attn
1020 |         )
1021 | 
1022 |         # https://arxiv.org/abs/2208.06061 proposes to add a residual for better gradients
1023 | 
1024 |         if exists(r):
1025 |             out = out * r + out
1026 | 
1027 |         # normformer scaling of heads
1028 | 
1029 |         if head_scale:
1030 |             out = out * self.head_scale_params
1031 | 
1032 |         # per head gating, from https://arxiv.org/abs/2306.12929
1033 | 
1034 |         if exists(self.to_v_head_gate):
1035 |             head_gate = self.to_v_head_gate(x)
1036 |             out = out * rearrange(head_gate, "b n h -> b h n 1").sigmoid()
1037 | 
1038 |         # merge heads
1039 | 
1040 |         out = rearrange(out, "b h n d -> b n (h d)")
1041 | 
1042 |         # alphafold2 styled gating of the values
1043 | 
1044 |         if exists(self.to_v_gate):
1045 |             gates = self.to_v_gate(x)
1046 |             out = out * gates.sigmoid()
1047 | 
1048 |         # combine the heads
1049 | 
1050 |         out = self.to_out(out)
1051 | 
1052 |         if exists(mask):
1053 |             mask = rearrange(mask, "b n -> b n 1")
1054 |             out = out.masked_fill(~mask, 0.0)
1055 | 
1056 |         if not return_intermediates:
1057 |             return out
1058 | 
1059 |         intermediates.cached_kv = cached_kv
1060 | 
1061 |         return out, intermediates
1062 | 
1063 | 
1064 | class AttentionLayers(nn.Module):
1065 |     def __init__(
1066 |         self,
1067 |         dim,
1068 |         depth,
1069 |         heads=8,
1070 |         causal=False,
1071 |         cross_attend=False,
1072 |         only_cross=False,
1073 |         use_scalenorm=False,
1074 |         use_rmsnorm=False,
1075 |         use_simple_rmsnorm=False,
1076 |         alibi_pos_bias=False,
1077 |         alibi_num_heads=None,
1078 |         rel_pos_bias=False,
1079 |         rel_pos_num_buckets=32,
1080 |         rel_pos_max_distance=128,
1081 |         dynamic_pos_bias=False,
1082 |         dynamic_pos_bias_log_distance=False,
1083 |         dynamic_pos_bias_mlp_depth=2,
1084 |         dynamic_pos_bias_norm=False,
1085 |         rotary_pos_emb=False,
1086 |         rotary_emb_dim=None,
1087 |         rotary_xpos=False,
1088 |         rotary_interpolation_factor=1.0,
1089 |         rotary_xpos_scale_base=512,
1090 |         rotary_base_rescale_factor=1.0,
1091 |         custom_layers=None,
1092 |         sandwich_coef=None,
1093 |         par_ratio=None,
1094 |         weight_tie_layers=False,  # Albert - https://arxiv.org/abs/1909.11942
1095 |         layers_execute_order=None,  # generalizes weight tying, can do arbitrary layer execution orders
1096 |         residual_attn=False,
1097 |         cross_residual_attn=False,
1098 |         macaron=False,
1099 |         pre_norm=True,
1100 |         pre_norm_has_final_norm=True,
1101 |         gate_residual=False,
1102 |         scale_residual=False,
1103 |         scale_residual_constant=1.0,
1104 |         shift_tokens=0,
1105 |         sandwich_norm=False,
1106 |         resi_dual=False,
1107 |         resi_dual_scale=1.0,
1108 |         zero_init_branch_output=False,
1109 |         layer_dropout=0.0,
1110 |         cross_attn_tokens_dropout=0.0,
1111 |         **kwargs,
1112 |     ):
1113 |         super().__init__()
1114 |         rotary_pos_emb = rotary_pos_emb or rotary_xpos
1115 | 
1116 |         ff_kwargs, kwargs = groupby_prefix_and_trim("ff_", kwargs)
1117 |         attn_kwargs, kwargs = groupby_prefix_and_trim("attn_", kwargs)
1118 | 
1119 |         dim_head = attn_kwargs.get("dim_head", DEFAULT_DIM_HEAD)
1120 | 
1121 |         self.dim = dim
1122 |         self.depth = depth
1123 |         self.causal = causal
1124 |         self.layers = nn.ModuleList([])
1125 | 
1126 |         self.has_pos_emb = rel_pos_bias or rotary_pos_emb
1127 | 
1128 |         rotary_emb_dim = max(default(rotary_emb_dim, dim_head // 2), 32)
1129 | 
1130 |         assert not (
1131 |             rotary_xpos and not causal
1132 |         ), "rotary xpos is not compatible with bidirectional attention"
1133 |         self.rotary_pos_emb = (
1134 |             RotaryEmbedding(
1135 |                 rotary_emb_dim,
1136 |                 use_xpos=rotary_xpos,
1137 |                 scale_base=rotary_xpos_scale_base,
1138 |                 interpolation_factor=rotary_interpolation_factor,
1139 |                 base_rescale_factor=rotary_base_rescale_factor,
1140 |             )
1141 |             if rotary_pos_emb
1142 |             else None
1143 |         )
1144 | 
1145 |         assert not (
1146 |             alibi_pos_bias and rel_pos_bias
1147 |         ), "you can only choose Alibi positional bias or T5 relative positional bias, not both"
1148 |         assert (
1149 |             rel_pos_num_buckets <= rel_pos_max_distance
1150 |         ), "number of relative position buckets must be less than the relative position max distance"
1151 | 
1152 |         # relative positional bias
1153 | 
1154 |         flash_attn = attn_kwargs.get("flash", False)
1155 |         assert (
1156 |             int(rel_pos_bias) + int(dynamic_pos_bias) + int(alibi_pos_bias)
1157 |         ) <= 1, "you can only choose up to one of t5, alibi, or dynamic positional bias"
1158 | 
1159 |         self.rel_pos = None
1160 |         if rel_pos_bias:
1161 |             assert (
1162 |                 not flash_attn
1163 |             ), "flash attention not compatible with t5 relative positional bias"
1164 |             self.rel_pos = RelativePositionBias(
1165 |                 scale=dim_head**0.5,
1166 |                 causal=causal,
1167 |                 heads=heads,
1168 |                 num_buckets=rel_pos_num_buckets,
1169 |                 max_distance=rel_pos_max_distance,
1170 |             )
1171 |         elif dynamic_pos_bias:
1172 |             assert (
1173 |                 not flash_attn
1174 |             ), "flash attention not compatible with dynamic positional bias"
1175 |             self.rel_pos = DynamicPositionBias(
1176 |                 dim=dim // 4,
1177 |                 heads=heads,
1178 |                 log_distance=dynamic_pos_bias_log_distance,
1179 |                 depth=dynamic_pos_bias_mlp_depth,
1180 |                 norm=dynamic_pos_bias_norm,
1181 |             )
1182 |         elif alibi_pos_bias:
1183 |             alibi_num_heads = default(alibi_num_heads, heads)
1184 |             assert (
1185 |                 alibi_num_heads <= heads
1186 |             ), "number of ALiBi heads must be less than the total number of heads"
1187 |             self.rel_pos = AlibiPositionalBias(heads=alibi_num_heads, total_heads=heads)
1188 | 
1189 |         assert (
1190 |             int(sandwich_norm) + int(resi_dual)
1191 |         ) <= 1, "either sandwich norm or resiDual is selected, but not both"
1192 |         assert not (
1193 |             not pre_norm and sandwich_norm
1194 |         ), "sandwich norm cannot be used when not using prenorm"
1195 | 
1196 |         if resi_dual:
1197 |             pre_norm = False
1198 | 
1199 |         self.pre_norm = pre_norm
1200 |         self.sandwich_norm = sandwich_norm
1201 | 
1202 |         self.resi_dual = resi_dual
1203 |         assert (
1204 |             0 < resi_dual_scale <= 1.0
1205 |         ), "resiDual prenorm residual must be scaled by a factor greater than 0 and less than or equal to 1."
1206 |         self.resi_dual_scale = resi_dual_scale
1207 | 
1208 |         self.residual_attn = residual_attn
1209 |         self.cross_residual_attn = cross_residual_attn
1210 |         assert not (
1211 |             flash_attn and (residual_attn or cross_residual_attn)
1212 |         ), "flash attention is not compatible with residual attention"
1213 | 
1214 |         self.cross_attend = cross_attend
1215 | 
1216 |         assert (
1217 |             int(use_scalenorm) + int(use_rmsnorm) + int(use_simple_rmsnorm)
1218 |         ) <= 1, "you can only use either scalenorm, rmsnorm, or simple rmsnorm"
1219 | 
1220 |         if use_scalenorm:
1221 |             norm_class = ScaleNorm
1222 |         elif use_rmsnorm:
1223 |             norm_class = RMSNorm
1224 |         elif use_simple_rmsnorm:
1225 |             norm_class = SimpleRMSNorm
1226 |         else:
1227 |             norm_class = nn.LayerNorm
1228 | 
1229 |         norm_fn = partial(norm_class, dim)
1230 | 
1231 |         if cross_attend and not only_cross:
1232 |             default_block = ("a", "c", "f")
1233 |         elif cross_attend and only_cross:
1234 |             default_block = ("c", "f")
1235 |         else:
1236 |             default_block = ("a", "f")
1237 | 
1238 |         if macaron:
1239 |             default_block = ("f",) + default_block
1240 | 
1241 |         # zero init
1242 | 
1243 |         if zero_init_branch_output:
1244 |             attn_kwargs = {**attn_kwargs, "zero_init_output": True}
1245 |             ff_kwargs = {**ff_kwargs, "zero_init_output": True}
1246 | 
1247 |         # setup weight tying, which is a special case of `layer_execute_order`
1248 | 
1249 |         assert not (
1250 |             weight_tie_layers
1251 |             and any([*map(exists, (custom_layers, par_ratio, sandwich_coef))])
1252 |         )
1253 | 
1254 |         if weight_tie_layers:
1255 |             assert not exists(layers_execute_order)
1256 |             layers_execute_order = tuple(range(len(default_block))) * depth
1257 |             depth = 1
1258 | 
1259 |         # calculate layer block order
1260 | 
1261 |         if exists(custom_layers):
1262 |             layer_types = custom_layers
1263 |         elif exists(par_ratio):
1264 |             par_depth = depth * len(default_block)
1265 |             assert 1 < par_ratio <= par_depth, "par ratio out of range"
1266 |             default_block = tuple(filter(not_equals("f"), default_block))
1267 |             par_attn = par_depth // par_ratio
1268 |             depth_cut = (
1269 |                 par_depth * 2 // 3
1270 |             )  # 2 / 3 attention layer cutoff suggested by PAR paper
1271 |             par_width = (depth_cut + depth_cut // par_attn) // par_attn
1272 |             assert (
1273 |                 len(default_block) <= par_width
1274 |             ), "default block is too large for par_ratio"
1275 |             par_block = default_block + ("f",) * (par_width - len(default_block))
1276 |             par_head = par_block * par_attn
1277 |             layer_types = par_head + ("f",) * (par_depth - len(par_head))
1278 |         elif exists(sandwich_coef):
1279 |             assert (
1280 |                 sandwich_coef > 0 and sandwich_coef <= depth
1281 |             ), "sandwich coefficient should be less than the depth"
1282 |             layer_types = (
1283 |                 ("a",) * sandwich_coef
1284 |                 + default_block * (depth - sandwich_coef)
1285 |                 + ("f",) * sandwich_coef
1286 |             )
1287 |         else:
1288 |             layer_types = default_block * depth
1289 | 
1290 |         self.layer_types = layer_types
1291 |         self.layers_execute_order = default(
1292 |             layers_execute_order, tuple(range(len(layer_types)))
1293 |         )
1294 | 
1295 |         assert all([i < len(self.layer_types) for i in self.layers_execute_order])
1296 | 
1297 |         self.num_attn_layers = len(list(filter(equals("a"), layer_types)))
1298 | 
1299 |         # stochastic depth
1300 | 
1301 |         self.layer_dropouts = cast_tuple(layer_dropout, len(layer_types))
1302 | 
1303 |         # structured dropout for cross attending
1304 | 
1305 |         self.cross_attn_tokens_dropout = cross_attn_tokens_dropout
1306 | 
1307 |         # calculate token shifting
1308 | 
1309 |         shift_tokens = cast_tuple(shift_tokens, len(layer_types))
1310 | 
1311 |         # whether it has post norm
1312 | 
1313 |         self.final_norm = norm_fn() if pre_norm or resi_dual else nn.Identity()
1314 | 
1315 |         # iterate and construct layers
1316 | 
1317 |         for ind, (layer_type, layer_shift_tokens) in enumerate(
1318 |             zip(self.layer_types, shift_tokens)
1319 |         ):
1320 |             ind == (len(self.layer_types) - 1)
1321 | 
1322 |             if layer_type == "a":
1323 |                 layer = Attention(dim, heads=heads, causal=causal, **attn_kwargs)
1324 |             elif layer_type == "c":
1325 |                 layer = Attention(dim, heads=heads, **attn_kwargs)
1326 |             elif layer_type == "f":
1327 |                 layer = FeedForward(dim, **ff_kwargs)
1328 |                 layer = layer if not macaron else Scale(0.5, layer)
1329 |             else:
1330 |                 raise Exception(f"invalid layer type {layer_type}")
1331 | 
1332 |             if layer_shift_tokens > 0:
1333 |                 shift_range_upper = layer_shift_tokens + 1
1334 |                 shift_range_lower = -layer_shift_tokens if not causal else 0
1335 |                 layer = ShiftTokens(range(shift_range_lower, shift_range_upper), layer)
1336 | 
1337 |             residual_fn = GRUGating if gate_residual else Residual
1338 |             residual = residual_fn(
1339 |                 dim,
1340 |                 scale_residual=scale_residual,
1341 |                 scale_residual_constant=scale_residual_constant,
1342 |             )
1343 | 
1344 |             pre_branch_norm = norm_fn() if pre_norm else None
1345 |             post_branch_norm = norm_fn() if sandwich_norm else None
1346 |             post_main_norm = norm_fn() if not pre_norm else None
1347 | 
1348 |             norms = nn.ModuleList([pre_branch_norm, post_branch_norm, post_main_norm])
1349 | 
1350 |             self.layers.append(nn.ModuleList([norms, layer, residual]))
1351 | 
1352 |     def forward(
1353 |         self,
1354 |         x,
1355 |         context=None,
1356 |         mask=None,
1357 |         context_mask=None,
1358 |         attn_mask=None,
1359 |         self_attn_kv_mask=None,
1360 |         mems=None,
1361 |         seq_start_pos: Optional[Tensor] = None,
1362 |         cache: Optional[LayerIntermediates] = None,
1363 |         cache_age=1,
1364 |         return_hiddens=False,
1365 |     ):
1366 |         assert not (
1367 |             self.cross_attend ^ exists(context)
1368 |         ), "context must be passed in if cross_attend is set to True"
1369 | 
1370 |         # initialize accums
1371 | 
1372 |         hiddens = []
1373 |         layer_hiddens = []
1374 |         intermediates = []
1375 | 
1376 |         prev_attn = None
1377 |         prev_cross_attn = None
1378 | 
1379 |         mems = mems.copy() if exists(mems) else [None] * self.num_attn_layers
1380 | 
1381 |         # handle left padded sequences
1382 | 
1383 |         if exists(seq_start_pos):
1384 |             seq_arange = torch.arange(x.shape[-2], device=x.device, dtype=torch.long)
1385 |             left_pad_mask = seq_arange >= seq_start_pos[..., None]
1386 | 
1387 |             if exists(self_attn_kv_mask):
1388 |                 self_attn_kv_mask = self_attn_kv_mask & left_pad_mask
1389 |             else:
1390 |                 self_attn_kv_mask = left_pad_mask
1391 | 
1392 |         # rotary positions
1393 | 
1394 |         rotary_pos_emb = None
1395 | 
1396 |         if exists(self.rotary_pos_emb):
1397 |             max_rotary_emb_length = max(
1398 |                 list(map(lambda m: (m.shape[1] if exists(m) else 0) + x.shape[1], mems))
1399 |             )
1400 |             rotary_pos_emb = self.rotary_pos_emb(max_rotary_emb_length)
1401 | 
1402 |         # assume cached key / values
1403 | 
1404 |         attn_cache = []
1405 | 
1406 |         if exists(cache):
1407 |             assert (
1408 |                 not self.training
1409 |                 and self.causal
1410 |                 and not any([*map(exists, (mask, attn_mask))])
1411 |             )
1412 | 
1413 |             if cache_age > 0:
1414 |                 x = x[:, -cache_age:]  # for spec decoding, may be greater than 1
1415 | 
1416 |             attn_cache = cache.attn_intermediates
1417 | 
1418 |         iter_attn_cache = iter(attn_cache)
1419 | 
1420 |         # outer residual - for resiDual paper
1421 | 
1422 |         outer_residual = x * self.resi_dual_scale
1423 | 
1424 |         # get layers to be executed
1425 | 
1426 |         layer_variables = (self.layer_types, self.layers, self.layer_dropouts)
1427 | 
1428 |         layer_variables = tuple(
1429 |             tuple(layer_variable[i] for i in self.layers_execute_order)
1430 |             for layer_variable in layer_variables
1431 |         )
1432 | 
1433 |         # go through the attention and feedforward layers
1434 | 
1435 |         for ind, (layer_type, (norm, block, residual_fn), layer_dropout) in enumerate(
1436 |             zip(*layer_variables)
1437 |         ):
1438 |             ind == (len(self.layers) - 1)
1439 | 
1440 |             if self.training and layer_dropout > 0.0 and random() < layer_dropout:
1441 |                 continue
1442 | 
1443 |             if layer_type == "a":
1444 |                 if return_hiddens:
1445 |                     hiddens.append(x)
1446 |                 layer_mem = mems.pop(0) if mems else None
1447 | 
1448 |             if layer_type == "c":
1449 |                 if self.training and self.cross_attn_tokens_dropout > 0.0:
1450 |                     context, context_mask = dropout_seq(
1451 |                         context, context_mask, self.cross_attn_tokens_dropout
1452 |                     )
1453 | 
1454 |             inner_residual = x
1455 | 
1456 |             if return_hiddens:
1457 |                 layer_hiddens.append(x)
1458 | 
1459 |             pre_norm, post_branch_norm, post_main_norm = norm
1460 | 
1461 |             if exists(pre_norm):
1462 |                 x = pre_norm(x)
1463 | 
1464 |             if layer_type == "a":
1465 |                 out, inter = block(
1466 |                     x,
1467 |                     mask=mask,
1468 |                     context_mask=self_attn_kv_mask,
1469 |                     attn_mask=attn_mask,
1470 |                     rel_pos=self.rel_pos,
1471 |                     rotary_pos_emb=rotary_pos_emb,
1472 |                     prev_attn=prev_attn,
1473 |                     cache=next(iter_attn_cache, None),
1474 |                     mem=layer_mem,
1475 |                     return_intermediates=True,
1476 |                 )
1477 |             elif layer_type == "c":
1478 |                 out, inter = block(
1479 |                     x,
1480 |                     context=context,
1481 |                     mask=mask,
1482 |                     context_mask=context_mask,
1483 |                     prev_attn=prev_cross_attn,
1484 |                     cache=next(iter_attn_cache, None),
1485 |                     return_intermediates=True,
1486 |                 )
1487 |             elif layer_type == "f":
1488 |                 out = block(x)
1489 | 
1490 |             if self.resi_dual:
1491 |                 outer_residual = outer_residual + out * self.resi_dual_scale
1492 | 
1493 |             if exists(post_branch_norm):
1494 |                 out = post_branch_norm(out)
1495 | 
1496 |             x = residual_fn(out, inner_residual)
1497 | 
1498 |             if layer_type in ("a", "c") and return_hiddens:
1499 |                 intermediates.append(inter)
1500 | 
1501 |             if layer_type == "a" and self.residual_attn:
1502 |                 prev_attn = inter.pre_softmax_attn
1503 |             elif layer_type == "c" and self.cross_residual_attn:
1504 |                 prev_cross_attn = inter.pre_softmax_attn
1505 | 
1506 |             if exists(post_main_norm):
1507 |                 x = post_main_norm(x)
1508 | 
1509 |         if return_hiddens:
1510 |             layer_hiddens.append(x)
1511 | 
1512 |         if self.resi_dual:
1513 |             x = x + self.final_norm(outer_residual)
1514 |         else:
1515 |             x = self.final_norm(x)
1516 | 
1517 |         if not return_hiddens:
1518 |             return x
1519 | 
1520 |         intermediates = LayerIntermediates(
1521 |             hiddens=hiddens,
1522 |             attn_intermediates=intermediates,
1523 |             layer_hiddens=layer_hiddens,
1524 |         )
1525 | 
1526 |         return x, intermediates
1527 | 
1528 | 
1529 | class Encoder(AttentionLayers):
1530 |     def __init__(self, **kwargs):
1531 |         assert "causal" not in kwargs, "cannot set causality on encoder"
1532 |         super().__init__(causal=False, **kwargs)
1533 | 
1534 | 
1535 | class Decoder(AttentionLayers):
1536 |     def __init__(self, **kwargs):
1537 |         assert "causal" not in kwargs, "cannot set causality on decoder"
1538 |         super().__init__(causal=True, **kwargs)
1539 | 
1540 | 
1541 | class CrossAttender(AttentionLayers):
1542 |     def __init__(self, **kwargs):
1543 |         super().__init__(cross_attend=True, only_cross=True, **kwargs)
1544 | 
1545 | 
1546 | class ViTransformerWrapper(nn.Module):
1547 |     def __init__(
1548 |         self,
1549 |         *,
1550 |         image_size,
1551 |         patch_size,
1552 |         attn_layers,
1553 |         channels=3,
1554 |         num_classes=None,
1555 |         post_emb_norm=False,
1556 |         num_register_tokens=0,
1557 |         emb_dropout=0.0,
1558 |     ):
1559 |         super().__init__()
1560 |         assert isinstance(attn_layers, Encoder), "attention layers must be an Encoder"
1561 |         assert divisible_by(
1562 |             image_size, patch_size
1563 |         ), "image dimensions must be divisible by the patch size"
1564 |         dim = attn_layers.dim
1565 |         num_patches = (image_size // patch_size) ** 2
1566 |         patch_dim = channels * patch_size**2
1567 | 
1568 |         self.patch_size = patch_size
1569 | 
1570 |         self.pos_embedding = nn.Parameter(torch.randn(1, num_patches, dim))
1571 | 
1572 |         has_register_tokens = num_register_tokens > 0
1573 |         self.has_register_tokens = has_register_tokens
1574 | 
1575 |         if has_register_tokens:
1576 |             self.register_tokens = nn.Parameter(torch.randn(num_register_tokens, dim))
1577 | 
1578 |         self.patch_to_embedding = nn.Sequential(
1579 |             nn.LayerNorm(patch_dim), nn.Linear(patch_dim, dim), nn.LayerNorm(dim)
1580 |         )
1581 | 
1582 |         self.post_emb_norm = nn.LayerNorm(dim) if post_emb_norm else nn.Identity()
1583 |         self.dropout = nn.Dropout(emb_dropout)
1584 | 
1585 |         self.attn_layers = attn_layers
1586 | 
1587 |         self.mlp_head = (
1588 |             nn.Linear(dim, num_classes) if exists(num_classes) else nn.Identity()
1589 |         )
1590 | 
1591 |     def forward(self, img, return_embeddings=False):
1592 |         b, p = img.shape[0], self.patch_size
1593 | 
1594 |         x = rearrange(img, "b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1=p, p2=p)
1595 |         x = self.patch_to_embedding(x)
1596 |         n = x.shape[1]
1597 | 
1598 |         x = x + self.pos_embedding[:, :n]
1599 | 
1600 |         x = self.post_emb_norm(x)
1601 |         x = self.dropout(x)
1602 | 
1603 |         if self.has_register_tokens:
1604 |             r = repeat(self.register_tokens, "n d -> b n d", b=b)
1605 |             x, ps = pack((x, r), "b * d")
1606 | 
1607 |         x = self.attn_layers(x)
1608 | 
1609 |         if self.has_register_tokens:
1610 |             x, _ = unpack(x, ps, "b * d")
1611 | 
1612 |         if not exists(self.mlp_head) or return_embeddings:
1613 |             return x
1614 | 
1615 |         x = x.mean(dim=-2)
1616 |         return self.mlp_head(x)
1617 | 
1618 | 
1619 | class TransformerWrapper(nn.Module):
1620 |     def __init__(
1621 |         self,
1622 |         *,
1623 |         num_tokens,
1624 |         max_seq_len,
1625 |         attn_layers,
1626 |         emb_dim=None,
1627 |         max_mem_len=0,
1628 |         shift_mem_down=0,
1629 |         emb_dropout=0.0,
1630 |         post_emb_norm=False,
1631 |         num_memory_tokens=None,
1632 |         memory_tokens_interspersed_every=None,
1633 |         tie_embedding=False,
1634 |         logits_dim=None,
1635 |         use_abs_pos_emb=True,
1636 |         scaled_sinu_pos_emb=False,
1637 |         l2norm_embed=False,
1638 |         emb_frac_gradient=1.0,  # GLM-130B and Cogview successfully used this, set at 0.1
1639 |         attn_z_loss_weight=1e-4,
1640 |     ):
1641 |         super().__init__()
1642 |         assert isinstance(
1643 |             attn_layers, AttentionLayers
1644 |         ), "attention layers must be one of Encoder or Decoder"
1645 | 
1646 |         dim = attn_layers.dim
1647 |         emb_dim = default(emb_dim, dim)
1648 |         self.emb_dim = emb_dim
1649 |         self.num_tokens = num_tokens
1650 | 
1651 |         self.max_seq_len = max_seq_len
1652 |         self.max_mem_len = max_mem_len
1653 |         self.shift_mem_down = shift_mem_down
1654 | 
1655 |         self.l2norm_embed = l2norm_embed
1656 |         self.token_emb = TokenEmbedding(emb_dim, num_tokens, l2norm_embed=l2norm_embed)
1657 | 
1658 |         if not (use_abs_pos_emb and not attn_layers.has_pos_emb):
1659 |             self.pos_emb = always(0)
1660 |         elif scaled_sinu_pos_emb:
1661 |             self.pos_emb = ScaledSinusoidalEmbedding(emb_dim)
1662 |         else:
1663 |             self.pos_emb = AbsolutePositionalEmbedding(
1664 |                 emb_dim, max_seq_len, l2norm_embed=l2norm_embed
1665 |             )
1666 | 
1667 |         self.emb_frac_gradient = emb_frac_gradient  # fraction of the gradient that should go to the embedding, https://arxiv.org/abs/2105.13290
1668 | 
1669 |         self.post_emb_norm = nn.LayerNorm(emb_dim) if post_emb_norm else nn.Identity()
1670 |         self.emb_dropout = nn.Dropout(emb_dropout)
1671 | 
1672 |         self.project_emb = nn.Linear(emb_dim, dim) if emb_dim != dim else nn.Identity()
1673 |         self.attn_layers = attn_layers
1674 | 
1675 |         self.init_()
1676 | 
1677 |         logits_dim = default(logits_dim, num_tokens)
1678 |         self.to_logits = (
1679 |             nn.Linear(dim, logits_dim)
1680 |             if not tie_embedding
1681 |             else lambda t: t @ self.token_emb.emb.weight.t()
1682 |         )
1683 | 
1684 |         # memory tokens (like [cls]) from Memory Transformers paper
1685 | 
1686 |         num_memory_tokens = default(num_memory_tokens, 0)
1687 |         self.num_memory_tokens = num_memory_tokens
1688 |         if num_memory_tokens > 0:
1689 |             self.memory_tokens = nn.Parameter(torch.randn(num_memory_tokens, dim))
1690 | 
1691 |         self.memory_tokens_interspersed_every = memory_tokens_interspersed_every
1692 | 
1693 |         # whether can do cached kv decoding
1694 | 
1695 |         self.can_cache_kv = self.num_memory_tokens == 0
1696 | 
1697 |     def init_(self):
1698 |         if self.l2norm_embed:
1699 |             nn.init.normal_(self.token_emb.emb.weight, std=1e-5)
1700 |             if not isinstance(self.pos_emb, always):
1701 |                 nn.init.normal_(self.pos_emb.emb.weight, std=1e-5)
1702 |             return
1703 | 
1704 |         nn.init.kaiming_normal_(self.token_emb.emb.weight)
1705 | 
1706 |     def forward(
1707 |         self,
1708 |         x,
1709 |         return_embeddings=False,
1710 |         return_logits_and_embeddings=False,
1711 |         return_intermediates=False,
1712 |         mask=None,
1713 |         return_mems=False,
1714 |         return_attn=False,
1715 |         mems=None,
1716 |         pos=None,
1717 |         prepend_embeds=None,
1718 |         sum_embeds=None,
1719 |         return_attn_z_loss=False,
1720 |         attn_z_loss_weight=1e-4,
1721 |         seq_start_pos=None,
1722 |         cache: Optional[LayerIntermediates] = None,
1723 |         **kwargs,
1724 |     ):
1725 |         b, n, device, num_mems, has_memory_tokens, emb_frac_gradient = (
1726 |             *x.shape,
1727 |             x.device,
1728 |             self.num_memory_tokens,
1729 |             self.num_memory_tokens > 0,
1730 |             self.emb_frac_gradient,
1731 |         )
1732 |         (return_mems | return_attn | return_intermediates | return_attn_z_loss)
1733 | 
1734 |         # absolute positional embedding
1735 | 
1736 |         external_pos_emb = exists(pos) and pos.dtype != torch.long
1737 |         pos_emb = (
1738 |             self.pos_emb(x, pos=pos, seq_start_pos=seq_start_pos)
1739 |             if not external_pos_emb
1740 |             else pos
1741 |         )
1742 |         x = self.token_emb(x) + pos_emb
1743 | 
1744 |         # for summing embeddings passed externally - needs this for self-conditioning in non-autoregressive training
1745 | 
1746 |         if exists(sum_embeds):
1747 |             x = x + sum_embeds
1748 | 
1749 |         # post embedding norm, purportedly leads to greater stabilization
1750 | 
1751 |         x = self.post_emb_norm(x)
1752 | 
1753 |         # whether to append embeds, as in PaLI, for image embeddings
1754 | 
1755 |         if exists(prepend_embeds):
1756 |             prepend_seq, prepend_dim = prepend_embeds.shape[1:]
1757 |             assert (
1758 |                 prepend_dim == x.shape[-1]
1759 |             ), "prepended embeddings need to have same dimensions as text model dimensions"
1760 | 
1761 |             x = torch.cat((prepend_embeds, x), dim=-2)
1762 | 
1763 |         # whether to reduce the gradient going to the embedding, from cogview paper, corroborated by GLM-130B model
1764 | 
1765 |         if emb_frac_gradient < 1:
1766 |             assert emb_frac_gradient > 0
1767 |             x = x * emb_frac_gradient + x.detach() * (1 - emb_frac_gradient)
1768 | 
1769 |         # embedding dropout
1770 | 
1771 |         x = self.emb_dropout(x)
1772 | 
1773 |         x = self.project_emb(x)
1774 | 
1775 |         if has_memory_tokens:
1776 |             mem_every = self.memory_tokens_interspersed_every
1777 | 
1778 |             if exists(mem_every):
1779 |                 assert mem_every > 0
1780 |                 assert isinstance(self.attn_layers, Decoder), "only for decoder"
1781 |                 next_seq_len = math.ceil(n / mem_every) * mem_every
1782 | 
1783 |                 x = pad_at_dim(x, (0, next_seq_len - n), dim=-2, value=0.0)
1784 |                 x = rearrange(x, "b (n m) d -> (b n) m d", m=mem_every)
1785 | 
1786 |             mem = repeat(self.memory_tokens, "n d -> b n d", b=x.shape[0])
1787 |             x, mem_packed_shape = pack((mem, x), "b * d")
1788 | 
1789 |             # auto-handle masking after appending memory tokens
1790 |             if not exists(mem_every) and exists(mask):
1791 |                 mask = pad_at_dim(mask, (num_mems, 0), dim=-1, value=True)
1792 | 
1793 |             if exists(mem_every):
1794 |                 x = rearrange(x, "(b n) m d -> b (n m) d", b=b)
1795 | 
1796 |         if self.shift_mem_down and exists(mems):
1797 |             mems_l, mems_r = mems[: self.shift_mem_down], mems[self.shift_mem_down :]
1798 |             mems = [*mems_r, *mems_l]
1799 | 
1800 |         x, intermediates = self.attn_layers(
1801 |             x,
1802 |             mask=mask,
1803 |             mems=mems,
1804 |             cache=cache,
1805 |             return_hiddens=True,
1806 |             seq_start_pos=seq_start_pos,
1807 |             **kwargs,
1808 |         )
1809 | 
1810 |         if has_memory_tokens:
1811 |             if exists(mem_every):
1812 |                 x = rearrange(x, "b (n m) d -> (b n) m d", m=(mem_every + num_mems))
1813 | 
1814 |             mem, x = unpack(x, mem_packed_shape, "b * d")
1815 | 
1816 |             if exists(mem_every):
1817 |                 x = rearrange(x, "(b n) m d -> b (n m) d", b=b)
1818 | 
1819 |             x = x[:, :n]
1820 | 
1821 |         if return_logits_and_embeddings:
1822 |             out = (self.to_logits(x), x)
1823 |         elif return_embeddings:
1824 |             out = x
1825 |         else:
1826 |             out = self.to_logits(x)
1827 | 
1828 |         if return_attn_z_loss:
1829 |             pre_softmax_attns = list(
1830 |                 map(lambda t: t.pre_softmax_attn, intermediates.attn_intermediates)
1831 |             )
1832 |             intermediates.attn_z_loss = calc_z_loss(
1833 |                 pre_softmax_attns, weight=attn_z_loss_weight
1834 |             )
1835 |             return_intermediates = True
1836 | 
1837 |         if return_mems:
1838 |             hiddens = intermediates.hiddens
1839 |             new_mems = (
1840 |                 list(map(lambda pair: torch.cat(pair, dim=-2), zip(mems, hiddens)))
1841 |                 if exists(mems)
1842 |                 else hiddens
1843 |             )
1844 |             new_mems = list(
1845 |                 map(lambda t: t[..., -self.max_mem_len :, :].detach(), new_mems)
1846 |             )
1847 | 
1848 |             if not return_intermediates:
1849 |                 return out, new_mems
1850 | 
1851 |             intermediates.mems = new_mems
1852 | 
1853 |         if return_intermediates:
1854 |             return out, intermediates
1855 | 
1856 |         if return_attn:
1857 |             attn_maps = list(
1858 |                 map(lambda t: t.post_softmax_attn, intermediates.attn_intermediates)
1859 |             )
1860 |             return out, attn_maps
1861 | 
1862 |         return out
1863 | 
1864 | 
1865 | class UL2(nn.Module):
1866 |     def __init__(
1867 |         self,
1868 |         *,
1869 |         dim,
1870 |         tie_token_emb=False,
1871 |         ignore_index=-100,
1872 |         pad_value=0,
1873 |         cross_attn_tokens_dropout=0.0,
1874 |         **kwargs,
1875 |     ):
1876 |         super().__init__()
1877 |         enc_kwargs, kwargs = groupby_prefix_and_trim("enc_", kwargs)
1878 |         dec_kwargs, kwargs = groupby_prefix_and_trim("dec_", kwargs)
1879 | 
1880 |         assert (
1881 |             "dim" not in enc_kwargs and "dim" not in dec_kwargs
1882 |         ), "dimension of either encoder or decoder must be set with `dim` keyword"
1883 |         enc_transformer_kwargs = pick_and_pop(["num_tokens", "max_seq_len"], enc_kwargs)
1884 |         enc_transformer_kwargs["emb_dropout"] = enc_kwargs.pop("emb_dropout", 0)
1885 |         enc_transformer_kwargs["num_memory_tokens"] = enc_kwargs.pop(
1886 |             "num_memory_tokens", None
1887 |         )
1888 |         enc_transformer_kwargs["scaled_sinu_pos_emb"] = enc_kwargs.pop(
1889 |             "scaled_sinu_pos_emb", False
1890 |         )
1891 |         enc_transformer_kwargs["use_abs_pos_emb"] = enc_kwargs.pop(
1892 |             "use_abs_pos_emb", True
1893 |         )
1894 | 
1895 |         dec_transformer_kwargs = pick_and_pop(["num_tokens", "max_seq_len"], dec_kwargs)
1896 |         dec_transformer_kwargs["emb_dropout"] = dec_kwargs.pop("emb_dropout", 0)
1897 |         dec_transformer_kwargs["scaled_sinu_pos_emb"] = dec_kwargs.pop(
1898 |             "scaled_sinu_pos_emb", False
1899 |         )
1900 |         dec_transformer_kwargs["use_abs_pos_emb"] = dec_kwargs.pop(
1901 |             "use_abs_pos_emb", True
1902 |         )
1903 | 
1904 |         self.cross_attn_tokens_dropout = cross_attn_tokens_dropout  # how many tokens from the encoder to dropout when cross attending from decoder - seen in a couple papers, including Perceiver AR - this will also be very effective regularization when cross attending to very long memories
1905 | 
1906 |         self.encoder = TransformerWrapper(
1907 |             **enc_transformer_kwargs, attn_layers=Encoder(dim=dim, **enc_kwargs)
1908 |         )
1909 | 
1910 |         self.decoder = TransformerWrapper(
1911 |             **dec_transformer_kwargs,
1912 |             attn_layers=Decoder(dim=dim, cross_attend=True, **dec_kwargs),
1913 |         )
1914 | 
1915 |         if tie_token_emb:
1916 |             self.decoder.token_emb = self.encoder.token_emb
1917 | 
1918 |         self.decoder = AutoregressiveWrapper(
1919 |             self.decoder, ignore_index=ignore_index, pad_value=pad_value
1920 |         )
1921 | 
1922 |     @torch.no_grad()
1923 |     def generate(
1924 |         self, seq_in, seq_out_start, seq_len, mask=None, attn_mask=None, **kwargs
1925 |     ):
1926 |         encodings = self.encoder(
1927 |             seq_in, mask=mask, attn_mask=attn_mask, return_embeddings=True
1928 |         )
1929 |         return self.decoder.generate(
1930 |             seq_out_start, seq_len, context=encodings, context_mask=mask, **kwargs
1931 |         )
1932 | 
1933 |     def forward(self, src, tgt, mask=None, attn_mask=None, src_prepend_embeds=None):
1934 |         if exists(src_prepend_embeds) and exists(mask):
1935 |             mask = pad_at_dim(
1936 |                 mask, (src_prepend_embeds.shape[-2], 0), dim=-1, value=True
1937 |             )
1938 | 
1939 |         enc = self.encoder(
1940 |             src,
1941 |             mask=mask,
1942 |             attn_mask=attn_mask,
1943 |             prepend_embeds=src_prepend_embeds,
1944 |             return_embeddings=True,
1945 |         )
1946 | 
1947 |         if self.training and self.cross_attn_tokens_dropout > 0:
1948 |             enc, mask = dropout_seq(enc, mask, self.cross_attn_tokens_dropout)
1949 | 
1950 |         out = self.decoder(tgt, context=enc, context_mask=mask)
1951 |         return out


--------------------------------------------------------------------------------
/pyproject.toml:
--------------------------------------------------------------------------------
 1 | [build-system]
 2 | requires = ["poetry-core>=1.0.0"]
 3 | build-backend = "poetry.core.masonry.api"
 4 | 
 5 | [tool.poetry]
 6 | name = "pali3"
 7 | version = "0.0.7"
 8 | description = "pali3 - Pytorch"
 9 | license = "MIT"
10 | authors = ["Kye Gomez <kye@apac.ai>"]
11 | homepage = "https://github.com/kyegomez/pali3"
12 | documentation = ""  # Add this if you have documentation.
13 | readme = "README.md"  # Assuming you have a README.md
14 | repository = "https://github.com/kyegomez/pali3"
15 | keywords = ["artificial intelligence", "deep learning", "optimizers", "Prompt Engineering"]
16 | classifiers = [
17 |     "Development Status :: 4 - Beta",
18 |     "Intended Audience :: Developers",
19 |     "Topic :: Scientific/Engineering :: Artificial Intelligence",
20 |     "License :: OSI Approved :: MIT License",
21 |     "Programming Language :: Python :: 3.6"
22 | ]
23 | 
24 | [tool.poetry.dependencies]
25 | python = "^3.6"
26 | torch = "*"
27 | einops = "*"
28 | horovod = "*"
29 | zetascale = "*"
30 | torchvision = "*"
31 | 
32 | 
33 | 
34 | 
35 | 
36 | [tool.poetry.dev-dependencies]
37 | # Add development dependencies here
38 | 
39 | 
40 | [tool.poetry.group.lint.dependencies]
41 | ruff = "^0.0.249"
42 | types-toml = "^0.10.8.1"
43 | types-redis = "^4.3.21.6"
44 | types-pytz = "^2023.3.0.0"
45 | black = "^23.1.0"
46 | types-chardet = "^5.0.4.6"
47 | mypy-protobuf = "^3.0.0"
48 | 
49 | 
50 | [tool.autopep8]
51 | max_line_length = 120
52 | ignore = "E501,W6"  # or ["E501", "W6"]
53 | in-place = true
54 | recursive = true
55 | aggressive = 3


--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | torch
2 | einops
3 | horovod
4 | zetascale
5 | torchvision


--------------------------------------------------------------------------------