File tree Expand file tree Collapse file tree
sub-packages/bionemo-evo2/src/bionemo/evo2/run Expand file tree Collapse file tree Original file line number Diff line number Diff line change 3737from torch import Tensor
3838
3939from bionemo .evo2 .data .fasta_dataset import SimpleFastaDataset
40+ from bionemo .evo2 .models .llama import LLAMA_MODEL_OPTIONS
4041
4142# Add import for Mamba models
4243from bionemo .evo2 .models .mamba import MAMBA_MODEL_OPTIONS , MambaModel
43- from bionemo .evo2 .models .llama import LLAMA_MODEL_OPTIONS
4444from bionemo .llm .lightning import LightningPassthroughPredictionMixin
4545from bionemo .llm .model .biobert .lightning import get_batch_on_this_context_parallel_rank
4646from bionemo .llm .utils .callbacks import PredictionWriter
@@ -82,7 +82,9 @@ def parse_args():
8282 "--model-size" ,
8383 type = str ,
8484 default = "7b" ,
85- choices = sorted (list (HYENA_MODEL_OPTIONS .keys ()) + list (MAMBA_MODEL_OPTIONS .keys ()) + list (LLAMA_MODEL_OPTIONS .keys ())),
85+ choices = sorted (
86+ list (HYENA_MODEL_OPTIONS .keys ()) + list (MAMBA_MODEL_OPTIONS .keys ()) + list (LLAMA_MODEL_OPTIONS .keys ())
87+ ),
8688 help = "Model size to use. Defaults to '7b'." ,
8789 )
8890 # output args:
You can’t perform that action at this time.
0 commit comments