File size: 455 Bytes
3a3b216
 
 
 
 
 
 
 
8c66169
 
3a3b216
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
from typing import List

import tiktoken

from graphgen.bases import BaseTokenizer


class TiktokenTokenizer(BaseTokenizer):
    def __init__(self, model_name: str = "cl100k_base"):
        super().__init__(model_name)
        self.enc = tiktoken.get_encoding(self.model_name)

    def encode(self, text: str) -> List[int]:
        return self.enc.encode(text)

    def decode(self, token_ids: List[int]) -> str:
        return self.enc.decode(token_ids)