| """Score a line of conversation. Needs torch, transformers, safetensors. |
| |
| python predict.py "I'll send the deck by Friday." |
| python predict.py --context "A: can someone own the budget deck?" --speaker B "sure, I'll do it." |
| echo '{"text": "...", "context": [{"speaker": "A", "text": "..."}], "speaker": "B"}' | python predict.py --jsonl |
| |
| --context is repeatable, oldest turn first; only the last 3 are used. --model is the model folder |
| (default: the folder this file is in). |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| from pathlib import Path |
|
|
| from cue import CueModel, parse_turn |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--model", default=str(Path(__file__).resolve().parent)) |
| ap.add_argument("--context", action="append", default=[], help='"SPEAKER: text", oldest first; repeatable') |
| ap.add_argument("--speaker", default=None, help="speaker of the target line (needed to tell same/other speaker)") |
| ap.add_argument("--jsonl", action="store_true", help="read {text, context?, speaker?} rows from stdin") |
| ap.add_argument("--device", default="cpu") |
| ap.add_argument("text", nargs="*") |
| a = ap.parse_args() |
|
|
| m = CueModel.from_pretrained(a.model, device=a.device) |
| if a.jsonl: |
| rows = [json.loads(line) for line in sys.stdin if line.strip()] |
| else: |
| ctx = [parse_turn(c) for c in a.context] |
| rows = [{"text": t, "context": ctx, "speaker": a.speaker} for t in a.text] |
| if not rows: |
| ap.error("give a text argument or --jsonl") |
| heads = m.product_heads |
| for r, s in zip(rows, m.score_many(rows)): |
| fired = m.decide(s) |
| print(json.dumps({"text": r["text"], **{h: round(s[h], 4) for h in heads}, |
| "fired": [h for h in heads if fired.get(h)]})) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
