File size: 1,658 Bytes
5e054a6 1fe108b 5e054a6 1fe108b 5e054a6 1fe108b 5e054a6 1fe108b 5e054a6 | 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 | """Download Reflex-1 and exercise single-request and batch inference."""
import argparse
import json
from dataclasses import asdict
import torch
from transformers import AutoModel
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--model", default="gai-labs/reflex-1")
parser.add_argument("--revision", help="Explicit Hub checkpoint branch, tag or commit")
parser.add_argument("--local-files-only", action="store_true")
args = parser.parse_args()
torch.set_num_threads(1)
model = AutoModel.from_pretrained(args.model, trust_remote_code=True,
revision=args.revision, local_files_only=args.local_files_only)
requests = [
{"state": "The lamp is red.", "question": "Which color is the lamp?",
"options": ["red", "blue", "green"]},
{"state": "The customer was charged twice for one card payment.",
"question": "Choose the matching issue.",
"options": ["duplicate charge", "lost card", "unknown fee", "cash withdrawal"]},
]
single = model.predict(requests[0])[0]
batch = model.predict_batch(requests)
assert single.choice == "red"
assert [decisions[0].choice for decisions in batch] == ["red", "duplicate charge"]
torch.testing.assert_close(torch.tensor(single.probabilities),
torch.tensor(batch[0][0].probabilities), atol=2e-5, rtol=1e-4)
print(json.dumps({"passed": True, "revision": getattr(model.config, "_commit_hash", None),
"decisions": [asdict(row[0]) for row in batch]}, indent=2))
if __name__ == "__main__":
main()
|