display model structure

backend/hf_model_utils.py

@@ -28,6 +28,8 @@ def module_hash(module):
     rep_bytes = str(rep).encode('utf-8')
     return hashlib.md5(rep_bytes).hexdigest()
 
+def is_number_string(value):
+    return isinstance(value, str) and value.isdigit()
 
 def hf_style_structural_dict(module):
     """
@@ -52,7 +54,7 @@ def hf_style_structural_dict(module):
             count = 1
             j = i + 1
             # Count consecutive children that are structurally identical
-            while j < len(children) and module_hash(children[j][1]) == current_hash:
+            while j < len(children) and is_number_string(name) and module_hash(children[j][1]) == current_hash:
                 count += 1
                 j += 1
 
@@ -93,6 +95,10 @@ def get_model_structure(model_name: str, model_type: str | None):
         "hidden_size": getattr(config, "hidden_size", None),
         "num_hidden_layers": getattr(config, "num_hidden_layers", None),
         "num_attention_heads": getattr(config, "num_attention_heads", None),
+        "image_size": getattr(config, "image_size", None),
+        "intermediate_size": getattr(config, "intermediate_size", None),
+        "patch_size": getattr(config, "patch_size", None),
+        "vocab_size": getattr(config, "vocab_size", None),
         "layers": hf_style_structural_dict(model)
     }
 

frontend/src/App.jsx

@@ -1,6 +1,7 @@
 import { useState } from "react";
 import ModelInputBar from "./components/ModelInputBar";
 import ModelLayersCard from "./components/ModelLayersCard";
+import ModelInfo from "./components/ModelInfo";
 
 export default function App() {
   const [structure, setStructure] = useState(null);
@@ -40,7 +41,7 @@ export default function App() {
       <div className="w-full max-w-3xl bg-white rounded-2xl shadow-lg p-6">
         {/* Header */}
         <h1 className="text-3xl font-bold text-center text-slate-800 mb-4">
-          Hugging Face Model Structure Viewer
+          Transformer Model Structure Viewer
         </h1>
         <p className="text-center text-slate-500 mb-6">
           Enter a model name (e.g. <code>deepseek-ai/deepseek-moe-16b-base</code>) to view its
@@ -56,19 +57,19 @@ export default function App() {
           <div className="text-red-600 text-center font-medium mb-4">{error}</div>
         )}
 
-        {/* Model Structure */}
-        {structure && (
-          <div className="max-h-[500px] overflow-y-auto bg-slate-50 rounded-xl p-4 border border-slate-200 shadow-inner">
-            <pre className="text-sm text-slate-800 whitespace-pre-wrap">
-              {structure.model_type}
-            </pre>
-          </div>
-        )}
+      {/* Model Info */}
+      {structure && (
+        <ModelInfo structure={
+          Object.fromEntries(
+            Object.entries(structure).filter(([key]) => key !== "layers")
+          )
+        } />
+      )}
       </div>
 
-      {/* Model Layers Card */}
+      {/* Model Layers */}
       {structure && (
-        <ModelLayersCard layers={structure?.layers?.children || {}} />
+        <ModelLayersCard name={structure?.layers?.class_name} layers={structure?.layers?.children || {}} />
       )}
 
       {/* Footer */}

frontend/src/components/ModelInfo.jsx

@@ -0,0 +1,39 @@
+import React from "react";
+
+export default function ModelStructure({ structure }) {
+  if (!structure) return null;
+
+  const items = {
+    model_type: "Model Type",
+    hidden_size: "Hidden Size",
+    num_attention_heads: "Number of Attention Heads",
+    num_hidden_layers: "Number of Hidden Layers",
+    image_size: "Image Size",
+    intermediate_size: "Intermediate Size",
+    patch_size: "Patch Size",
+    vocab_size: "Vocab Size",
+  }
+
+  // Find which keys in structure match our interest list
+  const matched = Object.keys(items).filter((k) => k in structure && !!structure[k]);
+
+  return (
+    <div className="max-h-[500px] overflow-y-auto bg-slate-50 rounded-xl p-4 border border-slate-200 shadow-inner">
+      {matched.length === 0 ? (
+        <pre className="text-sm text-slate-800 whitespace-pre-wrap">{structure.model_type}</pre>
+      ) : (
+        <div className="flex flex-wrap gap-2">
+          {matched.map((k) => (
+            <div
+              key={k}
+              className="bg-sky-100 border border-sky-300 text-sky-900 text-sm px-3 py-1 rounded-full flex items-center gap-3 shadow-sm"
+            >
+              <div className="font-medium">{items[k]}:</div>
+              <div className="text-sky-700">{String(structure[k])}</div>
+            </div>
+          ))}
+        </div>
+      )}
+    </div>
+  );
+}

frontend/src/components/ModelInputBar.jsx

@@ -31,51 +31,55 @@ export default function ModelInputBar({ loading, fetchModelStructure }) {
         }
     }
 
-    const [focused, setFocused] = useState(false);
-
     return (
-    <div className="flex flex-col gap-2 mb-4 w-full">
-        <div className="flex gap-2 w-full">
-            <input
-                type="text"
-                value={modelName}
-                onChange={(e) => setModelName(e.target.value)}
-                placeholder={placeholder}
-                className="flex-1 px-4 py-2 rounded-lg border border-slate-200 focus:outline-none focus:ring-2 focus:ring-sky-300"
-                onKeyDown={(e) => {
-                    if (e.key === "Enter") handleFetch();
-                }}
-                onFocus={() => setFocused(true)}
-                onBlur={() => setFocused(false)}
-                aria-label="Model name"
-            />
-
-            <button
-                onClick={handleFetch}
-                disabled={loading}
-                className={`px-4 py-2 rounded-lg font-medium text-white ${
-                    loading ? "bg-slate-400 cursor-wait" : "bg-sky-600 hover:bg-sky-700"
-                }`}
-            >
-                {loading ? "Loading..." : "Fetch"}
-            </button>
-        </div>
-
-        <div className="flex flex-wrap items-center gap-4 text-sm text-slate-700">
-            {options.map((option) => (
-                <label key={option.value} className="flex items-center gap-2">
+        <div className="flex flex-col gap-2 mb-4 w-full">
+            {/* Responsive flex row/col for title and input/button */}
+            <div className="flex flex-col sm:flex-row sm:items-center gap-2 w-full py-2">
+                <h3 className="text-sm font-semibold text-slate-800 mb-2 whitespace-nowrap sm:mb-0">
+                    Model Name:
+                </h3>
+                <div className="flex flex-1 gap-2">
                     <input
-                        type="radio"
-                        name="fetchOption"
-                        value={option.value}
-                        checked={selectedOption === option.value}
-                        onChange={handelRadioChange}
-                        className="w-4 h-4"
+                        type="text"
+                        value={modelName}
+                        onChange={(e) => setModelName(e.target.value)}
+                        placeholder={placeholder}
+                        className="flex-1 px-4 py-2 rounded-lg border border-slate-200 focus:outline-none focus:ring-2 focus:ring-sky-300"
+                        onKeyDown={(e) => {
+                            if (e.key === "Enter") handleFetch();
+                        }}
+                        aria-label="Model name"
                     />
-                    <span>{option.label}</span>
-                </label>
-            ))}
+                    <button
+                        onClick={handleFetch}
+                        disabled={loading}
+                        className={`model-name px-4 py-2 rounded-lg font-medium text-white ${
+                            loading ? "bg-slate-400 cursor-wait" : "bg-sky-600 hover:bg-sky-700"
+                        }`}
+                    >
+                        {loading ? "Loading..." : "View"}
+                    </button>
+                </div>
+            </div>
+
+            <div className="flex flex-col sm:flex-row sm:top gap-2 w-full py-2">
+                <h3 className="text-sm font-semibold text-slate-800 mb-2 whitespace-nowrap sm:mb-0">Model Type:</h3>
+                <div className="flex flex-wrap items-center gap-4 text-sm text-slate-700">
+                    {options.map((option) => (
+                        <label key={option.value} className="model-type flex items-center gap-2">
+                            <input
+                                type="radio"
+                                name="fetchOption"
+                                value={option.value}
+                                checked={selectedOption === option.value}
+                                onChange={handelRadioChange}
+                                className="w-4 h-4"
+                            />
+                            <span>{option.label}</span>
+                        </label>
+                    ))}
+                </div>
+            </div>
         </div>
-    </div>
-);
+    );
 }

frontend/src/components/ModelLayersCard.jsx

@@ -1,28 +1,53 @@
 import React from "react";
 
-function LayerNode({ node }) {
-  if (!node) return null;
+function isNumberString(value) {
+  return typeof value === "string" && /^[0-9]+$/.test(value);
+}
 
+function LayerNode({ node, nodeKey, last, depth = 0 }) {
+  if (!node) return null;
   const name = node.class_name || node.name || "<unknown>";
   const shape = node.params?.weight?.shape || [];
 
   // Normalize children: object map
   let children = [];
+  let keys = [];
   if (node.children && typeof node.children === "object") {
     children = Object.values(node.children);
+    keys = Object.keys(node.children);
   }
 
+  // choose background by depth for subtle alternation
+  const bgByDepth = ["bg-slate-50", "bg-slate-100", "bg-slate-200", "bg-slate-50"];
+  const bgClass = bgByDepth[depth % bgByDepth.length];
+  const isEmbed = String(name).toLowerCase().includes("embed");
+  const embedStyle = isEmbed ? { backgroundColor: "#fbdfe2" } : undefined;
+  const isAttn = String(name).toLowerCase().includes("attention") || String(name).toLowerCase().includes("attn");
+  const attnStyle = isAttn ? { backgroundColor: "#fddfba" } : undefined;
+  const isFFN = String(name).toLowerCase().includes("ffn") || String(name).toLowerCase().includes("mlp");
+  const ffnStyle = isFFN ? { backgroundColor: "#c2e6f8" } : undefined;
+  const isNorm = String(name).toLowerCase().includes("norm");
+  const normStyle = isNorm ? { backgroundColor: "#f3f7c3" } : undefined;
+  const isLastNorm = last && isNorm;
+  const lastNormStyle = isLastNorm ? { backgroundColor: "#DEDFF1" } : undefined;
+
+
   return (
-    <div className="pl-2">
-      <div className="flex items-center gap-2">
+    <div
+      className={`pl-2 ${!isEmbed ? bgClass : ""} rounded-md border border-slate-200 p-2`} 
+      style={embedStyle || attnStyle || ffnStyle || lastNormStyle || normStyle}
+    > 
+      <div className="flex items-center gap-3">
         <div className="text-sm text-slate-800 font-medium">{name}</div>
+        {nodeKey && !isNumberString(nodeKey) && <div className="text-sm text-slate-800 font-medium">({nodeKey})</div>}
         <div className="text-xs text-slate-500">{shape.join(" x ")}</div>
+        <div>{node.num_repeats && <span className="text-blue-600 font-bold tracking-wide">x {node.num_repeats}</span>}</div>
       </div>
 
       {children.length > 0 && (
-        <div className="pl-4 mt-2 border-l border-slate-100">
+        <div className="pl-4 mt-2">
           {children.map((child, idx) => (
-            <LayerNode key={idx} node={child} />
+            <LayerNode key={idx} node={child} nodeKey={keys[idx]} depth={depth + 1} />
           ))}
         </div>
       )}
@@ -30,15 +55,18 @@ function LayerNode({ node }) {
   );
 }
 
-export default function ModelLayersCard({ layers = {} }) {
+export default function ModelLayersCard({ layers = {}, name = "" }) {
   let rootNodes = [];
+  let rootKeys = [];
 
   if (layers && typeof layers === "object") {
     if (layers.children && typeof layers.children === "object") {
       rootNodes = Object.values(layers.children);
+      rootKeys = Object.keys(layers.children);
     } else {
       // Fallback: convert the top-level keyed object into an array
       rootNodes = Object.values(layers);
+      rootKeys = Object.keys(layers);
     }
   }
 
@@ -48,10 +76,10 @@ export default function ModelLayersCard({ layers = {} }) {
 
   return (
     <div className="w-full max-w-3xl mt-6 bg-white rounded-2xl shadow-lg p-4">
-      <h2 className="text-lg font-semibold text-slate-800 mb-3">Model Layers</h2>
+      <h2 className="text-lg font-semibold text-slate-800 mb-3">{name}</h2>
       <div className="space-y-2">
         {rootNodes.map((node, idx) => (
-          <LayerNode key={idx} node={node} />
+          <LayerNode key={idx} node={node} nodeKey={rootKeys[idx]} last={idx == rootNodes.length - 1} />
         ))}
       </div>
     </div>

frontend/src/index.css

@@ -1 +1,5 @@
 @import "tailwindcss";
+
+.model-name, .model-type {
+    cursor: pointer;
+}