File size: 2,405 Bytes
d669efd f2d0999 d669efd f2d0999 d669efd f2d0999 d669efd f2d0999 d669efd | 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 49 50 | """Standalone local decision inference (not an AIStudio Maya canvas component).
Inside the exported checkpoint this file can be run directly with --checkpoint .
The actual HF architecture lives in modeling_lemondecider.py/configuration_lemondecider.py.
"""
import argparse
import json
from pathlib import Path
import torch
from transformers import AutoModel
if Path(__file__).with_name('decision_utils.py').is_file():
from decision_utils import predict_batch, answer_codes
else:
from modeling.decision_utils import predict_batch, answer_codes
def load_model(checkpoint, device='cuda:0'):
torch.set_num_threads(8)
if torch.cuda.is_available():
torch.backends.cuda.enable_cudnn_sdp(False)
model = AutoModel.from_pretrained(str(checkpoint), trust_remote_code=True,
local_files_only=True, dtype=torch.bfloat16, device_map=device,
attn_implementation='sdpa').eval()
if Path(__file__).with_name('decision_utils.py').is_file():
from decision_utils import load_processor
else:
from modeling.decision_utils import load_processor
processor = load_processor(checkpoint)
codes, ids = answer_codes(processor.tokenizer, model.config.max_options,
getattr(model.config, 'code_scheme', 'letters-v1'))
if codes != model.config.codes or ids != model.config.token_ids:
raise ValueError('Checkpoint code/tokenizer mapping mismatch')
if tuple(model.readout.weight.shape) != (len(codes), model.config.text_config.hidden_size):
raise ValueError('Checkpoint readout/code-table shape mismatch')
return model, processor
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--checkpoint', default=str(Path(__file__).resolve().parent))
parser.add_argument('--device', default='cuda:0')
parser.add_argument('--input', help='JSON file containing a list of decision inputs')
args = parser.parse_args()
model, processor = load_model(args.checkpoint, args.device)
rows = json.loads(Path(args.input).read_text()) if args.input else [{
'state': '我想申请退款。', 'question': {'type': 'choice', 'instructions': '判断主要诉求',
'criteria': {'refund': '退款相关', 'delivery': '物流相关', 'account': '账户相关'}}}]
print(json.dumps(predict_batch(model, processor, rows), ensure_ascii=False, indent=2))
|