Update model.py

model.py

@@ -223,8 +223,14 @@ CHECKPOINT_PATH = hf_hub_download(
 
 VOCAB_PATH = os.path.join(script_dir, "vocab.pkl")
 
+class CustomUnpickler(pickle.Unpickler):
+    def find_class(self, module, name):
+        if name == "Vocabulary":
+            return Vocabulary
+        return super().find_class(module, name)
+
 with open(VOCAB_PATH, "rb") as f:
-    vocab = pickle.load(f)
+    vocab = CustomUnpickler(f).load()
 
 vocab_size = len(vocab)