[Update] translation_utils.py: Refactor weight handling and improve error logging

add translate model
- jncraton/m2m100_418M-ct2-int8
- jncraton/m2m100_1.2B-ct2-int8
- OpenNMT/nllb-200-3.3B-ct2-int8
- OpenNMT/nllb-200-distilled-1.3B-ct2-int8

src-python/models/translation/translation_utils.py

@@ -1,99 +1,88 @@
-import tempfile
-from zipfile import ZipFile
 from os import path as os_path
 from os import makedirs as os_makedirs
-from requests import get as requests_get
 from typing import Callable
-import hashlib
 import transformers
+import ctranslate2
+from huggingface_hub import hf_hub_url, list_repo_files
+from requests import get as requests_get
+try:
     from utils import errorLogging
+except Exception:
+    import traceback
+    def errorLogging():
+        print(traceback.format_exc())
 
 ctranslate2_weights = {
-    "small": { # M2M-100 418M-parameter model
+    "small": {
-        "url": "https://github.com/misyaguziya/VRCT-weights/releases/download/v1.0/m2m100_418m.zip",
+        # "hf_repo": "jncraton/m2m100_418M-ct2-int8",
-        "directory_name": "m2m100_418m",
+        # "directory_name": "m2m100_418M-ct2-int8",
-        "tokenizer": "facebook/m2m100_418M",
+        # "tokenizer": "facebook/m2m100_418M",
-        "hash": {
+        "hf_repo": "OpenNMT/nllb-200-distilled-1.3B-ct2-int8",
-            "model.bin": "e7c26a9abb5260abd0268fbe3040714070dec254a990b4d7fd3f74c5230e3acb",
+        "directory_name": "nllb-200-distilled-1.3B-ct2-int8",
-            "sentencepiece.model": "d8f7c76ed2a5e0822be39f0a4f95a55eb19c78f4593ce609e2edbc2aea4d380a",
+        "tokenizer": "facebook/nllb-200-distilled-1.3B",
-            "shared_vocabulary.txt": "bd440aa21b8ca3453fc792a0018a1f3fe68b3464aadddd4d16a4b72f73c86d8c",
-        }
     },
-    "large": { # M2M-100 1.2B-parameter model
+    "large": {
-        "url": "https://github.com/misyaguziya/VRCT-weights/releases/download/v1.0/m2m100_12b.zip",
+        "hf_repo": "jncraton/m2m100_1.2B-ct2-int8",
-        "directory_name": "m2m100_12b",
+        "directory_name": "m2m100_1.2B-ct2-int8",
-        "tokenizer": "facebook/m2m100_1.2b",
+        "tokenizer": "facebook/m2m100_1.2B",
-        "hash": {
+    },
-            "model.bin": "abb7bf4ba7e5e016b6e3ed480c752459b2f783ac8fca372e7587675e5bf3a919",
+    "nllb-200-3.3B-ct2-int8": {
-            "sentencepiece.model": "d8f7c76ed2a5e0822be39f0a4f95a55eb19c78f4593ce609e2edbc2aea4d380a",
+        "hf_repo": "OpenNMT/nllb-200-3.3B-ct2-int8",
-            "shared_vocabulary.txt": "bd440aa21b8ca3453fc792a0018a1f3fe68b3464aadddd4d16a4b72f73c86d8c",
+        "directory_name": "nllb-200-3.3B-ct2-int8",
-        }
+        "tokenizer": "facebook/nllb-200-3.3B",
+    },
+    "nllb-200-distilled-1.3B": {
+        "hf_repo": "OpenNMT/nllb-200-distilled-1.3B-ct2-int8",
+        "directory_name": "nllb-200-distilled-1.3B-ct2-int8",
+        "tokenizer": "facebook/nllb-200-distilled-1.3B",
     },
 }
 
-def calculate_file_hash(file_path, block_size=65536):
+def checkCTranslate2Weight(root: str, weight_type: str = "small"):
-    hash_object = hashlib.sha256()
-
-    with open(file_path, 'rb') as file:
-        for block in iter(lambda: file.read(block_size), b''):
-            hash_object.update(block)
-
-    return hash_object.hexdigest()
-
-def checkCTranslate2Weight(root, weight_type="small"):
     weight_directory_name = ctranslate2_weights[weight_type]["directory_name"]
-    hash_data = ctranslate2_weights[weight_type]["hash"]
+    path = os_path.join(root, "weights", "ctranslate2", weight_directory_name)
-    files = [
+    try:
-        "model.bin",
+        # モデルロード可能かどうかで判定
-        "sentencepiece.model",
+        ctranslate2.Translator(path)
-        "shared_vocabulary.txt"
+        return True
-    ]
+    except Exception:
-    path = os_path.join(root, "weights", "ctranslate2")
+        return False
 
-    # check already downloaded
+def downloadCTranslate2Weight(root: str, weight_type: str = "small", callback: Callable = None, end_callback: Callable = None):
-    already_downloaded = False
+    hf_repo = ctranslate2_weights[weight_type]["hf_repo"]
-    if all(os_path.exists(os_path.join(path, weight_directory_name, file)) for file in files):
+    files = list_repo_files(repo_id=hf_repo)
-        # check hash
+    path = os_path.join(root, "weights", "ctranslate2", ctranslate2_weights[weight_type]["directory_name"])
-        for file in files:
+    if checkCTranslate2Weight(root, weight_type):
-            original_hash = hash_data[file]
+        return True
-            current_hash = calculate_file_hash(os_path.join(path, weight_directory_name, file))
-            if original_hash != current_hash:
-                break
-        already_downloaded = True
-    return already_downloaded
-
-def downloadCTranslate2Weight(root, weight_type="small", callback=None, end_callback=None):
-    url = ctranslate2_weights[weight_type]["url"]
-    filename = "weight.zip"
-    path = os_path.join(root, "weights", "ctranslate2")
     os_makedirs(path, exist_ok=True)
 
-    if checkCTranslate2Weight(root, weight_type) is False:
+    def downloadFile(url: str, file_path: str, func: Callable = None):
         try:
-            with tempfile.TemporaryDirectory() as tmp_path:
             res = requests_get(url, stream=True)
+            res.raise_for_status()
             file_size = int(res.headers.get('content-length', 0))
             total_chunk = 0
-                with open(os_path.join(tmp_path, filename), 'wb') as file:
+            with open(file_path, 'wb') as file:
                 for chunk in res.iter_content(chunk_size=1024*2000):
                     file.write(chunk)
-                        if isinstance(callback, Callable):
+                    if func is not None:
                         total_chunk += len(chunk)
-                            callback(total_chunk/file_size)
+                        func(total_chunk/file_size)
-
-                with ZipFile(os_path.join(tmp_path, filename)) as zf:
-                    zf.extractall(path)
         except Exception:
             errorLogging()
 
-    if isinstance(end_callback, Callable):
+    for filename in files:
+        file_path = os_path.join(path, filename)
+        url = hf_hub_url(hf_repo, filename)
+        downloadFile(url, file_path, func=callback if filename == "model.bin" else None)
+
+    if end_callback is not None:
         end_callback()
 
-def downloadCTranslate2Tokenizer(path, weight_type="small"):
+def downloadCTranslate2Tokenizer(path: str, weight_type: str = "small"):
     directory_name = ctranslate2_weights[weight_type]["directory_name"]
     tokenizer = ctranslate2_weights[weight_type]["tokenizer"]
     tokenizer_path = os_path.join(path, "weights", "ctranslate2", directory_name, "tokenizer")
-
     try:
         os_makedirs(tokenizer_path, exist_ok=True)
         transformers.AutoTokenizer.from_pretrained(tokenizer, cache_dir=tokenizer_path)
@@ -101,3 +90,20 @@ def downloadCTranslate2Tokenizer(path, weight_type="small"):
         errorLogging()
         tokenizer_path = os_path.join("./weights", "ctranslate2", directory_name, "tokenizer")
         transformers.AutoTokenizer.from_pretrained(tokenizer, cache_dir=tokenizer_path)
+
+# テスト用コード（直接実行時のみ）
+if __name__ == "__main__":
+    def progress_callback(percent):
+        print(f"Download progress: {percent*100:.2f}%")
+
+    def end_callback():
+        print("Download finished.")
+
+    root = "./"  # 必要に応じてパスを変更
+    # for weight_type in ctranslate2_weights.keys():
+    #     print(f"Testing download for: {weight_type}")
+    #     downloadCTranslate2Weight(root, weight_type, callback=progress_callback, end_callback=end_callback)
+    #     result = checkCTranslate2Weight(root, weight_type)
+    #     print(f"Model loadable: {result}")
+    #     break
+    downloadCTranslate2Tokenizer(root, "small")