File size: 3,379 Bytes
c99f13f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
import gguf
import numpy as np
from constants import ARCH_FEATURES


def detect_architecture(tensors: dict) -> str:
    """Detect model architecture from tensor names."""
    names = list(tensors.keys())

    has_ssm = any("ssm_" in n for n in names)
    has_qkv = any("attn_qkv" in n for n in names)
    has_nextn = any("nextn" in n for n in names)
    has_moe = any("exps" in n for n in names)
    has_separate_qkv = any("attn_q.weight" in n for n in names)
    has_gemma_specific = any(
        t in n for t in ("layer_output_scale", "post_attention_norm", "post_ffw_norm")
        for n in names
    )

    if has_ssm and has_qkv:
        return "qwen35"
    if has_moe:
        return "mellum2"
    # Исправление: заменяем неопределённую has_blk_attn на уже существующий признак
    if has_separate_qkv or has_gemma_specific:
        return "gemma4"
    return "unknown"


def _detect_prefix(tensors: dict) -> str:
    for name in tensors:
        if name.startswith("BLK."):
            return "BLK"
    return "blk"


def _estimate_layers(tensors: dict) -> int:
    max_layer = 0
    for name in tensors:
        parts = name.split(".")
        if len(parts) >= 2 and parts[0] in ("blk", "BLK"):
            try:
                layer = int(parts[1])
                if layer > max_layer:
                    max_layer = layer
            except ValueError:
                pass
    return max_layer + 1  # layers are 0-indexed


def read_model(path: str) -> dict:
    """Parse BF16 GGUF, return model info."""
    r = gguf.GGUFReader(path)
    tensors = {}
    meta = {}

    for k, v in r.fields.items():
        # Безопасное извлечение данных (список/число, а не raw numpy)
        try:
            data = v.data
            if isinstance(data, np.ndarray):
                data = data.tolist()
            elif isinstance(data, (np.generic,)):
                data = data.item()
            meta[k] = data
        except Exception:
            meta[k] = str(v)

    for t in r.tensors:
        shape = list(t.shape)
        name = t.name
        n_elements = int(np.prod(shape))
        tensors[name] = {
            "shape": shape,
            "n_elements": n_elements,
            "size_mib": n_elements * 2 / 1024 / 1024,
        }

    arch = detect_architecture(tensors)
    arch_features = ARCH_FEATURES.get(arch, {}).copy()

    prefix = _detect_prefix(tensors)
    n_layers = _estimate_layers(tensors)

    if arch == "mellum2" and arch_features.get("moe_intermediate_size", 0) == 0:
        arch_features["moe_intermediate_size"] = 896

    arch_features["prefix"] = prefix
    if n_layers > 0:
        arch_features["n_layers"] = n_layers

    has_moe = any("exps" in n for n in tensors)
    if has_moe:
        arch_features["has_moe"] = True

    has_nextn = any("nextn" in n for n in tensors)
    # blk.32 is MTP only in ~32-layer models (Qwen). Skip for deeper models (Gemma4, 48 layers).
    has_blk32 = n_layers <= 33 and any(
        n.startswith("blk.32.") or n.startswith("BLK.32.") for n in tensors
    )
    arch_features["has_mtp"] = has_nextn or has_blk32

    return {
        "path": path,
        "architecture": arch,
        "features": arch_features,
        "tensors": tensors,
        "n_tensors": len(tensors),
        "meta": meta,
    }