translation demo for ascend 910b4

Dockerfile

@@ -0,0 +1,11 @@
+FROM swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.2.rc1-910b-ubuntu22.04-py3.11
+
+WORKDIR /workspace
+
+COPY requirements.txt /workspace
+RUN pip install --no-cache-dir -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
+
+ADD . /workspace
+
+EXPOSE 80
+CMD ["sh", "-c", "python3 fastapi_translate.py"]

README.md

@@ -1,3 +1,45 @@
 # enginex-ascend-910-translation
 
-行于【昇腾-910】系列算力卡的【翻译】引擎，基于 transformer 架构，支持 NLLB-200、SMaLL-100 等最新流行翻译模型
+行于【昇腾-910】系列算力卡的【翻译】引擎，基于 transformer 架构，支持 NLLB-200、SMaLL-100 等最新流行翻译模型
+
+# translation-transformers
+## Quickstart
+```shell
+#构建docker镜像
+docker build . -t ascend_910_translation
+
+#运行docker容器
+docker run -it -e ASCEND_VISIBLE_DEVICES=0 --device /dev/davinci0:/dev/davinci0 --device /dev/davinci_manager --device /dev/devmm_svm --device /dev/hisi_hdc -v /mnt:/mnt -v /usr/local/dcmi:/usr/local/dcmi -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info -v /etc/ascend_install.info:/etc/ascend_install.info --privileged -p 10078:80 --device=/dev/iluvatar0:/dev/iluvatar0 -e MODEL_TYPE=opus_mt -e MODEL_NAME=moxying/opus-mt-zh-en --name ascend_910_translation_test ascend_910_translation
+```
+等待模型下载完成，出现以下日志时，代表服务启动成功
+```shell
+INFO:     Application startup complete.
+INFO:     Uvicorn running on http://0.0.0.0:80 (Press CTRL+C to quit)
+```
+执行测试程序
+```shell
+python3 test.py
+```
+测试程序执行结果
+```
+Succeed!
+Response: [{'translations': [{'origin_text': '生活就像一块巧克力', 'translated': 'Life is like a piece of chocolate.'}, {'origin_text': '你来自哪里', 'translated': 'Where are you from?'}, {'origin_text': '你吃饭了吗', 'translated': 'Have you eaten yet?'}]}]
+```
+停止docker容器
+```
+docker stop ascend_910_translation_test 
+```
+## 模型支持
+在Quickstart中运行容器时，通过环境变量的方式，指定模型的类型和具体的模型名称，即：
+```
+-e MODEL_TYPE=opus_mt -e MODEL_NAME=moxying/opus-mt-zh-en
+```
+目前支持以下几种配置:
+| MODEL_TYPE |            MODEL_NAME                   |
+| ---------- | --------------------------------------- |
+| nllb200    | facebook/nllb-200-distilled-600M        |
+| small100   | aiyueqi/alirezamsh_small100             |
+| mbart      | facebook/mbart-large-50-many-to-many-mmt|
+| opus_mt    | moxying/opus-mt-zh-en                   |
+
+其中，MODEL_TYPE代表模型类型，必须为以上表格列举之一；MODEL_NAME是modelscope上面能够拉取到的模型名称，需要和MODEL_TYPE对应

fastapi_translate.py

@@ -0,0 +1,140 @@
+import os
+from fastapi import FastAPI, Query
+from fastapi.responses import PlainTextResponse
+from pydantic import BaseModel
+
+from typing import List, Any
+
+import uvicorn
+
+from modelscope import snapshot_download
+
+import torch_npu
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+import logger
+log = logger.get_logger(__file__)
+
+app = FastAPI()
+status = "Running"
+translator = None
+device = None
+model_type = None
+
+MODEL_TYPE = ("nllb200", "small100", "mbart", "opus_mt")
+MODEL_DIR = "/workspace/model"
+
+class TranslateRequest(BaseModel):
+    Text: str
+
+@app.on_event("startup")
+def load_model():
+    log.info("loading model")
+    global status, translator, device, model_type
+
+    model_type = extract_model_type()
+    log.info(f"model_type={model_type}")
+
+    fetch_model()
+
+    tokenizer, model = get_tokenizer_model()
+    #log.info(f"tokenizer={tokenizer}, model={model}")
+    model = model.to("npu")
+
+    translator = pipeline(task="translation", model=model, tokenizer=tokenizer, device="npu", use_cache=True)
+    warm_up()
+
+    status = "Success"
+    log.info("model loaded successfully")
+
+def fetch_model():
+    mn = os.environ.get("MODEL_NAME", "")
+    log.info(f"model_name={mn}")
+
+    os.makedirs(os.path.dirname(MODEL_DIR), exist_ok=True)
+    snapshot_download(mn, local_dir=MODEL_DIR)
+
+def translator_helper(text):
+    source_lang = "zh"
+    target_lang = "en"
+
+    if model_type == "nllb200":
+        source_lang = "zho_Hans"
+        target_lang = "eng_Latn"
+    if model_type == "mbart":
+        source_lang = "zh_CN"
+        target_lang = "en_XX"
+    if model_type == "opus_mt":
+        source_lang = "eng"
+        target_lang = "zho"
+
+    output = translator(text, src_lang=source_lang, tgt_lang=target_lang)
+    log.info(f"model_type={model_type}, src_lang={source_lang}, tgt_lang={target_lang}, output={output}")
+
+    return output
+
+def get_tokenizer_model():
+    if model_type == "small100":
+        from tokenization_small100 import SMALL100Tokenizer
+        tokenizer = SMALL100Tokenizer.from_pretrained(MODEL_DIR)
+        model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_DIR)
+    else:
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
+        model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_DIR)
+
+    return tokenizer, model
+
+def extract_model_type():
+    mt = os.environ.get("MODEL_TYPE", "")
+    log.info(f"model_type_input={mt}")
+
+    model = mt.lower()
+    if model not in MODEL_TYPE:
+        log.error(f"model_type {model} is not supported")
+        os._exit(1)
+
+    return model
+
+def warm_up():
+    log.info("warming up...")
+
+    warmup_test = translator_helper("今天的天气非常好")
+    log.info(f"warm up completed! model_type={model_type}, response={warmup_test}")
+
+    return warmup_test
+
+@app.get("/v1/get_status")
+async def get_status():
+    ret = {
+        "data": {
+            "status": status
+        }
+    }
+    return ret
+
+@app.post("/v1/translate")
+async def translate(
+    payload: List[TranslateRequest],
+):
+    if not payload:
+        return PlainTextResponse(text="Information missing", status_code=400)
+    results = []
+    texts = []
+    for trans_request in payload:
+        translations = []
+        texts.append(trans_request.Text)
+    
+    outputs = translator_helper(texts)
+    for i in range(0, len(texts)):
+        translations.append({
+            "origin_text": texts[i],
+            "translated": outputs[i]['translation_text']
+        })
+       
+    results.append({
+        "translations": translations
+    })
+    return results
+
+if __name__ == '__main__':
+    uvicorn.run("fastapi_translate:app", host="0.0.0.0", port=80, workers=1, access_log=False)
+

logger.py

@@ -0,0 +1,12 @@
+# -*- coding: utf-8 -*-
+import logging
+import os
+
+logging.basicConfig(
+    format="%(asctime)s %(name)-12s %(levelname)-4s %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+    level=os.environ.get("LOGLEVEL", "INFO"),
+)
+
+def get_logger(file):
+    return logging.getLogger(file)

requirements.txt

@@ -0,0 +1,9 @@
+fastapi
+uvicorn
+sentencepiece==0.2.0
+sacremoses==0.1.1
+protobuf==3.20.3
+torch==2.6.0
+torch-npu==2.6.0
+transformers==4.39.3
+modelscope

test.py

@@ -0,0 +1,30 @@
+import requests
+url = "http://127.0.0.1:10078/v1/translate"
+
+body = [
+    {"Text": "生活就像一块巧克力"},
+    {"Text": "你来自哪里"},
+    {"Text": "你吃饭了吗"},
+]
+
+headers = {
+    "Content-Type": "application/json",
+    "Accept": "application/json"
+}
+
+try:
+    response = requests.post(
+        url,
+        json=body,
+        headers=headers
+    )
+
+    if response.status_code == 200:
+        print("Succeed!")
+        print("Response:", response.json())
+    else:
+        print(f"Failed，code: {response.status_code}")
+        print("Error detail:", response.text)
+
+except requests.exceptions.RequestException as e:
+    print("request error:", str(e))

tokenization_small100.py

@@ -0,0 +1,366 @@
+# Copyright (c) 2022 Idiap Research Institute, http://www.idiap.ch/
+# Written by Alireza Mohammadshahi <alireza.mohammadshahi@idiap.ch>
+# This is a modified version of https://github.com/huggingface/transformers/blob/main/src/transformers/models/m2m_100/tokenization_m2m_100.py 
+# which owns by Fariseq Authors and The HuggingFace Inc. team.
+#
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tokenization classes for SMALL100."""
+import json
+import os
+from pathlib import Path
+from shutil import copyfile
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import sentencepiece
+
+from transformers.tokenization_utils import BatchEncoding, PreTrainedTokenizer
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+SPIECE_UNDERLINE = "▁"
+
+VOCAB_FILES_NAMES = {
+    "vocab_file": "vocab.json",
+    "spm_file": "sentencepiece.bpe.model",
+    "tokenizer_config_file": "tokenizer_config.json",
+}
+
+PRETRAINED_VOCAB_FILES_MAP = {
+    "vocab_file": {
+        "alirezamsh/small100": "https://huggingface.co/alirezamsh/small100/resolve/main/vocab.json",
+    },
+    "spm_file": {
+        "alirezamsh/small100": "https://huggingface.co/alirezamsh/small100/resolve/main/sentencepiece.bpe.model",
+    },
+    "tokenizer_config_file": {
+        "alirezamsh/small100": "https://huggingface.co/alirezamsh/small100/resolve/main/tokenizer_config.json",
+    },
+}
+
+PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
+    "alirezamsh/small100": 1024,
+}
+
+# fmt: off
+FAIRSEQ_LANGUAGE_CODES = {
+    "m2m100": ["af", "am", "ar", "ast", "az", "ba", "be", "bg", "bn", "br", "bs", "ca", "ceb", "cs", "cy", "da", "de", "el", "en", "es", "et", "fa", "ff", "fi", "fr", "fy", "ga", "gd", "gl", "gu", "ha", "he", "hi", "hr", "ht", "hu", "hy", "id", "ig", "ilo", "is", "it", "ja", "jv", "ka", "kk", "km", "kn", "ko", "lb", "lg", "ln", "lo", "lt", "lv", "mg", "mk", "ml", "mn", "mr", "ms", "my", "ne", "nl", "no", "ns", "oc", "or", "pa", "pl", "ps", "pt", "ro", "ru", "sd", "si", "sk", "sl", "so", "sq", "sr", "ss", "su", "sv", "sw", "ta", "th", "tl", "tn", "tr", "uk", "ur", "uz", "vi", "wo", "xh", "yi", "yo", "zh", "zu"]
+}
+# fmt: on
+
+
+class SMALL100Tokenizer(PreTrainedTokenizer):
+    """
+    Construct an SMALL100 tokenizer. Based on [SentencePiece](https://github.com/google/sentencepiece).
+    This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
+    this superclass for more information regarding those methods.
+    Args:
+        vocab_file (`str`):
+            Path to the vocabulary file.
+        spm_file (`str`):
+            Path to [SentencePiece](https://github.com/google/sentencepiece) file (generally has a .spm extension) that
+            contains the vocabulary.
+        tgt_lang (`str`, *optional*):
+            A string representing the target language.
+        eos_token (`str`, *optional*, defaults to `"</s>"`):
+            The end of sequence token.
+        sep_token (`str`, *optional*, defaults to `"</s>"`):
+            The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
+            sequence classification or for a text and a question for question answering. It is also used as the last
+            token of a sequence built with special tokens.
+        unk_token (`str`, *optional*, defaults to `"<unk>"`):
+            The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+            token instead.
+        pad_token (`str`, *optional*, defaults to `"<pad>"`):
+            The token used for padding, for example when batching sequences of different lengths.
+        language_codes (`str`, *optional*):
+            What language codes to use. Should be `"m2m100"`.
+        sp_model_kwargs (`dict`, *optional*):
+            Will be passed to the `SentencePieceProcessor.__init__()` method. The [Python wrapper for
+            SentencePiece](https://github.com/google/sentencepiece/tree/master/python) can be used, among other things,
+            to set:
+            - `enable_sampling`: Enable subword regularization.
+            - `nbest_size`: Sampling parameters for unigram. Invalid for BPE-Dropout.
+              - `nbest_size = {0,1}`: No sampling is performed.
+              - `nbest_size > 1`: samples from the nbest_size results.
+              - `nbest_size < 0`: assuming that nbest_size is infinite and samples from the all hypothesis (lattice)
+                using forward-filtering-and-backward-sampling algorithm.
+            - `alpha`: Smoothing parameter for unigram sampling, and dropout probability of merge operations for
+              BPE-dropout.
+    Examples:
+    ```python
+    >>> from tokenization_small100 import SMALL100Tokenizer
+    >>> tokenizer = SMALL100Tokenizer.from_pretrained("alirezamsh/small100", tgt_lang="ro")
+    >>> src_text = " UN Chief Says There Is No Military Solution in Syria"
+    >>> tgt_text = "Şeful ONU declară că nu există o soluţie militară în Siria"
+    >>> model_inputs = tokenizer(src_text, text_target=tgt_text, return_tensors="pt")
+    >>> model(**model_inputs)  # should work
+    ```"""
+
+    vocab_files_names = VOCAB_FILES_NAMES
+    max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
+    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
+    model_input_names = ["input_ids", "attention_mask"]
+
+    prefix_tokens: List[int] = []
+    suffix_tokens: List[int] = []
+
+    def __init__(
+        self,
+        vocab_file,
+        spm_file,
+        tgt_lang=None,
+        bos_token="<s>",
+        eos_token="</s>",
+        sep_token="</s>",
+        pad_token="<pad>",
+        unk_token="<unk>",
+        language_codes="m2m100",
+        sp_model_kwargs: Optional[Dict[str, Any]] = None,
+        num_madeup_words=8,
+        **kwargs,
+    ) -> None:
+        self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
+
+        self.language_codes = language_codes
+        fairseq_language_code = FAIRSEQ_LANGUAGE_CODES[language_codes]
+        self.lang_code_to_token = {lang_code: f"__{lang_code}__" for lang_code in fairseq_language_code}
+
+        kwargs["additional_special_tokens"] = kwargs.get("additional_special_tokens", [])
+        kwargs["additional_special_tokens"] += [
+            self.get_lang_token(lang_code)
+            for lang_code in fairseq_language_code
+            if self.get_lang_token(lang_code) not in kwargs["additional_special_tokens"]
+        ]
+
+        self.vocab_file = vocab_file
+        self.encoder = load_json(vocab_file)
+        self.decoder = {v: k for k, v in self.encoder.items()}
+        self.spm_file = spm_file
+        self.sp_model = load_spm(spm_file, self.sp_model_kwargs)
+
+        self.encoder_size = len(self.encoder)
+
+        self.lang_token_to_id = {
+            self.get_lang_token(lang_code): self.encoder_size + i for i, lang_code in enumerate(fairseq_language_code)
+        }
+        self.lang_code_to_id = {lang_code: self.encoder_size + i for i, lang_code in enumerate(fairseq_language_code)}
+        self.id_to_lang_token = {v: k for k, v in self.lang_token_to_id.items()}
+
+        self._tgt_lang = tgt_lang if tgt_lang is not None else "en"
+        self.cur_lang_id = self.get_lang_id(self._tgt_lang)
+        self.num_madeup_words = num_madeup_words
+        
+        super().__init__(
+            tgt_lang=tgt_lang,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            sep_token=sep_token,
+            unk_token=unk_token,
+            pad_token=pad_token,
+            language_codes=language_codes,
+            sp_model_kwargs=self.sp_model_kwargs,
+            num_madeup_words=num_madeup_words,
+            **kwargs,
+        )
+        
+        self.set_lang_special_tokens(self._tgt_lang)
+
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.encoder) + len(self.lang_token_to_id) + self.num_madeup_words
+
+    @property
+    def tgt_lang(self) -> str:
+        return self._tgt_lang
+
+    @tgt_lang.setter
+    def tgt_lang(self, new_tgt_lang: str) -> None:
+        self._tgt_lang = new_tgt_lang
+        self.set_lang_special_tokens(self._tgt_lang)
+
+    def _tokenize(self, text: str) -> List[str]:
+        return self.sp_model.encode(text, out_type=str)
+
+    def _convert_token_to_id(self, token):
+        if token in self.lang_token_to_id:
+            return self.lang_token_to_id[token]
+        return self.encoder.get(token, self.encoder[self.unk_token])
+
+    def _convert_id_to_token(self, index: int) -> str:
+        """Converts an index (integer) in a token (str) using the decoder."""
+        if index in self.id_to_lang_token:
+            return self.id_to_lang_token[index]
+        return self.decoder.get(index, self.unk_token)
+
+    def convert_tokens_to_string(self, tokens: List[str]) -> str:
+        """Converts a sequence of tokens (strings for sub-words) in a single string."""
+        return self.sp_model.decode(tokens)
+
+    def get_special_tokens_mask(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
+    ) -> List[int]:
+        """
+        Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
+        special tokens using the tokenizer `prepare_for_model` method.
+        Args:
+            token_ids_0 (`List[int]`):
+                List of IDs.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+            already_has_special_tokens (`bool`, *optional*, defaults to `False`):
+                Whether or not the token list is already formatted with special tokens for the model.
+        Returns:
+            `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
+        """
+
+        if already_has_special_tokens:
+            return super().get_special_tokens_mask(
+                token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
+            )
+
+        prefix_ones = [1] * len(self.prefix_tokens)
+        suffix_ones = [1] * len(self.suffix_tokens)
+        if token_ids_1 is None:
+            return prefix_ones + ([0] * len(token_ids_0)) + suffix_ones
+        return prefix_ones + ([0] * len(token_ids_0)) + ([0] * len(token_ids_1)) + suffix_ones
+
+    def build_inputs_with_special_tokens(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """
+        Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
+        adding special tokens. An MBART sequence has the following format, where `X` represents the sequence:
+        - `input_ids` (for encoder) `X [eos, src_lang_code]`
+        - `decoder_input_ids`: (for decoder) `X [eos, tgt_lang_code]`
+        BOS is never used. Pairs of sequences are not the expected use case, but they will be handled without a
+        separator.
+        Args:
+            token_ids_0 (`List[int]`):
+                List of IDs to which the special tokens will be added.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+        Returns:
+            `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
+        """
+        if token_ids_1 is None:
+            if self.prefix_tokens is None:
+                return token_ids_0 + self.suffix_tokens
+            else:
+                return self.prefix_tokens + token_ids_0 + self.suffix_tokens
+        # We don't expect to process pairs, but leave the pair logic for API consistency
+        if self.prefix_tokens is None:
+            return token_ids_0 + token_ids_1 + self.suffix_tokens
+        else:
+            return self.prefix_tokens + token_ids_0 + token_ids_1 + self.suffix_tokens
+
+    def get_vocab(self) -> Dict:
+        vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
+        vocab.update(self.added_tokens_encoder)
+        return vocab
+
+    def __getstate__(self) -> Dict:
+        state = self.__dict__.copy()
+        state["sp_model"] = None
+        return state
+
+    def __setstate__(self, d: Dict) -> None:
+        self.__dict__ = d
+
+        # for backward compatibility
+        if not hasattr(self, "sp_model_kwargs"):
+            self.sp_model_kwargs = {}
+
+        self.sp_model = load_spm(self.spm_file, self.sp_model_kwargs)
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        save_dir = Path(save_directory)
+        if not save_dir.is_dir():
+            raise OSError(f"{save_directory} should be a directory")
+        vocab_save_path = save_dir / (
+            (filename_prefix + "-" if filename_prefix else "") + self.vocab_files_names["vocab_file"]
+        )
+        spm_save_path = save_dir / (
+            (filename_prefix + "-" if filename_prefix else "") + self.vocab_files_names["spm_file"]
+        )
+
+        save_json(self.encoder, vocab_save_path)
+
+        if os.path.abspath(self.spm_file) != os.path.abspath(spm_save_path) and os.path.isfile(self.spm_file):
+            copyfile(self.spm_file, spm_save_path)
+        elif not os.path.isfile(self.spm_file):
+            with open(spm_save_path, "wb") as fi:
+                content_spiece_model = self.sp_model.serialized_model_proto()
+                fi.write(content_spiece_model)
+
+        return (str(vocab_save_path), str(spm_save_path))
+
+    def prepare_seq2seq_batch(
+        self,
+        src_texts: List[str],
+        tgt_texts: Optional[List[str]] = None,
+        tgt_lang: str = "ro",
+        **kwargs,
+    ) -> BatchEncoding:
+        self.tgt_lang = tgt_lang
+        self.set_lang_special_tokens(self.tgt_lang)
+        return super().prepare_seq2seq_batch(src_texts, tgt_texts, **kwargs)
+
+    def _build_translation_inputs(self, raw_inputs, tgt_lang: Optional[str], **extra_kwargs):
+        """Used by translation pipeline, to prepare inputs for the generate function"""
+        if tgt_lang is None:
+            raise ValueError("Translation requires a `tgt_lang` for this model")
+        self.tgt_lang = tgt_lang
+        inputs = self(raw_inputs, add_special_tokens=True, **extra_kwargs)
+        return inputs
+
+    def _switch_to_input_mode(self):
+        self.set_lang_special_tokens(self.tgt_lang)
+
+    def _switch_to_target_mode(self):
+        self.prefix_tokens = None
+        self.suffix_tokens = [self.eos_token_id]        
+
+    def set_lang_special_tokens(self, src_lang: str) -> None:
+        """Reset the special tokens to the tgt lang setting. No prefix and suffix=[eos, tgt_lang_code]."""
+        lang_token = self.get_lang_token(src_lang)
+        self.cur_lang_id = self.lang_token_to_id[lang_token]
+        self.prefix_tokens = [self.cur_lang_id]
+        self.suffix_tokens = [self.eos_token_id]
+
+    def get_lang_token(self, lang: str) -> str:
+        return self.lang_code_to_token[lang]
+
+    def get_lang_id(self, lang: str) -> int:
+        lang_token = self.get_lang_token(lang)
+        return self.lang_token_to_id[lang_token]
+
+
+def load_spm(path: str, sp_model_kwargs: Dict[str, Any]) -> sentencepiece.SentencePieceProcessor:
+    spm = sentencepiece.SentencePieceProcessor(**sp_model_kwargs)
+    spm.Load(str(path))
+    return spm
+
+
+def load_json(path: str) -> Union[Dict, List]:
+    with open(path, "r") as f:
+        return json.load(f)
+
+
+def save_json(data, path: str) -> None:
+    with open(path, "w") as f:
+        json.dump(data, f, indent=2)
+