Remove all connection to HF in CI (#1786)

Author: dakinggg

Diff for: .gitignore

@@ -157,3 +157,4 @@ notebooks/
 **/mlruns/*
 **/tokenizer-save-dir-*/**
 **/.downloaded_finetuning/
+tests/assets/tokenizers/

Diff for: tests/fixtures/models.py

@@ -2,11 +2,14 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import copy
+import hashlib
+import os
+import zipfile
 from typing import Any, Callable
 
 import pytest
+import requests
 from pytest import fixture
-from tenacity import retry, stop_after_attempt, wait_fixed
 from transformers import PreTrainedTokenizerBase
 
 from llmfoundry.models.hf.hf_causal_lm import ComposerHFCausalLM
@@ -195,109 +198,81 @@ def tiny_bert_config_helper():
     return config_object
 
 
-## TOKENIZER HELPERS ##
-@retry(
-    wait=wait_fixed(5),
-    stop=stop_after_attempt(1),
-)
-def tiny_gpt2_tokenizer_helper(add_pad: bool = False):
-    transformers = pytest.importorskip('transformers')
-
-    hf_tokenizer = transformers.AutoTokenizer.from_pretrained('gpt2')
-
-    if add_pad:
-        hf_tokenizer.add_special_tokens({'pad_token': '[PAD]'})
-    return hf_tokenizer
+def assets_path():
+    rank = os.environ.get('RANK', '0')
+    folder_name = 'tokenizers' + (f'_{rank}' if rank != '0' else '')
+    return os.path.join(
+        os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+        'assets',
+        folder_name,
+    )
 
 
-@retry(
-    wait=wait_fixed(5),
-    stop=stop_after_attempt(1),
-)
-def tiny_llama_tokenizer_helper():
-    transformers = pytest.importorskip('transformers')
+@pytest.fixture(scope='session')
+def tokenizers_assets():
+    download_tokenizers_files()
 
-    hf_tokenizer = transformers.AutoTokenizer.from_pretrained(
-        'huggyllama/llama-7b',
-        use_fast=False,
-    )
-    return hf_tokenizer
 
+def download_tokenizers_files():
+    """Download the tokenizers assets.
 
-@retry(
-    wait=wait_fixed(5),
-    stop=stop_after_attempt(1),
-)
-def tiny_codellama_tokenizer_helper():
-    transformers = pytest.importorskip('transformers')
+    We download from github, because downloading from HF directly is flaky and gets rate limited easily.
 
-    hf_tokenizer = transformers.AutoTokenizer.from_pretrained(
-        'codellama/CodeLlama-7b-hf',
-    )
-    return hf_tokenizer
+    Raises:
+        ValueError: If the checksum of the downloaded file does not match the expected checksum.
+    """
+    # Define paths
+    tokenizers_dir = assets_path()
 
+    if os.path.exists(tokenizers_dir):
+        return
 
-@retry(
-    wait=wait_fixed(5),
-    stop=stop_after_attempt(1),
-)
-def tiny_neox_tokenizer_helper():
-    transformers = pytest.importorskip('transformers')
+    # Create assets directory if it doesn't exist
+    os.makedirs(tokenizers_dir, exist_ok=True)
 
-    hf_tokenizer = transformers.AutoTokenizer.from_pretrained(
-        'EleutherAI/gpt-neox-20b',
-        model_max_length=2048,
-    )
-    return hf_tokenizer
+    # URL for the tokenizers.zip file
+    url = 'https://github.com/mosaicml/ci-testing/releases/download/tokenizers/tokenizers.zip'
+    expected_checksum = '12dc1f254270582f7806588f1f1d47945590c5b42dee28925e5dab95f2d08075'
 
+    # Download the zip file
+    response = requests.get(url, stream=True)
+    response.raise_for_status()
 
-@retry(
-    wait=wait_fixed(5),
-    stop=stop_after_attempt(1),
-)
-def tiny_t5_tokenizer_helper():
-    transformers = pytest.importorskip('transformers')
+    zip_path = os.path.join(tokenizers_dir, 'tokenizers.zip')
 
-    hf_tokenizer = transformers.AutoTokenizer.from_pretrained('t5-base',)
-    return hf_tokenizer
+    # Check the checksum
+    checksum = hashlib.sha256(response.content).hexdigest()
+    if checksum != expected_checksum:
+        raise ValueError(
+            f'Checksum mismatch: expected {expected_checksum}, got {checksum}',
+        )
 
+    with open(zip_path, 'wb') as f:
+        for chunk in response.iter_content(chunk_size=8192):
+            f.write(chunk)
 
-@retry(
-    wait=wait_fixed(5),
-    stop=stop_after_attempt(1),
-)
-def tiny_bert_tokenizer_helper():
-    transformers = pytest.importorskip('transformers')
+    # Extract the zip file
+    print(f'Extracting tokenizers.zip to {tokenizers_dir}')
+    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+        zip_ref.extractall(tokenizers_dir)
 
-    return transformers.AutoTokenizer.from_pretrained(
-        'google-bert/bert-base-uncased',
-    )
+    # Optionally remove the zip file after extraction
+    os.remove(zip_path)
 
 
-@retry(
-    wait=wait_fixed(5),
-    stop=stop_after_attempt(1),
-)
-def tiny_mpt_tokenizer_helper():
+## TOKENIZER HELPERS ##
+def assets_tokenizer_helper(name: str):
+    """Load a tokenizer from the assets directory."""
     transformers = pytest.importorskip('transformers')
 
-    return transformers.AutoTokenizer.from_pretrained(
-        'mosaicml/mpt-7b',
-        model_max_length=2048,
-    )
-
+    download_tokenizers_files()
 
-@retry(
-    wait=wait_fixed(5),
-    stop=stop_after_attempt(1),
-)
-def tiny_mpt_chat_tokenizer_helper():
-    transformers = pytest.importorskip('transformers')
+    assets_dir = assets_path()
+    tokenizer_path = os.path.join(assets_dir, name)
 
-    return transformers.AutoTokenizer.from_pretrained(
-        'mosaicml/mpt-7b-8k-chat',
-        model_max_length=2048,
-    )
+    # Load the tokenizer
+    hf_tokenizer = transformers.AutoTokenizer.from_pretrained(tokenizer_path)
+    return hf_tokenizer
 
 
 ## SESSION MODELS ##
@@ -336,48 +311,50 @@ def _session_tiny_bert_config():  # type: ignore
 
 ## SESSION TOKENIZERS ##
 @pytest.fixture(scope='session')
-def _session_tiny_gpt2_tokenizer():  # type: ignore
-    return tiny_gpt2_tokenizer_helper()
+def _session_tiny_gpt2_tokenizer(tokenizers_assets):  # type: ignore
+    return assets_tokenizer_helper('gpt2')
 
 
 @pytest.fixture(scope='session')
-def _session_tiny_gpt2_with_pad_tokenizer():  # type: ignore
-    return tiny_gpt2_tokenizer_helper(add_pad=True)
+def _session_tiny_gpt2_with_pad_tokenizer(tokenizers_assets):  # type: ignore
+    tokenizer = assets_tokenizer_helper('gpt2')
+    tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+    return tokenizer
 
 
 @pytest.fixture(scope='session')
-def _session_tiny_llama_tokenizer():  # type: ignore
-    return tiny_llama_tokenizer_helper()
+def _session_tiny_llama_tokenizer(tokenizers_assets):  # type: ignore
+    return assets_tokenizer_helper('llama')
 
 
 @pytest.fixture(scope='session')
-def _session_tiny_codellama_tokenizer():  # type: ignore
-    return tiny_codellama_tokenizer_helper()
+def _session_tiny_codellama_tokenizer(tokenizers_assets):  # type: ignore
+    return assets_tokenizer_helper('codellama')
 
 
 @pytest.fixture(scope='session')
-def _session_tiny_neox_tokenizer():  # type: ignore
-    return tiny_neox_tokenizer_helper()
+def _session_tiny_neox_tokenizer(tokenizers_assets):  # type: ignore
+    return assets_tokenizer_helper('neox')
 
 
 @pytest.fixture(scope='session')
-def _session_tiny_t5_tokenizer():  # type: ignore
-    return tiny_t5_tokenizer_helper()
+def _session_tiny_t5_tokenizer(tokenizers_assets):  # type: ignore
+    return assets_tokenizer_helper('t5')
 
 
 @pytest.fixture(scope='session')
-def _session_tiny_bert_tokenizer():  # type: ignore
-    return tiny_bert_tokenizer_helper()
+def _session_tiny_bert_tokenizer(tokenizers_assets):  # type: ignore
+    return assets_tokenizer_helper('bertt')
 
 
 @pytest.fixture(scope='session')
-def _session_tiny_mpt_tokenizer():  # type: ignore
-    return tiny_mpt_tokenizer_helper()
+def _session_tiny_mpt_tokenizer(tokenizers_assets):  # type: ignore
+    return assets_tokenizer_helper('mptt')
 
 
 @pytest.fixture(scope='session')
-def _session_tiny_mpt_chat_tokenizer():  # type: ignore
-    return tiny_mpt_chat_tokenizer_helper()
+def _session_tiny_mpt_chat_tokenizer(tokenizers_assets):  # type: ignore
+    return assets_tokenizer_helper('mptct')
 
 
 ## MODEL FIXTURES ##

Diff for: tests/tokenizers/test_tokenizer.py

@@ -4,7 +4,7 @@
 import pytest
 import torch
 from omegaconf import OmegaConf as om
-from transformers import AutoTokenizer, PreTrainedTokenizerBase
+from transformers import AutoTokenizer
 
 from llmfoundry.data.finetuning.tasks import _DEFAULT_CHAT_TEMPLATE
 from llmfoundry.tokenizers.utils import get_date_string
@@ -16,74 +16,6 @@ def get_config(conf_path: str = 'scripts/train/yamls/pretrain/mpt-125m.yaml'):
     return test_cfg
 
 
-def test_load_tokenizer(tiny_neox_tokenizer: PreTrainedTokenizerBase):
-    test_cfg = get_config(
-        conf_path='scripts/train/yamls/pretrain/mpt-125m.yaml',
-    )
-    truncation = True
-    padding = 'max_length'
-
-    resolved_om_tokenizer_config = om.to_container(
-        test_cfg.tokenizer,
-        resolve=True,
-    )
-    tokenizer_kwargs = resolved_om_tokenizer_config.get(  # type: ignore
-        'kwargs', {})
-    tokenizer_name = resolved_om_tokenizer_config['name']  # type: ignore
-    tokenizer = tiny_neox_tokenizer
-    tokenizer.pad_token = tokenizer.eos_token
-    assert tokenizer.vocab_size == 50254
-    assert tokenizer.name_or_path == 'EleutherAI/gpt-neox-20b'
-
-    # HuggingFace overrides model_max_length, so this check would fail. We explicitly reset the
-    # model_max_length in ComposerMPTCausalLM
-    # assert tokenizer.model_max_length == resolved_om_tokenizer_config['kwargs']['model_max_length']
-
-    in_str = 'hello\n\nhello'
-    out_token_key = [25521, 187, 187, 25521]
-
-    # test explicitly call tokenizer
-    out = tokenizer.encode(in_str)
-    assert out == out_token_key
-
-    # tokenizer __call__
-    out = tokenizer(in_str)['input_ids']
-    assert out == out_token_key
-
-    # tokenizer  __call__ with kwargs
-    padded_tokenize = tokenizer(
-        in_str,
-        truncation=truncation,
-        padding=padding,
-        max_length=tokenizer.model_max_length,
-    )['input_ids']
-    out_pad_tokens = out_token_key + [0] * (tokenizer.model_max_length - 4)
-    assert padded_tokenize == out_pad_tokens
-
-    # wrapper class __call__
-    out = tokenizer(in_str)['input_ids']
-    assert out == out_token_key
-
-    # wrapper class __call__ with kwargs
-    padded_tokenize = tokenizer(
-        in_str,
-        truncation=truncation,
-        padding=padding,
-        max_length=tokenizer.model_max_length,
-    )['input_ids']
-    assert padded_tokenize == out_pad_tokens
-
-    # check attn mask
-    attention_mask = tokenizer(
-        in_str,
-        truncation=truncation,
-        padding=padding,
-        max_length=tokenizer.model_max_length,
-    )['attention_mask']
-    attn_mask_key = [1, 1, 1, 1] + [0] * (tokenizer.model_max_length - 4)
-    assert attention_mask == attn_mask_key
-
-
 @pytest.mark.parametrize(
     'tokenizer_name',
     [