hanicker/pdfChatter

1import urllib.request2import fitz3import re4import numpy as np5import tensorflow_hub as hub6import openai7import gradio as gr8import os9from sklearn.neighbors import NearestNeighbors10 11def download_pdf(url, output_path):12    urllib.request.urlretrieve(url, output_path)13 14 15def preprocess(text):16    text = text.replace('\n', ' ')17    text = re.sub('\s+', ' ', text)18    return text19 20 21def pdf_to_text(path, start_page=1, end_page=None):22    doc = fitz.open(path)23    total_pages = doc.page_count24 25    if end_page is None:26        end_page = total_pages27 28    text_list = []29 30    for i in range(start_page-1, end_page):31        text = doc.load_page(i).get_text("text")32        text = preprocess(text)33        text_list.append(text)34 35    doc.close()36    return text_list37 38 39def text_to_chunks(texts, word_length=150, start_page=1):40    text_toks = [t.split(' ') for t in texts]41    page_nums = []42    chunks = []43    44    for idx, words in enumerate(text_toks):45        for i in range(0, len(words), word_length):46            chunk = words[i:i+word_length]47            if (i+word_length) > len(words) and (len(chunk) < word_length) and (48                len(text_toks) != (idx+1)):49                text_toks[idx+1] = chunk + text_toks[idx+1]50                continue51            chunk = ' '.join(chunk).strip()52            chunk = f'[Page no. {idx+start_page}]' + ' ' + '"' + chunk + '"'53            chunks.append(chunk)54    return chunks55 56class SemanticSearch:57    58    def __init__(self):59        self.use = hub.load('https://tfhub.dev/google/universal-sentence-encoder/4')60        self.fitted = False61    62    63    def fit(self, data, batch=1000, n_neighbors=5):64        self.data = data65        self.embeddings = self.get_text_embedding(data, batch=batch)66        n_neighbors = min(n_neighbors, len(self.embeddings))67        self.nn = NearestNeighbors(n_neighbors=n_neighbors)68        self.nn.fit(self.embeddings)69        self.fitted = True70    71    72    def __call__(self, text, return_data=True):73        inp_emb = self.use([text])74        neighbors = self.nn.kneighbors(inp_emb, return_distance=False)[0]75        76        if return_data:77            return [self.data[i] for i in neighbors]78        else:79            return neighbors80    81    82    def get_text_embedding(self, texts, batch=1000):83        embeddings = []84        for i in range(0, len(texts), batch):85            text_batch = texts[i:(i+batch)]86            emb_batch = self.use(text_batch)87            embeddings.append(emb_batch)88        embeddings = np.vstack(embeddings)89        return embeddings90 91 92 93def load_recommender(path, start_page=1):94    global recommender95    texts = pdf_to_text(path, start_page=start_page)96    chunks = text_to_chunks(texts, start_page=start_page)97    recommender.fit(chunks)98    return 'Corpus Loaded.'99 100def generate_text(openAI_key,prompt, engine="text-davinci-003"):101    openai.api_key = openAI_key102    completions = openai.Completion.create(103        engine=engine,104        prompt=prompt,105        max_tokens=512,106        n=1,107        stop=None,108        temperature=0.7,109    )110    message = completions.choices[0].text111    return message112 113def generate_answer(question,openAI_key):114    topn_chunks = recommender(question)115    prompt = ""116    prompt += 'search results:\n\n'117    for c in topn_chunks:118        prompt += c + '\n\n'119        120    prompt += "Instructions: Compose a comprehensive reply to the query using the search results given. "\121              "Cite each reference using [ Page Number] notation (every result has this number at the beginning). "\122              "Citation should be done at the end of each sentence. If the search results mention multiple subjects "\123              "with the same name, create separate answers for each. Only include information found in the results and "\124              "don't add any additional information. Make sure the answer is correct and don't output false content. "\125              "If the text does not relate to the query, simply state 'Text Not Found in PDF'. Ignore outlier "\126              "search results which has nothing to do with the question. Only answer what is asked. The "\127              "answer should be short and concise. Answer step-by-step. \n\nQuery: {question}\nAnswer: "128    129    prompt += f"Query: {question}\nAnswer:"130    answer = generate_text(openAI_key, prompt,"text-davinci-003")131    return answer132 133 134def question_answer(url, file, question,openAI_key):135    if openAI_key.strip()=='':136        return '[ERROR]: Please enter you Open AI Key. Get your key here : https://platform.openai.com/account/api-keys'137    if url.strip() == '' and file == None:138        return '[ERROR]: Both URL and PDF is empty. Provide atleast one.'139    140    if url.strip() != '' and file != None:141        return '[ERROR]: Both URL and PDF is provided. Please provide only one (eiter URL or PDF).'142 143    if url.strip() != '':144        glob_url = url145        download_pdf(glob_url, 'corpus.pdf')146        load_recommender('corpus.pdf')147 148    else:149        old_file_name = file.name150        file_name = file.name151        file_name = file_name[:-12] + file_name[-4:]152        os.rename(old_file_name, file_name)153        load_recommender(file_name)154 155    if question.strip() == '':156        return '[ERROR]: Question field is empty'157 158    return generate_answer(question,openAI_key)159 160 161recommender = SemanticSearch()162 163title = 'PDF GPT'164description = """ PDF GPT allows you to chat with your PDF file using Universal Sentence Encoder and Open AI. It gives hallucination free response than other tools as the embeddings are better than OpenAI. The returned response can even cite the page number in square brackets([]) where the information is located, adding credibility to the responses and helping to locate pertinent information quickly."""165 166with gr.Blocks() as demo:167 168    gr.Markdown(f'<center><h1>{title}</h1></center>')169    gr.Markdown(description)170 171    with gr.Row():172        173        with gr.Group():174            gr.Markdown(f'<p style="text-align:center">Get your Open AI API key <a href="https://platform.openai.com/account/api-keys">here</a></p>')175            openAI_key=gr.Textbox(label='Enter your OpenAI API key here')176            url = gr.Textbox(label='Enter PDF URL here')177            gr.Markdown("<center><h4>OR<h4></center>")178            file = gr.File(label='Upload your PDF/ Research Paper / Book here', file_types=['.pdf'])179            question = gr.Textbox(label='Enter your question here')180            btn = gr.Button(value='Submit')181            btn.style(full_width=True)182 183        with gr.Group():184            answer = gr.Textbox(label='The answer to your question is :')185 186        btn.click(question_answer, inputs=[url, file, question,openAI_key], outputs=[answer])187#openai.api_key = os.getenv('Your_Key_Here') 188demo.launch()