| 389 | |
| 390 | |
| 391 | class pipeline_runner: |
| 392 | def __init__(self, args, add_retrieval=False, process_id=0, server=False): |
| 393 | self.args = args |
| 394 | self.add_retrieval = add_retrieval |
| 395 | self.process_id = process_id |
| 396 | self.server = server |
| 397 | if not self.server: self.task_list = self.generate_task_list() |
| 398 | else: self.task_list = [] |
| 399 | |
| 400 | def get_backbone_model(self): |
| 401 | args = self.args |
| 402 | if args.backbone_model == "toolllama": |
| 403 | # ratio = 4 means the sequence length is expanded by 4, remember to change the model_max_length to 8192 (2048 * ratio) for ratio = 4 |
| 404 | ratio = int(args.max_sequence_length/args.max_source_sequence_length) |
| 405 | replace_llama_with_condense(ratio=ratio) |
| 406 | if args.lora: |
| 407 | backbone_model = ToolLLaMALoRA(base_name_or_path=args.model_path, model_name_or_path=args.lora_path, max_sequence_length=args.max_sequence_length) |
| 408 | else: |
| 409 | backbone_model = ToolLLaMA(model_name_or_path=args.model_path, max_sequence_length=args.max_sequence_length) |
| 410 | else: |
| 411 | backbone_model = args.backbone_model |
| 412 | return backbone_model |
| 413 | |
| 414 | def get_retriever(self): |
| 415 | return ToolRetriever(corpus_tsv_path=self.args.corpus_tsv_path, model_path=self.args.retrieval_model_path) |
| 416 | |
| 417 | def get_args(self): |
| 418 | return self.args |
| 419 | |
| 420 | def generate_task_list(self): |
| 421 | args = self.args |
| 422 | query_dir = args.input_query_file |
| 423 | answer_dir = args.output_answer_file |
| 424 | if not os.path.exists(answer_dir): |
| 425 | os.mkdir(answer_dir) |
| 426 | method = args.method |
| 427 | backbone_model = self.get_backbone_model() |
| 428 | white_list = get_white_list(args.tool_root_dir) |
| 429 | task_list = [] |
| 430 | querys = json.load(open(query_dir, "r")) |
| 431 | for query_id, data_dict in enumerate(querys): |
| 432 | if "query_id" in data_dict: |
| 433 | query_id = data_dict["query_id"] |
| 434 | if "api_list" in data_dict: |
| 435 | origin_tool_names = [standardize(cont["tool_name"]) for cont in data_dict["api_list"]] |
| 436 | tool_des = contain(origin_tool_names,white_list) |
| 437 | if tool_des == False: |
| 438 | continue |
| 439 | tool_des = [[cont["standard_tool_name"], cont["description"]] for cont in tool_des] |
| 440 | else: |
| 441 | tool_des = None |
| 442 | task_list.append((method, backbone_model, query_id, data_dict, args, answer_dir, tool_des)) |
| 443 | return task_list |
| 444 | |
| 445 | def method_converter(self, backbone_model, openai_key, method, env, process_id, single_chain_max_step=12, max_query_count=60, callbacks=None): |
| 446 | if callbacks is None: callbacks = [] |
| 447 | if backbone_model == "chatgpt_function": |
| 448 | model = "gpt-3.5-turbo-16k-0613" |
no outgoing calls
no test coverage detected