Repository: allenai/OLMo
File: hf_olmo/convert_olmo_to_hf.py (Python)

Code before:
        conf["model"]["eos_token_id"] = 50279  # pyright: ignore

    om.save(conf, path)


def download_s3_directory(bucket_name: str, prefix: str, local_dir: str, ignore: str | None = None):
    # Create S3 client
    s3_client = boto3.client("s3")

    re_ignore = re.compile(ignore) if ignore else None

    # List objects within the given prefix
    paginator = s3_client.get_paginator("list_objects_v2")
    pages = paginator.paginate(Bucket=bucket_name, Prefix=prefix)

    # Create a list to hold all the files to download
    files_to_download = []
    for page in pages:
        for obj in page.get("Contents", []):
            if re_ignore and re_ignore.search(obj["Key"]):
                continue
            files_to_download.append(obj["Key"])

    # Initialize the progress bar
    with tqdm(total=len(files_to_download), desc="Downloading files") as pbar:
        for s3_key in files_to_download:
            # Construct the full local path
            local_file_path = os.path.join(local_dir, os.path.relpath(s3_key, prefix))
            local_file_dir = os.path.dirname(local_file_path)

            # Ensure local directory exists
            if not os.path.exists(local_file_dir):
                os.makedirs(local_file_dir)

            # Download the file
            s3_client.download_file(bucket_name, s3_key, local_file_path)

            # Update the progress bar
            pbar.update(1)


@contextmanager
def make_local_checkpoint(checkpoint_dir: str) -> Generator[str, None, None]:
    parsed_dir = urlparse(checkpoint_dir)

    assert parsed_dir.scheme in ["s3", ""], "Only s3 and local paths are supported."

    if os.path.exists(checkpoint_dir):
        yield checkpoint_dir
        return


Review comment (suggestion):
Does something like the following not work? It would get rid of the need for `pbar.update(1)`, and remove 1 level of nesting
```
for s3_key in tqdm(files_to_download, desc="Downloading files"):
```

Code after:

    conf["tokenizer"]["identifier"] = tokenizer_name_or_path  # pyright: ignore

    if tokenizer_name_or_path == "allenai/gpt-neox-olmo-dolma-v1_5" or tokenizer_name_or_path.endswith(
        "allenai_eleuther-ai-gpt-neox-20b-pii-special.json"
    ):
        conf["model"]["eos_token_id"] = 50279  # pyright: ignore

    om.save(conf, path)


def download_s3_directory(bucket_name: str, prefix: str, local_dir: str, ignore: str | None = None):
    # Create S3 client
    s3_client = _get_s3_client("s3")

    re_ignore = re.compile(ignore) if ignore else None

    # List objects within the given prefix
    paginator = s3_client.get_paginator("list_objects_v2")
    pages = paginator.paginate(Bucket=bucket_name, Prefix=prefix)

    # Create a list to hold all the files to download
    files_to_download = []
    for page in pages:
        for obj in page.get("Contents", []):
            if re_ignore and re_ignore.search(obj["Key"]):
                continue
            files_to_download.append(obj["Key"])

    # Initialize the progress bar
    for s3_key in tqdm(files_to_download, desc="Downloading files"):
        # Construct the full local path
        local_file_path = os.path.join(local_dir, os.path.relpath(s3_key, prefix))
        local_file_dir = os.path.dirname(local_file_path)

        # Ensure local directory exists
        if not os.path.exists(local_file_dir):
            os.makedirs(local_file_dir)

        # Download the file
        s3_client.download_file(bucket_name, s3_key, local_file_path)


def make_local_checkpoint(checkpoint_dir: str) -> str:
    parsed_dir = urlparse(checkpoint_dir)

    assert parsed_dir.scheme in ["s3", ""], "Only s3 and local paths are supported."

    if os.path.exists(checkpoint_dir):
        return checkpoint_dir