scripts/tag.py

"""CLI: tag a financial sentence and print entity spans with confidence.

Usage:
    python scripts/tag.py "Apple reported $1.2B revenue in Q3 2024, up 12.5%."
    python scripts/tag.py --checkpoint checkpoints/best "SEC fined Goldman Sachs $500M."
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from rich.console import Console
from rich.table import Table

from finner.infer.predict import predict

console = Console()


def main():
    parser = argparse.ArgumentParser(description="Tag a financial sentence")
    parser.add_argument("text", help="Text to tag")
    parser.add_argument("--checkpoint", default=None, help="Checkpoint directory")
    args = parser.parse_args()

    result = predict(args.text, checkpoint=args.checkpoint)

    console.print(f"\n[bold]Input:[/bold] {args.text}\n")

    if result["entities"]:
        table = Table(title="Detected Entities", show_header=True)
        table.add_column("Entity", style="bold")
        table.add_column("Type", style="cyan")
        table.add_column("Confidence", justify="right")
        table.add_column("Tokens", justify="right")
        for ent in result["entities"]:
            table.add_row(
                ent["text"],
                ent["label"],
                f"{ent['confidence']:.3f}",
                f"{ent['start_token']}–{ent['end_token']}",
            )
        console.print(table)
    else:
        console.print("[yellow]No entities detected.[/yellow]")

    console.print("\n[dim]Token-level:[/dim]")
    for tok, lbl, conf in zip(result["tokens"], result["token_labels"], result["token_confidences"]):
        color = "green" if lbl != "O" else "dim"
        console.print(f"  [{color}]{tok:20s}  {lbl:15s}  {conf:.3f}[/{color}]")


if __name__ == "__main__":
    main()