File size: 1,854 Bytes
8a5e2d4 | 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 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | """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()
|