Update isoformer_tokenizer.py

isoformer_tokenizer.py

@@ -39,18 +39,22 @@ class IsoformerTokenizer(PreTrainedTokenizer):
         **kwargs
     ):
 
-        dna_hf_tokenizer = EsmTokenizer("dna_vocab_list.txt", model_max_length=196608)
+        pretrained_model_path = kwargs.get("pretrained_model_name_or_path", "")
+        dna_vocab_path = os.path.join(pretrained_model_path, "dna_vocab_list.txt")
+        rna_vocab_path = os.path.join(pretrained_model_path, "rna_vocab_list.txt")
+        protein_vocab_path = os.path.join(pretrained_model_path, "protein_vocab_list.txt")
+        
+        dna_hf_tokenizer = EsmTokenizer(dna_vocab_path, model_max_length=196608)
         dna_hf_tokenizer.eos_token = None  # Stops the tokenizer adding an EOS/SEP token at the end
         dna_hf_tokenizer.init_kwargs["eos_token"] = None  # Ensures it doesn't come back when reloading
         dna_hf_tokenizer.bos_token = None  # Stops the tokenizer adding an BOS/SEP token at the end
         dna_hf_tokenizer.init_kwargs["bos_token"] = None  # Ensures it doesn't come back when reloading
 
-
-        rna_hf_tokenizer = EsmTokenizer("rna_vocab_list.txt", model_max_length=1024)
+        rna_hf_tokenizer = EsmTokenizer(rna_vocab_path, model_max_length=1024)
         rna_hf_tokenizer.eos_token = None  # Stops the tokenizer adding an EOS/SEP token at the end
         rna_hf_tokenizer.init_kwargs["eos_token"] = None  # Ensures it doesn't come back when reloading
 
-        protein_hf_tokenizer = EsmTokenizer("protein_vocab_list.txt", model_max_length=1024)
+        protein_hf_tokenizer = EsmTokenizer(protein_vocab_path, model_max_length=1024)
         # protein_hf_tokenizer.eos_token = None  # Stops the tokenizer adding an EOS/SEP token at the end
         # protein_hf_tokenizer.init_kwargs["eos_token"] = None  # Ensures it doesn't come back when reloading
 
@@ -58,9 +62,9 @@ class IsoformerTokenizer(PreTrainedTokenizer):
         self.rna_tokenizer = rna_hf_tokenizer
         self.protein_tokenizer = protein_hf_tokenizer
 
-        self.dna_tokens = open("dna_vocab_list.txt", "r").read() .split("\n")
-        self.rna_tokens = open("rna_vocab_list.txt", "r").read() .split("\n")
-        self.protein_tokens = open("protein_vocab_list.txt", "r").read() .split("\n")
+        self.dna_tokens = open(dna_vocab_path, "r").read() .split("\n")
+        self.rna_tokens = open(rna_vocab_path, "r").read() .split("\n")
+        self.protein_tokens = open(protein_vocab_path, "r").read() .split("\n")
 
         super().__init__(**kwargs)