"""Convert jinaai/jina-reranker-v1-tiny-en's published ONNX export into a form usable by
Teradata's ONNXEmbeddings BYOM function. See README.md for why the concat workaround this
repo relies on is valid for this specific model, and the accuracy it was verified to.
"""

from pathlib import Path

import onnx
from huggingface_hub import hf_hub_download

UPSTREAM_REPO = "jinaai/jina-reranker-v1-tiny-en"
OUTPUT_TENSOR_NAME = "sentence_embedding"


def rename_graph_output(onnx_path: Path, new_name: str) -> None:
    """ONNXEmbeddings' model_output_tensor argument only accepts the literal values
    'sentence_embedding' or 'token_embeddings' -- confirmed against a live Vantage instance,
    which otherwise rejects the model load outright. Rename this model's native 'logits'
    output (and the node producing it) to satisfy that.
    """
    model = onnx.load(str(onnx_path))
    old_names = {out.name for out in model.graph.output}
    for output in model.graph.output:
        output.name = new_name
    for node in model.graph.node:
        node.output[:] = [new_name if o in old_names else o for o in node.output]
    onnx.save(model, str(onnx_path))


def main() -> None:
    repo_dir = Path(__file__).parent

    # This model already ships a ready-made ONNX export upstream (fp32, opset 11) --
    # no re-export via optimum/torch.onnx is needed here, only a rename of its output tensor.
    onnx_path = Path(hf_hub_download(UPSTREAM_REPO, "onnx/model.onnx", local_dir=repo_dir))
    hf_hub_download(UPSTREAM_REPO, "tokenizer.json", local_dir=repo_dir)
    hf_hub_download(UPSTREAM_REPO, "special_tokens_map.json", local_dir=repo_dir)
    hf_hub_download(UPSTREAM_REPO, "config.json", local_dir=repo_dir)

    rename_graph_output(onnx_path, OUTPUT_TENSOR_NAME)
    print(f"wrote {onnx_path} with output tensor renamed to '{OUTPUT_TENSOR_NAME}'")


if __name__ == "__main__":
    main()