(args)
| 55 | |
| 56 | |
| 57 | def init_model(args): |
| 58 | tokenizer = AutoTokenizer.from_pretrained(args.load_from) |
| 59 | if 'model' in args.load_from: |
| 60 | model = MiniMindForCausalLM(MiniMindConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe))) |
| 61 | moe_suffix = '_moe' if args.use_moe else '' |
| 62 | ckp = f'./{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth' |
| 63 | model.load_state_dict(torch.load(ckp, map_location=args.device), strict=True) |
| 64 | else: |
| 65 | model = AutoModelForCausalLM.from_pretrained(args.load_from, trust_remote_code=True) |
| 66 | get_model_params(model, model.config) |
| 67 | return model.half().eval().to(args.device), tokenizer |
| 68 | |
| 69 | |
| 70 | def parse_tool_calls(text): |
no test coverage detected