"""用法：python predict.py '商品标题'"""

import argparse
import json
from pathlib import Path

import torch
from transformers import AutoTokenizer

from model import BertTitleClassifier


def main():
    parser = argparse.ArgumentParser(description="商品标题分类")
    parser.add_argument("text", help="待分类的商品标题")
    args = parser.parse_args()
    model_dir = Path(__file__).resolve().parent
    settings = json.loads((model_dir / "inference_config.json").read_text(encoding="utf-8"))
    labels = json.loads((model_dir / "labels.json").read_text(encoding="utf-8"))
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = BertTitleClassifier.from_pretrained(model_dir, device=device)
    tokenizer = AutoTokenizer.from_pretrained(model_dir, local_files_only=True)
    inputs = tokenizer(
        args.text, return_tensors="pt", padding="max_length",
        truncation=True, max_length=settings["max_length"],
    )
    with torch.no_grad():
        logits = model(inputs["input_ids"].to(device), inputs["attention_mask"].to(device))
    pred_id = logits.argmax(dim=1).item()
    print(json.dumps({"id": pred_id, "label": labels["id2label"][str(pred_id)]}, ensure_ascii=False))


if __name__ == "__main__":
    main()