| from models import *
|
| from utils import *
|
| from .knowledge_base.case_repository import CaseRepositoryHandler
|
|
|
| class InformationExtractor:
|
| def __init__(self, llm: BaseEngine):
|
| self.llm = llm
|
|
|
| def extract_information(self, instruction="", text="", examples="", schema="", additional_info=""):
|
| examples = good_case_wrapper(examples)
|
| prompt = extract_instruction.format(instruction=instruction, examples=examples, text=text, additional_info=additional_info, schema=schema)
|
| response = self.llm.get_chat_response(prompt)
|
| response = extract_json_dict(response)
|
| return response
|
|
|
| def extract_information_compatible(self, task="", text="", constraint=""):
|
| instruction = instruction_mapper.get(task)
|
| prompt = extract_instruction_json.format(instruction=instruction, constraint=constraint, input=text)
|
| response = self.llm.get_chat_response(prompt)
|
| response = extract_json_dict(response)
|
| return response
|
|
|
| def summarize_answer(self, instruction="", answer_list="", schema="", additional_info=""):
|
| prompt = summarize_instruction.format(instruction=instruction, answer_list=answer_list, schema=schema, additional_info=additional_info)
|
| response = self.llm.get_chat_response(prompt)
|
| response = extract_json_dict(response)
|
| return response
|
|
|
| class ExtractionAgent:
|
| def __init__(self, llm: BaseEngine, case_repo: CaseRepositoryHandler):
|
| self.llm = llm
|
| self.module = InformationExtractor(llm = llm)
|
| self.case_repo = case_repo
|
| self.methods = ["extract_information_direct", "extract_information_with_case"]
|
|
|
| def __get_constraint(self, data: DataPoint):
|
| if data.constraint == "":
|
| return data
|
| if data.task == "NER":
|
| constraint = json.dumps(data.constraint)
|
| if "**Entity Type Constraint**" in constraint or self.llm.name == "OneKE":
|
| return data
|
| data.constraint = f"\n**Entity Type Constraint**: The type of entities must be chosen from the following list.\n{constraint}\n"
|
| elif data.task == "RE":
|
| constraint = json.dumps(data.constraint)
|
| if "**Relation Type Constraint**" in constraint or self.llm.name == "OneKE":
|
| return data
|
| data.constraint = f"\n**Relation Type Constraint**: The type of relations must be chosen from the following list.\n{constraint}\n"
|
| elif data.task == "EE":
|
| constraint = json.dumps(data.constraint)
|
| if "**Event Extraction Constraint**" in constraint:
|
| return data
|
| if self.llm.name != "OneKE":
|
| data.constraint = f"\n**Event Extraction Constraint**: The event type must be selected from the following dictionary keys, and its event arguments should be chosen from its corresponding dictionary values. \n{constraint}\n"
|
| else:
|
| try:
|
| result = [
|
| {
|
| "event_type": key,
|
| "trigger": True,
|
| "arguments": value
|
| }
|
| for key, value in data.constraint.items()
|
| ]
|
| data.constraint = json.dumps(result)
|
| except:
|
| print("Invalid Constraint: Event Extraction constraint must be a dictionary with event types as keys and lists of arguments as values.", data.constraint)
|
| elif data.task == "Triple":
|
| constraint = json.dumps(data.constraint)
|
| if "**Triple Extraction Constraint**" in constraint:
|
| return data
|
| if self.llm.name != "OneKE":
|
| if len(data.constraint) == 1:
|
| data.constraint = f"\n**Triple Extraction Constraint**: Entities type must chosen from following list:\n{constraint}\n"
|
| elif len(data.constraint) == 2:
|
| if data.constraint[0] == []:
|
| data.constraint = f"\n**Triple Extraction Constraint**: Relation type must chosen from following list:\n{data.constraint[1]}\n"
|
| elif data.constraint[1] == []:
|
| data.constraint = f"\n**Triple Extraction Constraint**: Entities type must chosen from following list:\n{data.constraint[0]}\n"
|
| else:
|
| data.constraint = f"\n**Triple Extraction Constraint**: Entities type must chosen from following list:\n{data.constraint[0]}\nRelation type must chosen from following list:\n{data.constraint[1]}\n"
|
| elif len(data.constraint) == 3:
|
| if data.constraint[0] == []:
|
| data.constraint = f"\n**Triple Extraction Constraint**: Relation type must chosen from following list:\n{data.constraint[1]}\nObject Entities must chosen from following list:\n{data.constraint[2]}\n"
|
| elif data.constraint[1] == []:
|
| data.constraint = f"\n**Triple Extraction Constraint**: Subject Entities must chosen from following list:\n{data.constraint[0]}\nObject Entities must chosen from following list:\n{data.constraint[2]}\n"
|
| elif data.constraint[2] == []:
|
| data.constraint = f"\n**Triple Extraction Constraint**: Subject Entities must chosen from following list:\n{data.constraint[0]}\nRelation type must chosen from following list:\n{data.constraint[1]}\n"
|
| else:
|
| data.constraint = f"\n**Triple Extraction Constraint**: Subject Entities must chosen from following list:\n{data.constraint[0]}\nRelation type must chosen from following list:\n{data.constraint[1]}\nObject Entities must chosen from following list:\n{data.constraint[2]}\n"
|
| else:
|
| data.constraint = f"\n**Triple Extraction Constraint**: The type of entities must be chosen from the following list:\n{constraint}\n"
|
| else:
|
| print("OneKE does not support Triple Extraction task now, please wait for the next version.")
|
|
|
| return data
|
|
|
| def extract_information_direct(self, data: DataPoint):
|
| data = self.__get_constraint(data)
|
| result_list = []
|
| for chunk_text in data.chunk_text_list:
|
| if self.llm.name != "OneKE":
|
| extract_direct_result = self.module.extract_information(instruction=data.instruction, text=chunk_text, schema=data.output_schema, examples="", additional_info=data.constraint)
|
| else:
|
| extract_direct_result = self.module.extract_information_compatible(task=data.task, text=chunk_text, constraint=data.constraint)
|
| result_list.append(extract_direct_result)
|
| function_name = current_function_name()
|
| data.set_result_list(result_list)
|
| data.update_trajectory(function_name, result_list)
|
| return data
|
|
|
| def extract_information_with_case(self, data: DataPoint):
|
| data = self.__get_constraint(data)
|
| result_list = []
|
| for chunk_text in data.chunk_text_list:
|
| examples = self.case_repo.query_good_case(data)
|
| extract_case_result = self.module.extract_information(instruction=data.instruction, text=chunk_text, schema=data.output_schema, examples=examples, additional_info=data.constraint)
|
| result_list.append(extract_case_result)
|
| function_name = current_function_name()
|
| data.set_result_list(result_list)
|
| data.update_trajectory(function_name, result_list)
|
| return data
|
|
|
| def summarize_answer(self, data: DataPoint):
|
| if len(data.result_list) == 0:
|
| return data
|
| if len(data.result_list) == 1:
|
| data.set_pred(data.result_list[0])
|
| return data
|
| summarized_result = self.module.summarize_answer(instruction=data.instruction, answer_list=data.result_list, schema=data.output_schema, additional_info=data.constraint)
|
| funtion_name = current_function_name()
|
| data.set_pred(summarized_result)
|
| data.update_trajectory(funtion_name, summarized_result)
|
| return data |
