File size: 3,019 Bytes
51bf86f | 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 | #!/usr/bin/env python3
"""Runs the original PyTorch pipeline (GecBERTModel from data/src) to get the outputs the mobile export must match.
Input: the model card's two examples, the unique FLEURS Vietnamese test transcriptions lower-cased with punctuation
removed (what an ASR model outputs), and paragraphs of 8 of them, long enough to be split into chunks. Writes out/reference.json: [{"text", "truth", "original"}], where "original" is
GecBERTModel's output with the settings of the model card (split_chunk=True, 3 iterations).
Usage:
.venv/bin/python -I reference.py
"""
import json
import sys
import time
from dataclasses import dataclass
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent)) # -I drops the script's own directory
from common import HERE, SRC, fleurs_paragraphs, fleurs_sentences # noqa: E402
from transformers import PreTrainedModel # noqa: E402
# Seq2LabelsModel resizes the embedding for its $START token while transformers initializes on the meta device,
# where mean resizing fails. The resized rows are overwritten by the checkpoint, so plain resizing is equivalent.
_resize = PreTrainedModel.resize_token_embeddings
PreTrainedModel.resize_token_embeddings = lambda self, n, *a, **k: _resize(self, n, *a, **{**k, "mean_resizing": False})
# The model repo's own inference code (read before use: plain PyTorch and file reads).
sys.path.insert(0, str(SRC))
import modeling_seq2labels # noqa: E402
# Written for transformers 4.18; current versions require ModelOutput subclasses to be dataclasses.
modeling_seq2labels.Seq2LabelsOutput = dataclass(modeling_seq2labels.Seq2LabelsOutput)
from gec_model import GecBERTModel # noqa: E402
EXAMPLES = [
"theo đó thủ tướng dự kiến tiếp bộ trưởng nông nghiệp mỹ tom wilsack bộ trưởng thương mại mỹ gina raimondo "
"bộ trưởng tài chính janet yellen gặp gỡ thượng nghị sĩ patrick leahy và một số nghị sĩ mỹ khác",
"những gói cước năm g mobifone sẽ mang đến cho bạn những trải nghiệm mới lạ trên cả tuyệt vời so với mạng "
"bốn g thì tốc độ truy cập mạng 5 g mobifone được nhận định là siêu đỉnh với mức truy cập nhanh gấp 10 lần",
]
def main() -> None:
model = GecBERTModel(vocab_path=str(SRC / "vocabulary"), model_paths=str(SRC), split_chunk=True, device="cpu")
for text in EXAMPLES:
print(model(text)[0])
rows = [{"text": text, "truth": truth} for text, truth in fleurs_sentences() + fleurs_paragraphs()]
rows = [{"text": text, "truth": ""} for text in EXAMPLES] + rows
start = time.time()
for row in rows:
row["original"] = model(row["text"])[0]
print(f"{len(rows)} sentences in {time.time() - start:.0f} s")
out = HERE / "out"
out.mkdir(exist_ok=True)
(out / "reference.json").write_text(json.dumps(rows, ensure_ascii=False, indent=1) + "\n", encoding="utf-8")
if __name__ == "__main__":
main()
|