|
|
import os |
|
|
import json |
|
|
from transformers import PreTrainedTokenizer |
|
|
from huggingface_hub import hf_hub_download |
|
|
|
|
|
|
|
|
class MedicalcodeTokenizer(PreTrainedTokenizer): |
|
|
def __init__(self, vocab_file, code2tokens_file, embedding_file, **kwargs): |
|
|
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 |
|
|
|
|
|
super().__init__(**kwargs) |
|
|
|
|
|
def _tokenize(self, text): |
|
|
if text in self.code2tok: |
|
|
return [f"token_{t}" for t in self.code2tok[text]] |
|
|
return self._infer_and_register(text) |
|
|
return self._infer_and_register(text) |
|
|
|
|
|
def _convert_token_to_id(self, token): |
|
|
return self.vocab.get(token, self.vocab.get(self.unk_token, 0)) |
|
|
|
|
|
def _convert_id_to_token(self, idx): |
|
|
return self.ids_to_tokens.get(idx, self.unk_token) |
|
|
|
|
|
def embed(self, text): |
|
|
tokens = self._tokenize(text) |
|
|
if text in self.code2emb: |
|
|
return self.code2emb[text] |
|
|
return tokens |
|
|
|
|
|
def get_vocab(self): |
|
|
return self.vocab |
|
|
|
|
|
def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None): |
|
|
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): |
|
|
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): |
|
|
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.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": ["tokenizer", "MedicalcodeTokenizer"] |
|
|
} |
|
|
} |
|
|
with open(os.path.join(save_directory, "tokenizer_config.json"), "w") as f: |
|
|
json.dump(tokenizer_config, f, indent=2) |
|
|
|
|
|
@classmethod |
|
|
def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): |
|
|
vocab_file = hf_hub_download(pretrained_model_name_or_path, "vocab.json") |
|
|
code2tokens_file = hf_hub_download(pretrained_model_name_or_path, "code2tokens.json") |
|
|
embedding_file = hf_hub_download(pretrained_model_name_or_path, "code2embeddings.json") |
|
|
return cls(vocab_file=vocab_file, code2tokens_file=code2tokens_file, embedding_file=embedding_file, **kwargs) |
|
|
|
