def run_infer_esm2()

in tutorials-and-examples/nvidia-bionemo/fine-tuning/job/finetuning.py [0:0]


def run_infer_esm2(work_dir, checkpoint_path, results_path):
    """Runs the infer_esm2 command using subprocess."""
    artificial_sequence_data = [
        "TLILGWSDKLGSLLNQLAIANESLGGGTIAVMAERDKEDMELDIGKMEFDFKGTSVI",
        "LYSGDHSTQGARFLRDLAENTGRAEYELLSLF",
        "GRFNVWLGGNESKIRQVLKAVKEIGVSPTLFAVYEKN",
        "DELTALGGLLHDIGKPVQRAGLYSGDHSTQGARFLRDLAENTGRAEYELLSLF",
        "KLGSLLNQLAIANESLGGGTIAVMAERDKEDMELDIGKMEFDFKGTSVI",
        "LFGAIGNAISAIHGQSAVEELVDAFVGGARISSAFPYSGDTYYLPKP",
        "LGGLLHDIGKPVQRAGLYSGDHSTQGARFLRDLAENTGRAEYELLSLF",
        "LYSGDHSTQGARFLRDLAENTGRAEYELLSLF",
        "ISAIHGQSAVEELVDAFVGGARISSAFPYSGDTYYLPKP",
        "SGSKASSDSQDANQCCTSCEDNAPATSYCVECSEPLCETCVEAHQRVKYTKDHTVRSTGPAKT",
    ]

    df = pd.DataFrame(artificial_sequence_data, columns=["sequences"])
    data_path = os.path.join(work_dir, "sequences.csv")
    df.to_csv(data_path, index=False)

    command = [
        "infer_esm2",
        "--checkpoint-path", checkpoint_path,
        "--data-path", data_path,
        "--results-path", results_path,
        "--config-class", "ESM2FineTuneSeqConfig",
    ]

    try:
        result = subprocess.run(command, check=True, capture_output=True, text=True)
        print("infer_esm2 output:\n", result.stdout)
        return True
    except subprocess.CalledProcessError as e:
        print(f"infer_esm2 failed with return code: {e.returncode}")
        print("stderr:\n", e.stderr)
        return False