Skip to content

Nemo2 to hf

main()

Convert a NeMo2 Evo2 model checkpoint to a Hugging Face model checkpoint.

Source code in bionemo/evo2/utils/checkpoint/nemo2_to_hf.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
def main():
    """Convert a NeMo2 Evo2 model checkpoint to a Hugging Face model checkpoint."""
    args = parse_args()
    model_type = args.model_type
    if model_type == "hyena":
        raise ValueError("Hyena models are not supported for conversion to Hugging Face yet.")
    elif model_type == "mamba":
        exporter = HFNemotronExporter(args.model_path)
    elif model_type == "llama":
        exporter = HFLlamaExporter(args.model_path)
    else:
        raise ValueError(f"Invalid model type: {model_type}.")
    exporter.apply(args.output_dir)

parse_args()

Parse command-line arguments.

Source code in bionemo/evo2/utils/checkpoint/nemo2_to_hf.py 
22
23
24
25
26
27
28
29
30
def parse_args():
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model-type", type=str, required=True, help="Model type to convert.", choices=["hyena", "mamba", "llama"]
    )
    parser.add_argument("--model-path", type=str, required=True, help="Model path to convert.")
    parser.add_argument("--output-dir", type=str, required=True, help="Output directory path for the converted model.")
    return parser.parse_args()