model_convert.py

import numpy as np
import h5py
import argparse
import shutil
from pathlib import Path
from typing import TypeVar

T = TypeVar("T")


def get_typed_item(parent: h5py.Group, key: str, expected_type: type[T]) -> T:
    item = parent[key]
    assert isinstance(item, expected_type), (
        f"Expected {expected_type.__name__}, got {type(item)}"
    )
    return item


def get_group(parent: h5py.Group, key: str) -> h5py.Group:
    return get_typed_item(parent, key, h5py.Group)


def get_dataset(parent: h5py.Group, key: str) -> h5py.Dataset:
    return get_typed_item(parent, key, h5py.Dataset)


def process_bias_dataset(group: h5py.Group) -> None:
    arr: np.ndarray = np.zeros(150)
    bias_dataset = get_dataset(group, "bias:0")
    bias_dataset.read_direct(arr)
    arr = arr.reshape(2, 75)
    del group["bias:0"]
    group.create_dataset("bias:0", data=arr, dtype="<f4")


def convert_file(file_path: Path) -> None:
    with h5py.File(file_path, "r+") as f:
        for layer_id in range(10, 15):
            process_bias_dataset(
                get_group(
                    get_group(
                        get_group(
                            get_group(f, "model_weights"), f"bidirectional_{layer_id}"
                        ),
                        f"bidirectional_{layer_id}",
                    ),
                    f"forward_cu_dnngru_{layer_id}",
                )
            )
            process_bias_dataset(
                get_group(
                    get_group(
                        get_group(
                            get_group(f, "model_weights"), f"bidirectional_{layer_id}"
                        ),
                        f"bidirectional_{layer_id}",
                    ),
                    f"backward_cu_dnngru_{layer_id}",
                )
            )


if __name__ == "__main__":
    parser: argparse.ArgumentParser = argparse.ArgumentParser()
    parser.add_argument("--file", type=Path, required=True)
    args: argparse.Namespace = parser.parse_args()

    original_file: Path = args.file
    copy_file: Path = original_file.with_suffix(".modified" + original_file.suffix)

    shutil.copy2(original_file, copy_file)
    convert_file(copy_file)