classify.py

"""
SentimentPulse CLI — classify a single headline.

Usage:
    python classify.py --text "Apple beats earnings by 8%"
    python classify.py --text "Goldman fell sharply on credit fears" --model zero-shot

Outputs JSON:
    {
      "text": "...",
      "sentiment": "positive",
      "sentimentScore": 0.79,
      "scores": {"positive": 0.87, "negative": 0.08, "neutral": 0.05},
      "modelVersion": "finbert-finetuned-v1.0"
    }
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import torch

ROOT = Path(__file__).parent
FINETUNED_DIR = ROOT / "models" / "finbert-finetuned"
ZEROSHOT_MODEL = "ProsusAI/finbert"


def classify(text: str, model_dir: str) -> dict:
    from transformers import AutoModelForSequenceClassification, AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(model_dir)
    model = AutoModelForSequenceClassification.from_pretrained(model_dir)
    model.eval()

    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        max_length=128,
        padding=True,
    )

    with torch.no_grad():
        logits = model(**inputs).logits

    probs = torch.softmax(logits, dim=-1).squeeze().tolist()

    # ProsusAI/finbert label order: positive=0, negative=1, neutral=2
    # Our fine-tuned model: negative=0, neutral=1, positive=2
    # Detect which ordering by checking model config.
    id2label = model.config.id2label
    label_map = {v.lower(): i for i, v in id2label.items()}

    pos_prob = probs[label_map["positive"]]
    neg_prob = probs[label_map["negative"]]
    neu_prob = probs[label_map["neutral"]]

    sentiment_score = round(pos_prob - neg_prob, 4)
    predicted_label = max(label_map, key=lambda k: probs[label_map[k]])

    model_version = (
        "finbert-finetuned-v1.0"
        if "finbert-finetuned" in str(model_dir)
        else "finbert-zero-shot"
    )

    return {
        "text": text,
        "sentiment": predicted_label,
        "sentimentScore": sentiment_score,
        "scores": {
            "positive": round(pos_prob, 4),
            "negative": round(neg_prob, 4),
            "neutral": round(neu_prob, 4),
        },
        "modelVersion": model_version,
    }


def main():
    parser = argparse.ArgumentParser(description="Classify a financial headline")
    parser.add_argument("--text", required=True, help="Headline text to classify")
    parser.add_argument(
        "--model",
        choices=["fine-tuned", "zero-shot"],
        default="fine-tuned",
        help="Model to use (default: fine-tuned; falls back to zero-shot if not available)",
    )
    args = parser.parse_args()

    if args.model == "fine-tuned" and (FINETUNED_DIR / "config.json").exists():
        model_dir = str(FINETUNED_DIR)
    else:
        if args.model == "fine-tuned":
            print(
                f"Warning: fine-tuned model not found at {FINETUNED_DIR}. "
                "Falling back to zero-shot FinBERT.",
                flush=True,
            )
        model_dir = ZEROSHOT_MODEL

    result = classify(args.text, model_dir)
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()