chore(ci): fix test path for scripts

Author: hspedro

.github/workflows/build.yml

@@ -35,20 +35,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
-      uses: actions/upload-artifact@v4
+    - name: Upload build artifacts
+      uses: actions/upload-artifact@v3
       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/test.yml

@@ -32,17 +32,25 @@ 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-${{ matrix.python-version }}-${{ hashFiles('**/poetry.lock') }}
+        restore-keys: |
+          ${{ runner.os }}-poetry-${{ matrix.python-version }}-
+
     - name: Install dependencies
       run: poetry install
 
     - name: Run tests
       run: make test
 
     - name: Check code quality
-      run: |
-        make lint
+      run: make lint
 
     - name: Upload coverage reports to Codecov
+      if: matrix.python-version == '3.9'
       uses: codecov/codecov-action@v3
       env:
         CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}

babeltron/scripts/__init__.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()