Merge pull request #29 from kyutai-labs/ci

CI [2/N]

.github/ISSUE_TEMPLATE/bug.md

@@ -0,0 +1,29 @@
+---
+name: 🐛 Bug Report
+about: Submit a bug report to help us improve
+labels: bug, triage
+---
+
+## 🐛 Bug Report
+
+<!-- A clear and concise description of what the bug is -->
+
+## To Reproduce
+
+<!-- How to reproduce the bug  -->
+
+## Your Environment
+
+<!-- Include as many relevant details about the environment you experienced the bug in.-->
+
+**Please fill this part, failure to do so will lead to your issue being directly closed.**
+
+- Backend impacted (PyTorch, MLX, Rust, or Other):
+- Operating system (OSX, Windows, Linux):
+- Operating system version:
+
+If the backend impacted is PyTorch:
+- Python version:
+- PyTorch version:
+- CUDA version (run `python -c 'import torch;  print(torch.version.cuda)'`):
+- GPU model and memory:

.github/ISSUE_TEMPLATE/question.md

@@ -0,0 +1,9 @@
+---
+name: "❓Questions/Help/Support"
+about: If you have a question about the paper, code or algorithm, please ask here!
+labels: question, triage
+---
+
+## ❓ Questions
+
+<!-- (Please ask your question here.) -->

.github/PULL_REQUEST_TEMPLATE.md

@@ -0,0 +1 @@
+# Test

.github/actions/moshi_build/action.yml

@@ -0,0 +1,27 @@
+name: moshi_build
+description: 'Build env.'
+runs:
+  using: "composite"
+  steps:
+  - uses: actions/setup-python@v2
+    with:
+      python-version: '3.10'
+  - uses: actions/cache@v3
+    id: cache
+    with:
+      path: env
+      key: env-${{ hashFiles('moshi/pyproject.toml') }}
+  - name: Install dependencies
+    if: steps.cache.outputs.cache-hit != 'true'
+    shell: bash
+    run: |
+      python3 -m venv env
+      .  env/bin/activate
+      python -m pip install --upgrade pip
+      pip install torch==2.4.0 --index-url https://download.pytorch.org/whl/cpu
+      pip install -e './moshi[dev]'
+  - name: Setup env
+    shell: bash
+    run: |
+      .  env/bin/activate
+      pre-commit install

.github/actions/rust_build/action.yml

@@ -0,0 +1,33 @@
+name: rust_build
+description: 'Setup rust env'
+inputs:
+  os:
+    default: ubuntu-latest
+  toolchain:
+    default: stable
+  target:
+    default: check
+runs:
+  using: "composite"
+  steps:
+    - uses: actions-rs/toolchain@v1
+      with:
+        profile: minimal
+        toolchain: ${{ inputs.toolchain }}
+        override: true
+    - name: cargo cache
+      uses: actions/cache@v3
+      with:
+        path: |
+          ~/.cargo/bin/
+          ~/.cargo/registry/index/
+          ~/.cargo/registry/cache/
+          ~/.cargo/git/db/
+          rust/target/
+        key: ${{ inputs.os }}-cargo-${{ inputs.target }}-${{ hashFiles('**/Cargo.lock') }}
+        restore-keys: ${{ inputs.os }}-cargo-
+    - name: install deps
+      shell: bash
+      run: |
+        sudo apt-get update
+        sudo apt-get install libasound2-dev

.github/workflows/precommit.yml

@@ -0,0 +1,17 @@
+name: precommit
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main, refacto ]
+
+jobs:
+  run_precommit:
+    name: Run precommit
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: ./.github/actions/moshi_build
+      - run: |
+          . env/bin/activate
+          bash .git/hooks/pre-commit

.github/workflows/rust-ci.yml

@@ -0,0 +1,53 @@
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main, refacto ]
+
+name: Rust CI
+
+jobs:
+  check:
+    name: Check
+    defaults:
+      run:
+        working-directory: ./rust
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-latest]
+        rust: [stable]
+    steps:
+      - uses: actions/checkout@v2
+      - uses: ./.github/actions/rust_build
+      - name: check
+        shell: bash
+        run: |
+          cargo check
+      - name: clippy
+        shell: bash
+        run: |
+          cargo clippy -- -D warnings
+      - name: fmt
+        shell: bash
+        run: |
+          cargo fmt --all -- --check
+  test:
+    name: Test
+    defaults:
+      run:
+        working-directory: ./rust
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        os: [ubuntu-latest]
+        rust: [stable]
+    steps:
+      - uses: actions/checkout@v2
+      - uses: ./.github/actions/rust_build
+        with:
+          target: test
+      - name: test
+        shell: bash
+        run: |
+          cargo test

.pre-commit-config.yaml

@@ -0,0 +1,24 @@
+repos:
+  - repo: local
+    hooks:
+    - id: flake8-moshi
+      name: flake8 on moshi package
+      language: system
+      entry: bash -c 'cd moshi && flake8'
+      pass_filenames: false
+    - id: pyright-moshi
+      name: pyright on moshi package
+      language: system
+      entry: bash -c 'cd moshi && pyright'
+      pass_filenames: false
+    - id: flake8-moshi_mlx
+      name: flake8 on moshi_mlx package
+      language: system
+      entry: bash -c 'cd moshi_mlx && flake8'
+      pass_filenames: false
+    - id: pyright-moshi_mlx
+      name: pyright on moshi_mlx package
+      language: system
+      entry: bash -c 'cd moshi_mlx && pyright'
+      pass_filenames: false
+

CONTRIBUTING.md

@@ -0,0 +1,30 @@
+# Contributing to Moshi
+
+## Pull Requests
+
+Moshi is the implementation of a research paper.
+Therefore, we do not plan on accepting many pull requests for new features.
+We certainly welcome them for bug fixes.
+
+1. Fork the repo and create your branch from `main`.
+2. If you've changed APIs, update the documentation.
+3. Ensure pre-commit hooks pass properly, in particular the linting and typing.
+4. Accept the Contributor License Agreement (see after).
+
+Note that in general we will not accept refactoring of the code.
+
+
+## Contributor License Agreement ("CLA")
+
+In order to accept your pull request, we need you to submit a Contributtor License Agreement.
+As this CLA is not ready yet, we will delay acceptance of your PR.
+
+## Issues
+
+Please submit issues on our Github repository.
+
+## License
+
+By contributing to Moshi, you agree that your contributions will be licensed
+under the LICENSE-* files in the root directory of this source tree.
+In particular, the rust code is licensed under APACHE, and the python code under MIT.

moshi/LICENSE

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

moshi/pyproject.toml

@@ -25,6 +25,9 @@ version = {attr = "moshi.__version__"}
 requires = ["setuptools"]
 build-backend = "setuptools.build_meta"
 
+[tool.setuptools.dynamic]
+version = {attr = "moshi.__version__"}
+
 [project.optional-dependencies]
 dev = [
     "pyright",

moshi_mlx/LICENSE

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

moshi_mlx/moshi_mlx/__init__.py

@@ -1,12 +1,12 @@
 # Copyright (c) Kyutai, all rights reserved.
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
+# flake8: noqa
 
 """
 moshi_mlx is the MLX inference codebase for Kyutai audio generation models.
 """
 
-# flake8: noqa
 from . import modules, models, utils
 
 __version__ = "0.0.1"

moshi_mlx/moshi_mlx/local.py

@@ -20,7 +20,7 @@
 
 from .client_utils import AnyPrinter, Printer, RawPrinter
 import rustymimi
-import moshi_mlx
+from moshi_mlx import models, utils
 
 import huggingface_hub
 
@@ -99,10 +99,10 @@ def log(s):
         printer_q.put_nowait((PrinterType.INFO, s))
 
     log(f"[SERVER] loading text tokenizer {tokenizer_file}")
-    text_tokenizer = sentencepiece.SentencePieceProcessor(tokenizer_file)
+    text_tokenizer = sentencepiece.SentencePieceProcessor(tokenizer_file)  # type: ignore
     mx.random.seed(299792458)
-    lm_config = moshi_mlx.models.config_v0_1()
-    model = moshi_mlx.models.Lm(lm_config)
+    lm_config = models.config_v0_1()
+    model = models.Lm(lm_config)
     model.set_dtype(mx.bfloat16)
     if args.quantized is not None:
         group_size = 32 if args.quantized == 4 else 64
@@ -114,11 +114,11 @@ def log(s):
 
     model.warmup()
     log("[SERVER] model warmed up")
-    gen = moshi_mlx.models.LmGen(
+    gen = models.LmGen(
         model=model,
         max_steps=steps + 5,
-        text_sampler=moshi_mlx.utils.Sampler(),
-        audio_sampler=moshi_mlx.utils.Sampler(),
+        text_sampler=utils.Sampler(),
+        audio_sampler=utils.Sampler(),
         check=False,
     )
 
@@ -137,7 +137,7 @@ def log(s):
             text_token = text_token[0].item()
             audio_tokens = gen.last_audio_tokens()
             if text_token not in (0, 3):
-                _text = text_tokenizer.id_to_piece(text_token)
+                _text = text_tokenizer.id_to_piece(text_token)  # type: ignore
                 _text = _text.replace("▁", " ")
                 printer_q.put_nowait((PrinterType.TOKEN, _text))
             else:
@@ -158,7 +158,7 @@ def client(printer_q, client_to_server, server_to_client, args):
         )
     input_queue = queue.Queue()
     output_queue = queue.Queue()
-    audio_tokenizer = rustymimi.StreamTokenizer(mimi_file)
+    audio_tokenizer = rustymimi.StreamTokenizer(mimi_file)  # type: ignore
     start = server_to_client.get()
     printer_q.put_nowait(
         (PrinterType.INFO, f"[CLIENT] received '{start}' from server, starting...")
@@ -255,9 +255,9 @@ def main(printer: AnyPrinter):
     parser.add_argument("--tokenizer", type=str)
     parser.add_argument("--model", type=str)
     parser.add_argument("--mimi", type=str)
-    parser.add_argument("--quantized", type=int)
+    parser.add_argument("-q", "--quantized", type=int)
     parser.add_argument("--steps", default=2500, type=int)
-    parser.add_argument("--hf-repo", type=str, default="")
+    parser.add_argument("--hf-repo", type=str, default="kmhf/msh-v0.1")
 
     args = parser.parse_args()