import chromadb from chromadb import Documents, EmbeddingFunction, Embeddings from transformers import AutoModel import json from numpy.linalg import norm import sqlite3 import urllib from django.conf import settings # this module act as a singleton class class JinaAIEmbeddingFunction(EmbeddingFunction): def __init__(self, model): super().__init__() self.model = model def __call__(self, input: Documents) -> Embeddings: embeddings = self.model.encode(input) return embeddings.tolist() # instance of embedding_model embedding_model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-en', trust_remote_code=True, cache_dir='models') # instance of JinaAIEmbeddingFunction ef = JinaAIEmbeddingFunction(embedding_model) # list of topics topic_descriptions = json.load(open("topic_descriptions.txt")) topics = list(dict.keys(topic_descriptions)) embeddings = [embedding_model.encode(topic_descriptions[key]) for key in topic_descriptions] cos_sim = lambda a,b: (a @ b.T) / (norm(a)*norm(b)) def choose_topic(summary): embed = embedding_model.encode(summary) topic = "" max_sim = 0. for i,key in enumerate(topics): sim = cos_sim(embed,embeddings[i]) if sim > max_sim: topic = key max_sim = sim return topic def authors_list_to_str(authors): """input a list of authors, return a string represent authors""" text = "" for author in authors: text+=author+", " return text[:-3] def authors_str_to_list(string): """input a string of authors, return a list of authors""" authors = [] list_auth = string.split("and") for author in list_auth: if author != "et al.": authors.append(author.strip()) return authors def chunk_texts(text, max_char=400): """ Chunk a long text into several chunks, with each chunk about 300-400 characters long, but make sure no word is cut in half. Args: text: The long text to be chunked. max_char: The maximum number of characters per chunk (default: 400). Returns: A list of chunks. """ chunks = [] current_chunk = "" words = text.split() for word in words: if len(current_chunk) + len(word) + 1 >= max_char: chunks.append(current_chunk) current_chunk = " " else: current_chunk += " " + word chunks.append(current_chunk.strip()) return chunks def trimming(txt): start = txt.find("{") end = txt.rfind("}") return txt[start:end+1].replace("\n"," ") # crawl data def extract_tag(txt,tagname): return txt[txt.find("<"+tagname+">")+len(tagname)+2:txt.find("")] def get_record(extract): id = extract_tag(extract,"id") updated = extract_tag(extract,"updated") published = extract_tag(extract,"published") title = extract_tag(extract,"title").replace("\n ","").strip() summary = extract_tag(extract,"summary").replace("\n","").strip() authors = [] while extract.find("")!=-1: author = extract_tag(extract,"name") extract = extract[extract.find("")+9:] authors.append(author) pattern = '") != -1: extract = xml[xml.find("")+7:xml.find("")] xml = xml[xml.find("")+8:] extract = get_record(extract) topic = choose_topic(extract[6]) records.append([topic,*extract]) return records except Exception as e: return "Error: "+str(e) def crawl_arxiv(keyword_list, max_results=100): baseurl = 'http://export.arxiv.org/api/query?search_query=' records = [] for i,keyword in enumerate(keyword_list): if i ==0: url = baseurl + 'all:' + keyword else: url = url + '+OR+' + 'all:' + keyword url = url+ '&max_results=' + str(max_results) url = url.replace(' ', '%20') try: arxiv_page = urllib.request.urlopen(url,timeout=100).read() xml = str(arxiv_page,encoding="utf-8") while xml.find("") != -1: extract = xml[xml.find("")+7:xml.find("")] xml = xml[xml.find("")+8:] extract = get_record(extract) topic = choose_topic(extract[6]) records.append([topic,*extract]) return records except Exception as e: return "Error: "+str(e) class ArxivSQL: def __init__(self, table="arxivsql", name="db.sqlite3"): self.con = sqlite3.connect(name) self.cur = self.con.cursor() self.table = table def query(self, title="", author=[]): if len(title)>0: query_title = 'title like "%{}%"'.format(title) else: query_title = "True" if len(author)>0: query_author = 'authors like ' for auth in author: query_author += "'%{}%' or ".format(auth) query_author = query_author[:-4] else: query_author = "True" query = "select * from {} where {} and {}".format(self.table,query_title,query_author) result = self.cur.execute(query) return result.fetchall() def query_id(self, ids=[]): try: if len(ids) == 0: return None query = "select * from {} where id in (".format(self.table) for id in ids: query+="'"+id+"'," query = query[:-1] + ")" result = self.cur.execute(query) return result.fetchall() except Exception as e: print(e) print("Error query: ",query) def add(self, crawl_records): """ Add crawl_records (list) obtained from arxiv_crawlers A record is a list of 8 columns: [topic, id, updated, published, title, author, link, summary] Return the final length of the database table """ results = "" for record in crawl_records: try: query = """insert into arxivsql values("{}","{}","{}","{}","{}","{}","{}")""".format( record[1][21:], record[0], record[4].replace('"',"'"), authors_list_to_str(record[5]), record[2][:10], record[3][:10], record[6] ) self.cur.execute(query) self.con.commit() except Exception as e: result+=str(e) result+="\n" + query + "\n" finally: return results # instance of ArxivSQL sqldb = ArxivSQL() class ArxivChroma: """ Create an interface to arxivdb, which only support query and addition. This interface do not support edition and deletion procedures. """ def __init__(self, table="arxiv_records", name="arxivdb/"): self.client = chromadb.PersistentClient(name) self.model = embedding_model self.collection = self.client.get_or_create_collection(table, embedding_function=JinaAIEmbeddingFunction( model = self.model )) def query_relevant(self, keywords, query_texts, n_results=3): """ Perform a query using a list of keywords (str), or using a relavant string """ contains = [] for keyword in keywords: contains.append({"$contains":keyword.lower()}) return self.collection.query( query_texts=query_texts, where_document={ "$or":contains }, n_results=n_results, ) def query_exact(self, id): ids = ["{}_{}".format(id,j) for j in range(0,10)] return self.collection.get(ids=ids) def add(self, crawl_records): """ Add crawl_records (list) obtained from arxiv_crawlers A record is a list of 8 columns: [topic, id, updated, published, title, author, link, summary] Return the final length of the database table """ for record in crawl_records: embed_text = """ Topic: {}, Title: {}, Summary: {} """.format(record[0],record[4],record[7]) chunks = chunk_texts(embed_text) ids = [record[1][21:]+"_"+str(j) for j in range(len(chunks))] paper_ids = [{"paper_id":record[1][21:]} for _ in range(len(chunks))] self.collection.add( documents = chunks, metadatas=paper_ids, ids = ids ) return self.collection.count() # instance of ArxivChroma db = ArxivChroma()