import json import logging from typing import Optional import re import requests from requests.adapters import HTTPAdapter, Retry import arxiv import PyPDF2 import requests import time from operator import itemgetter import fitz from tqdm import tqdm from langchain.text_splitter import CharacterTextSplitter import chromadb from chromadb.utils import embedding_functions import os import asyncio import uuid from langchain.text_splitter import CharacterTextSplitter emmbedding_model = "text-embedding-3-large" OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY") dir = "arxiv-data/2023" chunk_size = 20 openai_ef = embedding_functions.OpenAIEmbeddingFunction(model_name=emmbedding_model,api_key=OPENAI_API_KEY) # chroma_client = chromadb.HttpClient(host='localhost', port=8000) # use this if you are using chroma server Docker chroma_client = chromadb.PersistentClient(path="/home/ubuntu/research/data/emeddeings") text_splitter = CharacterTextSplitter( separator="\n", chunk_size=3000, chunk_overlap=0 ) collection_doc = chroma_client.get_or_create_collection(name="2024_main_document_lvl") collection_para = chroma_client.get_or_create_collection(name="2024_main_paragraph_lvl") def background(f): def wrapped(*args, **kwargs): return asyncio.get_event_loop().run_in_executor(None, f, *args, **kwargs) return wrapped paper_id_re = re.compile(r'https://arxiv.org/abs/(\d+\.\d+)') def retry_request_session(retries: Optional[int] = 5): # we setup retry strategy to retry on common errors retries = Retry( total=retries, backoff_factor=0.1, status_forcelist=[ 408, # request timeout 500, # internal server error 502, # bad gateway 503, # service unavailable 504 # gateway timeout ] ) # we setup a session with the retry strategy session = requests.Session() session.mount('https://', HTTPAdapter(max_retries=retries)) return session def get_paper_id(query: str, handle_not_found: bool = True): """Get the paper ID from a query. :param query: The query to search with :type query: str :param handle_not_found: Whether to return None if no paper is found, defaults to True :type handle_not_found: bool, optional :return: The paper ID :rtype: str """ special_chars = { ":": "%3A", "|": "%7C", ",": "%2C", " ": "+" } # create a translation table from the special_chars dictionary translation_table = query.maketrans(special_chars) # use the translate method to replace the special characters search_term = query.translate(translation_table) # init requests search session session = retry_request_session() # get the search results res = session.get(f"https://www.google.com/search?q={search_term}&sclient=gws-wiz-serp") try: # extract the paper id paper_id = paper_id_re.findall(res.text)[0] except IndexError: if handle_not_found: # if no paper is found, return None return None else: # if no paper is found, raise an error raise Exception(f'No paper found for query: {query}') return paper_id class Arxiv: refs_re = re.compile(r'\n(References|REFERENCES)\n') references = [] llm = None def __init__(self, paper_id: str): """Object to handle the extraction of an ArXiv paper and its relevant information. :param paper_id: The ID of the paper to extract :type paper_id: str """ self.id = paper_id self.url = f"https://export.arxiv.org/pdf/{paper_id}.pdf" # initialize the requests session self.session = requests.Session() def load(self, save: bool = False): """Load the paper from the ArXiv API or from a local file if it already exists. Stores the paper's text content and meta data in self.content and other attributes. :param save: Whether to save the paper to a local file, defaults to False :type save: bool, optional """ try: self._download_meta() if save: self.save() except Exception as e: print(f"Error while downloading paper {self.id}: {e}") raise e def get_refs(self, extractor, text_splitter): """Get the references for the paper. :param extractor: The LLMChain extractor model :type extractor: LLMChain :param text_splitter: The text splitter to use :type text_splitter: TokenTextSplitter :return: The references for the paper :rtype: list """ if len(self.references) == 0: self._download_refs(extractor, text_splitter) return self.references def _download_refs(self, extractor, text_splitter): """Download the references for the paper. Stores them in the self.references attribute. :param extractor: The LLMChain extractor model :type extractor: LLMChain :param text_splitter: The text splitter to use :type text_splitter: TokenTextSplitter """ # get references section of paper refs = self.refs_re.split(self.content)[-1] # we don't need the full thing, just the first page refs_page = text_splitter.split_text(refs)[0] # use LLM extractor to extract references out = extractor.run(refs=refs_page) out = out.split('\n') out = [o for o in out if o != ''] # with list of references, find the paper IDs ids = [get_paper_id(o) for o in out] # clean up into JSONL type format out = [o.split(' | ') for o in out] # in case we're missing some fields out = [o for o in out if len(o) == 3] meta = [{ 'id': _id, 'title': o[0], 'authors': o[1], 'year': o[2] } for o, _id in zip(out, ids) if _id is not None] logging.debug(f"Extracted {len(meta)} references") self.references = meta def _convert_pdf_to_text(self): """Convert the PDF to text and store it in the self.content attribute. """ text = [] with open("temp.pdf", 'rb') as f: # create a PDF object pdf = PyPDF2.PdfReader(f) # iterate over every page in the PDF for page in range(len(pdf.pages)): # get the page object page_obj = pdf.pages[page] # extract text from the page text.append(page_obj.extract_text()) text = "\n".join(text) self.content = text def _download_meta(self): """Download the meta information for the paper from the ArXiv API and store it in the self attributes. """ search = arxiv.Search( query=f'id:{self.id}', max_results=1, sort_by=arxiv.SortCriterion.SubmittedDate ) result = list(search.results()) if len(result) == 0: raise ValueError(f"No paper found for paper '{self.id}'") result = result[0] # remove 'v1', 'v2', etc. from the end of the pdf_url result.pdf_url = re.sub(r'v\d+$', '', result.pdf_url) self.authors = " ".join([author.name for author in result.authors]) self.journal_ref = result.journal_ref self.source = result.pdf_url self.summary = result.summary self.title = result.title logging.debug(f"Downloaded metadata for paper '{self.id}'") def save(self): """Save the paper to a local JSON file. """ with open(f'papers/{self.id}.json', 'w') as fp: json.dump(self.__dict__(), fp, indent=4) def save_chunks( self, include_metadata: bool = True, path: str = "chunks" ): """Save the paper's chunks to a local JSONL file. :param include_metadata: Whether to include the paper's metadata in the chunks, defaults to True :type include_metadata: bool, optional :param path: The path to save the file to, defaults to "papers" :type path: str, optional """ if not os.path.exists(path): os.makedirs(path) with open(f'{path}/{self.id}.jsonl', 'w') as fp: for chunk in self.dataset: if include_metadata: chunk.update(self.get_meta()) fp.write(json.dumps(chunk) + '\n') logging.debug(f"Saved paper to '{path}/{self.id}.jsonl'") def get_meta(self): """Returns the meta information for the paper. :return: The meta information for the paper :rtype: dict """ fields = self.__dict__() # drop content field because it's big # fields.pop('content') return fields def chunker(self, chunk_size=300): # Single Chunk is made for now clean_paper = self._clean_text(self.content) langchain_dataset = [] langchain_dataset.append({ 'doi': self.id, 'chunk-id': 1, 'chunk': clean_paper }) self.dataset = langchain_dataset def _clean_text(self, text): text = re.sub(r'-\n', '', text) return text def __dict__(self): return { 'id': self.id, 'title': self.title, 'summary': self.summary, 'source': self.source, 'authors': self.authors, 'references': " ".join(self.references) } def __repr__(self): return f"Arxiv(paper_id='{self.id}')" def fonts(doc, granularity=False): """Extracts fonts and their usage in PDF documents. :param doc: PDF document to iterate through :type doc: :param granularity: also use 'font', 'flags' and 'color' to discriminate text :type granularity: bool :rtype: [(font_size, count), (font_size, count}], dict :return: most used fonts sorted by count, font style information """ styles = {} font_counts = {} for page in doc: blocks = page.get_text("dict")["blocks"] for b in blocks: # iterate through the text blocks if b['type'] == 0: # block contains text for l in b["lines"]: # iterate through the text lines for s in l["spans"]: # iterate through the text spans if granularity: identifier = "{0}_{1}_{2}_{3}".format(s['size'], s['flags'], s['font'], s['color']) styles[identifier] = {'size': s['size'], 'flags': s['flags'], 'font': s['font'], 'color': s['color']} else: identifier = "{0}".format(s['size']) styles[identifier] = {'size': s['size'], 'font': s['font']} font_counts[identifier] = font_counts.get(identifier, 0) + 1 # count the fonts usage font_counts = sorted(font_counts.items(), key=itemgetter(1), reverse=True) if len(font_counts) < 1: raise ValueError("Zero discriminating fonts found!") return font_counts, styles def font_tags(font_counts, styles): """Returns dictionary with font sizes as keys and tags as value. :param font_counts: (font_size, count) for all fonts occuring in document :type font_counts: list :param styles: all styles found in the document :type styles: dict :rtype: dict :return: all element tags based on font-sizes """ p_style = styles[font_counts[0][0]] # get style for most used font by count (paragraph) p_size = p_style['size'] # get the paragraph's size # sorting the font sizes high to low, so that we can append the right integer to each tag font_sizes = [] for (font_size, count) in font_counts: font_sizes.append(float(font_size)) font_sizes.sort(reverse=True) # aggregating the tags for each font size idx = 0 size_tag = {} for size in font_sizes: idx += 1 if size == p_size: idx = 0 size_tag[size] = '

' if size > p_size: size_tag[size] = ''.format(idx) elif size < p_size: size_tag[size] = ''.format(idx) return size_tag def headers_para(doc, size_tag): """Scrapes headers & paragraphs from PDF and return texts with element tags. :param doc: PDF document to iterate through :type doc: :param size_tag: textual element tags for each size :type size_tag: dict :rtype: list :return: texts with pre-prended element tags """ paragraphs = [] # list with paragraphs first = True # boolean operator for first header previous_s = {} # previous span for page in doc: blocks = page.get_text("dict")["blocks"] for b in blocks: # iterate through the text blocks if b['type'] == 0: # this block contains text # REMEMBER: multiple fonts and sizes are possible IN one block block_string = "" # text found in block for l in b["lines"]: # iterate through the text lines for s in l["spans"]: # iterate through the text spans if s['text'].strip(): # removing whitespaces: if first: previous_s = s first = False block_string = s['text'] if size_tag[s['size']] == '

' else '' else: if s['size'] == previous_s['size']: if block_string: # in the same block, so concatenate strings block_string += " " + s['text'] else: if block_string: # new block has started, so append the paragraph paragraphs.append(block_string) block_string = s['text'] if size_tag[s['size']] == '

' else '' previous_s = s if block_string: # append the last paragraph in the block if len(block_string) > 80: # print(len(block_string), block_string,'\n') paragraphs.append(block_string) return paragraphs def get_pdf_info(document_path): docs = fitz.open(document_path) only_text = "" for page in docs: only_text += page.get_text() + " " font_counts, styles = fonts(docs, granularity=False) size_tag = font_tags(font_counts, styles) elements = headers_para(docs, size_tag) paragraphs = [] for element in elements: if len(element) > 100: paragraphs.append(element.lower()) pattern = r'\d+(\.\d+)?\n' cleaned_text = re.sub(pattern, '', only_text) return cleaned_text.lower(),paragraphs def generate_uuid(): return str(uuid.uuid4()) def add_document_chroma_collection(collection_object, document_list, embedding_list, metadata): metadata_list = [metadata for i in range(len(document_list))] ids_gen = [generate_uuid() for i in range(len(document_list))] collection_object.add(embeddings = embedding_list,documents = document_list,metadatas = metadata_list,ids = ids_gen) if collection_object: return True def chunk_list(input_list, chunk_size): return [input_list[i:i + chunk_size] for i in range(0, len(input_list), chunk_size)] directories_2023 = [ os.path.join(dir, o) for o in os.listdir(dir) if os.path.isdir(os.path.join(dir,o)) ] @background def task(chunk): chunked_files = [os.path.join(single_dir, file) for file in chunk] # print("len of chunked_files", len(chunked_files)) for file_path in (chunked_files): try: paper_id = file_path.split("/")[-1].split(".pdf")[0].replace("v1", "").replace("v2", "").replace("v3", "").replace("v4", "").replace("v5", "") # print("paper_id", paper_id) paper = Arxiv(paper_id) paper.load() metadata = paper.get_meta() only_text,paragraphs = get_pdf_info(file_path) texts_list = text_splitter.split_text(only_text) if len(texts_list) > 0: doc_embed = openai_ef(texts_list) add_document_chroma_collection(collection_doc, texts_list, doc_embed,metadata) if len(paragraphs) > 0: paragraphs_embed = openai_ef(paragraphs) add_document_chroma_collection(collection_para, paragraphs, paragraphs_embed,metadata) time.sleep(5) except Exception as e: print("Error", e) with open("error.txt", "a") as f: f.write(file_path) continue for single_dir in directories_2023: print("len of files", len(os.listdir(single_dir))) chunked_list = chunk_list(os.listdir(single_dir), 20) print("len of chunked_list", len(chunked_list)) for chunk in tqdm(chunked_list): task(chunk) time.sleep(5)