First Commit

app.py

@@ -0,0 +1,93 @@
+import gradio as gr
+import requests
+
+def get_model_api_info(model_name: str) -> dict:
+    """
+    Hugging Face Hub API から指定モデルの情報を取得する
+    """
+    api_url = f"https://huggingface.co/api/models/{model_name}"
+    response = requests.get(api_url)
+    if response.status_code != 200:
+        return None
+    return response.json()
+
+def sum_precision_sizes(parameters: dict) -> int:
+    """
+    parameters 辞書から、キーがあれば対応表に合わせて合算する
+    """
+    # 各データ型とその乗数の対応表
+    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())
+
+def estimate_gpu_memory(model_name: str) -> str:
+    """
+    指定したモデル名からAPIを呼び出し、safetensors内の各精度のサイズを合算して、
+    補正後に必要なGPUメモリ量を GB 単位で概算して返す。
+    """
+    model_info = get_model_api_info(model_name)
+    if model_info is None:
+        return f"エラー: モデル '{model_name}' の情報が取得できませんでした。"
+
+    # リポジトリ全体のファイルサイズ（参考）
+    usedStorage = model_info.get("usedStorage")
+    usedStorage_gb = usedStorage / (1024 ** 3) if usedStorage else None
+
+    safetensors = model_info.get("safetensors")
+
+    # safetensorsの情報が見つからない場合
+    if not safetensors or "parameters" not in safetensors:
+        result = "safetensors の情報が見つかりません。\n"
+        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
+
+    # safetensorsの情報がある場合
+    parameters = safetensors["parameters"]
+    estimated_bytes = sum_precision_sizes(parameters)
+    estimated_gb = estimated_bytes / (1024 ** 3)
+
+    result = f"モデル '{model_name}' のsafetensors情報より、各パラメータ数は\n"
+
+    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"
+
+    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")
+
+    if usedStorage_gb is not None:
+        result += f"参考: 該当リポジトリのファイルサイズは合計 {usedStorage_gb:.2f} GB"
+    return result
+
+# Gradio インターフェースの定義
+iface = gr.Interface(
+    fn=estimate_gpu_memory,
+    inputs=gr.Textbox(label="モデル名 (例: meta-llama/Llama-3.3-70B-Instruct)"),
+    outputs="text",
+    title="Model memory estimator",
+    description=("Hugging Face Hub APIの返り値に含まれるsafetensors情報から各パラメータサイズを合算し、概算GPUメモリ必要量をGB単位で返します。")
+)
+
+iface.launch()