#import csv import gradio as gr from langchain.embeddings.openai import OpenAIEmbeddings from langchain.text_splitter import CharacterTextSplitter from langchain.vectorstores import Chroma from langchain.chains import ConversationalRetrievalChain from langchain.chat_models import ChatOpenAI from langchain.document_loaders import PyPDFLoader from dotenv import load_dotenv import os import requests from io import BytesIO import fitz from PIL import Image #init load_dotenv() # Global variables COUNT, N = 0, 0 chat_history = [] chain = '' pdf_file = None pdf_url = None enable_box = gr.Textbox.update(value=None, placeholder='Upload your OpenAI API key', interactive=True) disable_box = gr.Textbox.update(value='OpenAI API key is Set', interactive=False) openai_api_key = os.getenv('OPENAI_API_KEY') os.environ['OPENAI_API_KEY'] = openai_api_key # Function to set the OpenAI API key def set_apikey(api_key): os.environ['OPENAI_API_KEY'] = api_key return disable_box # Function to enable the API key input box def enable_api_box(): return enable_box # Function to add text to the chat history def add_text(history, text): if not text: raise gr.Error('Enter text') history = history + [(text, '')] return history # Function to process the PDF file and create a conversation chain def process_file(file): if 'OPENAI_API_KEY' not in os.environ: raise gr.Error('Upload your OpenAI API key') # loader = PyPDFLoader(file.name) loader = PyPDFLoader(file) documents = loader.load() embeddings = OpenAIEmbeddings() pdfsearch = Chroma.from_documents(documents, embeddings) chain = ConversationalRetrievalChain.from_llm(ChatOpenAI(temperature=0.3), retriever=pdfsearch.as_retriever(search_kwargs={"k": 1}), return_source_documents=True) return chain # Function to generate a response based on the chat history and query # def generate_response(history, query, btn): # global COUNT, N, chat_history, chain # if not btn: # raise gr.Error(message='Upload a PDF') # if COUNT == 0: # chain = process_file(btn) # COUNT += 1 # result = chain({"question": query, 'chat_history': chat_history}, return_only_outputs=True) # chat_history += [(query, result["answer"])] # N = list(result['source_documents'][0])[1][1]['page'] # for char in result['answer']: # history[-1][-1] += char # yield history, '' def generate_response(history,query,btn): global COUNT, N, chat_history, chain, pdf_file # if not pdf_file and btn: # raise gr.Error(message='Add a url that ends with .pdf') if COUNT == 0: if not pdf_file: chain = process_file(btn.name) else: chain = process_file(pdf_url) result = chain({"question": query, 'chat_history': chat_history}, return_only_outputs=True) chat_history += [(query, result["answer"])] N = list(result['source_documents'][0])[1][1]['page'] for char in result['answer']: history[-1][-1] += char yield history, '' # Function to render a specific page of a PDF file as an image def render_file(file): set_gobals_to_none global N doc = fitz.open(file.name) page = doc[N] # Render the page as a PNG image with a resolution of 300 DPI pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72)) image = Image.frombytes('RGB', [pix.width, pix.height], pix.samples) return image def render_pdf_url(url): global N, pdf_url, pdf_file pdf_url = url response = requests.get(url) if response.status_code == 200: pdf_f = BytesIO(response.content) doc = fitz.open(stream=pdf_f, filetype="pdf") pdf_file = doc page = doc[N] pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72)) image = Image.frombytes('RGB', [pix.width, pix.height], pix.samples) return image # def render_pdf_url(url): # global N, pdf_url, pdf_file # pdf_url = url # response = requests.get(url) # if response.status_code == 200: # pdf_f = BytesIO(response.content) # doc = fitz.open(stream=pdf_f, filetype="pdf") # pdf_file = doc # # Check if N is within the valid range of pages # if N < 0 or N >= len(doc): # raise ValueError(f"Page number {N} is not within the valid range of pages (0 to {len(doc) - 1})") # page = doc[N] # pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72)) # # Use Image.frombuffer instead of Image.frombytes # image = Image.frombuffer('RGB', [pix.width, pix.height], pix.samples, 'raw', 'RGB', 0, 1) # return image def rerender_pdf(url): global N, pdf_file page = pdf_file[N] pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72)) image = Image.frombytes('RGB', [pix.width, pix.height], pix.samples) return image def choose_render(url): global pdf_file, pdf_url if not pdf_file: return render_file(url) else: return rerender_pdf(url) def set_gobals_to_none(): global pdf_file, pdf_url pdf_file,pdf_url = None # Gradio application setup with gr.Blocks(css= "footer {visibility: hidden}", title="DocDialogue", analytics_enabled=True) as demo: # Create a Gradio block gr.Markdown("Have dialogue with your pdf documents") with gr.Column(): # with gr.Row(): # with gr.Column(scale=0.8): # api_key = gr.Textbox( # placeholder='Enter OpenAI API key', # show_label=False, # interactive=True # ).style(container=False) # with gr.Column(scale=0.2): # change_api_key = gr.Button('Change Key') with gr.Row(): with gr.Column(scale=0.8): pdf_url = gr.Textbox(placeholder="Enter PDF url",show_label=False, interactive=True) with gr.Column(scale=0.2): pdf_url_btn = gr.Button(size='lg',value='Submit') with gr.Row(): chatbot = gr.Chatbot(value=[], elem_id='chatbot', height=680 ,label="Chat") show_img = gr.Image(label='Upload PDF', tool='select', height=680) with gr.Row(): with gr.Column(scale=0.70): txt = gr.Textbox( show_label=False, placeholder="Enter text and press enter", container=False ) with gr.Column(scale=0.15): submit_btn = gr.Button('Submit') with gr.Column(scale=0.15): btn = gr.UploadButton("📁 Upload a PDF", file_types=[".pdf"]) # Set up event handlers # Event handler for submitting the OpenAI API key # api_key.submit(fn=set_apikey, inputs=[api_key], outputs=[api_key]) # Event handler for changing the API key # change_api_key.click(fn=enable_api_box, outputs=[api_key]) # Event handler for uploading a PDF btn.upload(fn=render_file, inputs=[btn], outputs=[show_img]) pdf_url_btn.click(fn=render_pdf_url, inputs=[pdf_url], outputs=[show_img]) # Event handler for submitting text and generating response submit_btn.click( fn=add_text, inputs=[chatbot, txt], outputs=[chatbot], queue=False ).success( fn=generate_response, inputs=[chatbot, txt, btn], outputs=[chatbot, txt] ).success( fn=choose_render, inputs=[btn], outputs=[show_img] ) demo.queue() if __name__ == "__main__": demo.launch() # Run # gradio chatwithpdf.py