We have switched from the Regex tokenizer to the SentencePiece tokenizer. (The updated code was written by OpenSoftware-World and corrected by ChatGPT.)

SentencePiece is much more effective than the Regex tokenizer and is also used in large language models. At OpenSoftware-World, we have switched to the SentencePiece tokenizer for our OpenSoftware-World-OSW1 AI model so that it can produce high-quality and accurate results, just like other large language models.

model_init.py

@@ -7,6 +7,7 @@ import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+import sentencepiece as spm
 
 NUM_THREADS = os.cpu_count() or 4
 torch.set_num_threads(NUM_THREADS)
@@ -17,40 +18,62 @@ except RuntimeError:
 
 DEVICE = torch.device("cpu")
 print(f"🧵 Number of CPU threads : {NUM_THREADS}")
-TOKEN_RE = re.compile(r"\w+|[^\w\s]", re.UNICODE)
-
-def tokenize(text: str):
-    return TOKEN_RE.findall(text.lower())
 
 class Vocab:
-    PAD, UNK, BOS, EOS = "<pad>", "<unk>", "<bos>", "<eos>"
+    PAD = "<pad>"
+    UNK = "<unk>"
+    BOS = "<bos>"
+    EOS = "<eos>"
+
+    def __init__(self, model_path="opensoftware_world_osw1_tokenizer.model"):
+        self.sp = spm.SentencePieceProcessor()
+        self.sp.load(model_path)
 
-    def __init__(self):
-        self.stoi = {}
-        self.itos = []
+        self.pad_id = self.sp.pad_id()
+        self.unk_id = self.sp.unk_id()
+        self.bos_id = self.sp.bos_id()
+        self.eos_id = self.sp.eos_id()
 
     def encode(self, text, add_bos=False, add_eos=False):
-        ids = [self.stoi.get(t, self.stoi[Vocab.UNK]) for t in tokenize(text)]
+        ids = self.sp.encode(text, out_type=int)
+
         if add_bos:
-            ids = [self.stoi[Vocab.BOS]] + ids
+            ids = [self.bos_id] + ids
+
         if add_eos:
-            ids = ids + [self.stoi[Vocab.EOS]]
+            ids = ids + [self.eos_id]
+
         return ids
 
     def decode(self, ids):
-        toks = [self.itos[i] for i in ids if 0 <= i < len(self.itos)]
-        toks = [t for t in toks if t != Vocab.PAD and t != Vocab.BOS]
-        out = []
-        for t in toks:
-            if t == Vocab.EOS:
-                break
-            out.append(t)
-        text = " ".join(out)
-        text = re.sub(r"\s+([.,!?;:])", r"\1", text)
-        return text
+        ids = [
+            i for i in ids
+            if i not in (self.pad_id, self.bos_id)
+        ]
+
+        if self.eos_id in ids:
+            ids = ids[:ids.index(self.eos_id)]
+
+        return self.sp.decode(ids)
 
     def __len__(self):
-        return len(self.itos)
+        return self.sp.get_piece_size()
+
+    @property
+    def stoi(self):
+        return {
+            self.PAD: self.pad_id,
+            self.UNK: self.unk_id,
+            self.BOS: self.bos_id,
+            self.EOS: self.eos_id,
+        }
+
+    @property
+    def itos(self):
+        return [
+            self.sp.id_to_piece(i)
+            for i in range(self.sp.get_piece_size())
+        ]
 
 class CausalSelfAttention(nn.Module):
     def __init__(self, d_model, n_head, dropout):
@@ -155,10 +178,8 @@ def load_checkpoint(path: str):
     ckpt = torch.load(path, map_location="cpu")
 
     cfg = ckpt["config"]
-    vocab = Vocab()
-    vocab.stoi = ckpt["vocab_stoi"]
-    vocab.itos = ckpt["vocab_itos"]
-    pad_id = ckpt["pad_id"]
+    vocab = Vocab("opensoftware_world_osw1_tokenizer.model")
+    pad_id = vocab.sp.pad_id()
 
     model = OSW1Model(len(vocab), cfg, pad_id=pad_id).to(DEVICE)
     model.load_state_dict(ckpt["model_state_dict"])
@@ -187,7 +208,7 @@ def load_checkpoint(path: str):
 
 def chat_loop(model: OSW1Model, vocab: Vocab):
     print("=" * 64)
-    print("💬 OSW1 ready! You can start chatting. Type 'exit' to quit.")
+    print("💬 OpenSoftware-World-OSW1 ready! You can start chatting. Type 'exit' to quit.")
     print("=" * 64)
 
     eos_id = vocab.stoi[Vocab.EOS]
@@ -211,7 +232,7 @@ def chat_loop(model: OSW1Model, vocab: Vocab):
         out = model.generate(x, max_new_tokens=60, temperature=0.85, top_k=40, eos_id=eos_id)
         answer_ids = out[0, len(ids):].tolist()
         answer = vocab.decode(answer_ids)
-        print(f"OSW1: {answer if answer else '(...silence...)'}")
+        print(f"OpenSoftware-World-OSW1: {answer if answer else '(...silence...)'}")
 
 def main():
     if len(sys.argv) > 1:
@@ -224,8 +245,8 @@ def main():
         if ckpt_path is None:
             print(
                 "❌ No checkpoint files found in the directory.\n"
-                "   Please train a model using 'python train_osw1.py' or\n"
-                "   specify a checkpoint file using 'python model_init.py <file_path>'."
+                "   Please train a model using 'python3 model_training.py' or\n"
+                "   specify a checkpoint file using 'python3 model_init.py <file_path>'."
             )
             sys.exit(1)