Newer
Older
sentence-bert-base-ja-mean-tokens-v2-int8 / test / test_inference.py
@楽曲作りまくりおじさん 楽曲作りまくりおじさん 18 hours ago 768 bytes Initial commit: sentence-bert-base-ja-mean-tokens-v2 量子化 ONNX ビルド
"""target/ に配布物があるとき、推論スクリプトで shape を確認する。"""
from pathlib import Path

import pytest

ROOT = Path(__file__).resolve().parent.parent
TARGET_DIR = ROOT / "target"


@pytest.mark.skipif(
    not (TARGET_DIR / "model_quantized.onnx").exists(),
    reason="target/model_quantized.onnx がない(先にビルドすること)",
)
def test_run_inference_shape():
    """run_inference で 1 文をエンコードし、出力が (1, 768) であることを確認する。"""
    import sys
    sys.path.insert(0, str(ROOT / "scripts"))
    from run_inference import encode  # noqa: E402

    vec = encode("テスト文です。", target_dir=TARGET_DIR)
    assert vec.shape == (1, 768), f"expected (1, 768), got {vec.shape}"