refactoring

app.py

@@ -1,47 +1,70 @@
 import gradio as gr
 import requests
 
+
+# 定数定義
+API_BASE_URL = "https://huggingface.co/api/models/"
+MULTIPLIERS = {
+    "U8": 1,
+    "I8": 1,
+    "F8_E5M2": 1,
+    "F8_E4M3": 1,
+    "I16": 2,
+    "U16": 2,
+    "F16": 2,
+    "BF16": 2,
+    "I32": 4,
+    "U32": 4,
+    "F32": 4,
+    "F64": 8,
+    "I64": 8,
+    "U64": 8,
+}
+PRECISION_KEYS = list(MULTIPLIERS.keys())
+
 def get_model_api_info(model_name: str) -> dict:
     """
-    Hugging Face Hub API から指定モデルの情報を取得する
+    Hugging Face Hub APIから指定モデルの情報を取得する。
+
+    :param model_name: モデル名（例: "meta-llama/Llama-3.3-70B-Instruct"）
+    :return: モデル情報の辞書。取得に失敗した場合はNoneを返す。
     """
-    api_url = f"https://huggingface.co/api/models/{model_name}"
-    response = requests.get(api_url)
-    if response.status_code != 200:
+    api_url = f"{API_BASE_URL}{model_name}"
+    try:
+        response = requests.get(api_url)
+        response.raise_for_status()
+        return response.json()
+    except requests.RequestException:
         return None
-    return response.json()
 
 def sum_precision_sizes(parameters: dict) -> int:
     """
-    parameters 辞書から、キーがあれば対応表に合わせて合算する
+    safetensors内の各精度のパラメータ数に対して、定数の乗数をかけた合計バイト数を算出する。
+
+    :param parameters: 各精度のパラメータ数を含む辞書
+    :return: 合計バイト数
     """
     # 各データ型とその乗数の対応表
-    multipliers = {
-        "U8": 1,
-        "I8": 1,
-        "F8_E5M2": 1,
-        "F8_E4M3": 1,
-        "I16": 2,
-        "U16": 2,
-        "F16": 2,
-        "BF16": 2,
-        "I32": 4,
-        "U32": 4,
-        "F32": 4,
-        "F64": 8,
-        "I64": 8,
-        "U64": 8,
-    }
-    return sum(parameters.get(key, 0) * multiplier for key, multiplier in multipliers.items())
+    total_bytes = 0
+    for precision, multiplier in MULTIPLIERS.items():
+        count = parameters.get(precision, 0)
+        total_bytes += count * multiplier
+    return total_bytes
 
 def estimate_gpu_memory(model_name: str) -> str:
     """
-    指定したモデル名からAPIを呼び出し、safetensors内の各精度のサイズを合算して、
-    補正後に必要なGPUメモリ量を GB 単位で概算して返す。
+    指定したモデル名からAPI情報を取得し、safetensors内の各精度パラメータサイズの合算値から
+    GPUメモリ必要量を概算する。
+
+    :param model_name: モデル名
+    :return: GPUメモリ必要量などの情報を含むメッセージ文字列
     """
+    result_lines = []
+
     model_info = get_model_api_info(model_name)
     if model_info is None:
-        return f"エラー: モデル '{model_name}' の情報が取得できませんでした。"
+        result_lines.append(f"エラー: モデル '{model_name}' の情報が取得できませんでした。")
+        return "\n".join(result_lines)
 
     # リポジトリ全体のファイルサイズ（参考）
     usedStorage = model_info.get("usedStorage")
@@ -51,35 +74,40 @@ def estimate_gpu_memory(model_name: str) -> str:
 
     # safetensorsの情報が見つからない場合
     if not safetensors or "parameters" not in safetensors:
-        result = "safetensors の情報が見つかりません。\n"
+        result_lines.append("safetensorsの情報が見つかりません。")
         if usedStorage_gb is not None:
-          result += (f"\n参考: 該当リポジトリのファイルサイズは合計 {usedStorage_gb:.2f} GB \n"
-            f"仮にこれを全てモデルサイズとして推論時(1.5倍)、トレーニング時(4倍)の補正を適用すると概算GPUメモリ必要量は\n"
-            f"【推論】約 {usedStorage_gb * 1.5:.2f} GB \n"
-            f"【トレーニング】約 {usedStorage_gb* 4:.2f} GB \n"
-            f"となります。")
-        return result
+            result_lines.append("")
+            result_lines.append(f"参考までに、該当リポジトリのファイルサイズは合計 {usedStorage_gb:.2f} GBです。")
+            result_lines.append(f"仮にこれを全てモデルサイズとして推論時(1.5倍)、トレーニング時(4倍)の補正を適用すると概算GPUメモリ必要量は")
+            result_lines.append(f"【推論】約 {usedStorage_gb * 1.5:.2f} GB")
+            result_lines.append(f"【トレーニング】約 {usedStorage_gb * 4:.2f} GB")
+            result_lines.append("となります。")
+        return "\n".join(result_lines)
 
     # safetensorsの情報がある場合
     parameters = safetensors["parameters"]
     estimated_bytes = sum_precision_sizes(parameters)
     estimated_gb = estimated_bytes / (1024 ** 3)
 
-    result = f"モデル '{model_name}' のsafetensors情報より、各パラメータ数は\n"
+    result_lines.append(f"モデル '{model_name}' のsafetensors情報より、各パラメータ数は")
 
-    keys = ["U8", "I8", "F8_E5M2", "F8_E4M3", "I16", "U16", "F16", "BF16", "I32", "U32", "F32", "F64", "I64", "U64"]
-    lines = [f"【{key}】{parameters[key]:,}" for key in keys if key in parameters]
-    result += "\n".join(lines) + "\n"
+    # 各精度ごとのパラメータ数の出力
+    for precision in PRECISION_KEYS:
+        if precision in parameters:
+            result_lines.append(f"【{precision}】 {parameters[precision]:,}")
 
-    result += (f"これらを合算するとモデルのパラメータサイズは約 {estimated_gb:.2f} GB です。\n\n"
-                f"推論時(1.5倍)、トレーニング時(4倍)の補正を適用すると概算GPUメモリ必要量は\n"
-                f"【推論】約 {estimated_gb * 1.5:.2f} GB \n"
-                f"【トレーニング】約 {estimated_gb * 4:.2f} GB \n"
-                f"となります。\n")
+    result_lines.append(f"これらを合算するとモデルのパラメータサイズは約 {estimated_gb:.2f} GB です。")
+    result_lines.append("")
+    result_lines.append(f"推論時(1.5倍)、トレーニング時(4倍)の補正を適用すると概算GPUメモリ必要量は")
+    result_lines.append(f"【推論】約 {estimated_gb * 1.5:.2f} GB")
+    result_lines.append(f"【トレーニング】約 {estimated_gb * 4:.2f} GB")
+    result_lines.append("となります。")
 
     if usedStorage_gb is not None:
-        result += f"参考: 該当リポジトリのファイルサイズは合計 {usedStorage_gb:.2f} GB"
-    return result
+        result_lines.append("")
+        result_lines.append(f"参考: 該当リポジトリのファイルサイズは合計 {usedStorage_gb:.2f} GB")
+
+    return "\n".join(result_lines)
 
 # Gradio インターフェースの定義
 iface = gr.Interface(