1st

Document_QA.py

@@ -6,40 +6,43 @@ import pickle
 from tqdm import tqdm
 import argparse
 import os
+from PyPDF2 import PdfReader
 
-def create_embeddings(input):
+class Paper(object):
+
+    def __init__(self, pdf_path) -> None:
+        self._pdf_obj = PdfReader(pdf_path)
+        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 = []
-    # limit about 1000 tokens per request
-    # 记录文章每行的长度
-    # 0 [100]
-    # 1 [200]
-    # 2 [4100]
-    # 3 [999]
-    lens = [len(text) for text in input]
-    query_len = 0
-    start_index = 0
     tokens = 0
 
     def get_embedding(input_slice):
+        input_slice = [input_slice]
         embedding = openai.Embedding.create(model="text-embedding-ada-002", input=input_slice)
-        #返回了(文字，embedding)和文字的token
         return [(text, data.embedding) for text, data in zip(input_slice, embedding.data)], embedding.usage.total_tokens
-    #将文字的数量按照4096切分成多块，每一块去计算一次embedding，如果不足4096则一次计算所有文本的embedding
-    for index, l in tqdm(enumerate(lens)):
-        query_len += l
-        if query_len > 4096:
-            ebd, tk = get_embedding(input[start_index:index + 1])
-            query_len = 0
-            start_index = index + 1
-            tokens += tk
-            result.extend(ebd)
-
-    if query_len > 0:
-        ebd, tk = get_embedding(input[start_index:])
+    
+    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):
@@ -58,33 +61,35 @@ class QA():
         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], limit)
+        context = self.get_texts(embedding[1])
         #将用户的问题和涉及的文字告诉gpt，并将答案返回
         answer = self.completion(query,context)
         return answer,context
-    def get_texts(self,embeding,limit):
+    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+5])
+            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]
+        # 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
+        # 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(
@@ -100,24 +105,30 @@ class QA():
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description="Document QA")
-    parser.add_argument("--input_file", default="input.txt", 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("--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))
+    # 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

app.py

@@ -4,6 +4,8 @@ import openai
 # from gpt_reader.prompt import BASE_POINTS
 from Document_QA import QA
 from Document_QA import create_embeddings
+from Document_QA import Paper
+from PyPDF2 import PdfReader
 
 class GUI:
     def __init__(self):
@@ -14,27 +16,17 @@ class GUI:
     #load pdf and create all embedings
     def pdf_init(self, api_key, pdf_path):
         openai.api_key = api_key
-        print("--------------pdf_path is:",pdf_path)
-        with open(pdf_path,'r',encoding='utf-8') as f:
-            texts = f.readlines()
-            #按照行对文章进行切割
-            texts = [text.strip() for text in texts if text.strip()]
-            self.all_embedding,self.tokens = create_embeddings(texts)
+        pdf_reader = PdfReader(pdf_path)
+        paper = Paper(pdf_reader)
+        all_texts = paper.get_texts()
+        self.all_embedding, self.tokens = create_embeddings(all_texts)
+        print("全部文本消耗 {} tokens".format(self.tokens))
+
     def get_answer(self, question):
         qa = QA(self.all_embedding)
         answer,context = qa(question)
         return answer.strip()
 
-    # def analyse(self, api_key, pdf_file):
-    #     self.session = PaperReader(api_key, points_to_focus=BASE_POINTS)
-    #     return self.session.read_pdf_and_summarize(pdf_file)
-
-    # def ask_question(self, question):
-    #     if self.session == "":
-    #         return "Please upload PDF file first!"
-    #     return self.session.question(question)
-
-
 with gr.Blocks() as demo:
     gr.Markdown(
         """
@@ -57,4 +49,4 @@ with gr.Blocks() as demo:
 
 if __name__ == "__main__":
     demo.title = "CHATGPT-PAPER-READER"
-    demo.launch(debug=True,share=True)  # add "share=True" to share CHATGPT-PAPER-READER app on Internet.
+    demo.launch()  # add "share=True" to share CHATGPT-PAPER-READER app on Internet.

requirements.txt

@@ -2,3 +2,4 @@ numpy
 faiss-cpu
 tqdm
 openai
+PyPDF2