chore(ci): fix test path for scripts

.github/workflows/build.yml

@@ -6,10 +6,6 @@ on:
       artifact-name:
         description: "Name of the artifact containing the built package"
         value: ${{ jobs.build.outputs.artifact-name }}
-  pull_request:
-    branches: [ main ]
-  push:
-    branches: [ main ]
 
 jobs:
   build:
@@ -35,20 +31,26 @@ jobs:
         poetry config virtualenvs.create true --local
         poetry config virtualenvs.in-project true --local
 
+    # Add caching for Poetry dependencies
+    - name: Cache Poetry dependencies
+      uses: actions/cache@v3
+      with:
+        path: .venv
+        key: ${{ runner.os }}-poetry-3.9-${{ hashFiles('**/poetry.lock') }}
+        restore-keys: |
+          ${{ runner.os }}-poetry-3.9-
+
     - name: Install dependencies
       run: poetry install
 
     - name: Build package
       run: poetry build
 
-    - name: Upload artifact
+    - name: Upload build artifacts
       uses: actions/upload-artifact@v4
       with:
         name: dist
         path: dist/
-        retention-days: 1
 
     - name: Check code quality
-      run: |
-        poetry run black --check babeltron/
-        poetry run isort --check-only babeltron/
+      run: make lint

.github/workflows/main.yml

@@ -0,0 +1,19 @@
+name: Main Branch CI
+
+on:
+  push:
+    branches: [ main ]
+
+jobs:
+  build:
+    uses: ./.github/workflows/build.yml
+
+  test-matrix:
+    needs: build
+    strategy:
+      matrix:
+        python-version: ['3.9', '3.10', '3.11', '3.12', '3.13']
+    uses: ./.github/workflows/test.yml
+    with:
+      python-version: ${{ matrix.python-version }}
+      upload-coverage: ${{ matrix.python-version == '3.9' }}

.github/workflows/publish.yml

@@ -1,51 +1,31 @@
 name: Publish Package
 
 on:
-  push:
-    tags:
-      - 'v*'
+  release:
+    types: [created]
 
 jobs:
-  test:
-    uses: ./.github/workflows/test.yml
-
   build:
-    needs: test
     uses: ./.github/workflows/build.yml
 
-  publish:
+  test:
     needs: build
-    runs-on: ubuntu-latest
+    uses: ./.github/workflows/test.yml
+    with:
+      python-version: '3.9'
+      upload-coverage: true
 
+  publish:
+    needs: [build, test]
+    runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v4
-
-    - name: Set up Python
-      uses: actions/setup-python@v5
-      with:
-        python-version: '3.9'
-
-    - name: Install Poetry
-      uses: abatilo/actions-poetry@v4
-      with:
-        poetry-version: latest
-
-    - name: Setup virtualenv
-      run: |
-        poetry config virtualenvs.create true --local
-        poetry config virtualenvs.in-project true --local
-
-    - name: Configure Poetry
-      run: |
-        poetry config pypi-token.pypi ${{ secrets.PYPI_TOKEN }}
-
-#     - name: Download artifact
-#       uses: actions/download-artifact@v4
-#       with:
-#         name: ${{ needs.build.outputs.artifact-name }}
-#         path: dist
-
-#     - name: Publish to PyPI
-#       run: poetry publish
-#       env:
-#         POETRY_PYPI_TOKEN_PYPI: ${{ secrets.PYPI_TOKEN }}
+      - name: Download build artifacts
+        uses: actions/download-artifact@v3
+        with:
+          name: dist
+          path: dist/
+
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          password: ${{ secrets.PYPI_API_TOKEN }}

.github/workflows/pull_request.yml

@@ -0,0 +1,16 @@
+name: Pull Request Checks
+
+on:
+  pull_request:
+    branches: [ main ]
+
+jobs:
+  build:
+    uses: ./.github/workflows/build.yml
+
+  test:
+    needs: build
+    uses: ./.github/workflows/test.yml
+    with:
+      python-version: '3.9'
+      upload-coverage: true

.github/workflows/test.yml

@@ -2,25 +2,28 @@ name: Tests
 
 on:
   workflow_call:
-  pull_request:
-    branches: [ main ]
-  push:
-    branches: [ main ]
+    inputs:
+      python-version:
+        description: 'Python version to use'
+        required: false
+        default: '3.9'
+        type: string
+      upload-coverage:
+        description: 'Whether to upload coverage reports'
+        required: false
+        default: true
+        type: boolean
 
 jobs:
   test:
     runs-on: ubuntu-latest
-    strategy:
-      matrix:
-        python-version: ['3.9', '3.10', '3.11', '3.12', '3.13']
-
     steps:
     - uses: actions/checkout@v4
 
-    - name: Set up Python ${{ matrix.python-version }}
+    - name: Set up Python ${{ inputs.python-version }}
       uses: actions/setup-python@v5
       with:
-        python-version: ${{ matrix.python-version }}
+        python-version: ${{ inputs.python-version }}
 
     - name: Install Poetry
       uses: abatilo/actions-poetry@v4
@@ -32,17 +35,22 @@ jobs:
         poetry config virtualenvs.create true --local
         poetry config virtualenvs.in-project true --local
 
+    - name: Cache Poetry dependencies
+      uses: actions/cache@v3
+      with:
+        path: .venv
+        key: ${{ runner.os }}-poetry-${{ inputs.python-version }}-${{ hashFiles('**/poetry.lock') }}
+        restore-keys: |
+          ${{ runner.os }}-poetry-${{ inputs.python-version }}-
+
     - name: Install dependencies
       run: poetry install
 
     - name: Run tests
       run: make test
 
-    - name: Check code quality
-      run: |
-        make lint
-
     - name: Upload coverage reports to Codecov
+      if: inputs.upload-coverage
       uses: codecov/codecov-action@v3
       env:
         CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}

babeltron/scripts/download_models.py

@@ -0,0 +1,102 @@
+#!/usr/bin/env python3
+"""
+Script to download M2M-100 models for Babeltron.
+"""
+import argparse
+import os
+from pathlib import Path
+from typing import List, Optional, Union
+
+from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer
+
+VALID_MODEL_SIZES: List[str] = os.environ.get(
+    "BABELTRON_MODEL_SIZES", "418M,1.2B,12B"
+).split(",")
+DEFAULT_MODEL_SIZE: str = os.environ.get("BABELTRON_DEFAULT_MODEL_SIZE", "418M")
+DEFAULT_OUTPUT_DIR: Path = Path.home() / "models"
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Download M2M100 translation models")
+    parser.add_argument(
+        "--size",
+        choices=VALID_MODEL_SIZES,
+        default=DEFAULT_MODEL_SIZE,
+        help="Model size to download (418M, 1.2B, or 12B)",
+    )
+    parser.add_argument(
+        "--output-dir", default=DEFAULT_OUTPUT_DIR, help="Directory to save the model"
+    )
+    return parser.parse_args()
+
+
+def download_model(
+    model_size: str = DEFAULT_MODEL_SIZE,
+    output_dir: Optional[Union[str, Path]] = None,
+    show_progress: bool = True,
+) -> str:
+    """
+    Download M2M-100 model and tokenizer.
+
+    Args:
+        model_size (str): Size of the model to download (418M, 1.2B, or 12B)
+        output_dir (str or Path, optional): Directory to save the model to
+        show_progress (bool): Whether to show a progress bar
+
+    Returns:
+        str: Path to the downloaded model directory
+    """
+    if model_size not in VALID_MODEL_SIZES:
+        raise ValueError(f"Model size must be one of {VALID_MODEL_SIZES}")
+
+    model_name = f"facebook/m2m100_{model_size}"
+
+    if output_dir is None:
+        output_dir = DEFAULT_OUTPUT_DIR
+    else:
+        output_dir = Path(output_dir)
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    print(f"Downloading {model_name} model and tokenizer to {output_dir}...")
+
+    print("Downloading tokenizer...")
+    tokenizer = M2M100Tokenizer.from_pretrained(model_name)
+    tokenizer.save_pretrained(output_dir)
+
+    print("Downloading model (this may take a while)...")
+    model = M2M100ForConditionalGeneration.from_pretrained(model_name)
+    model.save_pretrained(output_dir)
+
+    print(f"Model and tokenizer successfully saved to {output_dir}")
+
+    return str(output_dir)
+
+
+def main():
+    args = parse_args()
+
+    try:
+        model_map = {"418M": "418M", "1.2B": "1.2B", "12B": "12B"}
+
+        model_size = model_map[args.size]
+        output_dir = args.output_dir
+
+        print(f"Downloading {args.size} model...")
+        print(
+            "This may take a while depending on your internet connection and the model size."
+        )
+
+        download_model(model_size=model_size, output_dir=output_dir)
+
+        print(f"Model successfully downloaded and saved to {output_dir}")
+
+    except Exception as e:
+        print(f"Error downloading model: {e}")
+        return 1
+
+    return 0
+
+
+if __name__ == "__main__":
+    main()