src/utils/common.py

import copy
from typing import List, Union, Dict

import PIL.Image
import torch
import numpy as np
import torchvision.transforms.functional as F
import transformers

from transformers import PreTrainedTokenizer

IGNORE_INDEX = -100

def print_trainable_params(model: torch.nn.Module) -> None:
    trainable_params, all_param = 0, 0
    for param in model.parameters():
        num_params = param.numel()
        # if using DS Zero 3 and the weights are initialized empty
        if num_params == 0 and hasattr(param, "ds_numel"):
            num_params = param.ds_numel
        all_param += num_params
        if param.requires_grad:
            trainable_params += num_params
    print("trainable params: {:d} || all params: {:d} || trainable%: {:.4f}".format(
        trainable_params, all_param, 100 * trainable_params / all_param))


def post_process_generate_ids(tokenizer: PreTrainedTokenizer, ids: torch.Tensor):
    ids = copy.deepcopy(ids)  # do not modify origin preds and targets
    ids[ids < 0] = tokenizer.pad_token_id
    # pad_to_multiof 开启后, 多余的部分没法解码, 这里暂时替换为 ','
    ids[ids >= len(tokenizer)] = tokenizer.convert_tokens_to_ids(',')
    return ids


def decode_generate_ids(tokenizer: PreTrainedTokenizer, ids: torch.Tensor) -> Union[List[str], str]:
    assert ids.ndim in [1, 2]
    only_one_sentence = ids.ndim == 1
    if only_one_sentence:
        ids = ids.unsqueeze(0)
    ids = post_process_generate_ids(tokenizer, ids)
    res = tokenizer.batch_decode(ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
    if only_one_sentence:
        return res[0]
    return res



def draw_bounding_boxes(
        image: Union[torch.Tensor, PIL.Image.Image],
        boxes: Union[torch.Tensor, List, np.ndarray],
        **kwargs,
):
    if isinstance(image, PIL.Image.Image):
        from torchvision.transforms import PILToTensor
        image = PILToTensor()(image)
    assert isinstance(image, torch.Tensor), ""

    if not isinstance(boxes, torch.Tensor):
        boxes = torch.as_tensor(boxes)
    assert isinstance(boxes, torch.Tensor)

    from torchvision.utils import draw_bounding_boxes as _draw_bounding_boxes
    return _draw_bounding_boxes(image, boxes, **kwargs)


# https://github.com/huggingface/tokenizers/issues/247#issuecomment-675458087
def smart_tokenizer_and_embedding_resize(
        special_tokens_dict: Dict,
        tokenizer: transformers.PreTrainedTokenizer,
        model: transformers.PreTrainedModel,
):
    """Resize tokenizer and embedding.

    Note: This is the unoptimized version that may make your embedding size not be divisible by 64.
    """
    num_new_tokens = tokenizer.add_special_tokens(special_tokens_dict)
    model.resize_token_embeddings(len(tokenizer))

    if num_new_tokens > 0:
        input_embeddings = model.get_input_embeddings().weight.data
        output_embeddings = model.get_output_embeddings().weight.data

        input_embeddings_avg = input_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)
        output_embeddings_avg = output_embeddings[:-num_new_tokens].mean(dim=0, keepdim=True)

        input_embeddings[-num_new_tokens:] = input_embeddings_avg
        output_embeddings[-num_new_tokens:] = output_embeddings_avg

def patch_transformer_logging():
    import logging
    import transformers
    def enable_explicit_format():
        handlers = transformers.utils.logging._get_library_root_logger().handlers

        for handler in handlers:
            formatter = logging.Formatter("[(%(levelname)s) %(pathname)s:%(lineno)s ] %(asctime)s >> %(message)s")
            handler.setFormatter(formatter)
    transformers.utils.logging.enable_explicit_format = enable_explicit_format

def print_model_stats(model):
    """
    通用模型参数统计工具。
    自动识别是 Dense 还是 MoE 模型，并计算 Total vs Active 参数量。
    """
    
    # 1. 计算物理总参数量 (Total Parameters)
    # 使用 set 避免计算共享参数 (Shared Weights)，例如 Embedding 和 lm_head 共享权重的情况
    unique_params = {p.data_ptr(): p for p in model.parameters()}.values()
    total_params = sum(p.numel() for p in unique_params)
    
    # 2. 初始化激活参数量 (Active Parameters)
    # 默认假设是 Dense 模型，所有参数都是激活的
    active_params = total_params
    
    moe_infos = [] # 用于存储发现的 MoE 层信息
    
    # 3. 遍历所有子模块，寻找 MoE 层特征
    # 我们不匹配类名，而是匹配"特征" (Duck Typing)
    for name, module in model.named_modules():
        # 特征判定：有 num_experts 属性，且有一个叫 experts 的 ModuleList
        if hasattr(module, 'num_experts') and hasattr(module, 'experts') and isinstance(module.experts, nn.ModuleList):
            
            # 获取关键超参
            num_experts = getattr(module, 'num_experts', 0)
            # 兼容不同的 top_k 命名 (top_k 或 num_experts_per_tok)
            top_k = getattr(module, 'top_k', getattr(module, 'num_experts_per_tok', 0))
            
            # 如果找不到 top_k，可能不是标准的 Sparse MoE，跳过
            if top_k == 0: 
                continue

            # --- 核心计算逻辑 ---
            # 1. 计算单个专家的参数量 (假设所有专家结构相同，取第一个)
            # 这里必须用 recursion=True 确保统计专家内部所有层
            single_expert_params = sum(p.numel() for p in module.experts[0].parameters())
            
            # 2. 计算"休眠"专家数量
            dormant_experts = num_experts - top_k
            
            # 3. 从激活总数中扣除休眠专家的参数
            # 注意：total_params 里已经包含了 N 个专家，我们只需要减去 (N-K) 个
            if dormant_experts > 0:
                deduction = dormant_experts * single_expert_params
                active_params -= deduction
                
                moe_infos.append({
                    "layer": name,
                    "experts": num_experts,
                    "active": top_k,
                    "expert_size": single_expert_params
                })

    # --- 4. 格式化输出 ---
    def format_num(num):
        if num >= 1e9: return f"{num/1e9:.2f}B"
        if num >= 1e6: return f"{num/1e6:.2f}M"
        if num >= 1e3: return f"{num/1e3:.2f}K"
        return str(num)

    print("=" * 50)
    print(f"Model Architecture Analysis")
    print("=" * 50)
    
    if len(moe_infos) > 0:
        print(f"👉 Detection: MoE Model (Sparse Mixture-of-Experts)")
        print(f"   - Found {len(moe_infos)} MoE layers")
        print(f"   - Config: {moe_infos[0]['experts']} Experts, Top-{moe_infos[0]['active']} Active")
    else:
        print(f"👉 Detection: Dense Model (Standard Transformer)")
        
    print("-" * 50)
    print(f"Total Parameters (VRAM):  {format_num(total_params)}")
    print(f"Active Parameters (FLOPs): {format_num(active_params)}")
    
    if len(moe_infos) > 0:
        sparsity = 1 - (active_params / total_params)
        print(f"Sparsity Ratio:            {sparsity:.2%}")
        # 计算相比 Dense 版本的倍数
        # 假设 Dense 版本就是 active_params 大小（不太严谨但直观）
        print(f"Upcycling Scale:           {total_params/active_params:.2f}x Larger than Dense Base")
    else:
        print(f"Sparsity Ratio:            0.00% (Dense)")
        
    print("=" * 50)
    
    return total_params, active_params