CoolFace
Modelpublic

proxy3d/multi-motions-28

sourceHugging Facemitupdated 4d agoView on Hugging Face
1likes38downloads
predict.py32 linesDownload Raw Back to root
1#!/usr/bin/env python32# Author: Ilya Zelenskiy (proxy3d)3# Telegram channel: https://t.me/greenruff4# Related project: https://github.com/iproxy3d/communication-styles-llm5from __future__ import annotations6 7import argparse8import json9 10from goemotions_en_ru import EmotionClassifier11 12 13def main() -> None:14    parser = argparse.ArgumentParser(description="EN/RU 28-label emotion classification")15    parser.add_argument("text", help="English or Russian text")16    parser.add_argument("--model", required=True, help="Local model directory or Hugging Face model ID")17    parser.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda", "mps"])18    parser.add_argument("--top-k", type=int, default=5)19    parser.add_argument("--thresholds", default=None, help="Optional thresholds.json")20    args = parser.parse_args()21 22    classifier = EmotionClassifier(23        args.model,24        device=args.device,25        thresholds_path=args.thresholds,26    )27    print(json.dumps(classifier.predict(args.text, top_k=args.top_k), ensure_ascii=False, indent=2))28 29 30if __name__ == "__main__":31    main()32