| """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() |
|
|
