Rename internal model references to CodonTranslatorModel

src/sampler.py

@@ -19,7 +19,7 @@ import torch.nn.functional as F
 import numpy as np
 from safetensors.torch import load_file
 
-from .models import CodonGPT
+from .models import CodonTranslatorModel
 from .tokenizer import CodonTokenizer
 
 logger = logging.getLogger(__name__)
@@ -130,7 +130,7 @@ class CodonSampler:
         # Model
         state = self._load_state_dict()
         arch = self._infer_arch_from_state_dict(state)
-        self.model = CodonGPT(
+        self.model = CodonTranslatorModel(
             vocab_size=self.V,
             hidden_size=int(arch["hidden_size"]),
             num_layers=int(arch["num_layers"]),