ztrain/io.py

# ztrain/io.py
# Copyright (c) 2024 Praxis Maldevide - cc-by-nc-4.0 granted

import os
from glob import glob

def flatten_index(model_paths : list[str], allow_list : list[str]):
    flat = []
    subtype = []
    index = {}
    ix = 0
    for g in sorted(model_paths):
        name = os.path.basename(g)
        if name in allow_list:
            index[name] = ix
            flat.append(name)
            if 'base' in g:
                subtype.append('base')
            elif 'instruct' in g:
                subtype.append('instruct')
            else:
                subtype.append('other')
        ix += 1
    return index, flat, subtype

def list_for_path(path: str, include_folders: list[str], search: str = "/**/*") -> tuple[list[str], list[str], list[str], dict[str, int]]:
    model_list = sorted([*[ f for f in glob(path + search)]])
    group_idx, model_names, subtypes = flatten_index(model_list, include_folders)
    groups = [[m] for m in model_names]
    return model_names, subtypes, model_list, group_idx