sem-spider/spider.py

import requests
import os
import signal
import argparse
import json
import pymongo
from queue import Queue
from threading import Thread
from urllib.parse import urlparse

# S2_API_KEY = os.getenv('S2_API_KEY')
QUERY_FIELDS1 = 'paperId,corpusId,title,authors,year,url,tldr,venue,externalIds,fieldsOfStudy,s2FieldsOfStudy,abstract,citationCount,referenceCount,publicationTypes,influentialCitationCount,publicationDate,journal'
QUERY_FIELDS2 = 'paperId,corpusId,title,authors,year,url,venue,externalIds,fieldsOfStudy,s2FieldsOfStudy,abstract,citationCount,referenceCount,publicationTypes,influentialCitationCount,publicationDate,journal'
QUERY_FIELDS3 = 'paperId,corpusId,title,authors'

# 读取配置文件中的数据库参数
with open("config.json", "r") as f:
    config = json.load(f)
db_url = config["db_url"]
db_name = config["db_name"]
db_collection = config["db_collection"]

NUM_THREADS = config["num_threads"]
TASK_QUEUE_LEN = config["task_queue_len"]
S2_API_KEY = config["s2_api_key"]

# 定义退出标志
quit_flag = False

# 连接数据库，创建 papers 集合
client = pymongo.MongoClient(db_url)
db = client[db_name]
papers = db[db_collection]
papers_data = db['{}_data'.format(db_collection)]


def read_file(filename):
    data_list = []
    with open(filename, 'r') as f:
        for line in f:
            line_dict = json.loads(line)
            data_list.append(line_dict)
        # 在这里可以对每个字典对象进行操作，例如：
        # print(data_dict['key'])
    return data_list


def add_paper(file_path):
    papers.create_index("corpusid", unique=True)
    # 读取 paper 文件，存入数据库
    # data_list = read_file(file_path)
    # 批量插入数据
    inserted_ids = 0

    try:
        sub_list = []
        with open(file_path, 'r') as f:
            for line in f:
                line_dict = json.loads(line)
                sub_list.append(line_dict)

                if len(sub_list) == 2000:
                    result = papers.insert_many(sub_list, ordered=False)
                    inserted_ids += len(result.inserted_ids)
                    sub_list = []

                print('-------process', inserted_ids, '/', '7318795')

            if sub_list:
                result = papers.insert_many(sub_list, ordered=False)
                inserted_ids += len(result.inserted_ids)
                sub_list = []

    except pymongo.errors.BulkWriteError as e:
        inserted_ids = e.details['nInserted']
    finally:
        # 输出插入结果
        print("总插入数据: {0}, 已插入数据: {1}, 已存在数据: {2}" .format(
            7318795, inserted_ids, papers.count_documents({})))


def crawl_data():
    # papers_data.create_index("corpusId", unique=True)

    # 创建任务队列和线程
    q = Queue(TASK_QUEUE_LEN)
    # num_threads = 4
    threads = []
    for i in range(NUM_THREADS):
        t = Thread(target=worker, args=(q,))
        t.daemon = True
        t.start()
        print("starting worker: {}".format(t.native_id))
        threads.append(t)

    # 从数据库中读取 URL，加入任务队列
    while True:
        try:
            for data in papers.find({'$or': [{'consumed': {'$exists': False}}, {'consumed': False}]}):
                if quit_flag:
                    break
                if 'consumed' in data and data['consumed']:
                    print(data['corpusid'], "already inserted")
                    continue
                print('add {} to the task queue'.format(data['corpusid']))
                q.put((data['url'], data['corpusid']))
            break
        except Exception as e:
            print('crawl_data error', e)
            continue

    #
    print("Waitting for the task queue to complete...")
    q.join()
    print("The task queue has been completed!")

    # 停止线程
    for i in range(NUM_THREADS):
        q.put(None)
    for t in threads:
        print("stoping worker: {}" . format(t.native_id))
        t.join()


def mark_data_as_consumed(corpus_id):
    result = papers.update_one({'corpusid': corpus_id}, {
                               '$set': {'consumed': True}})


def worker(q):
    while True:
        item = q.get()
        if item is None:
            break
        url = urlparse(item[0]).path
        paper_id = url.split('/')[-1]
        corpus_id = item[1]
        print('crawling {} data: {}'.format(corpus_id, url))

        try:
            data = fetch_data(paper_id)
            if data is not None:
                # papers_data.insert_one(data)
                filter = {'corpusId': corpus_id}
                update = {'$set': data}
                result = papers_data.update_one(filter, update, upsert=True)
                mark_data_as_consumed(corpus_id)
                print(result.upserted_id, "inserted successfully")
        except Exception as error:
            # handle the exception
            print("An exception occurred:", error)
        finally:
            q.task_done()


def get_paper(paper_id):
    rsp = requests.get(f'https://api.semanticscholar.org/graph/v1/paper/{paper_id}',
                       headers={'x-api-key': S2_API_KEY},
                       params={'fields': QUERY_FIELDS1})
    rsp.raise_for_status()
    return rsp.json()


def get_citations(paper_id):
    edges = get_citation_edges(url=f'https://api.semanticscholar.org/graph/v1/paper/{paper_id}/citations',
                               headers={'x-api-key': S2_API_KEY},
                               params={'fields': QUERY_FIELDS2})
    return list(edge['citingPaper'] for edge in edges)


def get_references(paper_id):
    edges = get_citation_edges(url=f'https://api.semanticscholar.org/graph/v1/paper/{paper_id}/references',
                               headers={'x-api-key': S2_API_KEY},
                               params={'fields': QUERY_FIELDS2})
    return list(edge['citedPaper'] for edge in edges)


# 接口存在人机验证
def get_related_papers(paper_id):
    rsp = requests.get(url=f'https://www.semanticscholar.org/api/1/paper/{paper_id}/related-papers?limit=10&recommenderType=relatedPapers',
                       headers={'x-api-key': S2_API_KEY},
                       params={'fields': QUERY_FIELDS3})
    rsp.raise_for_status()
    return rsp.json()['papers']


def get_recommended_papers(paper_id):
    rsp = requests.get(url=f'https://api.semanticscholar.org/recommendations/v1/papers/forpaper/{paper_id}',
                       headers={'x-api-key': S2_API_KEY},
                       params={'fields': QUERY_FIELDS2})
    rsp.raise_for_status()
    return rsp.json()['recommendedPapers']


def get_citation_edges(**req_kwargs):
    """This helps with API endpoints that involve paging."""
    page_size = 1000
    offset = 0
    while True:
        req_kwargs.setdefault('params', dict())
        req_kwargs['params']['limit'] = page_size
        req_kwargs['params']['offset'] = offset
        rsp = requests.get(**req_kwargs)
        rsp.raise_for_status()

        page = rsp.json()["data"]
        for element in page:
            yield element

        if len(page) < page_size:
            break  # no more pages
        offset += page_size


def fetch_data(paper_id):
    print("fetching data:", paper_id)
    data = get_paper(paper_id)
    # print(paper)
    data['citations'] = get_citations(paper_id)
    data['references'] = get_references(paper_id)
    data['recommendedPapers'] = get_recommended_papers(paper_id)
    print('>>> fetch data OK, citations: {0}, references: {1}, recommendedPapers: {2}'.format(
        len(data.get('citations', [])), len(data.get('references', [])), len(data.get('recommendedPapers', []))
    ))

    return data if isinstance(data, dict) else None


def onSigInt(signo, frame):
    global quit_flag
    quit_flag = True
    print('Ctrl C: Waiting for the process to exit...')


if __name__ == "__main__":
    # 主进程退出信号
    signal.signal(signal.SIGINT, onSigInt)

    parser = argparse.ArgumentParser(description="Crawl data from URLs")
    parser.add_argument(
        "command", choices=["add_paper", "crawl_data"], help="Command to execute"
    )
    parser.add_argument("--path", help="Path to add to papers")
    args = parser.parse_args()

    if args.command == "add_paper":
        add_paper(args.path)
    elif args.command == "crawl_data":
        crawl_data()