"""
## 封装

__all__ = ['icetk']
sentencepiece_model_pb2 给包装起来了，

## ss
ss
"""

import sentencepiece as spm
from icetk import sentencepiece_model_pb2 as model

def __init__(self, model_path):
    self.proto = model.ModelProto()
    with open(model_path, 'rb') as fin:
        proto_str = fin.read()
        self.proto.ParseFromString(proto_str)
    self.refresh()

def refresh(self):
    self.sp = spm.SentencePieceProcessor()
    self.sp.Load(model_proto=self.proto.SerializeToString())
    self.num_tokens = self.sp.vocab_size()