#!/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()