| 42 | self.response = completion.choices[0].message.content |
| 43 | |
| 44 | class TripleScorer: |
| 45 | def __init__(self, triple_path:str, triple_soure_path:str , output_path:str = None): |
| 46 | """ |
| 47 | 初始化打分类 |
| 48 | """ |
| 49 | self.triples = io_file.read(triple_path) |
| 50 | self.triple_sources = io_file.read(triple_soure_path) |
| 51 | self.triple_sources = {(item["page_idx"], item["paragraph_idx"]): item["text"] for item in self.triple_sources} |
| 52 | |
| 53 | logger.info(f"Loaded {len(self.triples)} triples from {triple_path}") |
| 54 | |
| 55 | self.prompt = prompt_kg_judge.score_triple_prompt |
| 56 | |
| 57 | llm_client_args = {"llm_model":"Qwen2.5-72B", "llm_url":"http://0.0.0.0:8001/v1", "llm_api_key":"EMPTY"} |
| 58 | self.client = llm_client(llm_client_args) |
| 59 | |
| 60 | self.output_path = output_path if output_path else triple_path.replace(".jsonl", ".scores.jsonl") |
| 61 | |
| 62 | def parse_triple(self, triple_str:str) -> dict[dict]: |
| 63 | """ |
| 64 | 解析三元组字符串为结构化字典 |
| 65 | :param triple_str: "Head|relation|Tail" 格式的字符串 |
| 66 | :return: {"head":..., "relation":..., "tail":...} |
| 67 | """ |
| 68 | |
| 69 | parts = triple_str.split("\t") |
| 70 | return { |
| 71 | "head": parts[0].strip().strip("<").strip(">"), |
| 72 | "relation": parts[1].strip().strip("<").strip(">"), |
| 73 | "tail": parts[2].strip().strip("<").strip(">") |
| 74 | } |
| 75 | |
| 76 | def score_triple(self, triple_data:dict) -> dict: |
| 77 | """ |
| 78 | :param jsonl_line: JSONL格式的输入数据 |
| 79 | :return: 包含所有打分结果的字典 |
| 80 | """ |
| 81 | # 解析输入数据 |
| 82 | triple_str = triple_data["triple"] |
| 83 | source_text = triple_data["source_text"] |
| 84 | triple = self.parse_triple(triple_str) |
| 85 | |
| 86 | prompt = self.prompt.format( |
| 87 | source_text=source_text, |
| 88 | head_entity=triple["head"], |
| 89 | relation=triple["relation"], |
| 90 | tail_entity=triple["tail"] |
| 91 | ) |
| 92 | |
| 93 | response = self._call_llm(prompt) |
| 94 | parsed_score = self.parse_result(response) |
| 95 | |
| 96 | return parsed_score |
| 97 | |
| 98 | def parse_result(self, response: str) -> dict: |
| 99 | try: |
| 100 | # 新增正则表达式提取核心字段 |
| 101 | score_match = re.search(r'"score"\s*:\s*(\d*\.?\d+)', response, re.DOTALL) |