| """Persistent JSONL inference for the v2 standalone candidate.""" | |
| import argparse | |
| import json | |
| from pathlib import Path | |
| import sys | |
| def main(): | |
| p = argparse.ArgumentParser(description=__doc__) | |
| p.add_argument('--model', default=str(Path(__file__).resolve().parent)) | |
| p.add_argument('--mode', choices=('default', 'flash'), default='default') | |
| p.add_argument('--input', type=Path) | |
| a = p.parse_args() | |
| import torch | |
| from transformers import AutoModel | |
| torch.set_num_threads(4); torch.cuda.set_device(0); torch.backends.cuda.enable_cudnn_sdp(False) | |
| model = AutoModel.from_pretrained(a.model, trust_remote_code=True, | |
| local_files_only=Path(a.model).is_dir(), dtype=torch.bfloat16, device_map={'': 'cuda:0'}).eval() | |
| _ = model.processor | |
| stream = a.input.open() if a.input else sys.stdin | |
| try: | |
| for line in stream: | |
| if line.strip(): print(json.dumps(model.decide(json.loads(line), a.mode), ensure_ascii=False, allow_nan=False), flush=True) | |
| finally: | |
| if a.input: stream.close() | |
| if __name__ == '__main__': main() | |