from __future__ import annotations

import argparse
from pathlib import Path

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "model_path",
        nargs="?",
        default=Path(__file__).resolve().parent,
        help="Local path or Hugging Face model id.",
    )
    args = parser.parse_args()

    tokenizer = AutoTokenizer.from_pretrained(
        args.model_path,
        trust_remote_code=True,
        use_fast=True,
    )
    model = AutoModelForSequenceClassification.from_pretrained(
        args.model_path,
        trust_remote_code=True,
    ).eval()

    texts = [
        "I hope you have a wonderful day.",
        "You are disgusting and should disappear.",
    ]
    inputs = tokenizer(
        texts,
        padding=True,
        truncation=True,
        max_length=512,
        return_tensors="pt",
    )

    cls_token_id = getattr(model.config, "student_cls_token_id", tokenizer.cls_token_id)
    if cls_token_id is not None and int(inputs.input_ids[0, 0]) != int(cls_token_id):
        raise RuntimeError("Tokenizer did not prepend the expected CLS token.")

    with torch.inference_mode():
        probs = torch.softmax(model(**inputs).logits, dim=-1)

    toxic_id = int(model.config.label2id.get("toxic", 1))
    for text, score in zip(texts, probs[:, toxic_id].tolist()):
        print(f"{score:.6f}\t{text}")


if __name__ == "__main__":
    main()