diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..6313b56 --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +* text=auto eol=lf diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..c06256f --- /dev/null +++ b/.gitignore @@ -0,0 +1,25 @@ +# OS +.DS_Store +Thumbs.db + +# Python +.venv/ +venv/ +__pycache__/ +*.py[cod] + +# Editor +.idea/ +.vscode/ +.cursor/ +*.swp +*.swo + +# 中間ファイル(量子化の入力など) +build/* +!build/sentences.txt +# 配布物(target/ には配布するものだけ入る) +target/ + +# 作業記録(ローカルのみ) +journals/ diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..1ed7219 --- /dev/null +++ b/LICENSE @@ -0,0 +1,15 @@ +Creative Commons Attribution-ShareAlike 4.0 International (CC BY-SA 4.0) + +本リポジトリで配布するモデル(sonoisa/sentence-bert-base-ja-mean-tokens-v2 の量子化 ONNX 派生)は、 +Creative Commons Attribution-ShareAlike 4.0 International の下に提供されます。 + +フルテキスト: https://creativecommons.org/licenses/by-sa/4.0/legalcode + +要約: +- 表示: 適切なクレジット(元モデル・著作者)を表示すること。 +- 継承: 本作品に基づく派生作品は、同一の CC BY-SA 4.0 で頒布すること。 +- 商用利用: 可能。 + +元モデル: +- sonoisa/sentence-bert-base-ja-mean-tokens-v2 +- cl-tohoku/bert-base-japanese-whole-word-masking (東北大学 乾・鈴木研究室) diff --git a/README.md b/README.md new file mode 100644 index 0000000..435e4ee --- /dev/null +++ b/README.md @@ -0,0 +1,92 @@ +# sentence-bert-base-ja-mean-tokens-v2-int8 + +sonoisa/sentence-bert-base-ja-mean-tokens-v2 の **量子化 ONNX モデル** を CC BY-SA 4.0 で配布するためのリポジトリです。 + +## 概要 + +- **元モデル**: [sonoisa/sentence-bert-base-ja-mean-tokens-v2](https://huggingface.co/sonoisa/sentence-bert-base-ja-mean-tokens-v2)(日本語 Sentence-BERT) +- **ベース**: [cl-tohoku/bert-base-japanese-whole-word-masking](https://huggingface.co/cl-tohoku/bert-base-japanese-whole-word-masking)(東北大 BERT) +- **形式**: ONNX(INT8 量子化済み)、768 次元の文ベクトル出力 +- **用途**: 埋め込み・意味検索(要前処理でトークナイズ) + +## 配布するもの(target/) + +**`target/` には配布するものだけ**が入ります。この一式をそのまま渡せば、Hugging Face にアクセスせずに推論できます。 + +| ファイル | 説明 | +|----------|------| +| `model_quantized.onnx` | 量子化済み ONNX モデル(推論用) | +| `vocab.txt` | トークナイザー用語彙 | +| `tokenizer_config.json` | トークナイザー設定 | +| `special_tokens_map.json` | 特殊トークン定義 | +| `config.json` | モデル設定 | + +中間ファイル(`model_fp32.onnx` など)は `build/` に出力され、配布しません。 + +**配布方針**: `target/` 一式をそのまま ZIP やリリースアセットで配布する想定。リポジトリに target を含めるかは任意(.gitignore 済みのため通常は含めず、ビルド後に手元で配布する運用でよい)。 + +## リポジトリ構成 + +``` +sentence-bert-base-ja-mean-tokens-v2-int8/ +├── README.md +├── LICENSE +├── requirements.txt +├── export_onnx.py … HF モデル → ONNX (FP32) エクスポート +├── quantize.py … 静的量子化 → model_quantized.onnx +├── setup_venv.bat … 仮想環境 .venv 作成 +├── run_quantize.bat … エクスポート〜量子化の一括実行 +├── scripts/ +│ └── run_inference.py … target/ での推論・動作確認用 +├── test/ … 単体・統合テスト(pytest) +├── journals/ … 作業記録(.gitignore 済み) +├── build/ … 中間ファイル(model_fp32.onnx 等)と sentences.txt(キャリブレーション用) +│ └── sentences.txt … 量子化用サンプル文(リポジトリに同梱。編集可) +├── target/ … 配布するものだけ(.gitignore 済み) +└── .venv/ … setup_venv.bat で作成(.gitignore 済み) +``` + +## ビルド手順(量子化モデルを作る) + +**必要環境:** Python 3.8 以上(3.12 推奨) + +1. **環境準備** + - `setup_venv.bat` で `.venv` を作成し、依存関係をインストール + - または `pip install -r requirements.txt` + +2. **エクスポート** + `python export_onnx.py` + → `build/model_fp32.onnx`(中間)、トークナイザー・config は `target/` に出力 + +3. **量子化** + `python quantize.py` + → `target/model_quantized.onnx` が出力される(配布用) + +**一括実行:** `run_quantize.bat`(venv があれば自動で有効化してから実行) + +キャリブレーションの精度を上げたい場合は、`build/sentences.txt` に日本語文を追加してから `quantize.py` を再実行してください。 + +**Lint・テスト(任意)**: `pip install -e ".[dev]"` で pytest と ruff を入れたうえで、`ruff check .` と `pytest test/` を実行できる。 + +## 推論の例 + +`target/` を配布先にコピーしたうえで、`python scripts/run_inference.py` で動作確認できる。 + +```python +# スクリプト利用例 +from scripts.run_inference import encode +vec = encode("今日は良い天気です。") +print(vec.shape) # (1, 768) +``` + +## クレジット(表示義務) + +- **Sentence-BERT 日本語 v2**: sonoisa([Hugging Face](https://huggingface.co/sonoisa/sentence-bert-base-ja-mean-tokens-v2)) +- **事前学習 BERT**: 東北大学 乾・鈴木研究室([cl-tohoku/bert-base-japanese-whole-word-masking](https://huggingface.co/cl-tohoku/bert-base-japanese-whole-word-masking)) + +## ライセンス + +本リポジトリで配布するモデルおよび派生物は **Creative Commons Attribution-ShareAlike 4.0 International (CC BY-SA 4.0)** に従います。 +詳細は [LICENSE](LICENSE) および [https://creativecommons.org/licenses/by-sa/4.0/](https://creativecommons.org/licenses/by-sa/4.0/) を参照してください。 + +商用利用可能です。利用・再配布の際は上記クレジットの表示と、同一ライセンス(CC BY-SA 4.0)での継承が必要です。 diff --git a/build/sentences.txt b/build/sentences.txt new file mode 100644 index 0000000..14fc506 --- /dev/null +++ b/build/sentences.txt @@ -0,0 +1,20 @@ +これはキャリブレーション用の日本語サンプル文です。 +今日は良い天気ですね。 +機械学習と深層学習について勉強しています。 +検索エンジンで関連する文書を探します。 +質問に対する回答を生成してください。 +東京は日本の首都です。 +自然言語処理の研究が進んでいます。 +製品の品質を向上させる取り組みを進めます。 +会議は午後三時から始まります。 +データの分析結果を報告します。 +新しいプロジェクトの企画書を作成しました。 +お客様のご要望にお応えします。 +システムの稼働状況を確認してください。 +明日の天気予報は晴れです。 +日本語の文法は難しいです。 +技術的な問題が発生しました。 +予算の見直しが必要です。 +進捗状況を共有しましょう。 +資料を添付いたします。 +ご確認のほどよろしくお願いいたします。 diff --git a/export_onnx.py b/export_onnx.py new file mode 100644 index 0000000..6b58cd4 --- /dev/null +++ b/export_onnx.py @@ -0,0 +1,89 @@ +""" +sonoisa/sentence-bert-base-ja-mean-tokens-v2 を ONNX (FP32) にエクスポートする。 +出力は mean pooling 済みの 768 次元文ベクトル。 +""" +from pathlib import Path + +import torch +from transformers import AutoModel, AutoTokenizer + +# 出力先: build/ = 中間ファイル, target/ = 配布するものだけ +ROOT = Path(__file__).resolve().parent +BUILD_DIR = ROOT / "build" +TARGET_DIR = ROOT / "target" +MODEL_ID = "sonoisa/sentence-bert-base-ja-mean-tokens-v2" +ONNX_PATH = BUILD_DIR / "model_fp32.onnx" + + +class BertWithMeanPooling(torch.nn.Module): + """BERT 出力に mean pooling を適用して文ベクトルを返すラッパー。""" + + def __init__(self, model): + super().__init__() + self.bert = model + + def forward(self, input_ids, attention_mask, token_type_ids): + outputs = self.bert( + input_ids=input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + ) + last_hidden = outputs.last_hidden_state # (batch, seq, 768) + # attention_mask: (batch, seq) → (batch, seq, 1) でブロードキャスト + mask = attention_mask.unsqueeze(-1).float() + sum_emb = (last_hidden * mask).sum(dim=1) + sum_mask = mask.sum(dim=1).clamp(min=1e-9) + return sum_emb / sum_mask # (batch, 768) + + +def main(): + BUILD_DIR.mkdir(parents=True, exist_ok=True) + TARGET_DIR.mkdir(parents=True, exist_ok=True) + print(f"Loading {MODEL_ID} ...") + tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) + model = AutoModel.from_pretrained(MODEL_ID) + wrapper = BertWithMeanPooling(model) + wrapper.eval() + + # ダミー入力(batch=2, seq=64 で動的軸を付与) + batch_size, seq_len = 2, 64 + dummy = tokenizer( + ["これはテストです。", "もう一つの文です。"], + padding="max_length", + max_length=seq_len, + return_tensors="pt", + truncation=True, + ) + input_ids = dummy["input_ids"] + attention_mask = dummy["attention_mask"] + token_type_ids = dummy.get("token_type_ids") + if token_type_ids is None: + token_type_ids = torch.zeros_like(input_ids, dtype=torch.long) + + with torch.no_grad(): + torch.onnx.export( + wrapper, + (input_ids, attention_mask, token_type_ids), + str(ONNX_PATH), + input_names=["input_ids", "attention_mask", "token_type_ids"], + output_names=["sentence_embedding"], + dynamic_axes={ + "input_ids": {0: "batch_size", 1: "sequence_length"}, + "attention_mask": {0: "batch_size", 1: "sequence_length"}, + "token_type_ids": {0: "batch_size", 1: "sequence_length"}, + "sentence_embedding": {0: "batch_size"}, + }, + opset_version=18, + do_constant_folding=False, + ) + + print(f"Exported: {ONNX_PATH}") + + # 配布用にトークナイザー・config だけ target/ へ(配布物に同梱する) + tokenizer.save_pretrained(TARGET_DIR) + model.config.save_pretrained(TARGET_DIR) + print(f"Saved tokenizer and config to {TARGET_DIR} (distribution)") + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..f7a39c6 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,24 @@ +[project] +name = "sentence-bert-base-ja-mean-tokens-v2-int8" +version = "0.1.0" +description = "Quantum ONNX build for sentence-bert-base-ja-mean-tokens-v2" +requires-python = ">=3.8" + +[project.optional-dependencies] +dev = ["pytest>=7.0", "ruff>=0.1"] + +[tool.ruff] +target-version = "py38" +line-length = 100 +exclude = [".venv", "build", "target"] + +[tool.ruff.lint] +select = ["E", "F", "I", "N", "W"] +ignore = ["E501"] + +[tool.ruff.lint.isort] +known-first-party = [] + +[tool.pytest.ini_options] +testpaths = ["test"] +pythonpath = ["."] diff --git a/quantize.py b/quantize.py new file mode 100644 index 0000000..310cf4a --- /dev/null +++ b/quantize.py @@ -0,0 +1,97 @@ +""" +model_fp32.onnx をキャリブレーションデータで静的量子化し、model_quantized.onnx を出力する。 +先に export_onnx.py で model_fp32.onnx を生成しておくこと。 +""" +from pathlib import Path + +import numpy as np +from transformers import AutoTokenizer + +# onnxruntime 1.16+ では onnxruntime.quantization に統合されている想定 +try: + from onnxruntime.quantization import ( + CalibrationDataReader, + CalibrationMethod, + QuantFormat, + QuantType, + quantize_static, + ) +except ImportError: + from onnxruntime.quantization.quantize import quantize_static + from onnxruntime.quantization.calibrate import CalibrationDataReader, CalibrationMethod + from onnxruntime.quantization.quant_utils import QuantFormat, QuantType + +ROOT = Path(__file__).resolve().parent +BUILD_DIR = ROOT / "build" +TARGET_DIR = ROOT / "target" +MODEL_ID = "sonoisa/sentence-bert-base-ja-mean-tokens-v2" +SENTENCES_PATH = BUILD_DIR / "sentences.txt" +MODEL_FP32 = BUILD_DIR / "model_fp32.onnx" +MODEL_QUANT = TARGET_DIR / "model_quantized.onnx" +MAX_LENGTH = 128 +BATCH_SIZE = 1 + + +class SentenceCalibrationDataReader(CalibrationDataReader): + """キャリブレーション用日本語文をトークナイズして ONNX 入力形式で返す。""" + + def __init__(self, tokenizer, sentences_path: Path, max_length: int = MAX_LENGTH): + super().__init__() + self.tokenizer = tokenizer + self.max_length = max_length + with open(sentences_path, "r", encoding="utf-8") as f: + self.sentences = [line.strip() for line in f if line.strip()] + self.index = 0 + + def __len__(self): + return len(self.sentences) + + def get_next(self): + if self.index >= len(self.sentences): + return None + text = self.sentences[self.index] + self.index += 1 + enc = self.tokenizer( + text, + padding="max_length", + max_length=self.max_length, + truncation=True, + return_tensors="np", + ) + return { + "input_ids": enc["input_ids"].astype(np.int64), + "attention_mask": enc["attention_mask"].astype(np.int64), + "token_type_ids": enc.get("token_type_ids", np.zeros_like(enc["input_ids"], dtype=np.int64)), + } + + +def main(): + if not MODEL_FP32.exists(): + raise FileNotFoundError( + f"{MODEL_FP32} がありません。先に python export_onnx.py を実行してください。" + ) + if not SENTENCES_PATH.exists(): + raise FileNotFoundError(f"キャリブレーション用文一覧がありません: {SENTENCES_PATH}") + + print(f"Tokenizer: {MODEL_ID}") + tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) + reader = SentenceCalibrationDataReader(tokenizer, SENTENCES_PATH, MAX_LENGTH) + print(f"Calibration samples: {len(reader)}") + + print("Quantizing (static, QDQ, S8S8) ...") + quantize_static( + str(MODEL_FP32), + str(MODEL_QUANT), + reader, + quant_format=QuantFormat.QDQ, + activation_type=QuantType.QInt8, + weight_type=QuantType.QInt8, + calibrate_method=CalibrationMethod.MinMax, + per_channel=False, + reduce_range=False, + ) + print(f"Saved: {MODEL_QUANT}") + + +if __name__ == "__main__": + main() diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..e21d24a --- /dev/null +++ b/requirements.txt @@ -0,0 +1,10 @@ +# 量子化パイプライン用(再現ビルド用にバージョン指定) +transformers>=4.30.0,<6.0.0 +torch>=2.0.0,<3.0.0 +onnx>=1.14.0,<2.0.0 +onnxruntime>=1.16.0,<2.0.0 +onnxscript>=0.2.0 + +# 日本語トークナイザ(東北大 BERT) +fugashi>=1.2.0 +unidic-lite>=1.0.8 diff --git a/run_quantize.bat b/run_quantize.bat new file mode 100644 index 0000000..a14b31f --- /dev/null +++ b/run_quantize.bat @@ -0,0 +1,19 @@ +@echo off +set PYTHONUTF8=1 +cd /d "%~dp0" +if exist .venv\Scripts\activate.bat ( + call .venv\Scripts\activate.bat +) else ( + echo Installing dependencies into current Python... + python -m pip install -r requirements.txt + if errorlevel 1 exit /b 1 +) +echo. +echo Running export_onnx.py... +python export_onnx.py +if errorlevel 1 exit /b 1 +echo. +echo Running quantize.py... +python quantize.py +if errorlevel 1 exit /b 1 +echo Done. diff --git a/scripts/__init__.py b/scripts/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/scripts/__init__.py diff --git a/scripts/run_inference.py b/scripts/run_inference.py new file mode 100644 index 0000000..48fb7e7 --- /dev/null +++ b/scripts/run_inference.py @@ -0,0 +1,56 @@ +""" +target/ の量子化 ONNX で推論する。target/ を配布先にコピーしたうえで実行する。 +""" +from pathlib import Path + +import numpy as np +import onnxruntime as ort +from transformers import AutoTokenizer + +ROOT = Path(__file__).resolve().parent.parent +TARGET_DIR = ROOT / "target" + + +def encode( + text: str, + target_dir: Path | None = None, + max_length: int = 128, +): + """文を 768 次元ベクトルにエンコードする。""" + target_dir = target_dir or TARGET_DIR + if not (target_dir / "model_quantized.onnx").exists(): + raise FileNotFoundError(f"target に model_quantized.onnx がありません: {target_dir}") + tokenizer = AutoTokenizer.from_pretrained(str(target_dir)) + session = ort.InferenceSession( + str(target_dir / "model_quantized.onnx"), + providers=["CPUExecutionProvider"], + ) + enc = tokenizer( + text, + padding="max_length", + max_length=max_length, + truncation=True, + return_tensors="np", + ) + token_type_ids = enc.get( + "token_type_ids", + np.zeros_like(enc["input_ids"], dtype=np.int64), + ) + out, = session.run( + None, + { + "input_ids": enc["input_ids"].astype(np.int64), + "attention_mask": enc["attention_mask"].astype(np.int64), + "token_type_ids": token_type_ids, + }, + ) + return out # (1, 768) + + +def main(): + vec = encode("今日は良い天気です。") + print(vec.shape) # (1, 768) + + +if __name__ == "__main__": + main() diff --git a/setup_venv.bat b/setup_venv.bat new file mode 100644 index 0000000..2a9c681 --- /dev/null +++ b/setup_venv.bat @@ -0,0 +1,18 @@ +@echo off +set PYTHONUTF8=1 +cd /d "%~dp0" +if exist .venv\Scripts\activate.bat ( + echo .venv already exists. Activate and run: pip install -r requirements.txt + exit /b 0 +) +echo Creating .venv... +python -m venv .venv +if errorlevel 1 ( + echo Failed to create venv. Ensure Python 3.8+ is installed. + exit /b 1 +) +call .venv\Scripts\activate.bat +echo Installing dependencies... +pip install -r requirements.txt +echo Done. Use: .venv\Scripts\activate then run export_onnx.py / quantize.py +exit /b 0 diff --git a/test/test_build_paths.py b/test/test_build_paths.py new file mode 100644 index 0000000..b302060 --- /dev/null +++ b/test/test_build_paths.py @@ -0,0 +1,35 @@ +"""ビルド用パス・必須ファイルの存在を確認する。""" +from pathlib import Path + +import pytest + +ROOT = Path(__file__).resolve().parent.parent +BUILD_DIR = ROOT / "build" +TARGET_DIR = ROOT / "target" + + +def test_build_sentences_txt_exists(): + """build/sentences.txt がリポジトリに同梱されている。""" + assert BUILD_DIR.exists(), "build/ が存在すること" + assert (BUILD_DIR / "sentences.txt").exists(), "build/sentences.txt が存在すること" + + +def test_export_onnx_has_build_and_target(): + """export_onnx が参照する BUILD_DIR / TARGET_DIR が定義されている。""" + import sys + sys.path.insert(0, str(ROOT)) + from export_onnx import BUILD_DIR as E_BUILD, TARGET_DIR as E_TARGET # noqa: E402 + assert E_BUILD == BUILD_DIR + assert E_TARGET == TARGET_DIR + + +def test_quantize_uses_build_and_target(): + """quantize が build/model_fp32.onnx を読み target に書く。""" + import sys + sys.path.insert(0, str(ROOT)) + from quantize import BUILD_DIR as Q_BUILD, TARGET_DIR as Q_TARGET # noqa: E402 + from quantize import MODEL_FP32, MODEL_QUANT # noqa: E402 + assert Q_BUILD == BUILD_DIR + assert Q_TARGET == TARGET_DIR + assert MODEL_FP32 == BUILD_DIR / "model_fp32.onnx" + assert MODEL_QUANT == TARGET_DIR / "model_quantized.onnx" diff --git a/test/test_inference.py b/test/test_inference.py new file mode 100644 index 0000000..261734d --- /dev/null +++ b/test/test_inference.py @@ -0,0 +1,21 @@ +"""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}"