predict.py

#!/usr/bin/env python3
# Author: Ilya Zelenskiy (proxy3d)
# Telegram channel: https://t.me/greenruff
# Related project: https://github.com/iproxy3d/communication-styles-llm
from __future__ import annotations

import argparse
import json

from goemotions_en_ru import EmotionClassifier


def main() -> None:
    parser = argparse.ArgumentParser(description="EN/RU 28-label emotion classification")
    parser.add_argument("text", help="English or Russian text")
    parser.add_argument("--model", required=True, help="Local model directory or Hugging Face model ID")
    parser.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda", "mps"])
    parser.add_argument("--top-k", type=int, default=5)
    parser.add_argument("--thresholds", default=None, help="Optional thresholds.json")
    args = parser.parse_args()

    classifier = EmotionClassifier(
        args.model,
        device=args.device,
        thresholds_path=args.thresholds,
    )
    print(json.dumps(classifier.predict(args.text, top_k=args.top_k), ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()