get NER of users' query, and classify the label of entity :param query: user input :return: llm output, consist of entity and its entity label within []
(self, query, dataset)
| 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: |
| 75 | res["limitation"] = res["entities"] |
| 76 | for i in ["limitation", "metric", "query"]: |
| 77 | if i not in res: |
| 78 | res[i] = [] |
| 79 | return res |