Update tokenizer.py

tokenizer.py

@@ -1,43 +1,24 @@
 import os
 import json
-import torch
 from transformers import PreTrainedTokenizer
 
 
 class MedicalcodeTokenizer(PreTrainedTokenizer):
     def __init__(self, vocab_file, code2tokens_file, embedding_file, **kwargs):
-        print(f"Loading vocab from: {vocab_file}")
-        print(f"Loading token map from: {code2tokens_file}")
-
         with open(vocab_file, "r") as f:
             self.vocab = json.load(f)
-
         with open(code2tokens_file, "r") as f:
             self.code2tok = json.load(f)
-        
         with open(embedding_file, 'r') as f:
             self.code2emb = json.load(f)
 
         self.ids_to_tokens = {v: k for k, v in self.vocab.items()}
-        self.tokens_to_ids = self.vocab  # alias
+        self.tokens_to_ids = self.vocab
 
         super().__init__(**kwargs)
 
-    # ---------- required interface ----------
     def _tokenize(self, text):
-        if text in self.code2tok:
-            return self.code2tok[text]
-        return self._infer_and_register(text)
-    
-    def embed(self, text):
-        tokens = self._tokenize(text)         # 先分词
-        if text in self.code2emb:
-            return self.code2emb[text]
-        return ids
-
-    def encode(self, text):
-        tokens = self._tokenize(text)         # 先分词
-        return tokens
+        return self.code2tok.get(text, [])
 
     def _convert_token_to_id(self, token):
         return self.vocab.get(token, self.vocab.get(self.unk_token, 0))
@@ -49,50 +30,29 @@ class MedicalcodeTokenizer(PreTrainedTokenizer):
         return self.vocab
 
     def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
-        """
-        Build model inputs from a sequence or a pair of sequences by adding special tokens.
-        token_ids_0: list[int] — input ids for the first sequence
-        token_ids_1: Optional[list[int]] — input ids for the second sequence (if any)
-        """
-        if token_ids_1 is None:
-            return token_ids_0
-        else:
-            return token_ids_0 + token_ids_1
-
+        return token_ids_0 if token_ids_1 is None else token_ids_0 + token_ids_1
 
     def get_special_tokens_mask(self, token_ids, already_has_special_tokens=False):
         return [0] * len(token_ids)
 
     def create_token_type_ids_from_sequences(self, token_ids_0, token_ids_1=None):
-        if token_ids_1 is None:
-            return [0] * len(token_ids_0)
-        return [0] * len(token_ids_0) + [1] * len(token_ids_1)
-
-    # ---------- dynamic extension interface (disabled) ----------
-    def _infer_and_register(self, code, code_desc="This is a medical code"):
-        raise NotImplementedError("Dynamic token generation is disabled in this version.")
-
-    # ---------- saving ----------
-    def save_updates(self, out_vocab="vocab.json", out_map="code2tokens.json"):
-        json.dump(self.vocab, open(out_vocab, "w"), indent=2)
-        json.dump(self.code2tok, open(out_map, "w"), indent=2)
+        return [0] * len(token_ids_0) if token_ids_1 is None else [0] * len(token_ids_0) + [1] * len(token_ids_1)
 
     def save_pretrained(self, save_directory):
-        import os, json
         os.makedirs(save_directory, exist_ok=True)
         with open(os.path.join(save_directory, "vocab.json"), "w") as f:
             json.dump(self.vocab, f, indent=2)
         with open(os.path.join(save_directory, "code2tokens.json"), "w") as f:
             json.dump(self.code2tok, f, indent=2)
         with open(os.path.join(save_directory, "code2embeddings.json"), "w") as f:
-            json.dump(self.code2embed, f, indent=2)
+            json.dump(self.code2emb, f, indent=2)
         tokenizer_config = {
             "tokenizer_class": "tokenizer.MedicalcodeTokenizer",
             "vocab_file": "vocab.json",
             "code2tokens_file": "code2tokens.json",
             "code2embedding_file": "code2embeddings.json",
             "auto_map": {
-                "AutoTokenizer": "MedTok"
+                "AutoTokenizer": ["tokenizer", "MedicalcodeTokenizer"]
             }
         }
         with open(os.path.join(save_directory, "tokenizer_config.json"), "w") as f:
@@ -100,8 +60,7 @@ class MedicalcodeTokenizer(PreTrainedTokenizer):
 
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
-        import os, json
         vocab_file = os.path.join(pretrained_model_name_or_path, "vocab.json")
         code2tokens_file = os.path.join(pretrained_model_name_or_path, "code2tokens.json")
         embedding_file = os.path.join(pretrained_model_name_or_path, "code2embeddings.json")
-        return cls(vocab_file=vocab_file, code2tokens_file=code2tokens_file, embedding_file=embedding_file, **kwargs)
+        return cls(vocab_file=vocab_file, code2tokens_file=code2tokens_file, embedding_file=embedding_file, **kwargs)