Spaces:
Running
Running
| # -*- coding: utf-8 -*- | |
| import json | |
| import concurrent.futures | |
| from typing import List, Dict | |
| import re | |
| import logging | |
| from config import APP_CONFIG | |
| class MetaExtractor: | |
| def __init__(self, agent): | |
| self.agent = agent | |
| self.max_workers = APP_CONFIG["lit_max_workers"] | |
| def extract_single_paper_meta(self, text: str) -> Dict: | |
| snippet = self.agent._truncate_context(text, max_tokens=3500) | |
| prompt = APP_CONFIG["prompt_templates"]["meta_extract"] | |
| try: | |
| # 修改点：显式将模型降级为 deepseek-chat (V3)，以获得更快的并行提取速度 | |
| response = self.agent.call_deepseek( | |
| prompt, | |
| snippet, | |
| max_tokens=1000, | |
| model_name="deepseek-chat", | |
| logical_node='literature_parsing', | |
| ) | |
| # 强制提取 JSON 部分，过滤模型可能的解释性前缀或后缀 | |
| json_match = re.search(r'\{.*\}', response, re.DOTALL) | |
| if json_match: | |
| return json.loads(json_match.group()) | |
| return {} | |
| except Exception as e: | |
| logging.error(f"Meta Extraction error: {e}") | |
| return {} | |
| def run_batch_extraction(self, paper_texts: List[str]) -> List[Dict]: | |
| golden_db = [] | |
| # 使用配置中的 lit_max_workers 进行多线程并行提取 | |
| with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers) as executor: | |
| future_to_paper = {executor.submit(self.extract_single_paper_meta, text): text for text in paper_texts} | |
| for future in concurrent.futures.as_completed(future_to_paper): | |
| try: | |
| data = future.result() | |
| if data: | |
| golden_db.append(data) | |
| except Exception as e: | |
| logging.error(f"Batch Future error: {e}") | |
| return golden_db | |
| def extract_golden_metadata(agent, paper_texts: List[str]): | |
| extractor = MetaExtractor(agent) | |
| return extractor.run_batch_extraction(paper_texts) | |