Update README.md

README.md

@@ -12,10 +12,23 @@ This model is copied version of DNABERT-2 which removes the FlashAttention integ
 ```
 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)
-model = AutoModel.from_pretrained("quietflamingo/dnabert2-fixed", trust_remote_code=True)
+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"]