Add download progress logging and compact UI with tabs

app.py

@@ -3,6 +3,10 @@ import tempfile
 import os
 import glob
 import shutil
+import logging
+import io
+import sys
+from contextlib import redirect_stdout, redirect_stderr
 
 import gradio as gr
 import torch
@@ -10,13 +14,27 @@ from huggingface_hub import hf_hub_download, scan_cache_dir
 from safetensors import safe_open
 
 
-def get_param(model_id: str, param_key: str):
+def get_param(model_id: str, param_key: str, log_buffer: io.StringIO, progress: gr.Progress):
     """
     Download and return a specific parameter tensor from a Hugging Face model.
     """
     # Try to download the index file (for sharded models)
     try:
-        index_path = hf_hub_download(model_id, "model.safetensors.index.json")
+        log_buffer.write(f"📥 Downloading index file for {model_id}...\n")
+        progress(0.1, desc="Downloading index...")
+
+        # Capture tqdm output from stderr
+        stderr_capture = io.StringIO()
+        with redirect_stderr(stderr_capture):
+            index_path = hf_hub_download(
+                model_id, "model.safetensors.index.json")
+
+        stderr_output = stderr_capture.getvalue()
+        if stderr_output:
+            log_buffer.write(stderr_output + "\n")
+
+        log_buffer.write(f"✓ Index file found: {index_path}\n")
+
         with open(index_path, "r", encoding="utf-8") as f:
             index = json.load(f)
         weight_map = index["weight_map"]
@@ -25,13 +43,34 @@ def get_param(model_id: str, param_key: str):
                 f"Parameter '{param_key}' not found in model. Available keys: {list(weight_map.keys())[:10]}..."
             )
         shard_file = weight_map[param_key]
-    except Exception:
-        shard_file = "model.safetensors"
+        log_buffer.write(f"✓ Parameter found in shard: {shard_file}\n")
+    except Exception as e:
+        if "404" in str(e) or "not found" in str(e).lower():
+            log_buffer.write("ℹ️ No index file, trying single model file...\n")
+            shard_file = "model.safetensors"
+        else:
+            raise
+
+    log_buffer.write(f"📥 Downloading shard: {shard_file}...\n")
+    progress(0.3, desc=f"Downloading {shard_file}...")
+
+    # Capture download progress
+    stderr_capture = io.StringIO()
+    with redirect_stderr(stderr_capture):
+        shard_path = hf_hub_download(model_id, shard_file)
 
-    shard_path = hf_hub_download(model_id, shard_file)
+    stderr_output = stderr_capture.getvalue()
+    if stderr_output:
+        log_buffer.write(stderr_output + "\n")
 
+    log_buffer.write(f"✓ Shard downloaded: {shard_path}\n")
+    progress(0.7, desc="Loading tensor...")
+
+    log_buffer.write(f"🔍 Loading tensor '{param_key}'...\n")
     with safe_open(shard_path, framework="pt") as f:
         tensor = f.get_tensor(param_key)
+    log_buffer.write(f"✓ Tensor loaded successfully\n")
+    progress(0.9, desc="Finalizing...")
 
     return tensor
 
@@ -67,16 +106,23 @@ def format_tensor_info(tensor: torch.Tensor) -> str:
     return "<br>".join(info)
 
 
-def fetch_param(model_id: str, param_key: str):
+def fetch_param(model_id: str, param_key: str, progress=gr.Progress()):
     """Fetch parameter and return formatted info and tensor preview."""
+    log_buffer = io.StringIO()
+
     if not model_id or not param_key:
-        return "Please provide both model ID and parameter key.", "", None
+        return "Please provide both model ID and parameter key.", "", None, "❌ Missing required inputs"
 
     try:
-        tensor = get_param(model_id, param_key)
+        log_buffer.write(f"🚀 Starting download for {model_id}\n")
+        log_buffer.write(f"🎯 Target parameter: {param_key}\n\n")
+        progress(0, desc="Initializing...")
+
+        tensor = get_param(model_id, param_key, log_buffer, progress)
         info = format_tensor_info(tensor)
 
         # Create tensor preview (first few elements)
+        log_buffer.write(f"\n📊 Creating preview...\n")
         flat = tensor.flatten()
         preview_size = min(100, flat.numel())
         preview = flat[:preview_size].tolist()
@@ -86,14 +132,19 @@ def fetch_param(model_id: str, param_key: str):
             preview_str += f"\n\n... and {flat.numel() - preview_size:,} more values"
 
         # Save tensor for download
+        log_buffer.write(f"💾 Saving tensor for download...\n")
         temp_dir = tempfile.gettempdir()
         safe_param_key = param_key.replace("/", "_").replace(".", "_")
         download_path = os.path.join(temp_dir, f"{safe_param_key}.pt")
         torch.save(tensor, download_path)
+        log_buffer.write(f"✓ Saved to: {download_path}\n")
 
-        return info, preview_str, download_path
+        progress(1.0, desc="Complete!")
+        log_buffer.write(f"\n✅ All operations completed successfully!\n")
+        return info, preview_str, download_path, log_buffer.getvalue()
     except Exception as e:
-        return f"**Error:** {str(e)}", "", None
+        log_buffer.write(f"\n❌ Error: {str(e)}\n")
+        return f"**Error:** {str(e)}", "", None, log_buffer.getvalue()
 
 
 def list_keys(model_id: str):
@@ -123,7 +174,7 @@ def clear_temp_files():
                 deleted_files.append(os.path.basename(file))
             except Exception:
                 pass
-        
+
         if deleted_files:
             files_list = "\n".join(deleted_files)
             return f"✅ Cleared {count} temporary file(s):\n\n{files_list}"
@@ -139,10 +190,10 @@ def clear_hf_cache():
         cache_info = scan_cache_dir()
         total_size = cache_info.size_on_disk
         total_repos = len(cache_info.repos)
-        
+
         if total_repos == 0:
             return "✅ HuggingFace cache is already empty"
-        
+
         # Get cache directory and clear it
         cache_dir = os.path.expanduser("~/.cache/huggingface/hub")
         if os.path.exists(cache_dir):
@@ -162,9 +213,10 @@ def get_cache_info():
         # Temp files
         temp_dir = tempfile.gettempdir()
         pt_files = glob.glob(os.path.join(temp_dir, "*.pt"))
-        temp_size = sum(os.path.getsize(f) for f in pt_files if os.path.exists(f))
+        temp_size = sum(os.path.getsize(f)
+                        for f in pt_files if os.path.exists(f))
         temp_size_mb = temp_size / (1024 * 1024)
-        
+
         # HF cache
         try:
             cache_info = scan_cache_dir()
@@ -173,7 +225,7 @@ def get_cache_info():
         except Exception:
             hf_size_mb = 0
             hf_repos = 0
-        
+
         info = f"📊 Cache Info:\n\n"
         info += f"Temp .pt files: {len(pt_files)} file(s), {temp_size_mb:.2f} MB\n"
         info += f"HuggingFace cache: {hf_repos} repo(s), {hf_size_mb:.2f} MB\n"
@@ -188,31 +240,27 @@ custom_css = """
 * {
     font-family: Consolas, Monaco, 'Courier New', monospace !important;
 }
+.compact-row {
+    gap: 0.5rem !important;
+}
 """
 
 with gr.Blocks(title="HuggingFace Model Weight Inspector") as demo:
-    gr.Markdown(
-        """
-        # 🔍 HuggingFace Model Weight Inspector
-        
-        Inspect specific parameter tensors from any HuggingFace model without downloading the entire model.
-        """
-    )
+    gr.Markdown("# 🔍 HuggingFace Model Weight Inspector")
 
     with gr.Row():
         model_id_input = gr.Textbox(
             label="Model ID",
             placeholder="e.g., meta-llama/Llama-2-7b-hf",
             value="zai-org/GLM-5",
+            scale=4,
         )
-
-    with gr.Row():
-        list_keys_btn = gr.Button("📋 List Available Keys", variant="secondary")
+        list_keys_btn = gr.Button("📋 List Keys", variant="secondary", scale=1)
 
     keys_output = gr.Textbox(
         label="Available Parameter Keys",
-        lines=5,
-        max_lines=10,
+        lines=3,
+        max_lines=8,
     )
 
     with gr.Row():
@@ -220,30 +268,28 @@ with gr.Blocks(title="HuggingFace Model Weight Inspector") as demo:
             label="Parameter Key",
             placeholder="e.g., model.embed_tokens.weight",
             value="model.layers.5.mlp.gate.e_score_correction_bias",
+            scale=4,
         )
-
-    with gr.Row():
-        fetch_btn = gr.Button("🔎 Fetch Parameter", variant="primary")
-
-    with gr.Row():
-        with gr.Column():
-            info_output = gr.Markdown(label="Tensor Info")
-        with gr.Column():
-            preview_output = gr.Markdown(label="Tensor Preview")
-    
-    with gr.Row():
-        download_output = gr.File(label="Download Tensor (.pt file)")
-    
-    with gr.Row():
-        with gr.Column():
-            clear_temp_btn = gr.Button("🗑️ Clear Temp Files", variant="secondary")
-        with gr.Column():
-            clear_hf_btn = gr.Button("🗑️ Clear HF Cache", variant="secondary")
-        with gr.Column():
-            get_info_btn = gr.Button("📊 Get Cache Info", variant="secondary")
-    
-    with gr.Row():
-        clear_status = gr.Textbox(label="Status", interactive=False, lines=5)
+        fetch_btn = gr.Button("🔎 Fetch", variant="primary", scale=1)
+
+    with gr.Tabs():
+        with gr.Tab("Tensor Info"):
+            with gr.Row():
+                with gr.Column():
+                    info_output = gr.Markdown()
+                with gr.Column():
+                    preview_output = gr.Markdown()
+        
+        with gr.Tab("Download & Logs"):
+            download_output = gr.File(label="Download Tensor (.pt file)")
+            log_output = gr.Textbox(label="📋 Download Log", lines=6, interactive=False)
+        
+        with gr.Tab("Cache Management"):
+            with gr.Row():
+                clear_temp_btn = gr.Button("🗑️ Temp", variant="secondary", scale=1)
+                clear_hf_btn = gr.Button("🗑️ HF Cache", variant="secondary", scale=1)
+                get_info_btn = gr.Button("📊 Info", variant="secondary", scale=1)
+            clear_status = gr.Textbox(label="Status", interactive=False, lines=4)
 
     # Event handlers
     list_keys_btn.click(
@@ -255,21 +301,21 @@ with gr.Blocks(title="HuggingFace Model Weight Inspector") as demo:
     fetch_btn.click(
         fn=fetch_param,
         inputs=[model_id_input, param_key_input],
-        outputs=[info_output, preview_output, download_output],
+        outputs=[info_output, preview_output, download_output, log_output],
     )
-    
+
     clear_temp_btn.click(
         fn=clear_temp_files,
         inputs=[],
         outputs=[clear_status],
     )
-    
+
     clear_hf_btn.click(
         fn=clear_hf_cache,
         inputs=[],
         outputs=[clear_status],
     )
-    
+
     get_info_btn.click(
         fn=get_cache_info,
         inputs=[],