chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,70 @@
|
||||
"""Unit tests for Gemma3 model architecture."""
|
||||
|
||||
import pytest
|
||||
|
||||
from mlc_llm.model import MODEL_PRESETS, MODELS
|
||||
|
||||
|
||||
def test_gemma3_model_registered():
|
||||
"""Verify Gemma3 model is in the registry."""
|
||||
assert "gemma3" in MODELS, "gemma3 should be registered in MODELS"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
[
|
||||
"gemma3_2b",
|
||||
"gemma3_9b",
|
||||
],
|
||||
)
|
||||
def test_gemma3_creation(model_name: str):
|
||||
"""Test Gemma3 model creation and export to TVM IR.
|
||||
|
||||
Verifies:
|
||||
- Config can be loaded from preset
|
||||
- Model instance can be created
|
||||
- Model exports to TVM IR successfully
|
||||
- Named parameters are extracted
|
||||
"""
|
||||
model_info = MODELS["gemma3"]
|
||||
config = model_info.config.from_dict(MODEL_PRESETS[model_name])
|
||||
model = model_info.model(config)
|
||||
mod, named_params = model.export_tvm(
|
||||
spec=model.get_default_spec(),
|
||||
)
|
||||
|
||||
# Verify export succeeded
|
||||
assert mod is not None
|
||||
assert len(named_params) > 0
|
||||
|
||||
# Optional: show module structure
|
||||
mod.show(black_format=False)
|
||||
|
||||
# Print parameters for debugging
|
||||
for name, param in named_params:
|
||||
print(name, param.shape, param.dtype)
|
||||
|
||||
|
||||
def test_gemma3_config_validation():
|
||||
"""Test Gemma3 configuration has required fields."""
|
||||
model_info = MODELS["gemma3"]
|
||||
config = model_info.config.from_dict(MODEL_PRESETS["gemma3_2b"])
|
||||
|
||||
# Check required config parameters
|
||||
assert hasattr(config, "hidden_size") and config.hidden_size > 0
|
||||
assert hasattr(config, "num_hidden_layers") and config.num_hidden_layers > 0
|
||||
assert hasattr(config, "num_attention_heads") and config.num_attention_heads > 0
|
||||
assert hasattr(config, "vocab_size") and config.vocab_size > 0
|
||||
|
||||
print(
|
||||
f"Gemma3 Config: hidden_size={config.hidden_size}, "
|
||||
f"layers={config.num_hidden_layers}, "
|
||||
f"heads={config.num_attention_heads}, "
|
||||
f"vocab={config.vocab_size}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Allow running tests directly
|
||||
test_gemma3_creation("gemma3_2b")
|
||||
test_gemma3_creation("gemma3_9b")
|
||||
@@ -0,0 +1,20 @@
|
||||
import pytest
|
||||
|
||||
from mlc_llm.model import MODEL_PRESETS, MODELS
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", ["gpt2"])
|
||||
def test_gpt2_creation(model_name: str):
|
||||
model_info = MODELS["gpt2"]
|
||||
config = model_info.config.from_dict(MODEL_PRESETS[model_name])
|
||||
model = model_info.model(config)
|
||||
mod, named_params = model.export_tvm(
|
||||
spec=model.get_default_spec(),
|
||||
)
|
||||
mod.show(black_format=False)
|
||||
for name, param in named_params:
|
||||
print(name, param.shape, param.dtype)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_gpt2_creation("gpt2")
|
||||
@@ -0,0 +1,20 @@
|
||||
import pytest
|
||||
|
||||
from mlc_llm.model import MODEL_PRESETS, MODELS
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", ["redpajama_3b_v1"])
|
||||
def test_mistral_creation(model_name: str):
|
||||
model_info = MODELS["gpt_neox"]
|
||||
config = model_info.config.from_dict(MODEL_PRESETS[model_name])
|
||||
model = model_info.model(config)
|
||||
mod, named_params = model.export_tvm(
|
||||
spec=model.get_default_spec(),
|
||||
)
|
||||
mod.show(black_format=False)
|
||||
for name, param in named_params:
|
||||
print(name, param.shape, param.dtype)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_mistral_creation("redpajama_3b_v1")
|
||||
@@ -0,0 +1,114 @@
|
||||
import tvm
|
||||
from tvm import tirx
|
||||
from tvm.relax.frontend.nn import core, modules, spec
|
||||
from tvm.script import ir as I
|
||||
from tvm.script import relax as R
|
||||
from tvm.script import tirx as T
|
||||
|
||||
from mlc_llm.nn.kv_cache import PagedKVCache, RopeMode
|
||||
|
||||
# mypy: disable-error-code="attr-defined"
|
||||
|
||||
|
||||
def test_nn_module_paged_kv_cache():
|
||||
# fmt: off
|
||||
@I.ir_module
|
||||
class Module:
|
||||
@R.function
|
||||
def create_paged_kv_cache(
|
||||
max_batch_size: R.Shape(["max_batch_size_1"]),
|
||||
max_total_seq_len: R.Shape(["max_total_seq_len_1"]),
|
||||
prefill_chunk_size: R.Shape(["prefill_chunk_size_1"]),
|
||||
page_size: R.Shape(["page_size_1"]),
|
||||
support_sliding_window: R.Shape(["support_sliding_window_1"]),
|
||||
) -> R.Object:
|
||||
max_batch_size_1 = T.int64()
|
||||
max_total_seq_len_1 = T.int64()
|
||||
prefill_chunk_size_1 = T.int64()
|
||||
page_size_1 = T.int64()
|
||||
support_sliding_window_1 = T.int64()
|
||||
R.func_attr({"num_input": 5})
|
||||
with R.dataflow():
|
||||
paged_kv_cache: R.Object = R.call_pure_packed("mlc.create_paged_kv_cache_generic", R.shape([max_batch_size_1, max_total_seq_len_1, prefill_chunk_size_1, page_size_1, support_sliding_window_1]), R.prim_value(32), R.prim_value(32), R.prim_value(32), R.prim_value(128), R.prim_value(1), R.prim_value(1), R.prim_value(10000), R.prim_value(128), R.dtype("float16"), sinfo_args=(R.Object,)) # noqa: E501
|
||||
gv1: R.Object = paged_kv_cache
|
||||
R.output(gv1)
|
||||
return gv1
|
||||
|
||||
@R.function
|
||||
def forward(
|
||||
cache: R.Object, qkv: R.Tensor((1, 100, 96, 128), dtype="float16")
|
||||
) -> R.Tensor((1, 100, 32, 128), dtype="float16"):
|
||||
R.func_attr({"num_input": 2})
|
||||
with R.dataflow():
|
||||
reshape: R.Tensor((100, 96, 128), dtype="float16") = R.reshape(
|
||||
qkv, R.shape([100, 96, 128])
|
||||
)
|
||||
lv = R.call_dps_packed(
|
||||
"vm.builtin.attention_kv_cache_attention_with_fused_qkv",
|
||||
(cache, R.prim_value(0), R.prim_value(T.float32(1)), reshape),
|
||||
out_sinfo=R.Tensor((100, 32, 128), dtype="float16"),
|
||||
)
|
||||
reshape1: R.Tensor((1, 100, 32, 128), dtype="float16") = R.reshape(
|
||||
lv, R.shape([1, 100, 32, 128])
|
||||
)
|
||||
gv: R.Tensor((1, 100, 32, 128), dtype="float16") = reshape1
|
||||
R.output(gv)
|
||||
return gv
|
||||
# fmt: on
|
||||
|
||||
class PagedKVCacheTest(modules.Module):
|
||||
def forward(
|
||||
self,
|
||||
cache: PagedKVCache,
|
||||
qkv: core.Tensor,
|
||||
) -> core.Tensor:
|
||||
return cache.attention_with_fused_qkv(0, qkv, num_qo_heads=32, sm_scale=128**-0.5)
|
||||
|
||||
def create_paged_kv_cache(
|
||||
self,
|
||||
max_batch_size: tirx.Var,
|
||||
max_total_seq_len: tirx.Var,
|
||||
prefill_chunk_size: tirx.Var,
|
||||
page_size: tirx.Var,
|
||||
support_sliding_window: tirx.Var,
|
||||
) -> PagedKVCache:
|
||||
return PagedKVCache.create_generic(
|
||||
attn_kind="mha",
|
||||
max_batch_size=max_batch_size,
|
||||
max_total_seq_len=max_total_seq_len,
|
||||
prefill_chunk_size=prefill_chunk_size,
|
||||
page_size=page_size,
|
||||
support_sliding_window=support_sliding_window,
|
||||
num_hidden_layers=32,
|
||||
num_attention_heads=32,
|
||||
num_key_value_heads=32,
|
||||
qk_head_dim=128,
|
||||
v_head_dim=128,
|
||||
rope_mode=RopeMode.NORMAL,
|
||||
rope_scale=1,
|
||||
rope_theta=10000,
|
||||
rotary_dim=128,
|
||||
dtype="float16",
|
||||
)
|
||||
|
||||
export_results = PagedKVCacheTest().export_tvm(
|
||||
spec={
|
||||
"forward": {
|
||||
"cache": spec.Object(object_type=PagedKVCache),
|
||||
"qkv": spec.Tensor((1, 100, 96, 128), "float16"),
|
||||
},
|
||||
"create_paged_kv_cache": {
|
||||
"max_batch_size": int,
|
||||
"max_total_seq_len": int,
|
||||
"prefill_chunk_size": int,
|
||||
"page_size": int,
|
||||
"support_sliding_window": int,
|
||||
},
|
||||
},
|
||||
)
|
||||
tvm_mod = export_results[0]
|
||||
tvm.ir.assert_structural_equal(tvm_mod, Module, True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_nn_module_paged_kv_cache()
|
||||
@@ -0,0 +1,25 @@
|
||||
import pytest
|
||||
|
||||
from mlc_llm.model import MODEL_PRESETS, MODELS
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_name", ["llama2_7b", "llama2_13b", "llama2_70b", "tinyllama_1b_chat_v1.0"]
|
||||
)
|
||||
def test_llama2_creation(model_name: str):
|
||||
model_info = MODELS["llama"]
|
||||
config = model_info.config.from_dict(MODEL_PRESETS[model_name])
|
||||
model = model_info.model(config)
|
||||
mod, named_params = model.export_tvm(
|
||||
spec=model.get_default_spec(),
|
||||
)
|
||||
mod.show(black_format=False)
|
||||
for name, param in named_params:
|
||||
print(name, param.shape, param.dtype)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_llama2_creation("llama2_7b")
|
||||
test_llama2_creation("llama2_13b")
|
||||
test_llama2_creation("llama2_70b")
|
||||
test_llama2_creation("tinyllama_1b_chat_v1")
|
||||
@@ -0,0 +1,72 @@
|
||||
import pytest
|
||||
|
||||
from mlc_llm.model import MODEL_PRESETS, MODELS
|
||||
from mlc_llm.quantization import QUANTIZATION
|
||||
from mlc_llm.quantization.group_quantization import (
|
||||
GroupQuantizeEmbedding,
|
||||
GroupQuantizeLinear,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
["llama2_7b", "llama2_13b", "llama2_70b"],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"quant_name",
|
||||
["q3f16_1", "q4f16_1", "q4f32_1"],
|
||||
)
|
||||
def test_llama2_group_quantization(model_name: str, quant_name: str):
|
||||
model_info = MODELS["llama"]
|
||||
config = model_info.config.from_dict(MODEL_PRESETS[model_name])
|
||||
model, quant_map = model_info.quantize["group-quant"](config, QUANTIZATION[quant_name])
|
||||
assert "model.embed_tokens.weight" in quant_map.param_map
|
||||
assert isinstance(
|
||||
model.model.embed_tokens,
|
||||
GroupQuantizeEmbedding,
|
||||
)
|
||||
assert "lm_head.weight" in quant_map.param_map
|
||||
assert isinstance(model.lm_head, GroupQuantizeLinear)
|
||||
for i in range(config.num_hidden_layers):
|
||||
assert f"model.layers.{i}.self_attn.qkv_proj.weight" in quant_map.param_map
|
||||
assert isinstance(
|
||||
model.model.layers[i].self_attn.qkv_proj,
|
||||
GroupQuantizeLinear,
|
||||
)
|
||||
assert f"model.layers.{i}.self_attn.o_proj.weight" in quant_map.param_map
|
||||
assert isinstance(
|
||||
model.model.layers[i].self_attn.o_proj,
|
||||
GroupQuantizeLinear,
|
||||
)
|
||||
assert f"model.layers.{i}.mlp.gate_up_proj.weight" in quant_map.param_map
|
||||
assert isinstance(
|
||||
model.model.layers[i].mlp.gate_up_proj,
|
||||
GroupQuantizeLinear,
|
||||
)
|
||||
assert f"model.layers.{i}.mlp.down_proj.weight" in quant_map.param_map
|
||||
assert isinstance(
|
||||
model.model.layers[i].mlp.down_proj,
|
||||
GroupQuantizeLinear,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_name",
|
||||
["llama2_7b", "llama2_13b", "llama2_70b"],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"quant_name",
|
||||
["q0f16", "q0f32"],
|
||||
)
|
||||
def test_llama2_no_quantization(model_name: str, quant_name: str):
|
||||
model_info = MODELS["llama"]
|
||||
config = model_info.config.from_dict(MODEL_PRESETS[model_name])
|
||||
_, quant_map = model_info.quantize["no-quant"](config, QUANTIZATION[quant_name])
|
||||
assert len(quant_map.param_map) == 0
|
||||
assert len(quant_map.map_func) == 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_llama2_group_quantization("llama2_7b", "q4f16_1")
|
||||
test_llama2_group_quantization("llama2_13b", "q4f16_1")
|
||||
test_llama2_group_quantization("llama2_70b", "q4f16_1")
|
||||
@@ -0,0 +1,20 @@
|
||||
import pytest
|
||||
|
||||
from mlc_llm.model import MODEL_PRESETS, MODELS
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", ["mistral_7b"])
|
||||
def test_mistral_creation(model_name: str):
|
||||
model_info = MODELS["mistral"]
|
||||
config = model_info.config.from_dict(MODEL_PRESETS[model_name])
|
||||
model = model_info.model(config)
|
||||
mod, named_params = model.export_tvm(
|
||||
spec=model.get_default_spec(),
|
||||
)
|
||||
mod.show(black_format=False)
|
||||
for name, param in named_params:
|
||||
print(name, param.shape, param.dtype)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_mistral_creation("mistral_7b")
|
||||
@@ -0,0 +1,21 @@
|
||||
import pytest
|
||||
|
||||
from mlc_llm.model import MODEL_PRESETS, MODELS
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name", ["phi-1_5", "phi-2"])
|
||||
def test_phi_creation(model_name: str):
|
||||
model_info = MODELS["phi-msft"]
|
||||
config = model_info.config.from_dict(MODEL_PRESETS[model_name])
|
||||
model = model_info.model(config)
|
||||
mod, named_params = model.export_tvm(
|
||||
spec=model.get_default_spec(),
|
||||
)
|
||||
mod.show(black_format=False)
|
||||
for name, param in named_params:
|
||||
print(name, param.shape, param.dtype)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_phi_creation("phi-1_5")
|
||||
test_phi_creation("phi-2")
|
||||
@@ -0,0 +1,147 @@
|
||||
import json
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
import tvm
|
||||
from safetensors import safe_open
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
from tvm import relax
|
||||
from tvm.contrib import tvmjs
|
||||
from tvm.runtime import ShapeTuple
|
||||
from tvm.runtime.vm import VirtualMachine
|
||||
|
||||
MLC_QWEN3_EMB_HF_DIR = os.environ.get("MLC_QWEN3_EMB_HF_DIR")
|
||||
MLC_QWEN3_EMB_MODEL_DIR = os.environ.get("MLC_QWEN3_EMB_MODEL_DIR")
|
||||
MLC_QWEN3_EMB_MODEL_LIB = os.environ.get("MLC_QWEN3_EMB_MODEL_LIB")
|
||||
MLC_QWEN3_EMB_DEVICE = os.environ.get("MLC_QWEN3_EMB_DEVICE", "cuda")
|
||||
|
||||
_skip = not all([MLC_QWEN3_EMB_HF_DIR, MLC_QWEN3_EMB_MODEL_DIR, MLC_QWEN3_EMB_MODEL_LIB])
|
||||
_skip_reason = (
|
||||
"Set MLC_QWEN3_EMB_HF_DIR, MLC_QWEN3_EMB_MODEL_DIR, MLC_QWEN3_EMB_MODEL_LIB to run this test"
|
||||
)
|
||||
|
||||
TEST_TEXTS = [
|
||||
"What is machine learning?",
|
||||
"CMU is Carnegie Mellon University",
|
||||
"机器学习是人工智能的一个分支",
|
||||
"量子コンピュータの基本原理を説明してください",
|
||||
"머신러닝은 인공지능의 한 분야입니다.",
|
||||
(
|
||||
"Instruct: Given a web search query, retrieve relevant passages "
|
||||
"that answer the query\nQuery: What is the capital of China?"
|
||||
),
|
||||
(
|
||||
"The Transformer architecture, introduced in the paper Attention Is All You Need, "
|
||||
"revolutionized natural language processing by replacing recurrent layers with "
|
||||
"self-attention mechanisms. This allows the model to process all positions in a "
|
||||
"sequence simultaneously rather than sequentially, leading to significant improvements "
|
||||
"in both training efficiency and the ability to capture long-range dependencies. "
|
||||
"The key innovation is the multi-head attention mechanism, which allows the model "
|
||||
"to jointly attend to information from different representation subspaces at "
|
||||
"different positions."
|
||||
),
|
||||
"Hello",
|
||||
"def fibonacci(n): return n if n <= 1 else fibonacci(n-1) + fibonacci(n-2)",
|
||||
]
|
||||
|
||||
|
||||
def _load_embed_weight(hf_dir):
|
||||
safetensor_files = [f for f in os.listdir(hf_dir) if f.endswith(".safetensors")]
|
||||
for sf in safetensor_files:
|
||||
with safe_open(os.path.join(hf_dir, sf), framework="pt", device="cpu") as f:
|
||||
if "embed_tokens.weight" in f.keys():
|
||||
return f.get_tensor("embed_tokens.weight")
|
||||
raise FileNotFoundError(f"embed_tokens.weight not found in {hf_dir}")
|
||||
|
||||
|
||||
def _hf_logits(text, tokenizer, hf_model, embed_weight):
|
||||
inputs = tokenizer(text, return_tensors="pt")
|
||||
with torch.no_grad():
|
||||
hidden = hf_model(**inputs).last_hidden_state.float()
|
||||
logits = hidden @ embed_weight.float().T
|
||||
return logits[0, -1, :].numpy()
|
||||
|
||||
|
||||
def _mlc_logits(text, tokenizer, mlc_module, params, metadata, dev, embed_weight):
|
||||
input_ids = tokenizer(text, return_tensors="pt")["input_ids"][0].numpy().astype(np.int32)
|
||||
seq_len = len(input_ids)
|
||||
|
||||
embed_func = mlc_module["embed"]
|
||||
prefill_func = mlc_module["prefill_to_last_hidden_states"]
|
||||
|
||||
if mlc_module.implements_function("create_flashinfer_paged_kv_cache"):
|
||||
create_kv = mlc_module["create_flashinfer_paged_kv_cache"]
|
||||
elif mlc_module.implements_function("create_tir_paged_kv_cache"):
|
||||
create_kv = mlc_module["create_tir_paged_kv_cache"]
|
||||
else:
|
||||
raise RuntimeError("Cannot find KV cache creation function")
|
||||
|
||||
sliding_window = metadata.get("sliding_window_size", -1)
|
||||
context_window = metadata.get("context_window_size", 32768)
|
||||
prefill_chunk = metadata.get("prefill_chunk_size", 2048)
|
||||
max_seq_len = sliding_window if context_window == -1 else context_window
|
||||
|
||||
kv_cache = create_kv(
|
||||
ShapeTuple([1]),
|
||||
ShapeTuple([max_seq_len]),
|
||||
ShapeTuple([prefill_chunk]),
|
||||
ShapeTuple([16]),
|
||||
ShapeTuple([int(sliding_window != -1)]),
|
||||
)
|
||||
|
||||
nd_view = tvm.get_global_func("vm.builtin.reshape")
|
||||
add_sequence = tvm.get_global_func("vm.builtin.kv_state_add_sequence")
|
||||
begin_forward = tvm.get_global_func("vm.builtin.kv_state_begin_forward")
|
||||
end_forward = tvm.get_global_func("vm.builtin.kv_state_end_forward")
|
||||
|
||||
tokens_tvm = tvm.runtime.tensor(input_ids, device=dev)
|
||||
embedding = embed_func(tokens_tvm, params)
|
||||
embedding = nd_view(embedding, ShapeTuple([1, seq_len, embedding.shape[-1]]))
|
||||
|
||||
add_sequence(kv_cache, 0)
|
||||
begin_forward(kv_cache, ShapeTuple([0]), ShapeTuple([seq_len]))
|
||||
hidden_states, _ = prefill_func(embedding, kv_cache, params)
|
||||
end_forward(kv_cache)
|
||||
|
||||
# Compute logits from hidden states using embed_tokens weight (tie_word_embeddings)
|
||||
hidden = hidden_states.numpy().astype(np.float32)
|
||||
logits = hidden @ embed_weight.float().numpy().T
|
||||
return logits[0, -1, :]
|
||||
|
||||
|
||||
@pytest.mark.skipif(_skip, reason=_skip_reason)
|
||||
def test_mlc_hf_logit_match():
|
||||
tokenizer = AutoTokenizer.from_pretrained(MLC_QWEN3_EMB_HF_DIR, padding_side="left")
|
||||
hf_model = AutoModel.from_pretrained(MLC_QWEN3_EMB_HF_DIR)
|
||||
embed_weight = _load_embed_weight(MLC_QWEN3_EMB_HF_DIR)
|
||||
|
||||
dev = tvm.runtime.device(MLC_QWEN3_EMB_DEVICE, 0)
|
||||
ex = tvm.runtime.load_module(MLC_QWEN3_EMB_MODEL_LIB)
|
||||
vm = relax.VirtualMachine(ex, dev)
|
||||
mlc_module = vm.module
|
||||
|
||||
metadata = json.loads(VirtualMachine(ex, tvm.runtime.device("cpu"))["_metadata"]())
|
||||
params_dict, _ = tvmjs.load_tensor_cache(MLC_QWEN3_EMB_MODEL_DIR, dev)
|
||||
param_names = [p["name"] for p in metadata["params"]]
|
||||
params = [params_dict[name] for name in param_names]
|
||||
|
||||
for text in TEST_TEXTS:
|
||||
hf = _hf_logits(text, tokenizer, hf_model, embed_weight)
|
||||
mlc = _mlc_logits(text, tokenizer, mlc_module, params, metadata, dev, embed_weight)
|
||||
|
||||
cos_sim = np.dot(hf, mlc) / (np.linalg.norm(hf) * np.linalg.norm(mlc))
|
||||
assert cos_sim > 0.99, f"[{text[:30]}] Cosine similarity {cos_sim:.6f} below 0.99"
|
||||
|
||||
max_diff = np.max(np.abs(hf - mlc))
|
||||
assert max_diff < 1.0, f"[{text[:30]}] Max absolute diff {max_diff:.6e} exceeds 1.0"
|
||||
|
||||
hf_top10 = set(np.argsort(hf)[-10:])
|
||||
mlc_top10 = set(np.argsort(mlc)[-10:])
|
||||
overlap = len(hf_top10 & mlc_top10)
|
||||
assert overlap >= 7, f"[{text[:30]}] Top-10 overlap {overlap}/10 below 7"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_mlc_hf_logit_match()
|
||||
Reference in New Issue
Block a user