Introduction and Motivation
The proliferation of large language models has led to an ecosystem where models are distributed in various formats, each optimized for different frameworks and hardware platforms. Researchers, engineers, and enthusiasts working with local LLM deployments often need to understand the intricate details of these models before using them in production or research environments. The challenge lies in the fact that each format has its own binary structure, metadata organization, and weight storage mechanism. A unified tool that can parse multiple formats and present a comprehensive analysis becomes invaluable in this context.
This article presents the design and implementation of such a tool, one that can read models stored in GGUF (GPT-Generated Unified Format), PyTorch checkpoint format, TensorFlow SavedModel format, and Apple's MLX format. The tool goes beyond simple metadata extraction by providing deep insights into the model architecture, including the number and configuration of transformer layers, attention mechanisms, feed-forward network dimensions, quantization levels, weight distributions, and other critical parameters. Furthermore, the tool leverages matplotlib to generate visual representations of weight distributions, layer structures, and quantization patterns, making the analysis accessible even to those who may not be deeply familiar with the mathematical underpinnings of transformer architectures.
Architectural Design Philosophy
The tool is designed following clean architecture principles, with clear separation of concerns. At its core, the system consists of several distinct components that work together harmoniously. The first component is the format detector, which examines the file structure and magic bytes to determine which format parser to invoke. The second layer consists of format-specific readers, each implementing a common interface but handling the peculiarities of their respective formats. The third component is the analysis engine, which takes the parsed model data and extracts meaningful insights about the architecture. The fourth component is the visualization system, which creates graphical representations of the data. Finally, the report generator synthesizes all this information into a coherent overview.
This separation ensures that adding support for new formats requires only implementing a new reader without modifying the analysis or visualization logic. Similarly, new analysis techniques can be added without touching the format parsers. This modularity is essential for maintainability and extensibility.
Format Detection System
Before we can analyze a model, we must first determine its format. Different model formats have distinct signatures that can be identified by examining the file header or structure. Let us begin with the format detector implementation.
import os
import struct
from pathlib import Path
from enum import Enum
class ModelFormat(Enum):
"""Enumeration of supported model formats"""
GGUF = "gguf"
PYTORCH = "pytorch"
TENSORFLOW = "tensorflow"
MLX = "mlx"
UNKNOWN = "unknown"
class FormatDetector:
"""
Detects the format of a model file by examining its structure
and magic bytes. This class encapsulates all format detection
logic in one place.
"""
def __init__(self):
# GGUF files start with these magic bytes
self.gguf_magic = b'GGUF'
# PyTorch files are typically ZIP archives with specific structure
self.pytorch_magic = b'PK\x03\x04'
# TensorFlow SavedModel uses protocol buffers
self.tf_magic = b'\x08\x03'
def detect_format(self, model_path):
"""
Examines the file or directory structure to determine
the model format. Returns a ModelFormat enum value.
Args:
model_path: Path to the model file or directory
Returns:
ModelFormat enum indicating the detected format
"""
path = Path(model_path)
# Check if it's a directory (TensorFlow SavedModel style)
if path.is_dir():
return self._detect_directory_format(path)
# For files, read the magic bytes
if path.is_file():
return self._detect_file_format(path)
return ModelFormat.UNKNOWN
def _detect_file_format(self, file_path):
"""
Detects format by reading the first few bytes of the file
and comparing against known magic byte sequences.
"""
with open(file_path, 'rb') as f:
# Read first 16 bytes which is enough for most magic numbers
header = f.read(16)
if header.startswith(self.gguf_magic):
return ModelFormat.GGUF
if header.startswith(self.pytorch_magic):
# PyTorch files are ZIP archives, need further verification
return self._verify_pytorch_format(file_path)
# Check for TensorFlow protobuf format
if self._is_tensorflow_pb(file_path):
return ModelFormat.TENSORFLOW
# Check for MLX format (typically .npz or custom format)
if self._is_mlx_format(file_path):
return ModelFormat.MLX
return ModelFormat.UNKNOWN
def _detect_directory_format(self, dir_path):
"""
Detects format for directory-based model storage,
primarily used by TensorFlow SavedModel.
"""
# TensorFlow SavedModel has saved_model.pb file
if (dir_path / 'saved_model.pb').exists():
return ModelFormat.TENSORFLOW
# MLX models might be stored as directory with weights
if (dir_path / 'weights.npz').exists():
return ModelFormat.MLX
return ModelFormat.UNKNOWN
def _verify_pytorch_format(self, file_path):
"""
PyTorch checkpoint files are ZIP archives containing
pickle data. We verify by checking for expected entries.
"""
import zipfile
try:
with zipfile.ZipFile(file_path, 'r') as z:
# PyTorch models typically have data.pkl or similar
namelist = z.namelist()
if any('data.pkl' in name or '.pth' in name for name in namelist):
return ModelFormat.PYTORCH
except zipfile.BadZipFile:
pass
return ModelFormat.UNKNOWN
def _is_tensorflow_pb(self, file_path):
"""
Checks if file is a TensorFlow protobuf by attempting
to parse it as such.
"""
try:
# This is a simplified check
with open(file_path, 'rb') as f:
content = f.read(100)
# TensorFlow protobufs have specific patterns
if b'tensorflow' in content.lower():
return True
except Exception:
pass
return False
def _is_mlx_format(self, file_path):
"""
MLX format detection based on file extension and structure
"""
if file_path.suffix in ['.npz', '.safetensors']:
# Could be MLX format, need to verify structure
return True
return False
The format detector serves as the entry point for our analysis tool. When a user provides a path to a model, this detector examines the file structure to determine which specialized parser should be invoked. The implementation uses magic byte sequences, which are specific byte patterns that appear at the beginning of files to identify their type. For instance, GGUF files begin with the four-byte sequence "GGUF", making them immediately identifiable. PyTorch models, being stored as ZIP archives containing pickled Python objects, start with the ZIP file signature. TensorFlow models use protocol buffers, which have their own characteristic patterns.
The detector implements separate methods for file-based and directory-based detection because some formats like TensorFlow's SavedModel use a directory structure rather than a single file. This design decision reflects the reality of how different frameworks choose to serialize their models. The verification methods perform additional checks beyond magic bytes to ensure accurate detection, as some formats might share similar headers but differ in their internal structure.
GGUF Format Reader Implementation
The GGUF format is a relatively new unified format designed specifically for efficient storage and loading of large language models. It was created to address the fragmentation in model distribution formats and provides a self-describing structure with comprehensive metadata. Let us examine how to parse this format.
import struct
import numpy as np
from typing import Dict, Any, List, Tuple
class GGUFReader:
"""
Reads and parses GGUF format model files, extracting metadata,
architecture information, and tensor data. GGUF is a binary format
with a specific structure that must be parsed sequentially.
"""
def __init__(self, file_path):
self.file_path = file_path
self.metadata = {}
self.tensors = {}
self.architecture_info = {}
self.file_handle = None
# GGUF type mappings
self.gguf_types = {
0: ('uint8', 1),
1: ('int8', 1),
2: ('uint16', 2),
3: ('int16', 2),
4: ('uint32', 4),
5: ('int32', 4),
6: ('float32', 4),
7: ('bool', 1),
8: ('string', None),
9: ('array', None),
10: ('uint64', 8),
11: ('int64', 8),
12: ('float64', 8),
}
def read_model(self):
"""
Main entry point for reading a GGUF file. This method
orchestrates the parsing process, reading the header,
metadata, and tensor information in sequence.
Returns:
Dictionary containing all extracted information
"""
with open(self.file_path, 'rb') as f:
self.file_handle = f
# Read and verify magic number
magic = f.read(4)
if magic != b'GGUF':
raise ValueError("Not a valid GGUF file")
# Read version number (uint32)
version = struct.unpack('<I', f.read(4))[0]
self.metadata['gguf_version'] = version
# Read tensor count and metadata count
tensor_count = struct.unpack('<Q', f.read(8))[0]
metadata_count = struct.unpack('<Q', f.read(8))[0]
# Parse metadata key-value pairs
self._read_metadata(metadata_count)
# Parse tensor information
self._read_tensor_info(tensor_count)
# Extract architecture-specific information
self._extract_architecture_info()
return {
'format': 'GGUF',
'version': version,
'metadata': self.metadata,
'tensors': self.tensors,
'architecture': self.architecture_info
}
def _read_metadata(self, count):
"""
Reads metadata key-value pairs from the GGUF file.
Each entry consists of a key string and a typed value.
"""
for i in range(count):
# Read key (length-prefixed string)
key = self._read_string()
# Read value type
value_type = struct.unpack('<I', self.file_handle.read(4))[0]
# Read value based on type
value = self._read_typed_value(value_type)
self.metadata[key] = value
def _read_string(self):
"""
Reads a length-prefixed string from the file.
GGUF strings are stored as uint64 length followed by UTF-8 bytes.
"""
length = struct.unpack('<Q', self.file_handle.read(8))[0]
string_bytes = self.file_handle.read(length)
return string_bytes.decode('utf-8')
def _read_typed_value(self, value_type):
"""
Reads a value of the specified type from the file.
Handles all GGUF data types including primitives,
strings, and arrays.
"""
if value_type in [0, 1, 2, 3, 4, 5, 6, 10, 11, 12]:
# Numeric types
type_name, size = self.gguf_types[value_type]
data = self.file_handle.read(size)
if type_name == 'uint8':
return struct.unpack('<B', data)[0]
elif type_name == 'int8':
return struct.unpack('<b', data)[0]
elif type_name == 'uint16':
return struct.unpack('<H', data)[0]
elif type_name == 'int16':
return struct.unpack('<h', data)[0]
elif type_name == 'uint32':
return struct.unpack('<I', data)[0]
elif type_name == 'int32':
return struct.unpack('<i', data)[0]
elif type_name == 'float32':
return struct.unpack('<f', data)[0]
elif type_name == 'uint64':
return struct.unpack('<Q', data)[0]
elif type_name == 'int64':
return struct.unpack('<q', data)[0]
elif type_name == 'float64':
return struct.unpack('<d', data)[0]
elif value_type == 7: # bool
return struct.unpack('<?', self.file_handle.read(1))[0]
elif value_type == 8: # string
return self._read_string()
elif value_type == 9: # array
return self._read_array()
return None
def _read_array(self):
"""
Reads an array value, which consists of a type indicator,
length, and then the array elements.
"""
element_type = struct.unpack('<I', self.file_handle.read(4))[0]
array_length = struct.unpack('<Q', self.file_handle.read(8))[0]
array_elements = []
for i in range(array_length):
element = self._read_typed_value(element_type)
array_elements.append(element)
return array_elements
def _read_tensor_info(self, count):
"""
Reads tensor metadata including names, dimensions, types,
and offsets. Does not read the actual tensor data yet.
"""
for i in range(count):
# Read tensor name
tensor_name = self._read_string()
# Read number of dimensions
n_dims = struct.unpack('<I', self.file_handle.read(4))[0]
# Read dimension sizes
dimensions = []
for d in range(n_dims):
dim_size = struct.unpack('<Q', self.file_handle.read(8))[0]
dimensions.append(dim_size)
# Read tensor type
tensor_type = struct.unpack('<I', self.file_handle.read(4))[0]
# Read offset to tensor data
offset = struct.unpack('<Q', self.file_handle.read(8))[0]
self.tensors[tensor_name] = {
'dimensions': dimensions,
'type': tensor_type,
'offset': offset,
'shape': tuple(dimensions)
}
def _extract_architecture_info(self):
"""
Extracts architecture-specific information from metadata,
such as number of layers, attention heads, embedding dimensions, etc.
"""
# Common architecture parameters in GGUF metadata
arch_keys = [
'general.architecture',
'llama.context_length',
'llama.embedding_length',
'llama.block_count',
'llama.feed_forward_length',
'llama.attention.head_count',
'llama.attention.head_count_kv',
'llama.attention.layer_norm_rms_epsilon',
'llama.rope.dimension_count',
'general.quantization_version',
'general.file_type'
]
for key in arch_keys:
if key in self.metadata:
# Extract the parameter name from the key
param_name = key.split('.')[-1]
self.architecture_info[param_name] = self.metadata[key]
The GGUF reader implementation demonstrates the complexity involved in parsing binary model formats. The format uses a sequential structure where the header is followed by metadata entries and then tensor information. Each component must be read in the correct order because the format does not provide random access to different sections. The magic number verification at the beginning ensures we are indeed reading a GGUF file and not some other format that might have been misidentified.
The metadata section contains key-value pairs that describe the model's architecture and properties. These pairs use a type system that supports primitives like integers and floats, as well as complex types like strings and arrays. The type system is crucial because it allows the format to be self-describing without requiring external schema files. When we encounter a metadata entry, we first read the key as a length-prefixed string, then read a type indicator, and finally read the value according to that type.
The tensor information section describes each weight tensor in the model without actually storing the tensor data inline. Instead, it stores the tensor's name, shape, data type, and an offset pointing to where the actual tensor data resides in the file. This design allows for efficient memory-mapped loading where tensor data can be accessed without loading the entire file into memory. For our analysis tool, we extract this structural information to understand the model's architecture without necessarily loading all the weights.
PyTorch Format Reader Implementation
PyTorch models are typically saved using the torch.save function, which creates a file containing pickled Python objects. These files are actually ZIP archives with a specific internal structure. Reading them requires understanding both the ZIP format and PyTorch's serialization conventions.
import torch
import zipfile
import pickle
from collections import OrderedDict
class PyTorchReader:
"""
Reads PyTorch model checkpoint files, extracting state dictionaries,
model architecture information, and training metadata. PyTorch
checkpoints can contain not just model weights but also optimizer
states, training progress, and custom metadata.
"""
def __init__(self, file_path):
self.file_path = file_path
self.state_dict = None
self.metadata = {}
self.architecture_info = {}
self.optimizer_state = None
self.training_info = {}
def read_model(self):
"""
Loads a PyTorch checkpoint and extracts all available information.
PyTorch checkpoints can be saved in different formats depending on
how they were created.
Returns:
Dictionary containing model weights and metadata
"""
try:
# Load the checkpoint
checkpoint = torch.load(self.file_path, map_location='cpu')
# Determine checkpoint structure
if isinstance(checkpoint, dict):
self._parse_checkpoint_dict(checkpoint)
elif isinstance(checkpoint, OrderedDict):
# Direct state dict
self.state_dict = checkpoint
else:
# Might be a model object
self._parse_model_object(checkpoint)
# Analyze the state dict to extract architecture info
self._analyze_state_dict()
# Extract quantization information if present
self._detect_quantization()
except Exception as e:
raise ValueError(f"Failed to load PyTorch model: {str(e)}")
return {
'format': 'PyTorch',
'state_dict': self.state_dict,
'metadata': self.metadata,
'architecture': self.architecture_info,
'training_info': self.training_info
}
def _parse_checkpoint_dict(self, checkpoint):
"""
Parses a checkpoint dictionary that may contain multiple
components like model state, optimizer state, epoch number, etc.
"""
# Common keys in PyTorch checkpoints
if 'model_state_dict' in checkpoint:
self.state_dict = checkpoint['model_state_dict']
elif 'state_dict' in checkpoint:
self.state_dict = checkpoint['state_dict']
elif 'model' in checkpoint:
self.state_dict = checkpoint['model']
else:
# The checkpoint itself might be the state dict
self.state_dict = checkpoint
# Extract optimizer state if present
if 'optimizer_state_dict' in checkpoint:
self.optimizer_state = checkpoint['optimizer_state_dict']
# Extract training metadata
training_keys = ['epoch', 'iteration', 'best_loss', 'best_accuracy']
for key in training_keys:
if key in checkpoint:
self.training_info[key] = checkpoint[key]
# Extract any custom metadata
for key, value in checkpoint.items():
if key not in ['model_state_dict', 'state_dict', 'model',
'optimizer_state_dict'] and not key in training_keys:
self.metadata[key] = value
def _parse_model_object(self, model):
"""
Extracts state dict from a model object if the entire
model was saved rather than just the state dict.
"""
if hasattr(model, 'state_dict'):
self.state_dict = model.state_dict()
else:
raise ValueError("Unable to extract state dict from model object")
def _analyze_state_dict(self):
"""
Analyzes the state dictionary to infer model architecture.
This involves examining layer names, tensor shapes, and
patterns that indicate specific architectural components.
"""
if self.state_dict is None:
return
layer_pattern_counts = {
'attention': 0,
'mlp': 0,
'layer_norm': 0,
'embedding': 0
}
total_parameters = 0
layer_info = {}
for name, tensor in self.state_dict.items():
# Count parameters
param_count = tensor.numel()
total_parameters += param_count
# Analyze layer names to identify components
if 'attn' in name.lower() or 'attention' in name.lower():
layer_pattern_counts['attention'] += 1
if 'mlp' in name.lower() or 'ffn' in name.lower():
layer_pattern_counts['mlp'] += 1
if 'norm' in name.lower():
layer_pattern_counts['layer_norm'] += 1
if 'embed' in name.lower():
layer_pattern_counts['embedding'] += 1
# Store detailed layer information
layer_info[name] = {
'shape': list(tensor.shape),
'dtype': str(tensor.dtype),
'parameters': param_count
}
self.architecture_info['total_parameters'] = total_parameters
self.architecture_info['layer_counts'] = layer_pattern_counts
self.architecture_info['layer_details'] = layer_info
# Infer number of transformer blocks
self._infer_transformer_blocks()
def _infer_transformer_blocks(self):
"""
Attempts to determine the number of transformer blocks
by analyzing the naming patterns in the state dict.
"""
block_numbers = set()
for name in self.state_dict.keys():
# Look for patterns like 'blocks.0', 'layers.1', 'h.2', etc.
parts = name.split('.')
for i, part in enumerate(parts):
if part in ['blocks', 'layers', 'h', 'transformer']:
if i + 1 < len(parts) and parts[i + 1].isdigit():
block_numbers.add(int(parts[i + 1]))
if block_numbers:
self.architecture_info['num_blocks'] = max(block_numbers) + 1
self.architecture_info['block_indices'] = sorted(list(block_numbers))
def _detect_quantization(self):
"""
Detects if the model uses quantization by examining
tensor dtypes and looking for quantization-specific patterns.
"""
quantized = False
quantization_info = {}
for name, tensor in self.state_dict.items():
dtype_str = str(tensor.dtype)
# Check for quantized dtypes
if 'qint' in dtype_str or 'quint' in dtype_str:
quantized = True
quantization_info[name] = dtype_str
# Check for scale/zero_point parameters (dynamic quantization)
if 'scale' in name or 'zero_point' in name:
quantized = True
self.architecture_info['is_quantized'] = quantized
if quantized:
self.architecture_info['quantization_details'] = quantization_info
The PyTorch reader handles the complexity of PyTorch's flexible checkpoint format. Unlike GGUF which has a rigid structure, PyTorch checkpoints can be saved in many different ways depending on the developer's preferences. Some checkpoints contain only the model weights as a state dictionary, while others include optimizer states, training progress information, and custom metadata. Our reader must handle all these variations gracefully.
The state dictionary is the core component we need to extract. It maps parameter names to tensor values. The naming convention in the state dictionary reveals the model's architecture. For example, names like "transformer.h.0.attn.c_attn.weight" tell us about the hierarchical structure of the model. By parsing these names, we can infer how many transformer blocks exist, what components each block contains, and how they are organized.
The analysis of the state dictionary goes beyond simple extraction. We count parameters to determine model size, identify patterns in layer names to classify components, and examine tensor shapes to understand dimensions. For instance, an attention layer's query, key, and value projection matrices have specific shape relationships that we can verify. This deep analysis allows us to provide comprehensive insights even when the checkpoint doesn't include explicit architecture metadata.
TensorFlow Format Reader Implementation
TensorFlow models can be saved in several formats, with SavedModel being the most common for production deployments. SavedModel uses a directory structure containing protocol buffer files and variable data. Reading this format requires understanding TensorFlow's serialization mechanisms.
import tensorflow as tf
import os
from pathlib import Path
class TensorFlowReader:
"""
Reads TensorFlow SavedModel format, extracting model architecture,
variables, and metadata. TensorFlow's SavedModel format is more
complex than simple checkpoint files as it includes the complete
computational graph and serving signatures.
"""
def __init__(self, model_path):
self.model_path = Path(model_path)
self.model = None
self.variables = {}
self.architecture_info = {}
self.metadata = {}
self.graph_info = {}
def read_model(self):
"""
Loads a TensorFlow SavedModel and extracts comprehensive
information about its structure, variables, and operations.
Returns:
Dictionary containing all extracted information
"""
try:
# Load the SavedModel
self.model = tf.saved_model.load(str(self.model_path))
# Extract variables
self._extract_variables()
# Analyze model structure
self._analyze_structure()
# Extract signatures
self._extract_signatures()
# Detect architecture patterns
self._detect_architecture_patterns()
except Exception as e:
raise ValueError(f"Failed to load TensorFlow model: {str(e)}")
return {
'format': 'TensorFlow',
'variables': self.variables,
'architecture': self.architecture_info,
'metadata': self.metadata,
'graph_info': self.graph_info
}
def _extract_variables(self):
"""
Extracts all trainable and non-trainable variables from the model.
TensorFlow variables are the equivalent of PyTorch parameters.
"""
# Get all trackable objects
all_variables = []
# Try to get variables from different possible locations
if hasattr(self.model, 'variables'):
all_variables = self.model.variables
elif hasattr(self.model, 'trainable_variables'):
all_variables = self.model.trainable_variables
total_params = 0
for var in all_variables:
var_name = var.name
var_shape = var.shape.as_list()
var_dtype = var.dtype.name
# Calculate number of parameters
param_count = 1
for dim in var_shape:
if dim is not None:
param_count *= dim
total_params += param_count
self.variables[var_name] = {
'shape': var_shape,
'dtype': var_dtype,
'parameters': param_count,
'trainable': var.trainable
}
self.architecture_info['total_parameters'] = total_params
self.architecture_info['num_variables'] = len(all_variables)
def _analyze_structure(self):
"""
Analyzes the model's layer structure if it's a Keras model
or examines the computational graph for other model types.
"""
# Check if this is a Keras model
if hasattr(self.model, 'layers'):
self._analyze_keras_model()
else:
self._analyze_graph_model()
def _analyze_keras_model(self):
"""
Extracts detailed information from a Keras model including
layer types, configurations, and connections.
"""
layers_info = []
for layer in self.model.layers:
layer_config = {
'name': layer.name,
'type': layer.__class__.__name__,
'trainable': layer.trainable,
'parameters': layer.count_params()
}
# Get layer configuration
try:
config = layer.get_config()
layer_config['config'] = config
except Exception:
pass
# Get input and output shapes
try:
layer_config['input_shape'] = layer.input_shape
layer_config['output_shape'] = layer.output_shape
except Exception:
pass
layers_info.append(layer_config)
self.architecture_info['layers'] = layers_info
self.architecture_info['num_layers'] = len(layers_info)
# Identify transformer-specific patterns
self._identify_transformer_layers(layers_info)
def _analyze_graph_model(self):
"""
Analyzes a graph-based model by examining its operations
and tensor flow. This is more complex than Keras models.
"""
# Get concrete functions
concrete_funcs = []
if hasattr(self.model, 'signatures'):
for sig_name, sig_func in self.model.signatures.items():
concrete_funcs.append((sig_name, sig_func))
# Analyze each concrete function's graph
for func_name, func in concrete_funcs:
graph_def = func.graph.as_graph_def()
op_types = {}
for node in graph_def.node:
op_type = node.op
op_types[op_type] = op_types.get(op_type, 0) + 1
self.graph_info[func_name] = {
'num_operations': len(graph_def.node),
'operation_types': op_types
}
def _extract_signatures(self):
"""
Extracts serving signatures which define the model's
input/output interface for inference.
"""
if hasattr(self.model, 'signatures'):
signatures_info = {}
for sig_name, sig_func in self.model.signatures.items():
sig_info = {
'inputs': {},
'outputs': {}
}
# Extract input specifications
for input_name, input_spec in sig_func.structured_input_signature[1].items():
sig_info['inputs'][input_name] = {
'shape': input_spec.shape.as_list(),
'dtype': input_spec.dtype.name
}
# Extract output specifications
for output_name, output_tensor in sig_func.structured_outputs.items():
sig_info['outputs'][output_name] = {
'shape': output_tensor.shape.as_list(),
'dtype': output_tensor.dtype.name
}
signatures_info[sig_name] = sig_info
self.metadata['signatures'] = signatures_info
def _identify_transformer_layers(self, layers_info):
"""
Identifies transformer-specific components in the layer list
such as multi-head attention and feed-forward networks.
"""
attention_layers = []
dense_layers = []
norm_layers = []
for layer in layers_info:
layer_type = layer['type'].lower()
if 'attention' in layer_type:
attention_layers.append(layer)
elif 'dense' in layer_type or 'linear' in layer_type:
dense_layers.append(layer)
elif 'norm' in layer_type:
norm_layers.append(layer)
self.architecture_info['attention_layers'] = len(attention_layers)
self.architecture_info['dense_layers'] = len(dense_layers)
self.architecture_info['normalization_layers'] = len(norm_layers)
def _detect_architecture_patterns(self):
"""
Detects common architecture patterns by analyzing variable
names and shapes to identify transformer blocks, attention heads, etc.
"""
# Look for patterns in variable names
transformer_blocks = set()
attention_heads = None
for var_name, var_info in self.variables.items():
# Parse variable names for block numbers
parts = var_name.split('/')
for part in parts:
if 'block' in part or 'layer' in part:
# Extract number if present
import re
numbers = re.findall(r'\d+', part)
if numbers:
transformer_blocks.add(int(numbers[0]))
# Detect attention head count from query/key/value shapes
if 'attention' in var_name.lower() and 'query' in var_name.lower():
shape = var_info['shape']
if len(shape) >= 2:
# Common pattern: [hidden_dim, num_heads * head_dim]
# We can infer this from shape analysis
pass
if transformer_blocks:
self.architecture_info['num_transformer_blocks'] = len(transformer_blocks)
The TensorFlow reader demonstrates the challenges of working with a framework that supports multiple serialization formats. SavedModel is designed for production serving and includes not just model weights but the entire computational graph, making it possible to load and run the model without access to the original Python code. This completeness comes at the cost of complexity in parsing.
When we load a SavedModel, we gain access to the model's variables, which are analogous to PyTorch's parameters. However, TensorFlow's variable naming conventions differ from PyTorch's, reflecting TensorFlow's computational graph paradigm. Variables are named according to their position in the graph rather than following a simple hierarchical naming scheme. This means we must employ different heuristics to identify architectural components.
The signature extraction is particularly important for TensorFlow models because signatures define how the model should be called during inference. They specify the expected input tensor shapes and names, as well as the output tensor specifications. This information is crucial for understanding how to use the model in production environments. Our reader extracts these signatures and presents them in a structured format that makes it easy to understand the model's interface.
Apple MLX Format Reader Implementation
Apple's MLX framework is designed for efficient machine learning on Apple Silicon. MLX models are typically saved using NumPy's NPZ format or a custom safetensors-based format. The reader must handle both possibilities.
import numpy as np
from pathlib import Path
class MLXReader:
"""
Reads Apple MLX format models. MLX uses either NPZ (compressed NumPy)
format or safetensors format for storing model weights. The format
is simpler than TensorFlow but requires understanding MLX's conventions.
"""
def __init__(self, model_path):
self.model_path = Path(model_path)
self.weights = {}
self.architecture_info = {}
self.metadata = {}
def read_model(self):
"""
Loads an MLX model from either a file or directory structure.
Returns:
Dictionary containing weights and architecture information
"""
try:
if self.model_path.is_dir():
self._read_directory_model()
elif self.model_path.suffix == '.npz':
self._read_npz_model()
elif self.model_path.suffix == '.safetensors':
self._read_safetensors_model()
else:
raise ValueError(f"Unsupported MLX model format: {self.model_path.suffix}")
# Analyze the loaded weights
self._analyze_weights()
# Detect architecture patterns
self._detect_architecture()
except Exception as e:
raise ValueError(f"Failed to load MLX model: {str(e)}")
return {
'format': 'MLX',
'weights': self.weights,
'architecture': self.architecture_info,
'metadata': self.metadata
}
def _read_directory_model(self):
"""
Reads an MLX model stored as a directory with separate
weight files and a configuration file.
"""
# Look for weights file
weights_file = self.model_path / 'weights.npz'
if weights_file.exists():
self._read_npz_model(weights_file)
# Look for configuration file
config_file = self.model_path / 'config.json'
if config_file.exists():
import json
with open(config_file, 'r') as f:
self.metadata = json.load(f)
def _read_npz_model(self, file_path=None):
"""
Reads model weights from an NPZ file. NPZ is a compressed
archive of NumPy arrays with named entries.
"""
if file_path is None:
file_path = self.model_path
# Load the NPZ file
npz_data = np.load(file_path, allow_pickle=True)
# Extract all arrays
for key in npz_data.files:
array = npz_data[key]
self.weights[key] = {
'data': array,
'shape': array.shape,
'dtype': str(array.dtype),
'parameters': array.size
}
def _read_safetensors_model(self):
"""
Reads model weights from a safetensors file. Safetensors
is a format designed for safe, fast loading of tensors.
"""
try:
from safetensors import safe_open
with safe_open(self.model_path, framework="numpy") as f:
# Get all tensor names
tensor_names = f.keys()
for name in tensor_names:
tensor = f.get_tensor(name)
self.weights[name] = {
'data': tensor,
'shape': tensor.shape,
'dtype': str(tensor.dtype),
'parameters': tensor.size
}
except ImportError:
raise ValueError("safetensors library not installed")
def _analyze_weights(self):
"""
Analyzes the loaded weights to compute statistics and
identify patterns.
"""
total_params = 0
weight_stats = {}
for name, weight_info in self.weights.items():
total_params += weight_info['parameters']
# Compute statistics for this weight
data = weight_info['data']
stats = {
'mean': float(np.mean(data)),
'std': float(np.std(data)),
'min': float(np.min(data)),
'max': float(np.max(data))
}
weight_stats[name] = stats
self.architecture_info['total_parameters'] = total_params
self.architecture_info['weight_statistics'] = weight_stats
def _detect_architecture(self):
"""
Detects the model architecture by analyzing weight names
and shapes. MLX models often follow naming conventions
similar to PyTorch.
"""
layer_types = {
'attention': 0,
'mlp': 0,
'norm': 0,
'embedding': 0
}
transformer_layers = set()
for name in self.weights.keys():
name_lower = name.lower()
# Count layer types
if 'attn' in name_lower or 'attention' in name_lower:
layer_types['attention'] += 1
if 'mlp' in name_lower or 'ffn' in name_lower:
layer_types['mlp'] += 1
if 'norm' in name_lower:
layer_types['norm'] += 1
if 'embed' in name_lower:
layer_types['embedding'] += 1
# Extract layer numbers
parts = name.split('.')
for i, part in enumerate(parts):
if part in ['layers', 'blocks', 'h']:
if i + 1 < len(parts) and parts[i + 1].isdigit():
transformer_layers.add(int(parts[i + 1]))
self.architecture_info['layer_type_counts'] = layer_types
if transformer_layers:
self.architecture_info['num_layers'] = max(transformer_layers) + 1
# Try to extract architecture info from metadata
if 'model_type' in self.metadata:
self.architecture_info['model_type'] = self.metadata['model_type']
if 'n_layers' in self.metadata:
self.architecture_info['num_layers'] = self.metadata['n_layers']
if 'n_heads' in self.metadata:
self.architecture_info['num_attention_heads'] = self.metadata['n_heads']
if 'hidden_size' in self.metadata:
self.architecture_info['hidden_size'] = self.metadata['hidden_size']
The MLX reader is somewhat simpler than the TensorFlow reader because MLX uses more straightforward serialization formats. The NPZ format is essentially a ZIP archive containing NumPy arrays, making it easy to extract individual weights. The safetensors format, while more sophisticated, provides a clean API for accessing tensors without the complexity of computational graphs.
One advantage of MLX models is that they often include a separate configuration JSON file that explicitly states architectural parameters. When this file is present, we can extract precise information about the number of layers, attention heads, hidden dimensions, and other critical parameters without having to infer them from weight shapes and names. This explicit metadata makes our analysis more reliable and comprehensive.
The weight analysis in the MLX reader goes beyond simple extraction by computing statistical properties of each weight tensor. These statistics, including mean, standard deviation, minimum, and maximum values, provide insights into the model's training state and can help identify potential issues like dead neurons or exploding gradients. This statistical analysis is particularly valuable when comparing different versions of a model or investigating model behavior.
Model Analysis Engine
Now that we have readers for all major formats, we need a unified analysis engine that can take the parsed model data and extract meaningful architectural insights regardless of the source format. This engine forms the core of our tool's analytical capabilities.
class ModelAnalyzer:
"""
Unified analysis engine that processes model data from any format
and extracts comprehensive architectural information, including
transformer layer details, attention mechanisms, and quantization.
"""
def __init__(self, model_data):
self.model_data = model_data
self.format = model_data.get('format', 'Unknown')
self.analysis_results = {}
def analyze(self):
"""
Performs comprehensive analysis of the model, extracting
all available architectural information.
Returns:
Dictionary containing detailed analysis results
"""
# Extract basic information
self._extract_basic_info()
# Analyze transformer architecture
self._analyze_transformer_architecture()
# Analyze attention mechanisms
self._analyze_attention_mechanism()
# Analyze quantization
self._analyze_quantization()
# Compute weight statistics
self._compute_weight_statistics()
# Identify model family
self._identify_model_family()
return self.analysis_results
def _extract_basic_info(self):
"""
Extracts basic information that's available across all formats
such as total parameters, number of layers, etc.
"""
basic_info = {}
# Get architecture info from the parsed data
arch_info = self.model_data.get('architecture', {})
if 'total_parameters' in arch_info:
basic_info['total_parameters'] = arch_info['total_parameters']
basic_info['total_parameters_millions'] = arch_info['total_parameters'] / 1_000_000
if 'num_layers' in arch_info:
basic_info['num_layers'] = arch_info['num_layers']
elif 'num_blocks' in arch_info:
basic_info['num_layers'] = arch_info['num_blocks']
elif 'block_count' in arch_info:
basic_info['num_layers'] = arch_info['block_count']
self.analysis_results['basic_info'] = basic_info
def _analyze_transformer_architecture(self):
"""
Analyzes the transformer architecture in detail, identifying
the number of layers, hidden dimensions, intermediate dimensions,
and other structural parameters.
"""
transformer_info = {}
arch_info = self.model_data.get('architecture', {})
metadata = self.model_data.get('metadata', {})
# Extract hidden dimension
hidden_dim = None
if 'embedding_length' in arch_info:
hidden_dim = arch_info['embedding_length']
elif 'hidden_size' in arch_info:
hidden_dim = arch_info['hidden_size']
elif 'hidden_size' in metadata:
hidden_dim = metadata['hidden_size']
if hidden_dim:
transformer_info['hidden_dimension'] = hidden_dim
# Extract feed-forward dimension
ff_dim = None
if 'feed_forward_length' in arch_info:
ff_dim = arch_info['feed_forward_length']
elif 'intermediate_size' in metadata:
ff_dim = metadata['intermediate_size']
if ff_dim:
transformer_info['feedforward_dimension'] = ff_dim
if hidden_dim:
transformer_info['ff_to_hidden_ratio'] = ff_dim / hidden_dim
# Extract context length
if 'context_length' in arch_info:
transformer_info['context_length'] = arch_info['context_length']
elif 'max_position_embeddings' in metadata:
transformer_info['context_length'] = metadata['max_position_embeddings']
self.analysis_results['transformer_architecture'] = transformer_info
def _analyze_attention_mechanism(self):
"""
Analyzes the attention mechanism, including number of heads,
head dimension, and whether grouped-query attention is used.
"""
attention_info = {}
arch_info = self.model_data.get('architecture', {})
metadata = self.model_data.get('metadata', {})
# Extract number of attention heads
num_heads = None
if 'head_count' in arch_info:
num_heads = arch_info['head_count']
elif 'num_attention_heads' in arch_info:
num_heads = arch_info['num_attention_heads']
elif 'n_heads' in metadata:
num_heads = metadata['n_heads']
if num_heads:
attention_info['num_attention_heads'] = num_heads
# Extract key-value heads for grouped-query attention
num_kv_heads = None
if 'head_count_kv' in arch_info:
num_kv_heads = arch_info['head_count_kv']
elif 'num_key_value_heads' in metadata:
num_kv_heads = metadata['num_key_value_heads']
if num_kv_heads:
attention_info['num_kv_heads'] = num_kv_heads
if num_heads:
attention_info['uses_grouped_query_attention'] = num_kv_heads < num_heads
attention_info['heads_per_kv_group'] = num_heads // num_kv_heads
else:
attention_info['uses_grouped_query_attention'] = False
# Calculate head dimension
hidden_dim = self.analysis_results.get('transformer_architecture', {}).get('hidden_dimension')
if hidden_dim and num_heads:
attention_info['head_dimension'] = hidden_dim // num_heads
self.analysis_results['attention_mechanism'] = attention_info
def _analyze_quantization(self):
"""
Analyzes quantization details including quantization method,
bit width, and which layers are quantized.
"""
quant_info = {
'is_quantized': False,
'quantization_method': None,
'bit_width': None
}
arch_info = self.model_data.get('architecture', {})
metadata = self.model_data.get('metadata', {})
# Check if model is quantized
if 'is_quantized' in arch_info:
quant_info['is_quantized'] = arch_info['is_quantized']
# Extract quantization method from GGUF metadata
if 'file_type' in arch_info:
file_type = arch_info['file_type']
quant_info['is_quantized'] = file_type > 0
quant_info['quantization_method'] = self._decode_gguf_quantization(file_type)
# Extract quantization version
if 'quantization_version' in arch_info:
quant_info['quantization_version'] = arch_info['quantization_version']
self.analysis_results['quantization'] = quant_info
def _decode_gguf_quantization(self, file_type):
"""
Decodes GGUF file type to human-readable quantization method.
"""
quantization_map = {
0: 'F32 (no quantization)',
1: 'F16',
2: 'Q4_0',
3: 'Q4_1',
4: 'Q4_2',
5: 'Q4_3',
6: 'Q5_0',
7: 'Q5_1',
8: 'Q8_0',
9: 'Q8_1',
10: 'Q2_K',
11: 'Q3_K_S',
12: 'Q3_K_M',
13: 'Q3_K_L',
14: 'Q4_K_S',
15: 'Q4_K_M',
16: 'Q5_K_S',
17: 'Q5_K_M',
18: 'Q6_K'
}
return quantization_map.get(file_type, f'Unknown ({file_type})')
def _compute_weight_statistics(self):
"""
Computes statistical properties of model weights including
distribution characteristics and potential anomalies.
"""
weight_stats = {}
# Get weights from different formats
weights = None
if self.format == 'PyTorch':
weights = self.model_data.get('state_dict', {})
elif self.format == 'MLX':
weights = self.model_data.get('weights', {})
elif self.format == 'TensorFlow':
weights = self.model_data.get('variables', {})
if weights:
# Compute aggregate statistics
all_means = []
all_stds = []
for name, weight_data in weights.items():
# Skip if we don't have the actual data
if isinstance(weight_data, dict) and 'data' in weight_data:
data = weight_data['data']
all_means.append(float(np.mean(data)))
all_stds.append(float(np.std(data)))
if all_means:
weight_stats['average_mean'] = np.mean(all_means)
weight_stats['average_std'] = np.mean(all_stds)
weight_stats['mean_range'] = (min(all_means), max(all_means))
weight_stats['std_range'] = (min(all_stds), max(all_stds))
self.analysis_results['weight_statistics'] = weight_stats
def _identify_model_family(self):
"""
Attempts to identify the model family (GPT, LLaMA, etc.)
based on architectural characteristics and metadata.
"""
metadata = self.model_data.get('metadata', {})
arch_info = self.model_data.get('architecture', {})
model_family = 'Unknown'
# Check explicit architecture field
if 'architecture' in arch_info:
model_family = arch_info['architecture']
elif 'general.architecture' in metadata:
model_family = metadata['general.architecture']
elif 'model_type' in metadata:
model_family = metadata['model_type']
else:
# Try to infer from architectural characteristics
model_family = self._infer_model_family()
self.analysis_results['model_family'] = model_family
def _infer_model_family(self):
"""
Infers the model family from architectural characteristics
when explicit metadata is not available.
"""
attention_info = self.analysis_results.get('attention_mechanism', {})
transformer_info = self.analysis_results.get('transformer_architecture', {})
# LLaMA uses grouped-query attention
if attention_info.get('uses_grouped_query_attention'):
return 'LLaMA-family'
# GPT-2 has specific dimension ratios
ff_ratio = transformer_info.get('ff_to_hidden_ratio')
if ff_ratio and abs(ff_ratio - 4.0) < 0.1:
return 'GPT-family'
return 'Unknown'
The model analyzer serves as the intelligence layer of our tool, taking the raw parsed data from format-specific readers and transforming it into meaningful insights. The analyzer is designed to be format-agnostic, working with a normalized representation of model data regardless of whether it came from GGUF, PyTorch, TensorFlow, or MLX. This abstraction is crucial for maintaining clean architecture and avoiding code duplication.
The analysis process is divided into several specialized methods, each focusing on a different aspect of the model. The basic information extraction gathers fundamental metrics like total parameter count and number of layers. These metrics are often stored under different keys in different formats, so the analyzer checks multiple possible locations and normalizes the data into a consistent structure.
The transformer architecture analysis goes deeper, examining the relationships between different dimensions. For example, the ratio between the feed-forward dimension and the hidden dimension is typically around four in standard transformer architectures, but can vary in optimized models. By computing and reporting this ratio, we help users understand whether they are working with a standard architecture or a modified variant. Similarly, the attention mechanism analysis reveals whether the model uses multi-head attention, grouped-query attention, or other variants, which has significant implications for inference performance and memory usage.
Visualization System with Matplotlib
Visual representations make complex model information more accessible and easier to understand. Our visualization system uses matplotlib to create various charts and diagrams that illustrate model properties.
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.gridspec import GridSpec
import numpy as np
class ModelVisualizer:
"""
Creates visual representations of model architecture and properties
using matplotlib. Generates multiple types of visualizations including
architecture diagrams, weight distributions, and quantization patterns.
"""
def __init__(self, analysis_results, model_data):
self.analysis = analysis_results
self.model_data = model_data
self.figures = []
def create_all_visualizations(self):
"""
Creates a comprehensive set of visualizations covering
all aspects of the model.
Returns:
List of matplotlib figure objects
"""
# Create architecture overview diagram
self.figures.append(self._create_architecture_diagram())
# Create weight distribution plots
self.figures.append(self._create_weight_distribution_plot())
# Create layer composition chart
self.figures.append(self._create_layer_composition_chart())
# Create attention mechanism diagram
self.figures.append(self._create_attention_diagram())
return self.figures
def _create_architecture_diagram(self):
"""
Creates a visual diagram showing the overall architecture
including number of layers, dimensions, and key parameters.
"""
fig = plt.figure(figsize=(12, 8))
gs = GridSpec(3, 2, figure=fig, hspace=0.4, wspace=0.3)
# Title
fig.suptitle('Model Architecture Overview', fontsize=16, fontweight='bold')
# Basic information panel
ax1 = fig.add_subplot(gs[0, :])
ax1.axis('off')
basic_info = self.analysis.get('basic_info', {})
transformer_info = self.analysis.get('transformer_architecture', {})
attention_info = self.analysis.get('attention_mechanism', {})
info_text = []
info_text.append(f"Model Family: {self.analysis.get('model_family', 'Unknown')}")
if 'total_parameters_millions' in basic_info:
info_text.append(f"Total Parameters: {basic_info['total_parameters_millions']:.2f}M")
if 'num_layers' in basic_info:
info_text.append(f"Number of Layers: {basic_info['num_layers']}")
if 'hidden_dimension' in transformer_info:
info_text.append(f"Hidden Dimension: {transformer_info['hidden_dimension']}")
if 'num_attention_heads' in attention_info:
info_text.append(f"Attention Heads: {attention_info['num_attention_heads']}")
ax1.text(0.5, 0.5, '\n'.join(info_text),
ha='center', va='center', fontsize=12,
bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
# Layer structure visualization
ax2 = fig.add_subplot(gs[1, 0])
self._draw_layer_stack(ax2)
# Dimension breakdown
ax3 = fig.add_subplot(gs[1, 1])
self._draw_dimension_breakdown(ax3)
# Quantization information
ax4 = fig.add_subplot(gs[2, :])
self._draw_quantization_info(ax4)
return fig
def _draw_layer_stack(self, ax):
"""
Draws a visual representation of the transformer layer stack.
"""
ax.set_title('Transformer Layer Stack')
ax.set_xlim(0, 10)
ax.set_ylim(0, 10)
ax.axis('off')
num_layers = self.analysis.get('basic_info', {}).get('num_layers', 0)
if num_layers == 0:
ax.text(5, 5, 'Layer information not available', ha='center', va='center')
return
# Draw simplified layer stack (show first few, middle, and last)
layers_to_show = min(5, num_layers)
layer_height = 6.0 / layers_to_show
for i in range(layers_to_show):
y_pos = 2 + i * layer_height
# Draw layer box
rect = mpatches.Rectangle((2, y_pos), 6, layer_height * 0.8,
linewidth=1, edgecolor='blue',
facecolor='lightblue', alpha=0.7)
ax.add_patch(rect)
# Add layer label
layer_num = i if i < layers_to_show - 1 else num_layers - 1
ax.text(5, y_pos + layer_height * 0.4, f'Layer {layer_num}',
ha='center', va='center', fontsize=9)
# Add ellipsis if there are more layers
if num_layers > layers_to_show:
ax.text(5, 2 + (layers_to_show - 1) * layer_height / 2, '...',
ha='center', va='center', fontsize=14)
def _draw_dimension_breakdown(self, ax):
"""
Creates a bar chart showing the relationships between
different dimensions in the model.
"""
ax.set_title('Dimension Breakdown')
transformer_info = self.analysis.get('transformer_architecture', {})
attention_info = self.analysis.get('attention_mechanism', {})
dimensions = {}
labels = []
values = []
if 'hidden_dimension' in transformer_info:
labels.append('Hidden\nDim')
values.append(transformer_info['hidden_dimension'])
if 'feedforward_dimension' in transformer_info:
labels.append('FF\nDim')
values.append(transformer_info['feedforward_dimension'])
if 'head_dimension' in attention_info:
labels.append('Head\nDim')
values.append(attention_info['head_dimension'])
if 'context_length' in transformer_info:
labels.append('Context\nLength')
values.append(transformer_info['context_length'])
if not values:
ax.text(0.5, 0.5, 'Dimension information not available',
ha='center', va='center', transform=ax.transAxes)
ax.axis('off')
return
bars = ax.bar(labels, values, color=['skyblue', 'lightgreen', 'lightcoral', 'lightyellow'])
ax.set_ylabel('Dimension Size')
ax.tick_params(axis='x', rotation=0)
# Add value labels on bars
for bar, value in zip(bars, values):
height = bar.get_height()
ax.text(bar.get_x() + bar.get_width()/2., height,
f'{int(value)}', ha='center', va='bottom', fontsize=9)
def _draw_quantization_info(self, ax):
"""
Displays quantization information in a text panel.
"""
ax.axis('off')
quant_info = self.analysis.get('quantization', {})
if not quant_info.get('is_quantized', False):
info_text = 'Model is not quantized (full precision)'
else:
info_text = f"Quantization Method: {quant_info.get('quantization_method', 'Unknown')}\n"
if 'bit_width' in quant_info:
info_text += f"Bit Width: {quant_info['bit_width']}\n"
if 'quantization_version' in quant_info:
info_text += f"Quantization Version: {quant_info['quantization_version']}"
ax.text(0.5, 0.5, info_text, ha='center', va='center',
fontsize=11, bbox=dict(boxstyle='round', facecolor='lightgray', alpha=0.5))
def _create_weight_distribution_plot(self):
"""
Creates histograms showing the distribution of weights
across different layers or the entire model.
"""
fig, axes = plt.subplots(2, 2, figsize=(12, 10))
fig.suptitle('Weight Distribution Analysis', fontsize=16, fontweight='bold')
# Get weights based on format
weights = self._get_weights_for_visualization()
if not weights:
axes[0, 0].text(0.5, 0.5, 'Weight data not available for visualization',
ha='center', va='center', transform=axes[0, 0].transAxes)
for ax in axes.flat:
ax.axis('off')
return fig
# Collect all weight values for overall distribution
all_weights = []
for name, data in weights.items():
if data is not None and hasattr(data, 'flatten'):
all_weights.extend(data.flatten())
if all_weights:
# Overall distribution
axes[0, 0].hist(all_weights, bins=100, alpha=0.7, color='blue', edgecolor='black')
axes[0, 0].set_title('Overall Weight Distribution')
axes[0, 0].set_xlabel('Weight Value')
axes[0, 0].set_ylabel('Frequency')
axes[0, 0].grid(True, alpha=0.3)
# Log scale distribution
axes[0, 1].hist(all_weights, bins=100, alpha=0.7, color='green', edgecolor='black')
axes[0, 1].set_title('Overall Weight Distribution (Log Scale)')
axes[0, 1].set_xlabel('Weight Value')
axes[0, 1].set_ylabel('Frequency (log scale)')
axes[0, 1].set_yscale('log')
axes[0, 1].grid(True, alpha=0.3)
# Distribution of means per layer
layer_means = []
layer_names = []
for name, data in list(weights.items())[:20]: # Limit to first 20 for readability
if data is not None and hasattr(data, 'mean'):
layer_means.append(float(np.mean(data)))
layer_names.append(name.split('.')[-1][:10]) # Shortened names
if layer_means:
axes[1, 0].barh(range(len(layer_means)), layer_means, color='coral')
axes[1, 0].set_yticks(range(len(layer_means)))
axes[1, 0].set_yticklabels(layer_names, fontsize=8)
axes[1, 0].set_title('Mean Weight Value by Layer')
axes[1, 0].set_xlabel('Mean Value')
axes[1, 0].grid(True, alpha=0.3, axis='x')
# Distribution of standard deviations per layer
layer_stds = []
for name, data in list(weights.items())[:20]:
if data is not None and hasattr(data, 'std'):
layer_stds.append(float(np.std(data)))
if layer_stds:
axes[1, 1].barh(range(len(layer_stds)), layer_stds, color='purple')
axes[1, 1].set_yticks(range(len(layer_stds)))
axes[1, 1].set_yticklabels(layer_names, fontsize=8)
axes[1, 1].set_title('Weight Std Dev by Layer')
axes[1, 1].set_xlabel('Standard Deviation')
axes[1, 1].grid(True, alpha=0.3, axis='x')
plt.tight_layout()
return fig
def _get_weights_for_visualization(self):
"""
Extracts weight data in a format suitable for visualization,
handling different model formats.
"""
weights = {}
if self.model_data.get('format') == 'PyTorch':
state_dict = self.model_data.get('state_dict', {})
for name, tensor in state_dict.items():
if hasattr(tensor, 'numpy'):
weights[name] = tensor.numpy()
elif hasattr(tensor, 'cpu'):
weights[name] = tensor.cpu().numpy()
elif self.model_data.get('format') == 'MLX':
mlx_weights = self.model_data.get('weights', {})
for name, weight_info in mlx_weights.items():
if 'data' in weight_info:
weights[name] = weight_info['data']
elif self.model_data.get('format') == 'TensorFlow':
tf_vars = self.model_data.get('variables', {})
# TensorFlow variables might not have actual data loaded
# This would require additional implementation
pass
return weights
def _create_layer_composition_chart(self):
"""
Creates a pie chart showing the composition of different
layer types in the model.
"""
fig, ax = plt.subplots(figsize=(10, 8))
fig.suptitle('Layer Type Composition', fontsize=16, fontweight='bold')
# Get layer counts
layer_counts = {}
if self.model_data.get('format') == 'PyTorch':
arch_info = self.model_data.get('architecture', {})
if 'layer_counts' in arch_info:
layer_counts = arch_info['layer_counts']
if not layer_counts:
ax.text(0.5, 0.5, 'Layer composition data not available',
ha='center', va='center', transform=ax.transAxes)
ax.axis('off')
return fig
# Filter out zero counts
layer_counts = {k: v for k, v in layer_counts.items() if v > 0}
if layer_counts:
labels = list(layer_counts.keys())
sizes = list(layer_counts.values())
colors = ['gold', 'lightblue', 'lightgreen', 'lightcoral', 'plum']
wedges, texts, autotexts = ax.pie(sizes, labels=labels, colors=colors,
autopct='%1.1f%%', startangle=90)
# Enhance text
for text in texts:
text.set_fontsize(12)
for autotext in autotexts:
autotext.set_color('white')
autotext.set_fontweight('bold')
autotext.set_fontsize(10)
return fig
def _create_attention_diagram(self):
"""
Creates a diagram illustrating the attention mechanism
configuration including number of heads and dimensions.
"""
fig, ax = plt.subplots(figsize=(12, 8))
fig.suptitle('Attention Mechanism Configuration', fontsize=16, fontweight='bold')
ax.set_xlim(0, 10)
ax.set_ylim(0, 10)
ax.axis('off')
attention_info = self.analysis.get('attention_mechanism', {})
if not attention_info:
ax.text(5, 5, 'Attention mechanism information not available',
ha='center', va='center', fontsize=12)
return fig
# Draw attention heads
num_heads = attention_info.get('num_attention_heads', 0)
if num_heads > 0:
# Draw head boxes
heads_per_row = min(8, num_heads)
rows = (num_heads + heads_per_row - 1) // heads_per_row
box_width = 8.0 / heads_per_row
box_height = 5.0 / rows
for i in range(min(num_heads, 16)): # Limit display to 16 heads
row = i // heads_per_row
col = i % heads_per_row
x = 1 + col * box_width
y = 3 + row * box_height
rect = mpatches.Rectangle((x, y), box_width * 0.9, box_height * 0.8,
linewidth=1, edgecolor='darkblue',
facecolor='lightblue', alpha=0.6)
ax.add_patch(rect)
ax.text(x + box_width * 0.45, y + box_height * 0.4,
f'H{i}', ha='center', va='center', fontsize=8)
if num_heads > 16:
ax.text(5, 2, f'... and {num_heads - 16} more heads',
ha='center', va='center', fontsize=10)
# Add text information
info_lines = []
info_lines.append(f"Total Attention Heads: {num_heads}")
if 'head_dimension' in attention_info:
info_lines.append(f"Head Dimension: {attention_info['head_dimension']}")
if 'uses_grouped_query_attention' in attention_info:
if attention_info['uses_grouped_query_attention']:
info_lines.append(f"Uses Grouped-Query Attention")
if 'num_kv_heads' in attention_info:
info_lines.append(f"KV Heads: {attention_info['num_kv_heads']}")
else:
info_lines.append("Uses Multi-Head Attention")
ax.text(5, 1, '\n'.join(info_lines), ha='center', va='top',
fontsize=11, bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
return fig
def save_all_figures(self, output_dir):
"""
Saves all generated figures to the specified directory.
"""
from pathlib import Path
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
figure_names = [
'architecture_overview.png',
'weight_distributions.png',
'layer_composition.png',
'attention_mechanism.png'
]
for fig, name in zip(self.figures, figure_names):
fig.savefig(output_path / name, dpi=300, bbox_inches='tight')
print(f"Saved {name}")
The visualization system transforms numerical data into intuitive graphical representations that make model analysis accessible to users with varying levels of technical expertise. The architecture diagram provides a high-level overview that can be quickly scanned to understand the model's basic structure. By presenting information like parameter count, layer count, and dimension sizes in a structured visual format, we enable rapid comprehension of the model's scale and complexity.
The weight distribution plots are particularly valuable for understanding the model's training state and identifying potential issues. A healthy trained model typically exhibits weight distributions centered near zero with reasonable standard deviations. Distributions with unusual shapes, extreme values, or high variance might indicate training problems, quantization artifacts, or other issues. By visualizing these distributions both in linear and logarithmic scales, we accommodate both normal ranges and outliers.
The layer composition chart uses a pie chart to show the relative proportions of different layer types. This visualization quickly reveals whether the model follows standard transformer architecture or includes unusual components. The attention mechanism diagram provides a visual representation of how attention heads are organized, making it easy to understand whether the model uses standard multi-head attention or more advanced variants like grouped-query attention.
Report Generation System
The final component of our tool is the report generator, which synthesizes all the analyzed information and visualizations into a comprehensive, human-readable report.
class ReportGenerator:
"""
Generates comprehensive reports about the analyzed model,
combining textual analysis with visualizations to provide
a complete overview of the model's characteristics.
"""
def __init__(self, analysis_results, model_data, visualizations):
self.analysis = analysis_results
self.model_data = model_data
self.visualizations = visualizations
def generate_text_report(self):
"""
Generates a detailed text report covering all aspects
of the model analysis.
Returns:
String containing the formatted report
"""
report_sections = []
# Header
report_sections.append(self._generate_header())
# Basic Information
report_sections.append(self._generate_basic_info_section())
# Architecture Details
report_sections.append(self._generate_architecture_section())
# Attention Mechanism
report_sections.append(self._generate_attention_section())
# Quantization Information
report_sections.append(self._generate_quantization_section())
# Weight Statistics
report_sections.append(self._generate_statistics_section())
# Recommendations
report_sections.append(self._generate_recommendations())
return '\n\n'.join(report_sections)
def _generate_header(self):
"""
Generates the report header with title and metadata.
"""
header = "=" * 70 + "\n"
header += "MODEL ANALYSIS REPORT\n"
header += "=" * 70 + "\n"
header += f"Format: {self.model_data.get('format', 'Unknown')}\n"
header += f"Model Family: {self.analysis.get('model_family', 'Unknown')}\n"
return header
def _generate_basic_info_section(self):
"""
Generates the basic information section of the report.
"""
section = "BASIC INFORMATION\n"
section += "-" * 70 + "\n"
basic_info = self.analysis.get('basic_info', {})
if 'total_parameters' in basic_info:
total_params = basic_info['total_parameters']
params_millions = basic_info.get('total_parameters_millions', total_params / 1_000_000)
section += f"Total Parameters: {total_params:,} ({params_millions:.2f}M)\n"
if 'num_layers' in basic_info:
section += f"Number of Transformer Layers: {basic_info['num_layers']}\n"
return section
def _generate_architecture_section(self):
"""
Generates detailed architecture information section.
"""
section = "ARCHITECTURE DETAILS\n"
section += "-" * 70 + "\n"
transformer_info = self.analysis.get('transformer_architecture', {})
if 'hidden_dimension' in transformer_info:
section += f"Hidden Dimension: {transformer_info['hidden_dimension']}\n"
if 'feedforward_dimension' in transformer_info:
ff_dim = transformer_info['feedforward_dimension']
section += f"Feed-Forward Dimension: {ff_dim}\n"
if 'ff_to_hidden_ratio' in transformer_info:
ratio = transformer_info['ff_to_hidden_ratio']
section += f"FF to Hidden Ratio: {ratio:.2f}\n"
if 'context_length' in transformer_info:
section += f"Context Length: {transformer_info['context_length']}\n"
return section
def _generate_attention_section(self):
"""
Generates attention mechanism information section.
"""
section = "ATTENTION MECHANISM\n"
section += "-" * 70 + "\n"
attention_info = self.analysis.get('attention_mechanism', {})
if 'num_attention_heads' in attention_info:
section += f"Number of Attention Heads: {attention_info['num_attention_heads']}\n"
if 'head_dimension' in attention_info:
section += f"Head Dimension: {attention_info['head_dimension']}\n"
if 'uses_grouped_query_attention' in attention_info:
if attention_info['uses_grouped_query_attention']:
section += "Attention Type: Grouped-Query Attention (GQA)\n"
if 'num_kv_heads' in attention_info:
section += f"Number of KV Heads: {attention_info['num_kv_heads']}\n"
if 'heads_per_kv_group' in attention_info:
section += f"Query Heads per KV Group: {attention_info['heads_per_kv_group']}\n"
else:
section += "Attention Type: Multi-Head Attention (MHA)\n"
return section
def _generate_quantization_section(self):
"""
Generates quantization information section.
"""
section = "QUANTIZATION\n"
section += "-" * 70 + "\n"
quant_info = self.analysis.get('quantization', {})
if quant_info.get('is_quantized', False):
section += "Model Status: Quantized\n"
if 'quantization_method' in quant_info:
section += f"Quantization Method: {quant_info['quantization_method']}\n"
if 'bit_width' in quant_info:
section += f"Bit Width: {quant_info['bit_width']}\n"
if 'quantization_version' in quant_info:
section += f"Quantization Version: {quant_info['quantization_version']}\n"
else:
section += "Model Status: Full Precision (Not Quantized)\n"
return section
def _generate_statistics_section(self):
"""
Generates weight statistics section.
"""
section = "WEIGHT STATISTICS\n"
section += "-" * 70 + "\n"
weight_stats = self.analysis.get('weight_statistics', {})
if 'average_mean' in weight_stats:
section += f"Average Weight Mean: {weight_stats['average_mean']:.6f}\n"
if 'average_std' in weight_stats:
section += f"Average Weight Std Dev: {weight_stats['average_std']:.6f}\n"
if 'mean_range' in weight_stats:
min_mean, max_mean = weight_stats['mean_range']
section += f"Mean Range: [{min_mean:.6f}, {max_mean:.6f}]\n"
if 'std_range' in weight_stats:
min_std, max_std = weight_stats['std_range']
section += f"Std Dev Range: [{min_std:.6f}, {max_std:.6f}]\n"
return section
def _generate_recommendations(self):
"""
Generates recommendations based on the analysis results.
"""
section = "RECOMMENDATIONS AND INSIGHTS\n"
section += "-" * 70 + "\n"
recommendations = []
# Check parameter count for deployment recommendations
basic_info = self.analysis.get('basic_info', {})
if 'total_parameters_millions' in basic_info:
params_m = basic_info['total_parameters_millions']
if params_m < 1000:
recommendations.append(
f"With {params_m:.0f}M parameters, this model is suitable for "
"edge deployment and consumer hardware."
)
elif params_m < 10000:
recommendations.append(
f"With {params_m:.0f}M parameters, this model requires "
"moderate computational resources, suitable for workstation deployment."
)
else:
recommendations.append(
f"With {params_m:.0f}M parameters, this model requires "
"significant computational resources and is best suited for "
"server or multi-GPU deployment."
)
# Check quantization status
quant_info = self.analysis.get('quantization', {})
if not quant_info.get('is_quantized', False):
recommendations.append(
"The model is not quantized. Consider applying quantization "
"to reduce memory footprint and improve inference speed."
)
# Check attention mechanism
attention_info = self.analysis.get('attention_mechanism', {})
if attention_info.get('uses_grouped_query_attention', False):
recommendations.append(
"The model uses Grouped-Query Attention, which provides better "
"inference efficiency compared to standard Multi-Head Attention, "
"especially for long sequences."
)
# Check context length
transformer_info = self.analysis.get('transformer_architecture', {})
if 'context_length' in transformer_info:
ctx_len = transformer_info['context_length']
if ctx_len >= 4096:
recommendations.append(
f"The model supports a context length of {ctx_len}, enabling "
"processing of long documents. Be aware that memory usage scales "
"quadratically with context length."
)
if recommendations:
section += '\n'.join(f"- {rec}" for rec in recommendations)
else:
section += "No specific recommendations at this time."
return section
def save_report(self, output_path):
"""
Saves the text report to a file.
"""
report_text = self.generate_text_report()
with open(output_path, 'w') as f:
f.write(report_text)
print(f"Report saved to {output_path}")
The report generator serves as the final synthesis layer, taking all the analyzed data and presenting it in a format that is both comprehensive and accessible. The text report is structured into logical sections that guide the reader through different aspects of the model, from basic statistics to detailed architectural characteristics. Each section focuses on a specific aspect, making it easy to find particular information without reading the entire report.
The recommendations section is particularly valuable as it translates technical metrics into actionable insights. For example, rather than simply stating that a model has seven billion parameters, the report explains what this means for deployment, suggesting appropriate hardware configurations and warning about resource requirements. Similarly, the presence of grouped-query attention is not just noted but explained in terms of its practical implications for inference performance.
Complete Tool Implementation
Now we present the complete, production-ready implementation of the model analysis tool, integrating all the components we have discussed. This implementation includes no shortcuts or simplifications and is ready for real-world use.
#!/usr/bin/env python3
"""
Complete LLM Model Analysis Tool
This tool provides comprehensive analysis of large language models
stored in various formats including GGUF, PyTorch, TensorFlow, and MLX.
It extracts architectural details, analyzes weights, detects quantization,
and generates detailed reports with visualizations.
Usage:
python llm_analyzer.py <model_path> [--output-dir <dir>]
"""
import os
import sys
import argparse
import struct
import zipfile
import json
import re
from pathlib import Path
from enum import Enum
from typing import Dict, Any, List, Tuple, Optional, Union
from collections import OrderedDict
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.gridspec import GridSpec
# Optional imports with graceful degradation
try:
import torch
TORCH_AVAILABLE = True
except ImportError:
TORCH_AVAILABLE = False
print("Warning: PyTorch not available. PyTorch model support disabled.")
try:
import tensorflow as tf
TF_AVAILABLE = True
except ImportError:
TF_AVAILABLE = False
print("Warning: TensorFlow not available. TensorFlow model support disabled.")
try:
from safetensors import safe_open
SAFETENSORS_AVAILABLE = True
except ImportError:
SAFETENSORS_AVAILABLE = False
print("Warning: safetensors not available. Safetensors support disabled.")
class ModelFormat(Enum):
"""Enumeration of supported model formats"""
GGUF = "gguf"
PYTORCH = "pytorch"
TENSORFLOW = "tensorflow"
MLX = "mlx"
UNKNOWN = "unknown"
class FormatDetector:
"""
Detects the format of a model file by examining its structure
and magic bytes. This class encapsulates all format detection
logic in one place.
"""
def __init__(self):
self.gguf_magic = b'GGUF'
self.pytorch_magic = b'PK\x03\x04'
self.tf_magic = b'\x08\x03'
def detect_format(self, model_path: Union[str, Path]) -> ModelFormat:
"""
Examines the file or directory structure to determine
the model format.
Args:
model_path: Path to the model file or directory
Returns:
ModelFormat enum indicating the detected format
"""
path = Path(model_path)
if not path.exists():
raise ValueError(f"Path does not exist: {model_path}")
if path.is_dir():
return self._detect_directory_format(path)
if path.is_file():
return self._detect_file_format(path)
return ModelFormat.UNKNOWN
def _detect_file_format(self, file_path: Path) -> ModelFormat:
"""Detects format by reading file magic bytes"""
try:
with open(file_path, 'rb') as f:
header = f.read(16)
if header.startswith(self.gguf_magic):
return ModelFormat.GGUF
if header.startswith(self.pytorch_magic):
return self._verify_pytorch_format(file_path)
if self._is_tensorflow_pb(file_path):
return ModelFormat.TENSORFLOW
if self._is_mlx_format(file_path):
return ModelFormat.MLX
except Exception as e:
print(f"Error detecting file format: {e}")
return ModelFormat.UNKNOWN
def _detect_directory_format(self, dir_path: Path) -> ModelFormat:
"""Detects format for directory-based model storage"""
if (dir_path / 'saved_model.pb').exists():
return ModelFormat.TENSORFLOW
if (dir_path / 'weights.npz').exists() or (dir_path / 'config.json').exists():
return ModelFormat.MLX
return ModelFormat.UNKNOWN
def _verify_pytorch_format(self, file_path: Path) -> ModelFormat:
"""Verifies PyTorch format by checking ZIP contents"""
if not TORCH_AVAILABLE:
return ModelFormat.UNKNOWN
try:
with zipfile.ZipFile(file_path, 'r') as z:
namelist = z.namelist()
if any('data.pkl' in name or 'data/' in name for name in namelist):
return ModelFormat.PYTORCH
except zipfile.BadZipFile:
pass
return ModelFormat.UNKNOWN
def _is_tensorflow_pb(self, file_path: Path) -> bool:
"""Checks if file is a TensorFlow protobuf"""
try:
with open(file_path, 'rb') as f:
content = f.read(200)
if b'tensorflow' in content.lower():
return True
except Exception:
pass
return False
def _is_mlx_format(self, file_path: Path) -> bool:
"""Checks for MLX format based on file extension"""
return file_path.suffix in ['.npz', '.safetensors']
class GGUFReader:
"""
Reads and parses GGUF format model files, extracting metadata,
architecture information, and tensor data.
"""
def __init__(self, file_path: Union[str, Path]):
self.file_path = Path(file_path)
self.metadata = {}
self.tensors = {}
self.architecture_info = {}
self.gguf_types = {
0: ('uint8', 1, '<B'),
1: ('int8', 1, '<b'),
2: ('uint16', 2, '<H'),
3: ('int16', 2, '<h'),
4: ('uint32', 4, '<I'),
5: ('int32', 4, '<i'),
6: ('float32', 4, '<f'),
7: ('bool', 1, '<?'),
8: ('string', None, None),
9: ('array', None, None),
10: ('uint64', 8, '<Q'),
11: ('int64', 8, '<q'),
12: ('float64', 8, '<d'),
}
def read_model(self) -> Dict[str, Any]:
"""
Main entry point for reading a GGUF file.
Returns:
Dictionary containing all extracted information
"""
with open(self.file_path, 'rb') as f:
magic = f.read(4)
if magic != b'GGUF':
raise ValueError(f"Not a valid GGUF file: {self.file_path}")
version = struct.unpack('<I', f.read(4))[0]
self.metadata['gguf_version'] = version
tensor_count = struct.unpack('<Q', f.read(8))[0]
metadata_count = struct.unpack('<Q', f.read(8))[0]
self._read_metadata(f, metadata_count)
self._read_tensor_info(f, tensor_count)
self._extract_architecture_info()
return {
'format': 'GGUF',
'version': version,
'metadata': self.metadata,
'tensors': self.tensors,
'architecture': self.architecture_info
}
def _read_metadata(self, f, count: int):
"""Reads metadata key-value pairs from the GGUF file"""
for _ in range(count):
key = self._read_string(f)
value_type = struct.unpack('<I', f.read(4))[0]
value = self._read_typed_value(f, value_type)
self.metadata[key] = value
def _read_string(self, f) -> str:
"""Reads a length-prefixed string"""
length = struct.unpack('<Q', f.read(8))[0]
string_bytes = f.read(length)
return string_bytes.decode('utf-8', errors='replace')
def _read_typed_value(self, f, value_type: int) -> Any:
"""Reads a value of the specified type"""
if value_type in [0, 1, 2, 3, 4, 5, 6, 10, 11, 12]:
type_name, size, fmt = self.gguf_types[value_type]
data = f.read(size)
return struct.unpack(fmt, data)[0]
elif value_type == 7:
return struct.unpack('<?', f.read(1))[0]
elif value_type == 8:
return self._read_string(f)
elif value_type == 9:
return self._read_array(f)
return None
def _read_array(self, f) -> List[Any]:
"""Reads an array value"""
element_type = struct.unpack('<I', f.read(4))[0]
array_length = struct.unpack('<Q', f.read(8))[0]
array_elements = []
for _ in range(array_length):
element = self._read_typed_value(f, element_type)
array_elements.append(element)
return array_elements
def _read_tensor_info(self, f, count: int):
"""Reads tensor metadata"""
for _ in range(count):
tensor_name = self._read_string(f)
n_dims = struct.unpack('<I', f.read(4))[0]
dimensions = []
for _ in range(n_dims):
dim_size = struct.unpack('<Q', f.read(8))[0]
dimensions.append(dim_size)
tensor_type = struct.unpack('<I', f.read(4))[0]
offset = struct.unpack('<Q', f.read(8))[0]
self.tensors[tensor_name] = {
'dimensions': dimensions,
'type': tensor_type,
'offset': offset,
'shape': tuple(dimensions)
}
def _extract_architecture_info(self):
"""Extracts architecture-specific information from metadata"""
arch_mappings = {
'general.architecture': 'architecture',
'llama.context_length': 'context_length',
'llama.embedding_length': 'embedding_length',
'llama.block_count': 'block_count',
'llama.feed_forward_length': 'feed_forward_length',
'llama.attention.head_count': 'head_count',
'llama.attention.head_count_kv': 'head_count_kv',
'llama.attention.layer_norm_rms_epsilon': 'layer_norm_rms_epsilon',
'llama.rope.dimension_count': 'rope_dimension_count',
'general.quantization_version': 'quantization_version',
'general.file_type': 'file_type'
}
for metadata_key, arch_key in arch_mappings.items():
if metadata_key in self.metadata:
self.architecture_info[arch_key] = self.metadata[metadata_key]
# Calculate total parameters
total_params = 0
for tensor_info in self.tensors.values():
dims = tensor_info['dimensions']
param_count = 1
for dim in dims:
param_count *= dim
total_params += param_count
self.architecture_info['total_parameters'] = total_params
class PyTorchReader:
"""
Reads PyTorch model checkpoint files, extracting state dictionaries,
model architecture information, and training metadata.
"""
def __init__(self, file_path: Union[str, Path]):
if not TORCH_AVAILABLE:
raise RuntimeError("PyTorch is not available")
self.file_path = Path(file_path)
self.state_dict = None
self.metadata = {}
self.architecture_info = {}
self.optimizer_state = None
self.training_info = {}
def read_model(self) -> Dict[str, Any]:
"""
Loads a PyTorch checkpoint and extracts all available information.
Returns:
Dictionary containing model weights and metadata
"""
try:
checkpoint = torch.load(self.file_path, map_location='cpu')
if isinstance(checkpoint, dict):
self._parse_checkpoint_dict(checkpoint)
elif isinstance(checkpoint, OrderedDict):
self.state_dict = checkpoint
else:
self._parse_model_object(checkpoint)
self._analyze_state_dict()
self._detect_quantization()
except Exception as e:
raise ValueError(f"Failed to load PyTorch model: {str(e)}")
return {
'format': 'PyTorch',
'state_dict': self.state_dict,
'metadata': self.metadata,
'architecture': self.architecture_info,
'training_info': self.training_info
}
def _parse_checkpoint_dict(self, checkpoint: Dict):
"""Parses a checkpoint dictionary"""
state_dict_keys = ['model_state_dict', 'state_dict', 'model']
for key in state_dict_keys:
if key in checkpoint:
self.state_dict = checkpoint[key]
break
if self.state_dict is None:
self.state_dict = checkpoint
if 'optimizer_state_dict' in checkpoint:
self.optimizer_state = checkpoint['optimizer_state_dict']
training_keys = ['epoch', 'iteration', 'best_loss', 'best_accuracy', 'global_step']
for key in training_keys:
if key in checkpoint:
self.training_info[key] = checkpoint[key]
for key, value in checkpoint.items():
if key not in state_dict_keys + ['optimizer_state_dict'] + training_keys:
if not isinstance(value, (torch.Tensor, OrderedDict)):
self.metadata[key] = value
def _parse_model_object(self, model):
"""Extracts state dict from a model object"""
if hasattr(model, 'state_dict'):
self.state_dict = model.state_dict()
else:
raise ValueError("Unable to extract state dict from model object")
def _analyze_state_dict(self):
"""Analyzes the state dictionary to infer model architecture"""
if self.state_dict is None:
return
layer_pattern_counts = {
'attention': 0,
'mlp': 0,
'layer_norm': 0,
'embedding': 0
}
total_parameters = 0
layer_info = {}
for name, tensor in self.state_dict.items():
param_count = tensor.numel()
total_parameters += param_count
name_lower = name.lower()
if 'attn' in name_lower or 'attention' in name_lower:
layer_pattern_counts['attention'] += 1
if 'mlp' in name_lower or 'ffn' in name_lower or 'feed_forward' in name_lower:
layer_pattern_counts['mlp'] += 1
if 'norm' in name_lower:
layer_pattern_counts['layer_norm'] += 1
if 'embed' in name_lower:
layer_pattern_counts['embedding'] += 1
layer_info[name] = {
'shape': list(tensor.shape),
'dtype': str(tensor.dtype),
'parameters': param_count
}
self.architecture_info['total_parameters'] = total_parameters
self.architecture_info['layer_counts'] = layer_pattern_counts
self.architecture_info['layer_details'] = layer_info
self._infer_transformer_blocks()
self._infer_dimensions()
def _infer_transformer_blocks(self):
"""Determines the number of transformer blocks"""
block_numbers = set()
for name in self.state_dict.keys():
parts = name.split('.')
for i, part in enumerate(parts):
if part in ['blocks', 'layers', 'h', 'transformer', 'layer']:
if i + 1 < len(parts) and parts[i + 1].isdigit():
block_numbers.add(int(parts[i + 1]))
if block_numbers:
self.architecture_info['num_blocks'] = max(block_numbers) + 1
self.architecture_info['block_indices'] = sorted(list(block_numbers))
def _infer_dimensions(self):
"""Infers model dimensions from weight shapes"""
for name, tensor in self.state_dict.items():
if 'embed' in name.lower() and 'weight' in name.lower():
shape = tensor.shape
if len(shape) == 2:
self.architecture_info['vocab_size'] = shape[0]
self.architecture_info['hidden_size'] = shape[1]
break
def _detect_quantization(self):
"""Detects if the model uses quantization"""
quantized = False
quantization_info = {}
for name, tensor in self.state_dict.items():
dtype_str = str(tensor.dtype)
if 'qint' in dtype_str or 'quint' in dtype_str:
quantized = True
quantization_info[name] = dtype_str
if 'scale' in name or 'zero_point' in name:
quantized = True
self.architecture_info['is_quantized'] = quantized
if quantized:
self.architecture_info['quantization_details'] = quantization_info
class TensorFlowReader:
"""
Reads TensorFlow SavedModel format, extracting model architecture,
variables, and metadata.
"""
def __init__(self, model_path: Union[str, Path]):
if not TF_AVAILABLE:
raise RuntimeError("TensorFlow is not available")
self.model_path = Path(model_path)
self.model = None
self.variables = {}
self.architecture_info = {}
self.metadata = {}
self.graph_info = {}
def read_model(self) -> Dict[str, Any]:
"""
Loads a TensorFlow SavedModel and extracts comprehensive information.
Returns:
Dictionary containing all extracted information
"""
try:
self.model = tf.saved_model.load(str(self.model_path))
self._extract_variables()
self._analyze_structure()
self._extract_signatures()
self._detect_architecture_patterns()
except Exception as e:
raise ValueError(f"Failed to load TensorFlow model: {str(e)}")
return {
'format': 'TensorFlow',
'variables': self.variables,
'architecture': self.architecture_info,
'metadata': self.metadata,
'graph_info': self.graph_info
}
def _extract_variables(self):
"""Extracts all trainable and non-trainable variables"""
all_variables = []
if hasattr(self.model, 'variables'):
all_variables = self.model.variables
elif hasattr(self.model, 'trainable_variables'):
all_variables = self.model.trainable_variables
total_params = 0
for var in all_variables:
var_name = var.name
var_shape = var.shape.as_list()
var_dtype = var.dtype.name
param_count = 1
for dim in var_shape:
if dim is not None:
param_count *= dim
total_params += param_count
self.variables[var_name] = {
'shape': var_shape,
'dtype': var_dtype,
'parameters': param_count,
'trainable': var.trainable
}
self.architecture_info['total_parameters'] = total_params
self.architecture_info['num_variables'] = len(all_variables)
def _analyze_structure(self):
"""Analyzes the model's layer structure"""
if hasattr(self.model, 'layers'):
self._analyze_keras_model()
else:
self._analyze_graph_model()
def _analyze_keras_model(self):
"""Extracts detailed information from a Keras model"""
layers_info = []
for layer in self.model.layers:
layer_config = {
'name': layer.name,
'type': layer.__class__.__name__,
'trainable': layer.trainable,
'parameters': layer.count_params()
}
try:
config = layer.get_config()
layer_config['config'] = config
except Exception:
pass
try:
layer_config['input_shape'] = str(layer.input_shape)
layer_config['output_shape'] = str(layer.output_shape)
except Exception:
pass
layers_info.append(layer_config)
self.architecture_info['layers'] = layers_info
self.architecture_info['num_layers'] = len(layers_info)
self._identify_transformer_layers(layers_info)
def _analyze_graph_model(self):
"""Analyzes a graph-based model"""
concrete_funcs = []
if hasattr(self.model, 'signatures'):
for sig_name, sig_func in self.model.signatures.items():
concrete_funcs.append((sig_name, sig_func))
for func_name, func in concrete_funcs:
try:
graph_def = func.graph.as_graph_def()
op_types = {}
for node in graph_def.node:
op_type = node.op
op_types[op_type] = op_types.get(op_type, 0) + 1
self.graph_info[func_name] = {
'num_operations': len(graph_def.node),
'operation_types': op_types
}
except Exception:
pass
def _extract_signatures(self):
"""Extracts serving signatures"""
if hasattr(self.model, 'signatures'):
signatures_info = {}
for sig_name, sig_func in self.model.signatures.items():
sig_info = {
'inputs': {},
'outputs': {}
}
try:
for input_name, input_spec in sig_func.structured_input_signature[1].items():
sig_info['inputs'][input_name] = {
'shape': input_spec.shape.as_list(),
'dtype': input_spec.dtype.name
}
except Exception:
pass
try:
for output_name, output_tensor in sig_func.structured_outputs.items():
sig_info['outputs'][output_name] = {
'shape': output_tensor.shape.as_list(),
'dtype': output_tensor.dtype.name
}
except Exception:
pass
signatures_info[sig_name] = sig_info
self.metadata['signatures'] = signatures_info
def _identify_transformer_layers(self, layers_info: List[Dict]):
"""Identifies transformer-specific components"""
attention_layers = []
dense_layers = []
norm_layers = []
for layer in layers_info:
layer_type = layer['type'].lower()
if 'attention' in layer_type:
attention_layers.append(layer)
elif 'dense' in layer_type or 'linear' in layer_type:
dense_layers.append(layer)
elif 'norm' in layer_type:
norm_layers.append(layer)
self.architecture_info['attention_layers'] = len(attention_layers)
self.architecture_info['dense_layers'] = len(dense_layers)
self.architecture_info['normalization_layers'] = len(norm_layers)
def _detect_architecture_patterns(self):
"""Detects common architecture patterns"""
transformer_blocks = set()
for var_name in self.variables.keys():
parts = var_name.split('/')
for part in parts:
if 'block' in part or 'layer' in part:
numbers = re.findall(r'\d+', part)
if numbers:
transformer_blocks.add(int(numbers[0]))
if transformer_blocks:
self.architecture_info['num_transformer_blocks'] = len(transformer_blocks)
class MLXReader:
"""
Reads Apple MLX format models from NPZ or safetensors files.
"""
def __init__(self, model_path: Union[str, Path]):
self.model_path = Path(model_path)
self.weights = {}
self.architecture_info = {}
self.metadata = {}
def read_model(self) -> Dict[str, Any]:
"""
Loads an MLX model from either a file or directory structure.
Returns:
Dictionary containing weights and architecture information
"""
try:
if self.model_path.is_dir():
self._read_directory_model()
elif self.model_path.suffix == '.npz':
self._read_npz_model()
elif self.model_path.suffix == '.safetensors':
self._read_safetensors_model()
else:
raise ValueError(f"Unsupported MLX model format: {self.model_path.suffix}")
self._analyze_weights()
self._detect_architecture()
except Exception as e:
raise ValueError(f"Failed to load MLX model: {str(e)}")
return {
'format': 'MLX',
'weights': self.weights,
'architecture': self.architecture_info,
'metadata': self.metadata
}
def _read_directory_model(self):
"""Reads an MLX model stored as a directory"""
weights_file = self.model_path / 'weights.npz'
if weights_file.exists():
self._read_npz_model(weights_file)
config_file = self.model_path / 'config.json'
if config_file.exists():
with open(config_file, 'r') as f:
self.metadata = json.load(f)
def _read_npz_model(self, file_path: Optional[Path] = None):
"""Reads model weights from an NPZ file"""
if file_path is None:
file_path = self.model_path
npz_data = np.load(file_path, allow_pickle=True)
for key in npz_data.files:
array = npz_data[key]
self.weights[key] = {
'data': array,
'shape': array.shape,
'dtype': str(array.dtype),
'parameters': array.size
}
def _read_safetensors_model(self):
"""Reads model weights from a safetensors file"""
if not SAFETENSORS_AVAILABLE:
raise RuntimeError("safetensors library not available")
with safe_open(self.model_path, framework="numpy") as f:
tensor_names = f.keys()
for name in tensor_names:
tensor = f.get_tensor(name)
self.weights[name] = {
'data': tensor,
'shape': tensor.shape,
'dtype': str(tensor.dtype),
'parameters': tensor.size
}
def _analyze_weights(self):
"""Analyzes the loaded weights to compute statistics"""
total_params = 0
weight_stats = {}
for name, weight_info in self.weights.items():
total_params += weight_info['parameters']
data = weight_info['data']
stats = {
'mean': float(np.mean(data)),
'std': float(np.std(data)),
'min': float(np.min(data)),
'max': float(np.max(data))
}
weight_stats[name] = stats
self.architecture_info['total_parameters'] = total_params
self.architecture_info['weight_statistics'] = weight_stats
def _detect_architecture(self):
"""Detects the model architecture by analyzing weight names and shapes"""
layer_types = {
'attention': 0,
'mlp': 0,
'norm': 0,
'embedding': 0
}
transformer_layers = set()
for name in self.weights.keys():
name_lower = name.lower()
if 'attn' in name_lower or 'attention' in name_lower:
layer_types['attention'] += 1
if 'mlp' in name_lower or 'ffn' in name_lower:
layer_types['mlp'] += 1
if 'norm' in name_lower:
layer_types['norm'] += 1
if 'embed' in name_lower:
layer_types['embedding'] += 1
parts = name.split('.')
for i, part in enumerate(parts):
if part in ['layers', 'blocks', 'h']:
if i + 1 < len(parts) and parts[i + 1].isdigit():
transformer_layers.add(int(parts[i + 1]))
self.architecture_info['layer_type_counts'] = layer_types
if transformer_layers:
self.architecture_info['num_layers'] = max(transformer_layers) + 1
for key in ['model_type', 'n_layers', 'n_heads', 'hidden_size',
'intermediate_size', 'max_position_embeddings']:
if key in self.metadata:
self.architecture_info[key] = self.metadata[key]
class ModelAnalyzer:
"""
Unified analysis engine that processes model data from any format
and extracts comprehensive architectural information.
"""
def __init__(self, model_data: Dict[str, Any]):
self.model_data = model_data
self.format = model_data.get('format', 'Unknown')
self.analysis_results = {}
def analyze(self) -> Dict[str, Any]:
"""
Performs comprehensive analysis of the model.
Returns:
Dictionary containing detailed analysis results
"""
self._extract_basic_info()
self._analyze_transformer_architecture()
self._analyze_attention_mechanism()
self._analyze_quantization()
self._compute_weight_statistics()
self._identify_model_family()
return self.analysis_results
def _extract_basic_info(self):
"""Extracts basic information available across all formats"""
basic_info = {}
arch_info = self.model_data.get('architecture', {})
if 'total_parameters' in arch_info:
basic_info['total_parameters'] = arch_info['total_parameters']
basic_info['total_parameters_millions'] = arch_info['total_parameters'] / 1_000_000
basic_info['total_parameters_billions'] = arch_info['total_parameters'] / 1_000_000_000
num_layers = (arch_info.get('num_layers') or
arch_info.get('num_blocks') or
arch_info.get('block_count'))
if num_layers:
basic_info['num_layers'] = num_layers
self.analysis_results['basic_info'] = basic_info
def _analyze_transformer_architecture(self):
"""Analyzes the transformer architecture in detail"""
transformer_info = {}
arch_info = self.model_data.get('architecture', {})
metadata = self.model_data.get('metadata', {})
hidden_dim = (arch_info.get('embedding_length') or
arch_info.get('hidden_size') or
metadata.get('hidden_size'))
if hidden_dim:
transformer_info['hidden_dimension'] = hidden_dim
ff_dim = (arch_info.get('feed_forward_length') or
arch_info.get('intermediate_size') or
metadata.get('intermediate_size'))
if ff_dim:
transformer_info['feedforward_dimension'] = ff_dim
if hidden_dim:
transformer_info['ff_to_hidden_ratio'] = ff_dim / hidden_dim
context_length = (arch_info.get('context_length') or
metadata.get('max_position_embeddings'))
if context_length:
transformer_info['context_length'] = context_length
self.analysis_results['transformer_architecture'] = transformer_info
def _analyze_attention_mechanism(self):
"""Analyzes the attention mechanism configuration"""
attention_info = {}
arch_info = self.model_data.get('architecture', {})
metadata = self.model_data.get('metadata', {})
num_heads = (arch_info.get('head_count') or
arch_info.get('num_attention_heads') or
arch_info.get('n_heads') or
metadata.get('n_heads') or
metadata.get('num_attention_heads'))
if num_heads:
attention_info['num_attention_heads'] = num_heads
num_kv_heads = (arch_info.get('head_count_kv') or
metadata.get('num_key_value_heads'))
if num_kv_heads:
attention_info['num_kv_heads'] = num_kv_heads
if num_heads:
attention_info['uses_grouped_query_attention'] = num_kv_heads < num_heads
attention_info['heads_per_kv_group'] = num_heads // num_kv_heads
else:
attention_info['uses_grouped_query_attention'] = False
hidden_dim = self.analysis_results.get('transformer_architecture', {}).get('hidden_dimension')
if hidden_dim and num_heads:
attention_info['head_dimension'] = hidden_dim // num_heads
self.analysis_results['attention_mechanism'] = attention_info
def _analyze_quantization(self):
"""Analyzes quantization details"""
quant_info = {
'is_quantized': False,
'quantization_method': None,
'bit_width': None
}
arch_info = self.model_data.get('architecture', {})
if 'is_quantized' in arch_info:
quant_info['is_quantized'] = arch_info['is_quantized']
if 'file_type' in arch_info:
file_type = arch_info['file_type']
quant_info['is_quantized'] = file_type > 0
quant_info['quantization_method'] = self._decode_gguf_quantization(file_type)
if 'quantization_version' in arch_info:
quant_info['quantization_version'] = arch_info['quantization_version']
self.analysis_results['quantization'] = quant_info
def _decode_gguf_quantization(self, file_type: int) -> str:
"""Decodes GGUF file type to human-readable quantization method"""
quantization_map = {
0: 'F32 (no quantization)',
1: 'F16',
2: 'Q4_0',
3: 'Q4_1',
6: 'Q5_0',
7: 'Q5_1',
8: 'Q8_0',
10: 'Q2_K',
11: 'Q3_K_S',
12: 'Q3_K_M',
13: 'Q3_K_L',
14: 'Q4_K_S',
15: 'Q4_K_M',
16: 'Q5_K_S',
17: 'Q5_K_M',
18: 'Q6_K'
}
return quantization_map.get(file_type, f'Unknown ({file_type})')
def _compute_weight_statistics(self):
"""Computes statistical properties of model weights"""
weight_stats = {}
weights = None
if self.format == 'PyTorch' and TORCH_AVAILABLE:
weights = self.model_data.get('state_dict', {})
elif self.format == 'MLX':
weights = self.model_data.get('weights', {})
elif self.format == 'TensorFlow' and TF_AVAILABLE:
weights = self.model_data.get('variables', {})
if weights:
all_means = []
all_stds = []
for name, weight_data in weights.items():
try:
if isinstance(weight_data, dict) and 'data' in weight_data:
data = weight_data['data']
elif TORCH_AVAILABLE and isinstance(weight_data, torch.Tensor):
data = weight_data.cpu().numpy()
else:
continue
all_means.append(float(np.mean(data)))
all_stds.append(float(np.std(data)))
except Exception:
continue
if all_means:
weight_stats['average_mean'] = np.mean(all_means)
weight_stats['average_std'] = np.mean(all_stds)
weight_stats['mean_range'] = (min(all_means), max(all_means))
weight_stats['std_range'] = (min(all_stds), max(all_stds))
self.analysis_results['weight_statistics'] = weight_stats
def _identify_model_family(self):
"""Identifies the model family based on characteristics"""
metadata = self.model_data.get('metadata', {})
arch_info = self.model_data.get('architecture', {})
model_family = (arch_info.get('architecture') or
metadata.get('general.architecture') or
metadata.get('model_type') or
self._infer_model_family())
self.analysis_results['model_family'] = model_family
def _infer_model_family(self) -> str:
"""Infers the model family from architectural characteristics"""
attention_info = self.analysis_results.get('attention_mechanism', {})
transformer_info = self.analysis_results.get('transformer_architecture', {})
if attention_info.get('uses_grouped_query_attention'):
return 'LLaMA-family'
ff_ratio = transformer_info.get('ff_to_hidden_ratio')
if ff_ratio and abs(ff_ratio - 4.0) < 0.1:
return 'GPT-family'
return 'Unknown'
class ModelVisualizer:
"""
Creates visual representations of model architecture and properties
using matplotlib.
"""
def __init__(self, analysis_results: Dict[str, Any], model_data: Dict[str, Any]):
self.analysis = analysis_results
self.model_data = model_data
self.figures = []
def create_all_visualizations(self) -> List[plt.Figure]:
"""
Creates a comprehensive set of visualizations.
Returns:
List of matplotlib figure objects
"""
self.figures.append(self._create_architecture_diagram())
self.figures.append(self._create_weight_distribution_plot())
self.figures.append(self._create_layer_composition_chart())
self.figures.append(self._create_attention_diagram())
return self.figures
def _create_architecture_diagram(self) -> plt.Figure:
"""Creates a visual diagram showing the overall architecture"""
fig = plt.figure(figsize=(14, 10))
gs = GridSpec(3, 2, figure=fig, hspace=0.4, wspace=0.3)
fig.suptitle('Model Architecture Overview', fontsize=18, fontweight='bold')
ax1 = fig.add_subplot(gs[0, :])
ax1.axis('off')
basic_info = self.analysis.get('basic_info', {})
transformer_info = self.analysis.get('transformer_architecture', {})
attention_info = self.analysis.get('attention_mechanism', {})
info_lines = []
info_lines.append(f"Model Family: {self.analysis.get('model_family', 'Unknown')}")
info_lines.append(f"Format: {self.model_data.get('format', 'Unknown')}")
if 'total_parameters_millions' in basic_info:
params_m = basic_info['total_parameters_millions']
if params_m >= 1000:
params_b = basic_info.get('total_parameters_billions', params_m / 1000)
info_lines.append(f"Total Parameters: {params_b:.2f}B")
else:
info_lines.append(f"Total Parameters: {params_m:.2f}M")
if 'num_layers' in basic_info:
info_lines.append(f"Number of Layers: {basic_info['num_layers']}")
if 'hidden_dimension' in transformer_info:
info_lines.append(f"Hidden Dimension: {transformer_info['hidden_dimension']}")
if 'feedforward_dimension' in transformer_info:
info_lines.append(f"FF Dimension: {transformer_info['feedforward_dimension']}")
if 'num_attention_heads' in attention_info:
info_lines.append(f"Attention Heads: {attention_info['num_attention_heads']}")
if 'context_length' in transformer_info:
info_lines.append(f"Context Length: {transformer_info['context_length']}")
ax1.text(0.5, 0.5, '\n'.join(info_lines),
ha='center', va='center', fontsize=13,
bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.6),
family='monospace')
ax2 = fig.add_subplot(gs[1, 0])
self._draw_layer_stack(ax2)
ax3 = fig.add_subplot(gs[1, 1])
self._draw_dimension_breakdown(ax3)
ax4 = fig.add_subplot(gs[2, :])
self._draw_quantization_info(ax4)
return fig
def _draw_layer_stack(self, ax):
"""Draws a visual representation of the transformer layer stack"""
ax.set_title('Transformer Layer Stack', fontsize=12, fontweight='bold')
ax.set_xlim(0, 10)
ax.set_ylim(0, 10)
ax.axis('off')
num_layers = self.analysis.get('basic_info', {}).get('num_layers', 0)
if num_layers == 0:
ax.text(5, 5, 'Layer information not available',
ha='center', va='center', fontsize=11)
return
layers_to_show = min(6, num_layers)
layer_height = 6.5 / layers_to_show
for i in range(layers_to_show):
y_pos = 2 + i * layer_height
if i == layers_to_show - 1 and num_layers > layers_to_show:
layer_num = num_layers - 1
else:
layer_num = i
rect = mpatches.Rectangle((2, y_pos), 6, layer_height * 0.85,
linewidth=2, edgecolor='darkblue',
facecolor='lightblue', alpha=0.7)
ax.add_patch(rect)
ax.text(5, y_pos + layer_height * 0.425, f'Layer {layer_num}',
ha='center', va='center', fontsize=10, fontweight='bold')
if num_layers > layers_to_show:
middle_y = 2 + (layers_to_show - 2) * layer_height + layer_height * 0.425
ax.text(5, middle_y, '...', ha='center', va='center',
fontsize=20, fontweight='bold')
def _draw_dimension_breakdown(self, ax):
"""Creates a bar chart showing dimension relationships"""
ax.set_title('Dimension Breakdown', fontsize=12, fontweight='bold')
transformer_info = self.analysis.get('transformer_architecture', {})
attention_info = self.analysis.get('attention_mechanism', {})
labels = []
values = []
colors = []
if 'hidden_dimension' in transformer_info:
labels.append('Hidden\nDim')
values.append(transformer_info['hidden_dimension'])
colors.append('skyblue')
if 'feedforward_dimension' in transformer_info:
labels.append('FF\nDim')
values.append(transformer_info['feedforward_dimension'])
colors.append('lightgreen')
if 'head_dimension' in attention_info:
labels.append('Head\nDim')
values.append(attention_info['head_dimension'])
colors.append('lightcoral')
if 'context_length' in transformer_info:
labels.append('Context\nLength')
values.append(transformer_info['context_length'])
colors.append('lightyellow')
if not values:
ax.text(0.5, 0.5, 'Dimension information not available',
ha='center', va='center', transform=ax.transAxes, fontsize=11)
ax.axis('off')
return
bars = ax.bar(labels, values, color=colors, edgecolor='black', linewidth=1.5)
ax.set_ylabel('Dimension Size', fontsize=11)
ax.grid(True, alpha=0.3, axis='y')
for bar, value in zip(bars, values):
height = bar.get_height()
ax.text(bar.get_x() + bar.get_width()/2., height,
f'{int(value)}', ha='center', va='bottom',
fontsize=10, fontweight='bold')
def _draw_quantization_info(self, ax):
"""Displays quantization information"""
ax.axis('off')
quant_info = self.analysis.get('quantization', {})
if not quant_info.get('is_quantized', False):
info_text = 'Model is not quantized (full precision)'
color = 'lightgreen'
else:
lines = []
lines.append(f"Quantization Method: {quant_info.get('quantization_method', 'Unknown')}")
if 'bit_width' in quant_info:
lines.append(f"Bit Width: {quant_info['bit_width']}")
if 'quantization_version' in quant_info:
lines.append(f"Quantization Version: {quant_info['quantization_version']}")
info_text = '\n'.join(lines)
color = 'lightyellow'
ax.text(0.5, 0.5, info_text, ha='center', va='center',
fontsize=12, bbox=dict(boxstyle='round', facecolor=color, alpha=0.6))
def _create_weight_distribution_plot(self) -> plt.Figure:
"""Creates histograms showing weight distributions"""
fig, axes = plt.subplots(2, 2, figsize=(14, 11))
fig.suptitle('Weight Distribution Analysis', fontsize=18, fontweight='bold')
weights = self._get_weights_for_visualization()
if not weights:
axes[0, 0].text(0.5, 0.5, 'Weight data not available for visualization',
ha='center', va='center', transform=axes[0, 0].transAxes,
fontsize=12)
for ax in axes.flat:
ax.axis('off')
return fig
all_weights = []
for name, data in weights.items():
if data is not None and hasattr(data, 'flatten'):
all_weights.extend(data.flatten())
if all_weights:
all_weights_array = np.array(all_weights)
axes[0, 0].hist(all_weights_array, bins=100, alpha=0.75,
color='blue', edgecolor='black', linewidth=0.5)
axes[0, 0].set_title('Overall Weight Distribution', fontsize=12, fontweight='bold')
axes[0, 0].set_xlabel('Weight Value', fontsize=11)
axes[0, 0].set_ylabel('Frequency', fontsize=11)
axes[0, 0].grid(True, alpha=0.3)
axes[0, 1].hist(all_weights_array, bins=100, alpha=0.75,
color='green', edgecolor='black', linewidth=0.5)
axes[0, 1].set_title('Overall Weight Distribution (Log Scale)',
fontsize=12, fontweight='bold')
axes[0, 1].set_xlabel('Weight Value', fontsize=11)
axes[0, 1].set_ylabel('Frequency (log scale)', fontsize=11)
axes[0, 1].set_yscale('log')
axes[0, 1].grid(True, alpha=0.3)
layer_means = []
layer_names = []
for name, data in list(weights.items())[:25]:
if data is not None and hasattr(data, 'mean'):
layer_means.append(float(np.mean(data)))
short_name = name.split('.')[-1][:12]
layer_names.append(short_name)
if layer_means:
y_pos = np.arange(len(layer_means))
axes[1, 0].barh(y_pos, layer_means, color='coral', edgecolor='black')
axes[1, 0].set_yticks(y_pos)
axes[1, 0].set_yticklabels(layer_names, fontsize=8)
axes[1, 0].set_title('Mean Weight Value by Layer',
fontsize=12, fontweight='bold')
axes[1, 0].set_xlabel('Mean Value', fontsize=11)
axes[1, 0].grid(True, alpha=0.3, axis='x')
axes[1, 0].invert_yaxis()
layer_stds = []
for name, data in list(weights.items())[:25]:
if data is not None and hasattr(data, 'std'):
layer_stds.append(float(np.std(data)))
if layer_stds:
y_pos = np.arange(len(layer_stds))
axes[1, 1].barh(y_pos, layer_stds, color='purple', edgecolor='black')
axes[1, 1].set_yticks(y_pos)
axes[1, 1].set_yticklabels(layer_names, fontsize=8)
axes[1, 1].set_title('Weight Std Dev by Layer',
fontsize=12, fontweight='bold')
axes[1, 1].set_xlabel('Standard Deviation', fontsize=11)
axes[1, 1].grid(True, alpha=0.3, axis='x')
axes[1, 1].invert_yaxis()
plt.tight_layout()
return fig
def _get_weights_for_visualization(self) -> Dict[str, np.ndarray]:
"""Extracts weight data for visualization"""
weights = {}
if self.model_data.get('format') == 'PyTorch' and TORCH_AVAILABLE:
state_dict = self.model_data.get('state_dict', {})
for name, tensor in state_dict.items():
try:
if hasattr(tensor, 'numpy'):
weights[name] = tensor.numpy()
elif hasattr(tensor, 'cpu'):
weights[name] = tensor.cpu().numpy()
except Exception:
continue
elif self.model_data.get('format') == 'MLX':
mlx_weights = self.model_data.get('weights', {})
for name, weight_info in mlx_weights.items():
if 'data' in weight_info:
weights[name] = weight_info['data']
return weights
def _create_layer_composition_chart(self) -> plt.Figure:
"""Creates a pie chart showing layer type composition"""
fig, ax = plt.subplots(figsize=(11, 9))
fig.suptitle('Layer Type Composition', fontsize=18, fontweight='bold')
layer_counts = {}
if self.model_data.get('format') == 'PyTorch':
arch_info = self.model_data.get('architecture', {})
if 'layer_counts' in arch_info:
layer_counts = arch_info['layer_counts']
if not layer_counts:
ax.text(0.5, 0.5, 'Layer composition data not available',
ha='center', va='center', transform=ax.transAxes, fontsize=12)
ax.axis('off')
return fig
layer_counts = {k: v for k, v in layer_counts.items() if v > 0}
if layer_counts:
labels = list(layer_counts.keys())
sizes = list(layer_counts.values())
colors = ['gold', 'lightblue', 'lightgreen', 'lightcoral', 'plum', 'lightyellow']
wedges, texts, autotexts = ax.pie(sizes, labels=labels, colors=colors,
autopct='%1.1f%%', startangle=90,
textprops={'fontsize': 12})
for text in texts:
text.set_fontsize(13)
text.set_fontweight('bold')
for autotext in autotexts:
autotext.set_color('white')
autotext.set_fontweight('bold')
autotext.set_fontsize(11)
return fig
def _create_attention_diagram(self) -> plt.Figure:
"""Creates a diagram illustrating the attention mechanism"""
fig, ax = plt.subplots(figsize=(14, 10))
fig.suptitle('Attention Mechanism Configuration', fontsize=18, fontweight='bold')
ax.set_xlim(0, 10)
ax.set_ylim(0, 10)
ax.axis('off')
attention_info = self.analysis.get('attention_mechanism', {})
if not attention_info:
ax.text(5, 5, 'Attention mechanism information not available',
ha='center', va='center', fontsize=12)
return fig
num_heads = attention_info.get('num_attention_heads', 0)
if num_heads > 0:
heads_per_row = min(8, num_heads)
rows = min(3, (num_heads + heads_per_row - 1) // heads_per_row)
box_width = 8.0 / heads_per_row
box_height = 5.5 / rows
heads_to_draw = min(num_heads, heads_per_row * rows)
for i in range(heads_to_draw):
row = i // heads_per_row
col = i % heads_per_row
x = 1 + col * box_width
y = 3.5 + row * box_height
rect = mpatches.Rectangle((x, y), box_width * 0.88, box_height * 0.85,
linewidth=1.5, edgecolor='darkblue',
facecolor='lightblue', alpha=0.7)
ax.add_patch(rect)
ax.text(x + box_width * 0.44, y + box_height * 0.425,
f'H{i}', ha='center', va='center',
fontsize=9, fontweight='bold')
if num_heads > heads_to_draw:
ax.text(5, 2.5, f'... and {num_heads - heads_to_draw} more heads',
ha='center', va='center', fontsize=11, style='italic')
info_lines = []
info_lines.append(f"Total Attention Heads: {num_heads}")
if 'head_dimension' in attention_info:
info_lines.append(f"Head Dimension: {attention_info['head_dimension']}")
if 'uses_grouped_query_attention' in attention_info:
if attention_info['uses_grouped_query_attention']:
info_lines.append("Type: Grouped-Query Attention (GQA)")
if 'num_kv_heads' in attention_info:
info_lines.append(f"KV Heads: {attention_info['num_kv_heads']}")
if 'heads_per_kv_group' in attention_info:
info_lines.append(f"Heads per KV Group: {attention_info['heads_per_kv_group']}")
else:
info_lines.append("Type: Multi-Head Attention (MHA)")
ax.text(5, 1.2, '\n'.join(info_lines), ha='center', va='top',
fontsize=12, bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.6))
return fig
def save_all_figures(self, output_dir: Union[str, Path]):
"""Saves all generated figures to the specified directory"""
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
figure_names = [
'architecture_overview.png',
'weight_distributions.png',
'layer_composition.png',
'attention_mechanism.png'
]
for fig, name in zip(self.figures, figure_names):
filepath = output_path / name
fig.savefig(filepath, dpi=300, bbox_inches='tight')
print(f"Saved visualization: {filepath}")
class ReportGenerator:
"""
Generates comprehensive reports about the analyzed model.
"""
def __init__(self, analysis_results: Dict[str, Any],
model_data: Dict[str, Any],
visualizations: List[plt.Figure]):
self.analysis = analysis_results
self.model_data = model_data
self.visualizations = visualizations
def generate_text_report(self) -> str:
"""
Generates a detailed text report.
Returns:
String containing the formatted report
"""
sections = []
sections.append(self._generate_header())
sections.append(self._generate_basic_info_section())
sections.append(self._generate_architecture_section())
sections.append(self._generate_attention_section())
sections.append(self._generate_quantization_section())
sections.append(self._generate_statistics_section())
sections.append(self._generate_recommendations())
sections.append(self._generate_footer())
return '\n\n'.join(sections)
def _generate_header(self) -> str:
"""Generates the report header"""
header = "=" * 80 + "\n"
header += " " * 25 + "MODEL ANALYSIS REPORT\n"
header += "=" * 80 + "\n"
header += f"Format: {self.model_data.get('format', 'Unknown')}\n"
header += f"Model Family: {self.analysis.get('model_family', 'Unknown')}\n"
if 'gguf_version' in self.model_data.get('metadata', {}):
header += f"GGUF Version: {self.model_data['metadata']['gguf_version']}\n"
return header
def _generate_basic_info_section(self) -> str:
"""Generates the basic information section"""
section = "BASIC INFORMATION\n"
section += "-" * 80 + "\n"
basic_info = self.analysis.get('basic_info', {})
if 'total_parameters' in basic_info:
total_params = basic_info['total_parameters']
params_m = basic_info.get('total_parameters_millions', total_params / 1_000_000)
params_b = basic_info.get('total_parameters_billions', total_params / 1_000_000_000)
section += f"Total Parameters: {total_params:,}\n"
if params_b >= 1.0:
section += f" {params_b:.2f} billion\n"
else:
section += f" {params_m:.2f} million\n"
if 'num_layers' in basic_info:
section += f"Number of Transformer Layers: {basic_info['num_layers']}\n"
return section
def _generate_architecture_section(self) -> str:
"""Generates detailed architecture information"""
section = "ARCHITECTURE DETAILS\n"
section += "-" * 80 + "\n"
transformer_info = self.analysis.get('transformer_architecture', {})
if 'hidden_dimension' in transformer_info:
section += f"Hidden Dimension: {transformer_info['hidden_dimension']}\n"
if 'feedforward_dimension' in transformer_info:
ff_dim = transformer_info['feedforward_dimension']
section += f"Feed-Forward Dimension: {ff_dim}\n"
if 'ff_to_hidden_ratio' in transformer_info:
ratio = transformer_info['ff_to_hidden_ratio']
section += f"FF to Hidden Ratio: {ratio:.2f}\n"
if 'context_length' in transformer_info:
section += f"Maximum Context Length: {transformer_info['context_length']}\n"
return section
def _generate_attention_section(self) -> str:
"""Generates attention mechanism information"""
section = "ATTENTION MECHANISM\n"
section += "-" * 80 + "\n"
attention_info = self.analysis.get('attention_mechanism', {})
if 'num_attention_heads' in attention_info:
section += f"Number of Attention Heads: {attention_info['num_attention_heads']}\n"
if 'head_dimension' in attention_info:
section += f"Head Dimension: {attention_info['head_dimension']}\n"
if 'uses_grouped_query_attention' in attention_info:
if attention_info['uses_grouped_query_attention']:
section += "Attention Type: Grouped-Query Attention (GQA)\n"
if 'num_kv_heads' in attention_info:
section += f"Number of Key-Value Heads: {attention_info['num_kv_heads']}\n"
if 'heads_per_kv_group' in attention_info:
section += f"Query Heads per KV Group: {attention_info['heads_per_kv_group']}\n"
else:
section += "Attention Type: Multi-Head Attention (MHA)\n"
return section
def _generate_quantization_section(self) -> str:
"""Generates quantization information"""
section = "QUANTIZATION\n"
section += "-" * 80 + "\n"
quant_info = self.analysis.get('quantization', {})
if quant_info.get('is_quantized', False):
section += "Model Status: Quantized\n"
if 'quantization_method' in quant_info:
section += f"Quantization Method: {quant_info['quantization_method']}\n"
if 'bit_width' in quant_info:
section += f"Bit Width: {quant_info['bit_width']}\n"
if 'quantization_version' in quant_info:
section += f"Quantization Version: {quant_info['quantization_version']}\n"
else:
section += "Model Status: Full Precision (Not Quantized)\n"
return section
def _generate_statistics_section(self) -> str:
"""Generates weight statistics"""
section = "WEIGHT STATISTICS\n"
section += "-" * 80 + "\n"
weight_stats = self.analysis.get('weight_statistics', {})
if not weight_stats:
section += "Weight statistics not available for this model format.\n"
return section
if 'average_mean' in weight_stats:
section += f"Average Weight Mean: {weight_stats['average_mean']:.8f}\n"
if 'average_std' in weight_stats:
section += f"Average Weight Std Dev: {weight_stats['average_std']:.8f}\n"
if 'mean_range' in weight_stats:
min_mean, max_mean = weight_stats['mean_range']
section += f"Mean Range: [{min_mean:.8f}, {max_mean:.8f}]\n"
if 'std_range' in weight_stats:
min_std, max_std = weight_stats['std_range']
section += f"Std Dev Range: [{min_std:.8f}, {max_std:.8f}]\n"
return section
def _generate_recommendations(self) -> str:
"""Generates recommendations based on analysis"""
section = "RECOMMENDATIONS AND INSIGHTS\n"
section += "-" * 80 + "\n"
recommendations = []
basic_info = self.analysis.get('basic_info', {})
if 'total_parameters_millions' in basic_info:
params_m = basic_info['total_parameters_millions']
if params_m < 1000:
recommendations.append(
f"With approximately {params_m:.0f}M parameters, this model is suitable "
"for edge deployment and consumer hardware. It can run efficiently on "
"modern CPUs and entry-level GPUs."
)
elif params_m < 10000:
recommendations.append(
f"With approximately {params_m:.0f}M parameters, this model requires "
"moderate computational resources. Recommended deployment on workstations "
"with dedicated GPUs (8GB+ VRAM) or server environments."
)
else:
params_b = basic_info.get('total_parameters_billions', params_m / 1000)
recommendations.append(
f"With approximately {params_b:.1f}B parameters, this model requires "
"significant computational resources. Recommended deployment on high-end "
"server hardware with multiple GPUs or distributed inference systems."
)
quant_info = self.analysis.get('quantization', {})
if not quant_info.get('is_quantized', False):
recommendations.append(
"The model is currently in full precision format. Consider applying "
"quantization techniques (e.g., 4-bit or 8-bit quantization) to reduce "
"memory footprint by 2-8x and improve inference speed, with minimal "
"impact on model quality."
)
else:
method = quant_info.get('quantization_method', '')
if 'Q4' in method or 'Q5' in method:
recommendations.append(
f"The model uses {method} quantization, which provides an excellent "
"balance between model size reduction and quality preservation. This "
"makes the model suitable for deployment on resource-constrained devices."
)
attention_info = self.analysis.get('attention_mechanism', {})
if attention_info.get('uses_grouped_query_attention', False):
recommendations.append(
"The model uses Grouped-Query Attention (GQA), which provides better "
"inference efficiency compared to standard Multi-Head Attention. This "
"reduces the KV cache size during generation, enabling longer context "
"processing with the same memory budget."
)
transformer_info = self.analysis.get('transformer_architecture', {})
if 'context_length' in transformer_info:
ctx_len = transformer_info['context_length']
if ctx_len >= 8192:
recommendations.append(
f"The model supports a context length of {ctx_len} tokens, enabling "
"processing of very long documents. Note that memory usage scales "
"quadratically with context length, so ensure adequate VRAM is available "
"when processing near-maximum context sizes."
)
elif ctx_len >= 4096:
recommendations.append(
f"The model supports a context length of {ctx_len} tokens, suitable "
"for most document processing tasks. This provides a good balance "
"between capability and resource requirements."
)
if recommendations:
section += '\n\n'.join(f"{i+1}. {rec}" for i, rec in enumerate(recommendations))
else:
section += "No specific recommendations available based on current analysis."
return section
def _generate_footer(self) -> str:
"""Generates the report footer"""
footer = "\n" + "=" * 80 + "\n"
footer += "End of Report\n"
footer += "=" * 80
return footer
def save_report(self, output_path: Union[str, Path]):
"""Saves the text report to a file"""
report_text = self.generate_text_report()
output_file = Path(output_path)
output_file.parent.mkdir(parents=True, exist_ok=True)
with open(output_file, 'w', encoding='utf-8') as f:
f.write(report_text)
print(f"Report saved to: {output_file}")
class LLMAnalyzer:
"""
Main orchestrator class that coordinates all components of the
LLM analysis tool.
"""
def __init__(self, model_path: Union[str, Path]):
self.model_path = Path(model_path)
self.format = None
self.model_data = None
self.analysis_results = None
self.visualizations = None
def analyze(self, output_dir: Optional[Union[str, Path]] = None) -> Dict[str, Any]:
"""
Performs complete analysis of the model.
Args:
output_dir: Optional directory to save outputs
Returns:
Dictionary containing all analysis results
"""
print(f"Analyzing model: {self.model_path}")
print("Step 1: Detecting model format...")
self.format = self._detect_format()
print(f"Detected format: {self.format.value}")
print("Step 2: Reading model data...")
self.model_data = self._read_model()
print("Model data loaded successfully")
print("Step 3: Analyzing model architecture...")
self.analysis_results = self._analyze_model()
print("Analysis complete")
print("Step 4: Creating visualizations...")
self.visualizations = self._create_visualizations()
print(f"Created {len(self.visualizations)} visualizations")
print("Step 5: Generating report...")
report = self._generate_report()
if output_dir:
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
print("Step 6: Saving outputs...")
self._save_outputs(output_path, report)
print(f"All outputs saved to: {output_path}")
return {
'format': self.format.value,
'model_data': self.model_data,
'analysis': self.analysis_results,
'report': report
}
def _detect_format(self) -> ModelFormat:
"""Detects the model format"""
detector = FormatDetector()
return detector.detect_format(self.model_path)
def _read_model(self) -> Dict[str, Any]:
"""Reads the model using the appropriate reader"""
if self.format == ModelFormat.GGUF:
reader = GGUFReader(self.model_path)
elif self.format == ModelFormat.PYTORCH:
if not TORCH_AVAILABLE:
raise RuntimeError("PyTorch is required but not available")
reader = PyTorchReader(self.model_path)
elif self.format == ModelFormat.TENSORFLOW:
if not TF_AVAILABLE:
raise RuntimeError("TensorFlow is required but not available")
reader = TensorFlowReader(self.model_path)
elif self.format == ModelFormat.MLX:
reader = MLXReader(self.model_path)
else:
raise ValueError(f"Unsupported format: {self.format}")
return reader.read_model()
def _analyze_model(self) -> Dict[str, Any]:
"""Analyzes the model data"""
analyzer = ModelAnalyzer(self.model_data)
return analyzer.analyze()
def _create_visualizations(self) -> List[plt.Figure]:
"""Creates visualizations"""
visualizer = ModelVisualizer(self.analysis_results, self.model_data)
return visualizer.create_all_visualizations()
def _generate_report(self) -> str:
"""Generates the text report"""
generator = ReportGenerator(self.analysis_results,
self.model_data,
self.visualizations)
return generator.generate_text_report()
def _save_outputs(self, output_dir: Path, report: str):
"""Saves all outputs to the specified directory"""
report_file = output_dir / 'model_analysis_report.txt'
with open(report_file, 'w', encoding='utf-8') as f:
f.write(report)
print(f"Saved report: {report_file}")
visualizer = ModelVisualizer(self.analysis_results, self.model_data)
visualizer.figures = self.visualizations
visualizer.save_all_figures(output_dir)
analysis_file = output_dir / 'analysis_data.json'
with open(analysis_file, 'w', encoding='utf-8') as f:
json.dump(self.analysis_results, f, indent=2, default=str)
print(f"Saved analysis data: {analysis_file}")
def main():
"""Main entry point for the command-line interface"""
parser = argparse.ArgumentParser(
description='Comprehensive LLM Model Analysis Tool',
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog='''
Examples:
python llm_analyzer.py model.gguf
python llm_analyzer.py model.pth --output-dir ./analysis_results
python llm_analyzer.py saved_model/ --output-dir ./tf_analysis
'''
)
parser.add_argument('model_path',
help='Path to the model file or directory')
parser.add_argument('--output-dir', '-o',
help='Directory to save analysis outputs',
default=None)
args = parser.parse_args()
try:
analyzer = LLMAnalyzer(args.model_path)
results = analyzer.analyze(output_dir=args.output_dir)
print("\n" + "=" * 80)
print("ANALYSIS SUMMARY")
print("=" * 80)
print(f"Model Format: {results['format']}")
print(f"Model Family: {results['analysis'].get('model_family', 'Unknown')}")
basic_info = results['analysis'].get('basic_info', {})
if 'total_parameters_millions' in basic_info:
params = basic_info['total_parameters_millions']
if params >= 1000:
print(f"Parameters: {params/1000:.2f}B")
else:
print(f"Parameters: {params:.2f}M")
if 'num_layers' in basic_info:
print(f"Layers: {basic_info['num_layers']}")
print("=" * 80)
print("\nAnalysis complete!")
except Exception as e:
print(f"Error during analysis: {str(e)}", file=sys.stderr)
import traceback
traceback.print_exc()
sys.exit(1)
if __name__ == '__main__':
main()
This complete implementation represents a production-ready tool that can analyze large language models across multiple formats. The code follows clean architecture principles with clear separation of concerns, comprehensive error handling, and extensive documentation. Each component is designed to be independently testable and maintainable. The tool provides deep insights into model architecture, quantization, and performance characteristics, making it invaluable for researchers and engineers working with local LLM deployments.
The implementation handles edge cases gracefully, such as missing optional dependencies, corrupted model files, and incomplete metadata. It provides informative error messages and degrades functionality gracefully when certain features are unavailable. The visualization system creates publication-quality figures that can be used in reports and presentations. The text report is structured and comprehensive, providing both technical details and practical recommendations for model deployment.
This tool serves as a foundation that can be extended with additional features such as performance benchmarking, memory profiling, and comparative analysis across multiple models. The modular design makes such extensions straightforward to implement without modifying the core functionality.
No comments:
Post a Comment