make_calib_data.py Source File

make_calib_data.py Source File#

Mobilint SDK qb Compiler: make_calib_data.py Source File
Mobilint SDK qb Compiler v0.12.0.0
MCS002-EN
 
 
from typing import Union, Callable, List, Optional
from shutil import rmtree
 
import os
import numpy as np
import random
import yaml
 
from tqdm import tqdm
from pathlib import Path
from .preprocess import preproc_builder
from .utils_calib import (
    list_np_files_in_txt,
    list_np_files_in_json,
    get_balanced_ids_coco,
    read_structure,
    make_calib_dict,
    save_calib_to_msgpack,
)
 
 
def get_yaml(yaml_path: str):
    assert (
        isinstance(yaml_path, str)
        and os.path.isfile(yaml_path)
        and yaml_path.endswith(".yaml")
    ), f"Got wrong yaml path: {yaml_path}."
    with open(yaml_path) as f:
        cfg = yaml.load(f, Loader=yaml.FullLoader)
    return cfg
 
 
def make_calib_man(
    pre_ftn,
    data_dir: str,
    save_dir: str,
    save_name: str,
    anno_json: str = None,
    file_format: str = "%012d.jpg",
    max_size: int = -1,
    remove_npy: bool = False,
    seed: int = 2023,
    save_calib_msg: bool = False,
    msg_path: str = None,
):
    """
    @brief From images and a user-provided preprocessing function, generate calibration NumPy files and a manifest.
 
    @param pre_ftn Callable. Pre-processing function that accepts an image path and returns a NumPy array.
    @param data_dir string. Directory containing the raw calibration images.
    @param save_dir string. Directory in which to store the generated NumPy files and metadata text file.
    @param save_name string. Basename for the saved assets. NumPy files are written to
    @c {save_dir}/{save_name}_npy and the manifest to @c {save_dir}/{save_name}.txt. Defaults to the basename of
    @p data_dir when omitted.
    @param anno_json string. Optional COCO-style annotation file. When provided, samples are randomly selected while
    maintaining class balance.
    @param file_format string. Filename format that uses @c image_idx (for example "%012d.jpg").
    @param max_size int. Maximum number of calibration samples. A value of -1 disables the limit.
    @param remove_npy bool. Remove existing NumPy outputs before generating new data.
    @param seed int. Random seed applied to sampling and class-balanced selection. Defaults to 2023.
    @param save_calib_msg bool. Persist the calibration metadata as an MSGpack file in addition to the manifest.
    @param msg_path string. Explicit MSGpack output path. When omitted, the path is auto-generated by dataset name and
    sample count.
    @return string. Path to the generated text file that lists the preprocessed NumPy samples.
 
    @note Calibration samples are written as float32 NumPy arrays derived from the @p pre_ftn output. When
    @p max_size is positive and the dataset is larger, a random subset is used. COCO annotations override this sampling
    to balance categories.
    """
    assert isinstance(
        pre_ftn, Callable
    ), f"Expected a function for pre_ftn, but got type={type(pre_ftn)}."
 
    if save_name is None:
        save_name = os.path.basename(data_dir)
    save_npy_dir = os.path.join(save_dir, f"{save_name}")
    save_txt_dir = os.path.join(save_dir, f"{save_name}.txt")
 
    # Remove pre-existing numpy files if remove_npy = True.
    if (remove_npy) and (os.path.isdir(save_dir)) and (os.path.isdir(save_npy_dir)):
        tmp_np_list = [x for x in os.listdir(save_npy_dir) if x.endswith("npy")]
        print(f"There are {len(tmp_np_list)} npy files in {save_npy_dir}.")
        rmtree(save_dir)
    os.makedirs(save_dir, exist_ok=True)
    os.makedirs(save_npy_dir, exist_ok=True)
 
    img_path_list = read_structure(data_dir)
    fname_list = []
    npy_list = []
    ids_list = []
 
    # Limit maximum size of calibration dataset.
    random.seed(seed)
    if anno_json is not None:
        img_path_list, ids_list = get_balanced_ids_coco(
            anno_json, img_path_list, max_size, fname_format=file_format, seed=seed
        )
        dataname = "coco"
    elif max_size > 0 and len(img_path_list) > max_size:
        img_path_list = random.sample(img_path_list, max_size)
        dataname = "any_sampled"
    else:
        dataname = "any_all"
 
    # Collect samples and make/save numpy files.
    for img_path in tqdm(img_path_list, desc="Making Numpy File"):
        out = pre_ftn(img_path)
        if out is None:
            continue
        if not isinstance(out, np.ndarray):
            raise AssertionError(f"Got unexpected tensor type: {type(out)}.", out)
 
        save_npy = out.astype(np.float32)
        fname = os.path.basename(img_path)
        fname_list.append(fname)
        npy_list.append(save_npy)
        img_name = Path(fname).stem
        save_path = os.path.join(save_npy_dir, img_name + ".npy")
        np.save(save_path, save_npy)
 
    if save_calib_msg:
        calib_dict = make_calib_dict(
            img_path_list, fname_list, npy_list, dataname, ids_list
        )
        save_calib_to_msgpack(calib_dict, msg_path)
 
    # make metadata text for numpy files.
    list_np_files_in_txt(save_npy_dir, save_txt_dir)
    return save_txt_dir
 
 
def make_calib(
    args_pre: Union[str, dict],
    data_dir: str,
    save_dir: str,
    save_name: str = None,
    anno_json: str = None,
    file_format: str = "%012d.jpg",
    max_size: int = -1,
    remove_npy: bool = False,
    seed: int = 2023,
    save_calib_msg: bool = False,
    msg_path: str = None,
) -> str:
    """
    @brief Create calibration NumPy files and a manifest using a YAML or dict preprocessing configuration.
 
    @param args_pre string or dict. Path to a YAML file or an in-memory dictionary describing the preprocessing
    pipeline. See the preproc_builder class for all supported operators and parameters.
    @param data_dir string. Directory of calibration images.
    @param save_dir string. Directory where the generated NumPy files and manifest text file are stored.
    @param save_name string. Optional basename for outputs. Defaults to the basename of @p data_dir.
    @param anno_json string. Optional COCO-format annotation file used to sample images with class balance.
    @param file_format string. Filename template built from @c image_idx (defaults to "%012d.jpg").
    @param max_size int. Maximum number of calibration samples to generate; -1 keeps all available items.
    @param remove_npy bool. Remove existing NumPy outputs before running preprocessing.
    @param seed int. Random seed used for sampling. Defaults to 2023.
    @param save_calib_msg bool. Save calibration metadata as MSGpack alongside NumPy and text outputs.
    @param msg_path string. Explicit MSGpack output path. Auto-generated from dataset name and sample count when
    omitted.
    @return string. Path to the manifest text file that lists the generated NumPy samples.
 
    @note The preprocessing pipeline is constructed through @c preproc_builder and executed via @c make_calib_man. Use
    the YAML schema to combine primitives such as GetImage, Resize, Normalize, Pad, CenterCrop, and SetOrder in a
    Pre-Order list.
    """
    if isinstance(args_pre, str):
        args_pre = get_yaml(args_pre)
    elif isinstance(args_pre, dict):
        pass
    else:
        raise TypeError(
            f"Expected yaml path (str) or cfg dictionary, but got type={type(args_pre)}."
        )
 
    pre_process_ftn = preproc_builder(args_pre, to_tensor=False)
    save_txt_dir = make_calib_man(
        pre_process_ftn,
        data_dir,
        save_dir,
        save_name,
        anno_json,
        file_format,
        max_size,
        remove_npy,
        seed,
        save_calib_msg,
        msg_path,
    )
    return save_txt_dir
 
 
def make_calib_llm(
    args_pre: Union[str, dict],
    dataset_name: str,
    subset_list: Union[List[str], str],
    text_field: str,
    split: str,
    save_dir: str,
    max_calib: int = 512,
    save_prefix: Optional[str] = None,
) -> List[str]:
    """Generate calibration embeddings for LLM models using Hugging Face datasets.
 
    Args:
        args_pre: YAML path or preprocessing configuration dictionary. Must set
            Datatype to "LLM".
        dataset_name: Name of the Hugging Face dataset to load.
        subset_list: List of dataset subsets (e.g. language splits) to iterate.
        split: HF dataset split to use.
        text_field: Key containing text in dataset entries.
        save_dir: Root directory where calibration samples are stored.
        max_calib: Maximum number of calibration samples per subset.
        save_prefix: Override prefix used in output directory naming. If None,
            derives from tokenizer directory name when available.
 
    Returns:
        List of paths to generated calibration txt files, one per subset.
    """
 
    from datasets import load_dataset  # Delay import to avoid optional dep cost
 
    if isinstance(args_pre, str):
        args_pre = get_yaml(args_pre)
    elif not isinstance(args_pre, dict):
        raise TypeError(
            f"Expected yaml path (str) or cfg dictionary, but got type={type(args_pre)}."
        )
 
    pre_process_ftn = preproc_builder(args_pre, to_tensor=False)
    if getattr(pre_process_ftn, "datatype", None) != "LLM":
        raise ValueError(
            "LLM calibration requires Datatype='LLM' in preprocessing config."
        )
 
    embedding_stage = next(
        (fnc for fnc in pre_process_ftn.preprocess if hasattr(fnc, "model_name")),
        None,
    )
    model_name = getattr(embedding_stage, "model_name", None)
    if save_prefix is None:
        save_prefix = model_name or "llm"
 
    os.makedirs(save_dir, exist_ok=True)
    generated_txt_paths: List[str] = []
 
    if isinstance(subset_list, str):
        subset_list = [subset_list]
 
    for subset_name in subset_list:
        dataset = load_dataset(dataset_name, subset_name, split=split)
        lang_suffix = subset_name.split(".")[-1] if "." in subset_name else subset_name
        target_dir_name = (
            f"{save_prefix}-{dataset_name.replace('/', '_')}-{lang_suffix}"
        )
        target_dir = os.path.join(save_dir, target_dir_name)
        os.makedirs(target_dir, exist_ok=True)
 
        cur_num_calib = 0
        for idx, sample in enumerate(dataset):
            if isinstance(sample, dict):
                if text_field not in sample:
                    raise KeyError(
                        f"Key '{text_field}' not found in dataset sample for subset {subset_name}."
                    )
                sample_input = sample[text_field]
            else:
                sample_input = sample
 
            try:
                embedded_text = pre_process_ftn(sample_input)
            except Exception as err:
                print(f"Skipping dataset item {idx} due to error: {err}")
                continue
 
            if embedded_text is None:
                continue
 
            if not isinstance(embedded_text, np.ndarray):
                raise AssertionError(
                    f"LLM preprocessing should return numpy.ndarray, got {type(embedded_text)}."
                )
 
            save_path = os.path.join(target_dir, f"inputs_embeds_{cur_num_calib}.npy")
            np.save(save_path, embedded_text)
            cur_num_calib += 1
 
            if cur_num_calib >= max_calib:
                break
 
        if cur_num_calib == 0:
            print(
                f"No calibration samples were generated for subset {subset_name}; check preprocessing constraints."
            )
            continue
 
        txt_path = os.path.join(save_dir, f"{target_dir_name}.txt")
        list_np_files_in_txt(target_dir, txt_path)
        generated_txt_paths.append(txt_path)
 
    return generated_txt_paths
 
 