Update README.md

README.md

@@ -22,24 +22,17 @@ ProkBERT-mini (kmer=6, shift=1) is part of the ProkBERT family of genomic langua
 The following example demonstrates how to use the ProkBERT-mini model for processing a DNA sequence:
 
 ```python
-from transformers import MegatronBertForMaskedLM
-from prokbert.prokbert_tokenizer import ProkBERTTokenizer
+from transformers import AutoTokenizer, AutoModelForMaskedLM
 
-# Tokenization parameters
-tokenization_parameters = {
-    'kmer': 6,
-    'shift': 1
-}
-# Initialize the tokenizer and model
-tokenizer = ProkBERTTokenizer(tokenization_params=tokenization_parameters, operation_space='sequence')
-model = MegatronBertForMaskedLM.from_pretrained("neuralbioinfo/prokbert-mini")
-# Example DNA sequence
-sequence = 'ATGTCCGCGGGACCT'
-# Tokenize the sequence
-inputs = tokenizer(sequence, return_tensors="pt")
-# Ensure that inputs have a batch dimension
-inputs = {key: value.unsqueeze(0) for key, value in inputs.items()}
-# Generate outputs from the model
+tokenizer = AutoTokenizer.from_pretrained("neuralbioinfo/prokbert-mini", trust_remote_code=True)
+model = AutoModel.from_pretrained("neuralbioinfo/prokbert-mini", trust_remote_code=True)
+
+segment = "ATGTCCGCGGGACCT"
+
+# Tokenize the input and return as PyTorch tensors
+inputs = tokenizer(segment, return_tensors="pt")
+
+# Pass the tokenized input to the model
 outputs = model(**inputs)
 ```