nucleotide finetuning about hyena-dna HOT 5 CLOSED

# experiment settings:
num_epochs = 80  # ~100 seems fine
max_length = 500  # max len of sequence of dataset (of what you want)
use_padding = True
data_path = './DNABERT_2/eval/GUE/EMP/H3'

batch_size = 256
learning_rate = 6e-4  # good default for Hyena
rc_aug = True  # reverse complement augmentation
add_eos = False  # add end of sentence token
weight_decay = 0.1

# for fine-tuning, only the 'tiny' model can fit on colab
pretrained_model_name = 'hyenadna-tiny-1k256d-seqlen'  # use None if training from scratch

# we need these for the decoder head, if using
use_head = True
n_classes = 1

# you can override with your own backbone config here if you want,
# otherwise we'll load the HF one by default
backbone_cfg = None

device = 'cuda' if torch.cuda.is_available() else 'cpu'
print("Using device:", device)

# instantiate the model (pretrained here)
if pretrained_model_name in ['hyenadna-tiny-1k256d-seqlen']:
    # use the pretrained Huggingface wrapper instead
    model = HyenaDNAPreTrainedModel.from_pretrained(
        './checkpoints',
        pretrained_model_name,
        download=False,
        config=backbone_cfg,
        device=device,
        use_head=use_head,
        n_classes=n_classes,
    )

# from scratch
else:
    model = HyenaDNAModel(**backbone_cfg, use_head=use_head, n_classes=n_classes)
 # create tokenizer
tokenizer = CharacterTokenizer(
    characters=['A', 'C', 'G', 'T', 'N'],  # add DNA characters, N is uncertain
    model_max_length=max_length + 2,  # to account for special tokens, like EOS
    add_special_tokens=False,  # we handle special tokens elsewhere
    padding_side='left', # since HyenaDNA is causal, we pad on the left
)

# create datasets
ds_train = SupervisedDataset(tokenizer=tokenizer,
                             data_path = os.path.join(data_path, "train.csv"),
                             kmer = -1)
ds_test = SupervisedDataset(tokenizer=tokenizer,
                            data_path= os.path.join(data_path,'test.csv'),
                            kmer = -1)
train_loader = DataLoader(ds_train, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(ds_test, batch_size=batch_size, shuffle=False)

# loss function
loss_fn = nn.BCEWithLogitsLoss()
#loss_fn = nn.MSELoss()
# create optimizer
optimizer = optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay)

model.to(device)

for epoch in range(num_epochs):
    train(model, device, train_loader, optimizer, epoch, loss_fn)
    test(model, device, test_loader, loss_fn)
    optimizer.step()

docker pull hyenadna/hyena-dna-nt6:latest 
docker run --gpus all -it -p80:3000 hyenadna/hyena-dna-nt6 /bin/bash