Fix loading instructions: use AutoModelForMaskedLM and trust_remote_code=True

README.md

@@ -16,15 +16,17 @@ This repository contains jointly trained Nucleotide Transformer (NT) and ESM2 mo
 
 ## Usage
 
+**IMPORTANT**: Both models are masked language models. The DNA model uses the Nucleotide Transformer architecture which requires `trust_remote_code=True`.
+
 ```python
-from transformers import AutoModel, AutoTokenizer
+from transformers import AutoModelForMaskedLM, AutoTokenizer
 
-# Load DNA model
-dna_model = AutoModel.from_pretrained("vsubasri/joint-nt-esm2-transcript-coding", subfolder="dna")
-dna_tokenizer = AutoTokenizer.from_pretrained("vsubasri/joint-nt-esm2-transcript-coding", subfolder="dna")
+# Load DNA model - requires trust_remote_code for custom NT architecture
+dna_model = AutoModelForMaskedLM.from_pretrained("vsubasri/joint-nt-esm2-transcript-coding", subfolder="dna", trust_remote_code=True)
+dna_tokenizer = AutoTokenizer.from_pretrained("vsubasri/joint-nt-esm2-transcript-coding", subfolder="dna", trust_remote_code=True)
 
 # Load protein model
-protein_model = AutoModel.from_pretrained("vsubasri/joint-nt-esm2-transcript-coding", subfolder="protein")  
+protein_model = AutoModelForMaskedLM.from_pretrained("vsubasri/joint-nt-esm2-transcript-coding", subfolder="protein")
 protein_tokenizer = AutoTokenizer.from_pretrained("vsubasri/joint-nt-esm2-transcript-coding", subfolder="protein")
 
 # Example joint usage