File size: 2,057 Bytes
60fd3f2
 
 
 
 
967ac28
 
60fd3f2
 
 
967ac28
 
60fd3f2
 
967ac28
 
60fd3f2
 
39921dc
 
 
 
 
 
 
 
 
60fd3f2
 
 
 
 
967ac28
60fd3f2
 
 
 
39921dc
967ac28
60fd3f2
 
967ac28
 
 
 
 
 
60fd3f2
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
# -*- 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"
            )
            
            # 强制提取 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)