import logging import copy import pdb import math import os import json import yaml import time import re from typing import List, Dict from factool.utils.base.pipeline import pipeline from factool.code.helper.postprocess import PostProcessor from factool.code.helper.execution import evaluate_test_cases_multi_solution from factool.utils.utils_json import CustomJSONEncoder class code_pipeline(pipeline): def __init__(self, foundation_model, multi_solution_cnt, testcases_input_cnt): super().__init__('code', foundation_model) self.multi_solution_cnt = multi_solution_cnt self.testcases_input_cnt = testcases_input_cnt with open(os.path.join(self.prompts_path, "query_generation.yaml"), 'r') as file: data = yaml.load(file, Loader=yaml.FullLoader) self.query_generation_prompt = data['code'] async def _testcases_input_generation(self, batch, testcases_input_cnt): messages_list = [] if self.company == 'openai': messages_list = [ [ {"role": "system", "content": self.query_generation_prompt['system']}, {"role": "user", "content": self.query_generation_prompt[ 'user_testcases_' + str(testcases_input_cnt) ].format(input_question=sample['prompt'], entry_point=sample['entry_point']) }, ] for sample in batch ] elif self.company == 'anthropic': messages_list = [self.query_generation_prompt[ 'user_testcases_' + str(testcases_input_cnt) ].format(input_question=sample['prompt'], entry_point=sample['entry_point']) for sample in batch] return await self.chat.async_run(messages_list, Dict) async def _multi_solution_generation(self, batch, multi_solution_cnt): bsize = 15 messages_list = [ [ {"role": "system", "content": self.query_generation_prompt['system']}, {"role": "user", "content": self.query_generation_prompt[ 'user_solutions'].format(input_question=sample['prompt'], entry_point=sample['entry_point'])}, ] for sample in batch ] final_messages_list = [copy.deepcopy(messages) for messages in messages_list for _ in range(multi_solution_cnt) ] responses = [] for i in range(0, len(final_messages_list), bsize): batch = final_messages_list[i:i + bsize] responses += await self.chat.async_run(batch, Dict) # Split the list into lists of length of multi_solution_cnt responses_split = [responses[i:i + multi_solution_cnt] for i in range(0, len(responses), multi_solution_cnt)] # Transform each element in each list multi_solutions = [] for solutions in responses_split: key_names = [f"python_solution_{i}" for i in range(1, multi_solution_cnt + 1)] new_element = {key: solutions[i]['python_solution'] if solutions[i] != None else "None" for i, key in enumerate(key_names)} multi_solutions.append(new_element) return multi_solutions async def run_with_tool_live(self, batch, batch_size): testcases_input = await self._testcases_input_generation(batch, self.testcases_input_cnt) multi_solutions = await self._multi_solution_generation(batch, self.multi_solution_cnt) if testcases_input == None or multi_solutions == None: return None responses = [] for i in range(batch_size): response = {'testcases_input': [], 'multi_solutions': [], 'with_tool_classification': "None"} try: response['testcases_input'] = list(testcases_input[i].values()) # Append the solution to be verified to the LAST element # of multi_solutions response['multi_solutions']\ = [multi_solutions[i][f'python_solution_{j}'] for j in range(1, self.multi_solution_cnt + 1)] +\ [batch[i]['completion']] except: response['testcases_input'] = ["None"] * self.testcases_input_cnt response['multi_solutions'] = ["None"] * (self.multi_solution_cnt + 1) exec_result = evaluate_test_cases_multi_solution( batch[i]['prompt'], response['testcases_input'], response['multi_solutions'], timeout=0.1) response['exec_result'] = exec_result response['with_tool_classification'] = True # must pass all testcases to be classified as "True" for testcase_result in exec_result: # syntax or timeout error happening on the potential solution if isinstance(testcase_result[-1], str) \ and testcase_result[-1].startswith('FAILURE'): response['with_tool_classification'] = False # majority voting. Note that the last element # is the solution to be verified. Also, multi solutions that return "FAILURE" are not counted and removed. else: failure_indices = [ i for i, res in enumerate(testcase_result[:-1]) if isinstance(res, str) and res.startswith('FAILURE')] testcase_result = [ res for i, res in enumerate(testcase_result) if i not in failure_indices] try: if testcase_result[:-1].count(testcase_result[-1]) \ < math.ceil(len(testcase_result) / 2): response['with_tool_classification'] = False # sometimes numpy array is included in testcase_result, so this error will be raised except: response['with_tool_classification'] = False responses.append(response) return responses async def run_with_tool_api_call(self, prompts, responses, entry_points): # response preprocessing to extract the code snippet: claims = [] for i, response in enumerate(responses): if "```python" in response: match = re.search(r"```python\n(.*?)\n```", response, re.DOTALL) if match: claims.append(match.group(1)) else: claims.append("") elif "```" in response: match = re.search(r"```\n(.*?)\n```", response, re.DOTALL) if match: claims.append(match.group(1)) else: claims.append("") else: claims.append(response) batch_size = 5 num_batches = math.ceil(len(prompts) / batch_size) self.sample_list = [ {"prompt": prompt, "response": response, "entry_point": entry_point, "completion": claim, "category": 'code'} for prompt, response, entry_point, claim in zip(prompts, responses, entry_points, claims)] for i in range(num_batches): print(i) batch_start = i * batch_size batch_end = min((i + 1) * batch_size, len(responses)) responses_returned = await self.run_with_tool_live(self.sample_list[batch_start: batch_end], batch_end - batch_start) for j, response_returned in enumerate(responses_returned): index = batch_start + j self.sample_list[index].update({ 'claim': self.sample_list[index]['completion'], 'testcases_queries': response_returned['testcases_input'], 'potential_solutions_queries': response_returned['multi_solutions'], 'exec_results': response_returned['exec_result'], 'claim_level_factuality': response_returned['with_tool_classification'], 'response_level_factuality': response_returned['with_tool_classification'] }) del self.sample_list[index]["completion"] return self.sample_list async def run_with_tool_dataset(self, annotated_dataset_path: str, with_tool_classified_dataset_path: str, rerun: bool = False, rerun_indices: list = []): data_path = with_tool_classified_dataset_path if rerun else annotated_dataset_path with open(data_path, 'r') as f: data = [json.loads(line) for line in f] self.sample_list = data rerun_elements = self.sample_list if not rerun else [self.sample_list[i] for i in rerun_indices] batch_size = 5 num_batches = math.ceil(len(rerun_elements) / batch_size) # 5 for i in range(num_batches): print(i) batch_start = i * batch_size batch_end = min((i + 1) * batch_size, len(rerun_elements)) responses = await self.run_with_tool_live(rerun_elements[batch_start:batch_end], batch_end - batch_start) for j, response in enumerate(responses): index = batch_start + j if not rerun else rerun_indices[batch_start + j] self.sample_list[index]['with_tool_classification'] = response['with_tool_classification'] if response is not None else 'None' if response is not None: self.sample_list[index].update({ 'testcases_input': response['testcases_input'], 'multi_solutions': response['multi_solutions'], 'exec_result': response['exec_result'] }) # save everything after each batch to prevent data loss with open(with_tool_classified_dataset_path, 'w') as f: for item in self.sample_list: try: json_str = json.dumps(item, cls=CustomJSONEncoder) except: continue f.write(json_str + '\n') async def run_self_check_live(self, fewshot, batch): user_prompt_key = 'user_3_shot_CoT' if fewshot else 'user_zero_shot_CoT' messages_list = [ [ {"role": "system", "content": self.self_check_prompt['system']}, {"role": "user", "content": self.self_check_prompt[user_prompt_key].format(input_question=response['prompt'], input_solution=response['completion'])}, ] for response in batch ] return await self.chat.async_run(messages_list, Dict) async def run_self_check_dataset(self, annotated_dataset_path: str, self_check_classified_dataset_path: str, fewshot: bool = False, rerun: bool = False, rerun_indices: list = []): if rerun == False: with open(annotated_dataset_path, 'r') as f: self.sample_list = [json.loads(line) for line in f] rerun_elements = self.sample_list else: with open(self_check_classified_dataset_path, 'r') as f: self.sample_list = [json.loads(line) for line in f] rerun_elements = [self.sample_list[i] for i in rerun_indices] batch_size = 5 num_batches = math.ceil(len(rerun_elements) / batch_size) for i in range(num_batches): print(i) batch_start = i * batch_size batch_end = (i + 1) * batch_size batch = rerun_elements[batch_start:batch_end] responses = await self.run_self_check_live(fewshot, batch) for j, response in enumerate(responses): index = batch_start + j if rerun == False else rerun_indices[batch_start + j] self.sample_list[index]['self_check_classification'] = response.get('factuality', 'None') if response is not None else 'None' self.sample_list[index]['self_check_reasoning'] = response.get('reasoning', 'None') if response is not None else 'None' # save everything after each batch to prevent data loss with open(self_check_classified_dataset_path, 'w') as f: for item in self.sample_list: json_str = json.dumps(item) f.write(json_str + '\n')