File size: 1,105 Bytes
299bfda | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 | #!/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()
|