import openai import faiss import numpy as np import pickle from tqdm import tqdm import argparse import os from PyPDF2 import PdfReader class Paper(object): def __init__(self, pdf_obj: PdfReader) -> None: self._pdf_obj = pdf_obj self._paper_meta = self._pdf_obj.metadata self.texts = [] def iter_pages(self, iter_text_len: int = 1000): page_idx = 0 for page in self._pdf_obj.pages: txt = page.extract_text() for i in range((len(txt) // iter_text_len) + 1): yield page_idx, i, txt[i * iter_text_len:(i + 1) * iter_text_len] page_idx += 1 def get_texts(self): for (page_idx, part_idx, text) in self.iter_pages(): self.texts.append(text.strip()) return self.texts def create_embeddings(inputs): """Create embeddings for the provided input.""" # input = ['ddd','aaa','ccccccccccccccc','ddddd'] result = [] tokens = 0 def get_embedding(input_slice): input_slice = [input_slice] embedding = openai.Embedding.create(model="text-embedding-ada-002", input=input_slice) return [(text, data.embedding) for text, data in zip(input_slice, embedding.data)], embedding.usage.total_tokens for i in range(0,len(inputs)): ebd, tk = get_embedding(inputs[i]) tokens += tk result.extend(ebd) return result, tokens def create_embedding(text): """Create an embedding for the provided text.""" embedding = openai.Embedding.create(model="text-embedding-ada-002", input=text) return text, embedding.data[0].embedding class QA(): def __init__(self,data_embe) -> None: d = 1536 index = faiss.IndexFlatL2(d) embe = np.array([emm[1] for emm in data_embe]) data = [emm[0] for emm in data_embe] index.add(embe) #所有emdding self.index = index #所有文字 self.data = data print("now all data is:\n",self.data) def __call__(self, query): embedding = create_embedding(query) #输出与用户的问题相关的文字 context = self.get_texts(embedding[1]) #将用户的问题和涉及的文字告诉gpt,并将答案返回 answer = self.completion(query,context) return answer,context def get_texts(self,embeding,limit=5): _,text_index = self.index.search(np.array([embeding]),limit) context = [] for i in list(text_index[0]): context.extend(self.data[i:i+2]) # context = [self.data[i] for i in list(text_index[0])] #输出与用户的问题相关的文字 return context def completion(self,query, context): """Create a completion.""" # lens = [len(text) for text in context] # maximum = 3000 # for index, l in enumerate(lens): # maximum -= l # if maximum < 0: # context = context[:index + 1] # print("超过最大长度,截断到前", index + 1, "个片段") # break text = "\n".join(f"{index}. {text}" for index, text in enumerate(context)) response = openai.ChatCompletion.create( model="gpt-3.5-turbo", messages=[ {'role': 'system', 'content': f'你是一个有帮助的AI文章助手,从下文中提取有用的内容进行回答,不能回答不在下文提到的内容,相关性从高到底排序:\n\n{text}'}, {'role': 'user', 'content': query}, ], ) print("使用的tokens:", response.usage.total_tokens) return response.choices[0].message.content if __name__ == '__main__': parser = argparse.ArgumentParser(description="Document QA") parser.add_argument("--input_file", default="slimming-pages-1.pdf", dest="input_file", type=str,help="输入文件路径") # parser.add_argument("--file_embeding", default="input_embed.pkl", dest="file_embeding", type=str,help="文件embeding文件路径") parser.add_argument("--print_context", action='store_true',help="是否打印上下文") args = parser.parse_args() # if os.path.isfile(args.file_embeding): # data_embe = pickle.load(open(args.file_embeding,'rb')) # else: # with open(args.input_file,'r',encoding='utf-8') as f: # texts = f.readlines() # #按照行对文章进行切割 # texts = [text.strip() for text in texts if text.strip()] # data_embe,tokens = create_embeddings(texts) # pickle.dump(data_embe,open(args.file_embeding,'wb')) # print("文本消耗 {} tokens".format(tokens)) paper = Paper(args.input_file) all_texts = paper.get_texts() data_embe, tokens = create_embeddings(all_texts) print("全部文本消耗 {} tokens".format(tokens)) qa =QA(data_embe) limit = 10 while True: query = input("请输入查询(help可查看指令):") if query == "quit": break elif query.startswith("limit"): try: limit = int(query.split(" ")[1]) print("已设置limit为", limit) except Exception as e: print("设置limit失败", e) continue elif query == "help": print("输入limit [数字]设置limit") print("输入quit退出") continue answer,context = qa(query) if args.print_context: print("已找到相关片段:") for text in context: print('\t', text) print("=====================================") print("回答如下\n\n") print(answer.strip()) print("=====================================")