From 80b24e66b745bb303cc90afaf7b12cfac1eff9ae Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=BA=90=E6=96=87=E9=9B=A8?=
 <41315874+fumiama@users.noreply.github.com>
Date: Mon, 5 Aug 2024 21:06:52 +0800
Subject: [PATCH] fix(tokenizer): weights_only warning by dropping `.pt`

---
 ChatTTS/config/config.py    |  2 +-
 ChatTTS/core.py             |  2 +-
 ChatTTS/model/tokenizer.py  |  6 +++++-
 ChatTTS/res/sha256_map.json |  7 +++++--
 ChatTTS/utils/dl.py         | 15 ++++++++++++++-
 tools/checksum/tmpl.go      | 12 +++++++++---
 6 files changed, 35 insertions(+), 9 deletions(-)

diff --git a/ChatTTS/config/config.py b/ChatTTS/config/config.py
index b58fc26..661407c 100644
--- a/ChatTTS/config/config.py
+++ b/ChatTTS/config/config.py
@@ -7,7 +7,7 @@ class Path:
     dvae_ckpt_path: str = "asset/DVAE_full.pt"
     gpt_ckpt_path: str = "asset/GPT.pt"
     decoder_ckpt_path: str = "asset/Decoder.pt"
-    tokenizer_path: str = "asset/tokenizer.pt"
+    tokenizer_path: str = "asset/tokenizer"
 
 
 @dataclass(repr=False, eq=False)
diff --git a/ChatTTS/core.py b/ChatTTS/core.py
index 2a51d12..cb90bde 100644
--- a/ChatTTS/core.py
+++ b/ChatTTS/core.py
@@ -332,7 +332,7 @@ class Chat:
         self.logger.log(logging.INFO, "decoder loaded.")
 
         if tokenizer_path:
-            self.tokenizer = Tokenizer(tokenizer_path, device)
+            self.tokenizer = Tokenizer(tokenizer_path)
             self.logger.log(logging.INFO, "tokenizer loaded.")
 
         self.coef = coef
diff --git a/ChatTTS/model/tokenizer.py b/ChatTTS/model/tokenizer.py
index 372bb5d..b98bcd8 100644
--- a/ChatTTS/model/tokenizer.py
+++ b/ChatTTS/model/tokenizer.py
@@ -19,11 +19,15 @@ from ..utils import del_all
 
 class Tokenizer:
     def __init__(
-        self, tokenizer_path: torch.serialization.FILE_LIKE, device: torch.device
+        self, tokenizer_path: torch.serialization.FILE_LIKE,
     ):
+        """
         tokenizer: BertTokenizerFast = torch.load(
             tokenizer_path, map_location=device, mmap=True
         )
+        # tokenizer.save_pretrained("asset/tokenizer", legacy_format=False)
+        """
+        tokenizer: BertTokenizerFast = BertTokenizerFast.from_pretrained(tokenizer_path)
         self._tokenizer = tokenizer
 
         self.len = len(tokenizer)
diff --git a/ChatTTS/res/sha256_map.json b/ChatTTS/res/sha256_map.json
index 32f0141..b64b89e 100644
--- a/ChatTTS/res/sha256_map.json
+++ b/ChatTTS/res/sha256_map.json
@@ -3,6 +3,9 @@
 	"sha256_asset_DVAE_full_pt"      : "553eb75763511e23f3e5f86303e2163c5ca775489d637fb635d979c8ae58bbe5",
 	"sha256_asset_GPT_pt"       : "d7d4ee6461ea097a2be23eb40d73fb94ad3b3d39cb64fbb50cb3357fd466cadb",
 	"sha256_asset_spk_stat_pt"  : "3228d8a4cbbf349d107a1b76d2f47820865bd3c9928c4bdfe1cefd5c7071105f",
-	"sha256_asset_tokenizer_pt" : "e911ae7c6a7c27953433f35c44227a67838fe229a1f428503bdb6cd3d1bcc69c",
-	"sha256_asset_Vocos_pt"     : "09a670eda1c08b740013679c7a90ebb7f1a97646ea7673069a6838e6b51d6c58"
+	"sha256_asset_Vocos_pt"     : "09a670eda1c08b740013679c7a90ebb7f1a97646ea7673069a6838e6b51d6c58",
+
+	"sha256_asset_tokenizer_special_tokens_map_json": "bd0ac9d9bb1657996b5c5fbcaa7d80f8de530d01a283da97f89deae5b1b8d011",
+	"sha256_asset_tokenizer_tokenizer_config_json"  : "43e9d658b554fa5ee8d8e1d763349323bfef1ed7a89c0794220ab8861387d421",
+	"sha256_asset_tokenizer_tokenizer_json"         : "843838a64e121e23e774cc75874c6fe862198d9f7dd43747914633a8fd89c20e"
 }
diff --git a/ChatTTS/utils/dl.py b/ChatTTS/utils/dl.py
index d2dff9c..338721c 100644
--- a/ChatTTS/utils/dl.py
+++ b/ChatTTS/utils/dl.py
@@ -45,13 +45,13 @@ def check_model(
 
 def check_all_assets(base_dir: Path, sha256_map: Dict[str, str], update=False) -> bool:
     logger.get_logger().info("checking assets...")
+
     current_dir = base_dir / "asset"
     names = [
         "Decoder.pt",
         "DVAE_full.pt",
         "GPT.pt",
         "spk_stat.pt",
-        "tokenizer.pt",
         "Vocos.pt",
     ]
     for model in names:
@@ -61,6 +61,19 @@ def check_all_assets(base_dir: Path, sha256_map: Dict[str, str], update=False) -
         ):
             return False
 
+    current_dir = base_dir / "asset" / "tokenizer"
+    names = [
+        "special_tokens_map.json",
+        "tokenizer_config.json",
+        "tokenizer.json",
+    ]
+    for model in names:
+        menv = model.replace(".", "_")
+        if not check_model(
+            current_dir, model, sha256_map[f"sha256_asset_tokenizer_{menv}"], update
+        ):
+            return False
+
     logger.get_logger().info("all assets are already latest.")
     return True
 
diff --git a/tools/checksum/tmpl.go b/tools/checksum/tmpl.go
index 04533e1..fe5617e 100644
--- a/tools/checksum/tmpl.go
+++ b/tools/checksum/tmpl.go
@@ -5,8 +5,11 @@ var files = [...]string{
 	"asset/DVAE_full.pt",
 	"asset/GPT.pt",
 	"asset/spk_stat.pt",
-	"asset/tokenizer.pt",
 	"asset/Vocos.pt",
+
+	"asset/tokenizer/special_tokens_map.json",
+	"asset/tokenizer/tokenizer_config.json",
+	"asset/tokenizer/tokenizer.json",
 }
 
 const jsontmpl = `{
@@ -14,7 +17,10 @@ const jsontmpl = `{
 	"sha256_asset_DVAE_full_pt"      : "%s",
 	"sha256_asset_GPT_pt"       : "%s",
 	"sha256_asset_spk_stat_pt"  : "%s",
-	"sha256_asset_tokenizer_pt" : "%s",
-	"sha256_asset_Vocos_pt"     : "%s"
+	"sha256_asset_Vocos_pt"     : "%s",
+
+	"sha256_asset_tokenizer_special_tokens_map_json": "%s",
+	"sha256_asset_tokenizer_tokenizer_config_json"  : "%s",
+	"sha256_asset_tokenizer_tokenizer_json"         : "%s"
 }
 `