niqun123123 commited on
Commit
d669efd
·
verified ·
1 Parent(s): 410ca9f

Upload model.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. model.py +43 -0
model.py ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Standalone local decision inference (not an AIStudio Maya canvas component).
2
+
3
+ Inside the exported checkpoint this file can be run directly with --checkpoint .
4
+ The actual HF architecture lives in modeling_autojev.py/configuration_autojev.py.
5
+ """
6
+ import argparse
7
+ import json
8
+ from pathlib import Path
9
+ import torch
10
+ from transformers import AutoModel, AutoProcessor
11
+
12
+ if Path(__file__).with_name('decision_utils.py').is_file():
13
+ from decision_utils import predict_batch, answer_codes
14
+ else:
15
+ from modeling.decision_utils import predict_batch, answer_codes
16
+
17
+
18
+ def load_model(checkpoint, device='cuda:0'):
19
+ torch.set_num_threads(8)
20
+ if torch.cuda.is_available():
21
+ torch.backends.cuda.enable_cudnn_sdp(False)
22
+ model = AutoModel.from_pretrained(str(checkpoint), trust_remote_code=True,
23
+ local_files_only=True, dtype=torch.bfloat16, device_map=device,
24
+ attn_implementation='sdpa').eval()
25
+ processor = AutoProcessor.from_pretrained(str(checkpoint), local_files_only=True, trust_remote_code=True)
26
+ processor.tokenizer.padding_side = 'left'
27
+ codes, ids = answer_codes(processor.tokenizer)
28
+ if codes != model.config.codes or ids != model.config.token_ids:
29
+ raise ValueError('Checkpoint code/tokenizer mapping mismatch')
30
+ return model, processor
31
+
32
+
33
+ if __name__ == '__main__':
34
+ parser = argparse.ArgumentParser()
35
+ parser.add_argument('--checkpoint', default=str(Path(__file__).resolve().parent))
36
+ parser.add_argument('--device', default='cuda:0')
37
+ parser.add_argument('--input', help='JSON file containing a list of decision inputs')
38
+ args = parser.parse_args()
39
+ model, processor = load_model(args.checkpoint, args.device)
40
+ rows = json.loads(Path(args.input).read_text()) if args.input else [{
41
+ 'state': '我想申请退款。', 'question': {'type': 'choice', 'instructions': '判断主要诉求',
42
+ 'criteria': {'refund': '退款相关', 'delivery': '物流相关', 'account': '账户相关'}}}]
43
+ print(json.dumps(predict_batch(model, processor, rows), ensure_ascii=False, indent=2))