extract Main Entity from the query
| 15 | |
| 16 | |
| 17 | class NER(object): |
| 18 | """ |
| 19 | extract Main Entity from the query |
| 20 | """ |
| 21 | |
| 22 | def __init__(self, engine=CONFIG["methods"]["engine"], model_name=CONFIG["methods"]["model_name"], |
| 23 | temperature=CONFIG["methods"]["temperature"]): |
| 24 | self.engine = engine |
| 25 | self.max_tokens = CONFIG["max_tokens"] |
| 26 | self.model_name = model_name |
| 27 | self.temperature = temperature |
| 28 | self.today = datetime.date.today().strftime("%Y%m%d") |
| 29 | self.yesterday = (datetime.date.today() + datetime.timedelta(days=-1)).strftime("%Y%m%d") |
| 30 | |
| 31 | def parse_custom_string(self, input_str): |
| 32 | input_str = input_str.replace(' ', '').replace('\n', '').replace('"', '') |
| 33 | key_value_pattern = r"(\w+):\s*({.*?}|[^,{}]+)" |
| 34 | matches = re.findall(key_value_pattern, input_str) |
| 35 | |
| 36 | result = {} |
| 37 | for key, value in matches: |
| 38 | if value.startswith("{"): |
| 39 | result[key] = self.parse_custom_string(value) |
| 40 | else: |
| 41 | result[key] = value |
| 42 | return result |
| 43 | |
| 44 | def parse_json(self, input_str): |
| 45 | res = {} |
| 46 | try: |
| 47 | pattern = r"\{.*?\}" |
| 48 | match = re.search(pattern, input_str) |
| 49 | if match: |
| 50 | res = eval(match.group()) |
| 51 | except: |
| 52 | try: |
| 53 | res = self.parse_custom_string(input_str) |
| 54 | except Exception as e: |
| 55 | print(f"=========eval fail==========={e}", exc_info=True) |
| 56 | return res |
| 57 | |
| 58 | def run(self, query, dataset): |
| 59 | """ |
| 60 | get NER of users' query, and classify the label of entity |
| 61 | :param query: user input |
| 62 | :return: llm output, consist of entity and its entity label within [] |
| 63 | """ |
| 64 | |
| 65 | generate_sql_prompt = get_ner_prompt(dataset) |
| 66 | |
| 67 | prompt_dict = { |
| 68 | "query": query |
| 69 | } |
| 70 | |
| 71 | prompt = get_prompt_content(generate_sql_prompt, prompt_dict) |
| 72 | res = ask_llm(prompt) |
| 73 | res = self.parse_json(res) |
| 74 | if "entities" in res: |