fetch_abs.py

import time
import feedparser
import requests
from requests.adapters import HTTPAdapter, Retry
from typing import Dict, List, Optional
import json
from thefuzz import fuzz
from diskcache import Cache


cache = Cache('cache/semantic_scholar_references')
abs_cache = Cache('cache/abstract')


class ReferenceAbstractFetcher:
    def __init__(self):
        self.base_url = "https://api.semanticscholar.org/graph/v1/paper"
        
    def get_paper_references_cached(self, title: str) -> List[Dict]:
        references = cache.get(title, None)
        if references is not None:
            return references
        
        references = self.get_paper_references(title)
        
        if references is not None:
            cache[title] = references
        return references

    def get_paper_references(self, title: str) -> Optional[List[Dict]]:
        while True:
            response = requests.get(
                f'{self.base_url}/search/match',
                params={
                    'query': title,
                    'fields': 'references.title,references.abstract,references.year,references.externalIds'
                },
            )
            time.sleep(1)
            
            if response.status_code == 200:
                data = response.json()
                if data.get('data') and len(data['data']) > 0:
                    return data['data'][0].get('references', None)
            elif response.status_code == 429:
                continue
            else:
                response.raise_for_status()


def match_reference(ref_title: str, bib_entries: List[Dict], threshold: int = 85) -> str:
    best_match = None
    best_score = 0
    
    for entry in bib_entries:
        if 'title' not in entry:
            continue
        
        titles = entry['title']
        if isinstance(titles, str):
            titles = [titles]
        for title in titles:
            score = fuzz.partial_ratio(ref_title.lower(), title.lower())
            score_full = fuzz.ratio(ref_title.lower(), title.lower())
            score = score + score_full / 128
            if score >= threshold and score > best_score:
                best_score = score
                best_match = getattr(entry, 'key', None) or entry['ID']
    
    return best_match


def get_arxiv_api(id_list: str):
    session = requests.Session()
    session.mount('https://export.arxiv.org', HTTPAdapter(max_retries=Retry(read=5)))

    url = 'https://export.arxiv.org/api/query'
    entries = []
    num_max_results = len(entries) + 1
    start = 0
    num_items_per_page = 1000
    
    while len(entries) < num_max_results:
        param = dict(
            id_list=id_list,
            start=start,
            max_results=num_items_per_page,
            sortBy='submittedDate',
            sortOrder='ascending',
        )
        result = session.post(url, data=param)
        result.raise_for_status()
        data = feedparser.parse(result.text)
        num_max_results = int(data['feed']['opensearch_totalresults'])
        
        entries += data['entries']
        start += len(data['entries'])
        time.sleep(3)
    
    return entries


def update_abstract_cache(references: list[dict]) -> None:
    arxiv_ids = {}
    for ref in references:
        if ref['paperId'] in abs_cache:
            pass
        elif 'abstract' in ref and ref['abstract'] is not None:
            abs_cache[ref['paperId']] = ref['abstract']
        elif 'ArXiv' in (ref['externalIds'] or {}):
            arxiv_ids[ref['externalIds']['ArXiv']] = ref['paperId']
    
    if len(arxiv_ids) == 0:
        return
    
    entries = get_arxiv_api(','.join(arxiv_ids.keys()))
    for entry in entries:
        arxiv_id: str = entry['id'].split('/')[-1].split('v')[0]
        try:
            abs_cache[arxiv_ids[arxiv_id]] = entry['summary']
        except KeyError:
            pass


def process_paper_references(title: str, bib_entries: List[Dict]) -> Dict[str, Dict]:
    fetcher = ReferenceAbstractFetcher()
    citations_map = {}
    
    print(f"Fetching references for paper '{title}'...")
    references = fetcher.get_paper_references_cached(title)
    
    if not references:
        print("No reference information found")
        return {}
        
    print(f"Found {len(references)} references")
    print("Updating abstract cache...")
    update_abstract_cache(references)
    print("Matching bib entries...")
    
    for ref in references:
        ref['abstract'] = abs_cache.get(ref['paperId'], ref['abstract'])
    
    for ref in references:
        if not ref.get('title'):
            print("[Warning] Semantic Scholar returned a reference without a title")
            continue
        if not ref.get('abstract'):
            print(f"[Warning] Semantic Scholar returned a reference without an abstract: {ref['title']}")
            continue
        
        bib_id = match_reference(ref['title'], bib_entries)
        if bib_id:
            citations_map[bib_id] = {
                'title': ref['title'],
                'abstract': ref['abstract'],
            }
    
    print(f"Matched {len(citations_map)} references")
    for k, v in citations_map.items():
        print(k, v['title'])
    
    return citations_map


def save_citations_map(citations_map: Dict[str, Dict], output_file: str):
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(citations_map, f, ensure_ascii=False, indent=2)