[fix] modify huggingface url

app/utils.py

@@ -1,3 +1,10 @@
+import re
+
+hf_pattern = re.compile(
+    r"https://huggingface.co/(?P<account>[^/]+)/(?P<repo_id>[^/]+)/.*/(?P<branch>[^/]+)/(?P<filename>[^?]+).*"
+)
+
+
 def human_readable_size(size_in_bytes: int) -> str:
     # Convert file size to a human-readable format
     for unit in ["B", "KB", "MB", "GB", "TB", "PB"]:
@@ -13,3 +20,15 @@ def abbreviate_number(number: int) -> str:
         if number >= threshold:
             return f"{number/threshold:.2f} {unit}"
     return str(number)
+
+
+def cleanup_url(url: str) -> str:
+    match = hf_pattern.match(url)
+    if match:
+        account = match.group("account")
+        repo_id = match.group("repo_id")
+        branch = match.group("branch")
+        filename = match.group("filename")
+        return f"https://huggingface.co/{account}/{repo_id}/resolve/{branch}/{filename}"
+
+    return url.strip()

main.py

@@ -3,11 +3,11 @@ import os
 from pathlib import Path
 
 import gradio as gr
-import pandas as pd
 
 from app.devices import Device
 from app.models import GgufParser
 from app.tables import get_estimate_df, get_gpus_df, get_model_info_df
+from app.utils import cleanup_url
 
 GGUF_PARSER_VERSION = os.getenv("GGUF_PARSER_VERSION", "v0.12.0")
 gguf_parser = Path("gguf-parser-linux-amd64")
@@ -28,6 +28,7 @@ def process_url(url, context_length, device_selection):
     try:
         device_name = device_selection.split(" (")[0]
         selected_device = devices[device_name]
+        url = cleanup_url(url)
         res = os.popen(
             f'./{gguf_parser} --ctx-size={context_length} -url {url} --device-metric "{selected_device.FLOPS};{selected_device.memory_bandwidth}GBps" --json'
         ).read()
@@ -57,13 +58,13 @@ if __name__ == "__main__":
         url_input = gr.Textbox(
             label="GGUF File URL", placeholder="Enter GGUF URL", value=DEFAULT_URL
         )
-        context_length = gr.Number(label="Context Length", value=8192)
+        context_length_input = gr.Number(label="Context Length", value=8192)
         device_dropdown = gr.Dropdown(label="Select Device", choices=device_options)
         submit_btn = gr.Button("Send")
 
         submit_btn.click(
             fn=process_url,
-            inputs=[url_input, context_length, device_dropdown],
+            inputs=[url_input, context_length_input, device_dropdown],
             outputs=[
                 gr.DataFrame(label="Model Info"),
                 gr.DataFrame(label="ESTIMATE"),

tests/test_url.py

@@ -0,0 +1,19 @@
+from app.utils import cleanup_url
+
+# https://huggingface.co/phate334/Llama-3.1-8B-Instruct-Q4_K_M-GGUF/resolve/main/llama-3.1-8b-instruct-q4_k_m.gguf?download=true
+# https://huggingface.co/phate334/Llama-3.1-8B-Instruct-Q4_K_M-GGUF/resolve/main/llama-3.1-8b-instruct-q4_k_m.gguf
+# https://huggingface.co/phate334/Llama-3.1-8B-Instruct-Q4_K_M-GGUF/blob/main/llama-3.1-8b-instruct-q4_k_m.gguf
+# cleanup_url 輸出都應該要是 https://huggingface.co/phate334/Llama-3.1-8B-Instruct-Q4_K_M-GGUF/resolve/main/llama-3.1-8b-instruct-q4_k_m.gguf
+# 其餘非 huggingface.co 的 url 只要前後沒有空白就好
+
+resolve_url = "https://huggingface.co/phate334/Llama-3.1-8B-Instruct-Q4_K_M-GGUF/resolve/main/llama-3.1-8b-instruct-q4_k_m.gguf"
+resolve_url_download = "https://huggingface.co/phate334/Llama-3.1-8B-Instruct-Q4_K_M-GGUF/resolve/main/llama-3.1-8b-instruct-q4_k_m.gguf?download=true"
+blob_url = "https://huggingface.co/phate334/Llama-3.1-8B-Instruct-Q4_K_M-GGUF/blob/main/llama-3.1-8b-instruct-q4_k_m.gguf"
+other_url = "https://git.gss.com.tw/phate_wang/llm/llama-3.1-8b-instruct-q4_k_m.gguf"
+
+
+def test_cleanup_url():
+    assert cleanup_url(resolve_url) == resolve_url
+    assert cleanup_url(resolve_url_download) == resolve_url
+    assert cleanup_url(blob_url) == resolve_url
+    assert cleanup_url(other_url) == other_url