Function load_predictor_model

convert.py:1117–1134

Load MLP models for sparse FFN inference from directory.

(path: Path)

    return model_plus

def load_predictor_model(path: Path) -> ModelPlus:
    '''Load MLP models for sparse FFN inference from directory.'''
    assert path.is_dir(), f"MLP model path {path} is not a directory"
    
    first_model_path = path / "model_0.pt"
    assert first_model_path.resolve(), f"MLP model path {path} does not contain model_0.pt"

    model_paths = find_multifile_paths(first_model_path)
    models_plus: list[ModelPlus] = []
    for model_path in model_paths:
        # find number in model_path
        model_layer = int(re.search(r'model_(\d+).pt', str(model_path)).group(1))
        print(f"Loading MLP model file {model_path}")
        mlp_model = lazy_load_file(model_path)
        mlp_model.model = {f"model.layers.{model_layer}.{name}": tensor for name, tensor in mlp_model.model.items()}
        models_plus.append(mlp_model)

    return merge_multifile_models(models_plus)


def load_vocab(path: Path, vocabtype: str | None) -> Vocab: