import os TOGETHER_API_KEY = os.environ.get("TOGETHER_API_KEY") SEMANTIC_SCHOLAR_API_KEY = os.environ.get("SEMANTIC_SCHOLAR_API_KEY") import re import time import json import shutil import requests import spacy #!python -m spacy download en_core_web_lg from openai import OpenAI, APIError from llama_index import ( VectorStoreIndex, SimpleDirectoryReader, ServiceContext, load_index_from_storage ) from llama_index.embeddings import HuggingFaceEmbedding, TogetherEmbedding from llama_index.storage.storage_context import StorageContext # 处理原始.text数据抹去citation，或者直接从用户出获得没有citation的introduct def remove_citation(text): # Regular expression to match \cite{...} pattern = r'\\cite\{[^}]*\}' # Replace \cite{...} with an empty string text = re.sub(pattern, '', text) # Replace multiple spaces with a single space text = re.sub(r' +', ' ', text) # Replace spaces before punctuation marks with just the punctuation marks text = re.sub(r"\s+([,.!?;:()\[\]{}])", r"\1", text) return text def get_chat_completion(client, prompt, llm_model, max_tokens): messages = [ { "role": "system", "content": "You are an AI assistant", }, { "role": "user", "content": prompt, } ] try: chat_completion = client.chat.completions.create( messages=messages, model=llm_model, max_tokens=max_tokens ) return chat_completion.choices[0].message.content except APIError as e: # Handle specific API errors print(f"API Error: {e}") except Exception as e: # Handle other exceptions print(f"Error: {e}") def get_relevant_papers(search_query, sort=True, count=10): """ search_query (str): the required query parameter and its value (in this case, the keyword we want to search for) count (int): the number of relevant papers to return for each query Semantic Scholar Rate limit: 1 request per second for the following endpoints: /paper/batch /paper/search /recommendations 10 requests / second for all other calls """ # Define the paper search endpoint URL; All keywords in the search query are matched against the paper’s title and abstract. url = 'https://api.semanticscholar.org/graph/v1/paper/search' # Define headers with API key headers = {'x-api-key': SEMANTIC_SCHOLAR_API_KEY} query_params = { 'query': search_query, 'fields': 'url,title,year,abstract,authors.name,journal,citationStyles,tldr,referenceCount,citationCount', 'limit': 20, } # Send the API request response = requests.get(url, params=query_params, headers=headers) # Check response status if response.status_code == 200: json_response = response.json() if json_response['total'] != 0: papers = json_response['data'] else: papers = [] # Sort the papers based on citationCount in descending order if sort: papers = sorted(papers, key=lambda x: x['citationCount'], reverse=True) return papers[:count] else: print(f"Request failed with status code {response.status_code}: {response.text}") def save_papers(unique_dir, papers): os.makedirs(unique_dir, exist_ok=True) # Save each dictionary to a separate JSON file for i, dictionary in enumerate(papers): filename = os.path.join(unique_dir, f"{dictionary['paperId']}.json") with open(filename, 'w') as json_file: json.dump(dictionary, json_file, indent=4) print(f"{len(papers)} papers saved as JSON files successfully at {unique_dir}.") def get_index(service_context, docs_dir, persist_dir): documents = SimpleDirectoryReader(docs_dir, filename_as_id=True).load_data() # check if storage already exists PERSIST_DIR = persist_dir if not os.path.exists(PERSIST_DIR): print('create new index') index = VectorStoreIndex.from_documents( documents, service_context=service_context, show_progress=False ) # store it for later index.storage_context.persist(persist_dir=PERSIST_DIR) else: print('load the existing index') # load the existing index storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR) index = load_index_from_storage(storage_context, service_context=service_context) # refresh the index refreshed_docs = index.refresh_ref_docs(documents, update_kwargs={"delete_kwargs": {"delete_from_docstore": True}}) print(f'refreshed_docs:\n{refreshed_docs}') return index def get_paper_data(text): """text = node.text """ dictionary_from_json = json.loads(text) bibtex = dictionary_from_json['citationStyles']['bibtex'] bibtex = bibtex.replace('&', 'and') citation_label = re.findall(r"@(\w+){([\w'-]+)", bibtex)[0][1] citationCount = dictionary_from_json['citationCount'] if dictionary_from_json['tldr'] is not None: tldr = dictionary_from_json['tldr']['text'] else: tldr = 'No tldr available' url = dictionary_from_json['url'] return citation_label, (bibtex, citationCount, tldr, url) def move_cite_inside_sentence(sent, ez_citation): if sent[-1]!='\n': character = sent[-1] sent_new = sent[:-1] + ' ' + character else: count = sent.count('\n') sent = sent.strip() character = sent[-1] sent_new = sent[:-1] + ' ' + character + '\n'*count return sent_new.replace('', ez_citation) def write_bib_file(bib_file_content, data): bibtex, citationCount, tldr, url = data bib_file_content = bib_file_content + f'\n%citationCount: {citationCount}\n%tldr: {tldr}\n%url: {url}\n' + bibtex return bib_file_content def write_citation(sent, bib_file_content, retrieved_nodes, sim_threshold=0.75): labels = [] for node in retrieved_nodes: citation_label, data = get_paper_data(node.text) print('relevant paper id (node.id_):', node.id_, 'match score (node.score):', node.score) print('relevant paper data:', *data) print('-'*30) if node.score > sim_threshold and citation_label != "None": labels.append(citation_label) if not (citation_label in bib_file_content): bib_file_content = write_bib_file(bib_file_content, data) else: continue labels = ', '.join(labels) if labels: ez_citation = f'\cite{{{labels}}}' sent_new = move_cite_inside_sentence(sent, ez_citation) else: sent_new = sent return sent_new, bib_file_content get_prompt = lambda sentence: f""" I want to use semantic scholar paper search api to find the relevant papers, can you read the following text then suggest me an suitable search query for this task? Here is an example for using the api: ```python import requests # Define the paper search endpoint URL url = 'https://api.semanticscholar.org/graph/v1/paper/search' # Define the required query parameter and its value (in this case, the keyword we want to search for) query_params = {{ 'query': 'semantic scholar platform', 'limit': 3 }} # Make the GET request with the URL and query parameters searchResponse = requests.get(url, params=query_params) ``` Here is the text: {sentence} """ # main block def main(sentences, count, client, llm_model, max_tokens, service_context): """count (int): the number of relevant papers to return for each query""" sentences_new = [] bib_file_content = '' for sentence in sentences: prompt = get_prompt(sentence) response = get_chat_completion(client, prompt, llm_model, max_tokens) # Define a regular expression pattern to find the value of 'query' pattern = r"'query': '(.*?)'" matches = re.findall(pattern, response) if matches: search_query = matches[0] else: search_query = sentence[:2] # use the first two words as the search query relevant_papers = get_relevant_papers(search_query, sort=True, count=count) if relevant_papers: # save papers to json files and build index unique_dir = os.path.join("papers", f"{int(time.time())}") persist_dir = os.path.join("index", f"{int(time.time())}") save_papers(unique_dir, relevant_papers) index = get_index(service_context, unique_dir, persist_dir) # get sentence's most similar papers retriever = index.as_retriever(service_context=service_context, similarity_top_k=5) retrieved_nodes = retriever.retrieve(sentence) sent_new, bib_file_content = write_citation(sentence, bib_file_content, retrieved_nodes, sim_threshold=0.7) sentences_new.append(sent_new) else: sentences_new.append(sentence) print('sentence:', sentence.strip()) print('search_query:', search_query) print('='*30) return sentences_new, bib_file_content def ez_cite(introduction, debug=False): nlp = spacy.load("en_core_web_lg") doc = nlp(introduction) sentences = [sentence.text for sentence in doc.sents] sentences = [ remove_citation(sentence) for sentence in sentences] client = OpenAI(api_key=TOGETHER_API_KEY, base_url='https://api.together.xyz', ) llm_model = "Qwen/Qwen1.5-72B-Chat" max_tokens = 1000 embed_model = TogetherEmbedding(model_name="togethercomputer/m2-bert-80M-8k-retrieval", api_key=TOGETHER_API_KEY) service_context = ServiceContext.from_defaults( llm=None, embed_model=embed_model, chunk_size=8192, # chunk_size must be bigger than the whole .json so that all info is preserved, in this case, one doc is one node ) if debug: sentences = sentences[:2] sentences_new, bib_file_content = main(sentences, count=10, client=client, llm_model=llm_model, max_tokens=max_tokens, service_context=service_context) with open('intro.bib', 'w') as bib_file: bib_file.write(bib_file_content) final_intro = ' '.join(sentences_new) print(final_intro) print('='*30) dir_path = "index" try: # Delete the directory and its contents shutil.rmtree(dir_path) print(f"Directory '{dir_path}' deleted successfully.") except Exception as e: print(f"Error deleting directory '{dir_path}': {e}") dir_path = "papers" try: # Delete the directory and its contents shutil.rmtree(dir_path) print(f"Directory '{dir_path}' deleted successfully.") except Exception as e: print(f"Error deleting directory '{dir_path}': {e}") return final_intro, bib_file_content example1 = r"""In the current Noisy Intermediate-Scale Quantum (NISQ) era, a few methods have been proposed to construct useful quantum algorithms that are compatible with mild hardware restrictions. Most of these methods involve the specification of a quantum circuit Ansatz, optimized in a classical fashion to solve specific computational tasks. Next to variational quantum eigensolvers in chemistry and variants of the quantum approximate optimization algorithm, machine learning approaches based on such parametrized quantum circuits stand as some of the most promising practical applications to yield quantum advantages.""" if __name__ == "__main__": final_intro, bib_file_content = ez_cite(example1, debug=True)