DevGuptaa commited on
Commit
1c807a8
·
1 Parent(s): e8fa815

Upload tokenizer.py

Browse files
Files changed (1) hide show
  1. tokenizer.py +56 -0
tokenizer.py ADDED
@@ -0,0 +1,56 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ input_file = "/content/bothcan.txt" # Replace with actual input file path
2
+ model_prefix = "botchan" # Replace with desired model save path
3
+ import sentencepiece as spm
4
+ spm.SentencePieceTrainer.train(
5
+ input=input_file,
6
+ model_prefix=model_prefix,
7
+ vocab_size=1000, # Adjust as needed, this is just an example value
8
+ model_type="unigram", # You can use different models like unigram or bpe
9
+ )
10
+
11
+ from sentencepiece import SentencePieceProcessor
12
+
13
+ model_path = "botchan.model" # Replace with the actual path
14
+ sp_model = SentencePieceProcessor(model_file=model_path)
15
+ vocab_size = 4000
16
+
17
+ import os
18
+ from logging import getLogger
19
+ from typing import List
20
+
21
+ from sentencepiece import SentencePieceProcessor
22
+
23
+
24
+ logger = getLogger()
25
+
26
+
27
+ class Tokenizer:
28
+ def __init__(self, model_path: str):
29
+ # reload tokenizer
30
+ assert os.path.isfile(model_path), model_path
31
+ self.sp_model = SentencePieceProcessor(model_file=model_path)
32
+ logger.info(f"Reloaded SentencePiece model from {model_path}")
33
+
34
+ # BOS / EOS token IDs
35
+ self.n_words: int = self.sp_model.vocab_size()
36
+ self.bos_id: int = self.sp_model.bos_id()
37
+ self.eos_id: int = self.sp_model.eos_id()
38
+ self.pad_id: int = self.sp_model.pad_id()
39
+ logger.info(
40
+ f"#words: {self.n_words} - BOS ID: {self.bos_id} - EOS ID: {self.eos_id}"
41
+ )
42
+ assert self.sp_model.vocab_size() == self.sp_model.get_piece_size()
43
+
44
+ def encode(self, s: str, bos: bool, eos: bool) -> List[int]:
45
+ assert type(s) is str
46
+ t = self.sp_model.encode(s)
47
+ if bos:
48
+ t = [self.bos_id] + t
49
+ if eos:
50
+ t = t + [self.eos_id]
51
+ return t
52
+
53
+ def decode(self, t: List[int]) -> str:
54
+ return self.sp_model.decode(t)
55
+
56
+ tokenizer = Tokenizer(model_path="botchan.model") # Replace with actual model path