Spaces:
Running
Running
import requests | |
import json | |
import yaml | |
import scipdf | |
import os | |
import time | |
import aiohttp | |
import asyncio | |
import numpy as np | |
import random | |
def get_content_between_a_b(start_tag, end_tag, text): | |
extracted_text = "" | |
start_index = text.find(start_tag) | |
while start_index != -1: | |
end_index = text.find(end_tag, start_index + len(start_tag)) | |
if end_index != -1: | |
extracted_text += text[start_index + len(start_tag) : end_index] + " " | |
start_index = text.find(start_tag, end_index + len(end_tag)) | |
else: | |
break | |
return extracted_text.strip() | |
def extract(text, type): | |
if text: | |
target_str = get_content_between_a_b(f"<{type}>", f"</{type}>", text) | |
if target_str: | |
return target_str | |
else: | |
return text | |
else: | |
return "" | |
def download(url): | |
try: | |
response = requests.get(url) | |
if response.status_code == 200: | |
return response.content | |
else: | |
print(f"Failed to download the file from the URL: {url}") | |
return None | |
except requests.RequestException as e: | |
print(f"An error occurred while downloading the file from the URL: {url}") | |
print(e) | |
return None | |
except Exception as e: | |
print(f"An unexpected error occurred while downloading the file from the URL: {url}") | |
print(e) | |
return None | |
class Result: | |
def __init__(self,title="",abstract="",article = "",citations_conut = 0,year = None) -> None: | |
self.title = title | |
self.abstract = abstract | |
self.article = article | |
self.citations_conut = citations_conut | |
self.year = year | |
# Define the API endpoint URL | |
semantic_fields = ["title", "abstract", "year", "authors.name", "authors.paperCount", "authors.citationCount","authors.hIndex","url","referenceCount","citationCount","influentialCitationCount","isOpenAccess","openAccessPdf","fieldsOfStudy","s2FieldsOfStudy","embedding.specter_v1","embedding.specter_v2","publicationDate","citations"] | |
fieldsOfStudy = ["Computer Science","Medicine","Chemistry","Biology","Materials Science","Physics","Geology","Art","History","Geography","Sociology","Business","Political Science","Philosophy","Art","Literature","Music","Economics","Philosophy","Mathematics","Engineering","Environmental Science","Agricultural and Food Sciences","Education","Law","Linguistics"] | |
# citations.paperId, citations.title, citations.year, citations.authors.name, citations.authors.paperCount, citations.authors.citationCount, citations.authors.hIndex, citations.url, citations.referenceCount, citations.citationCount, citations.influentialCitationCount, citations.isOpenAccess, citations.openAccessPdf, citations.fieldsOfStudy, citations.s2FieldsOfStudy, citations.publicationDate | |
# publicationDateOrYear: 2019-03-05 ; 2019-03 ; 2019 ; 2016-03-05:2020-06-06 ; 1981-08-25: ; :2020-06-06 ; 1981:2020 | |
# publicationTypes: Review ; JournalArticle CaseReport ; ClinicalTrial ; Dataset ; Editorial ; LettersAndComments ; MetaAnalysis ; News ; Study ; Book ; BookSection | |
def process_fields(fields): | |
return ",".join(fields) | |
class SementicSearcher: | |
def __init__(self, ban_paper = []) -> None: | |
self.ban_paper = ban_paper | |
def search_papers(self, query, limit=5, offset=0, fields=["title", "paperId", "abstract", "isOpenAccess", 'openAccessPdf', "year","publicationDate","citations.title","citations.abstract","citations.isOpenAccess","citations.openAccessPdf","citations.citationCount","citationCount","citations.year"], | |
publicationDate=None, minCitationCount=0, year=None, | |
publicationTypes=None, fieldsOfStudy=None): | |
url = 'https://api.semanticscholar.org/graph/v1/paper/search' | |
fields = process_fields(fields) if isinstance(fields, list) else fields | |
# More specific query parameter | |
query_params = { | |
'query': query, | |
"limit": limit, | |
"offset": offset, | |
'fields': fields, | |
'publicationDateOrYear': publicationDate, | |
'minCitationCount': minCitationCount, | |
'year': year, | |
'publicationTypes': publicationTypes, | |
'fieldsOfStudy': fieldsOfStudy | |
} | |
# Load the API key from the configuration file | |
api_key = os.environ.get('SEMENTIC_SEARCH_API_KEY',None) | |
headers = {'x-api-key': api_key} if api_key else None | |
try: | |
filtered_query_params = {key: value for key, value in query_params.items() if value is not None} | |
response = requests.get(url, params=filtered_query_params, headers=headers) | |
if response.status_code == 200: | |
response_data = response.json() | |
return response_data | |
elif response.status_code == 429: | |
time.sleep(1) | |
print(f"Request failed with status code {response.status_code}: begin to retry") | |
return self.search_papers(query, limit, offset, fields, publicationDate, minCitationCount, year, publicationTypes, fieldsOfStudy) | |
else: | |
print(f"Request failed with status code {response.status_code}: {response.text}") | |
return None | |
except requests.RequestException as e: | |
print(f"An error occurred: {e}") | |
return None | |
def cal_cosine_similarity(self, vec1, vec2): | |
return np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2)) | |
def cal_cosine_similarity_matric(self,matric1, matric2): | |
if isinstance(matric1, list): | |
matric1 = np.array(matric1) | |
if isinstance(matric2, list): | |
matric2 = np.array(matric2) | |
if len(matric1.shape) == 1: | |
matric1 = matric1.reshape(1, -1) | |
if len(matric2.shape) == 1: | |
matric2 = matric2.reshape(1, -1) | |
dot_product = np.dot(matric1, matric2.T) | |
norm1 = np.linalg.norm(matric1, axis=1) | |
norm2 = np.linalg.norm(matric2, axis=1) | |
cos_sim = dot_product / np.outer(norm1, norm2) | |
scores = cos_sim.flatten() | |
return scores.tolist() | |
def read_arxiv_from_path(self, pdf_path): | |
def is_pdf(binary_data): | |
pdf_header = b'%PDF-' | |
return binary_data.startswith(pdf_header) | |
try: | |
flag = is_pdf(pdf_path) | |
if not flag: | |
return None | |
except Exception as e: | |
pass | |
try: | |
article_dict = scipdf.parse_pdf_to_dict(pdf_path) | |
except Exception as e: | |
print(f"Failed to parse the PDF file: {pdf_path}") | |
return None | |
return article_dict | |
def get_paper_embbeding_and_score(self,query_embedding, paper,llm): | |
paper_content = f""" | |
Title: {paper['title']} | |
Abstract: {paper['abstract']} | |
""" | |
paper_embbeding = llm.get_embbeding(paper_content) | |
paper_embbeding = np.array(paper_embbeding) | |
score = self.cal_cosine_similarity(query_embedding,paper_embbeding) | |
return [paper,score] | |
def rerank_papers(self, query_embedding, paper_list,llm): | |
if len(paper_list) == 0: | |
return [] | |
paper_list = [paper for paper in paper_list if paper] | |
if len(paper_list) >= 50: | |
paper_list = random.sample(paper_list,50) | |
paper_contents = [] | |
for paper in paper_list: | |
paper_content = f""" | |
Title: {paper['title']} | |
Abstract: {paper['abstract']} | |
""" | |
paper_contents.append(paper_content) | |
paper_contents_embbeding = llm.get_embbeding(paper_contents) | |
paper_contents_embbeding = np.array(paper_contents_embbeding) | |
scores = self.cal_cosine_similarity_matric(query_embedding,paper_contents_embbeding) | |
# 根据score对paper_list进行排序 | |
paper_list = sorted(zip(paper_list,scores),key = lambda x: x[1],reverse = True) | |
paper_list = [paper[0] for paper in paper_list] | |
return paper_list | |
def search(self,query,max_results = 5 ,paper_list = None ,rerank_query = None,llm = None,year = None,publicationDate = None,need_download = True,fields = ["title", "paperId", "abstract", "isOpenAccess", 'openAccessPdf', "year","publicationDate","citationCount"]): | |
if rerank_query: | |
rerank_query_embbeding = llm.get_embbeding(rerank_query) | |
rerank_query_embbeding = np.array(rerank_query_embbeding) | |
readed_papers = [] | |
if paper_list: | |
if isinstance(paper_list,set): | |
paper_list = list(paper_list) | |
if len(paper_list) == 0 : | |
pass | |
elif isinstance(paper_list[0], str): | |
readed_papers = paper_list | |
elif isinstance(paper_list[0], Result): | |
readed_papers = [paper.title for paper in paper_list] | |
print(f"Searching for papers related to the query: <{query}>") | |
results = self.search_papers(query,limit = 10 * max_results,year=year,publicationDate = publicationDate,fields = fields) | |
if not results or "data" not in results: | |
return [] | |
new_results = [] | |
for result in results['data']: | |
if result['title'] in self.ban_paper: | |
continue | |
new_results.append(result) | |
results = new_results | |
final_results = [] | |
if need_download: | |
paper_candidates = [] | |
for result in results: | |
if not result['isOpenAccess'] or not result['openAccessPdf'] or result['title'] in readed_papers: | |
continue | |
else: | |
paper_candidates.append(result) | |
else: | |
paper_candidates = results | |
if llm and rerank_query: | |
paper_candidates = self.rerank_papers(rerank_query_embbeding, paper_candidates,llm) | |
if need_download: | |
for result in paper_candidates: | |
pdf_link = result['openAccessPdf']["url"] | |
try: | |
content = self.download_pdf(pdf_link) | |
if not content: | |
continue | |
except Exception as e: | |
continue | |
title = result['title'] | |
abstract = result['abstract'] | |
citationCount = result['citationCount'] | |
year = result['year'] | |
article = self.read_arxiv_from_path(content) | |
if not article: | |
continue | |
final_results.append(Result(title,abstract,article,citationCount,year)) | |
if len(final_results) >= max_results: | |
break | |
else: | |
for result in paper_candidates: | |
title = result['title'] | |
abstract = result['abstract'] | |
citationCount = result['citationCount'] | |
year = result['year'] | |
final_results.append(Result(title,abstract,None,citationCount,year)) | |
if len(final_results) >= max_results: | |
break | |
return final_results | |
def search_related_paper(self,title,need_citation = True,need_reference = True,rerank_query = None,llm = None,paper_list = []): | |
print(f"Searching for the related papers of <{title}>, need_citation: {need_citation}, need_reference: {need_reference}") | |
fileds = ["title","abstract","citations.title","citations.abstract","citations.citationCount","references.title","references.abstract","references.citationCount","citations.isOpenAccess","citations.openAccessPdf","references.isOpenAccess","references.openAccessPdf","citations.year","references.year"] | |
results = self.search_papers(title,limit = 3,fields=fileds) | |
related_papers = [] | |
related_papers_title = [] | |
if not results or "data" not in results: | |
return None | |
for result in results["data"]: | |
if not result: | |
continue | |
if need_citation: | |
for citation in result["citations"]: | |
if "openAccessPdf" not in citation or not citation["openAccessPdf"]: | |
continue | |
elif citation["title"] in related_papers_title or citation["title"] in self.ban_paper or citation["title"] in paper_list: | |
continue | |
elif citation["isOpenAccess"] == False or citation["openAccessPdf"] == None: | |
continue | |
else: | |
related_papers.append(citation) | |
related_papers_title.append(citation["title"]) | |
if need_reference: | |
for reference in result["references"]: | |
if "openAccessPdf" not in reference or not reference["openAccessPdf"]: | |
continue | |
elif reference["title"] in related_papers_title or reference["title"] in self.ban_paper or reference["title"] in paper_list: | |
continue | |
elif reference["isOpenAccess"] == False or reference["openAccessPdf"] == None: | |
continue | |
else: | |
related_papers.append(reference) | |
related_papers_title.append(reference["title"]) | |
if result: | |
break | |
if len(related_papers) >= 200: | |
related_papers = related_papers[:200] | |
if rerank_query and llm: | |
rerank_query_embbeding = llm.get_embbeding(rerank_query) | |
rerank_query_embbeding = np.array(rerank_query_embbeding) | |
related_papers = self.rerank_papers(rerank_query_embbeding, related_papers,llm) | |
related_papers = [[paper["title"],paper["abstract"],paper["openAccessPdf"]["url"],paper["citationCount"],paper['year']] for paper in related_papers] | |
else: | |
related_papers = [[paper["title"],paper["abstract"],paper["openAccessPdf"]["url"],paper["citationCount"],paper['year']] for paper in related_papers] | |
related_papers = sorted(related_papers,key = lambda x: x[3],reverse = True) | |
print(f"Found {len(related_papers)} related papers") | |
for paper in related_papers: | |
url = paper[2] | |
content = self.download_pdf(url) | |
if content: | |
article = self.read_arxiv_from_path(content) | |
if not article: | |
continue | |
result = Result(paper[0],paper[1],article,paper[3],paper[4]) | |
return result | |
return None | |
def download_pdf(self, pdf_link): | |
content = download(pdf_link) | |
return content | |
def read_paper_title_abstract(self,article): | |
title = article["title"] | |
abstract = article["abstract"] | |
paper_content = f""" | |
Title: {title} | |
Abstract: {abstract} | |
""" | |
return paper_content | |
def read_paper_content(self,article): | |
paper_content = self.read_paper_title_abstract(article) | |
for section in article["sections"]: | |
paper_content += f"section: {section['heading']}\n content: {section['text']}\n ref_ids: {section['publication_ref']}\n" | |
return paper_content | |
def read_paper_content_with_ref(self,article): | |
paper_content = self.read_paper_content(article) | |
paper_content += "<References>\n" | |
i = 1 | |
for refer in article["references"]: | |
ref_id = refer["ref_id"] | |
title = refer["title"] | |
year = refer["year"] | |
paper_content += f"Ref_id:{ref_id} Title: {title} Year: ({year})\n" | |
i += 1 | |
paper_content += "</References>\n" | |
return paper_content | |