quietflamingo/

dnabert2-no-flashattention

import torch
from transformers import AutoTokenizer, AutoModel
from transformers.models.bert.configuration_bert import BertConfig

device = torch.device("cuda")

tokenizer = AutoTokenizer.from_pretrained(
    "quietflamingo/dnabert2-fixed",
    trust_remote_code=True,
)

config = BertConfig.from_pretrained(
    "quietflamingo/dnabert2-fixed",
)

self.model = AutoModel.from_pretrained(
    "quietflamingo/dnabert2-fixed",
    config=config
)

dna = "ACGTAGCATCGGATCTATCTATCGACACTTGGTTATCGATCTACGAGCATCTCGTTAGC"
inputs = tokenizer(dna, return_tensors = 'pt')["input_ids"]

inputs = inputs.to(device)
model = model.to(device)

hidden_states = model(inputs)[0] # [1, sequence_length, 768]

embedding_mean = torch.mean(hidden_states[0], dim=0)
print(torch.mean(embedding_mean) # Outputs 0.0045, matches DNABERT2

embedding_max = torch.max(hidden_states[0], dim=0)[0]
print(torch.mean(embedding_max) # Outputs 0.2840, matches DNABERT2

@misc{zhou2023dnabert2,
      title={DNABERT-2: Efficient Foundation Model and Benchmark For Multi-Species Genome}, 
      author={Zhihan Zhou and Yanrong Ji and Weijian Li and Pratik Dutta and Ramana Davuluri and Han Liu},
      year={2023},
      eprint={2306.15006},
      archivePrefix={arXiv},
      primaryClass={q-bio.GN}
}

import torch
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("zhihan1996/DNABERT-2-117M", trust_remote_code=True)
model = AutoModel.from_pretrained("zhihan1996/DNABERT-2-117M", trust_remote_code=True)

dna = "ACGTAGCATCGGATCTATCTATCGACACTTGGTTATCGATCTACGAGCATCTCGTTAGC"
inputs = tokenizer(dna, return_tensors = 'pt')["input_ids"]
hidden_states = model(inputs)[0] # [1, sequence_length, 768]

# embedding with mean pooling
embedding_mean = torch.mean(hidden_states[0], dim=0)
print(embedding_mean.shape) # expect to be 768

# embedding with max pooling
embedding_max = torch.max(hidden_states[0], dim=0)[0]
print(embedding_max.shape) # expect to be 768