Spaces:

thatsean
/

chatpdf

Runtime error

Create app.py

9be6285 verified about 2 years ago

7.07 kB

import os
import gradio as gr
import tempfile
from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.chat_models import ChatOpenAI
from langchain.chains import ConversationalRetrievalChain

# Set your OpenAI API key
os.environ["OPENAI_API_KEY"] = "your-api-key-here"

class RAGChatbot:
    def __init__(self):
        self.vector_store = None
        self.qa_chain = None

    def process_pdf(self, pdf_file):
        # Save the uploaded file temporarily
        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
            temp_file.write(pdf_file.read())
            temp_path = temp_file.name

        # Load the PDF
        loader = PyPDFLoader(temp_path)
        documents = loader.load()

        # Split the documents into chunks
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
        splits = text_splitter.split_documents(documents)

        # Create embeddings and vector store
        embeddings = HuggingFaceEmbeddings()
        self.vector_store = FAISS.from_documents(splits, embeddings)

        # Create the conversational chain
        llm = ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo")
        self.qa_chain = ConversationalRetrievalChain.from_llm(
            llm, retriever=self.vector_store.as_retriever(), return_source_documents=True
        )

        # Clean up the temporary file
        os.unlink(temp_path)

        return "PDF processed successfully. You can now ask questions about the document."

    def ask_question(self, question, history):
        if not self.qa_chain:
            return "Please upload a PDF document first.", []

        # Prepare the chat history
        chat_history = [(h[0], h[1]) for h in history]

        # Get the answer from the QA chain
        result = self.qa_chain({"question": question, "chat_history": chat_history})
        answer = result["answer"]
        source_docs = result["source_documents"]

        # Extract page numbers and content for highlighting
        highlights = []
        for doc in source_docs:
            page_number = doc.metadata.get("page", 0) + 1  # PyPDFLoader uses 0-based indexing
            content = doc.page_content
            highlights.append({"page": page_number, "content": content})

        return answer, highlights

chatbot = RAGChatbot()

def user(message, history):
    return "", history + [[message, None]]

def bot(history):
    question = history[-1][0]
    answer, highlights = chatbot.ask_question(question, history[:-1])
    history[-1][1] = answer
    return history, highlights

with gr.Blocks() as demo:
    gr.Markdown("# RAG Chatbot for PDF Documents")
    
    with gr.Row():
        with gr.Column(scale=1):
            chatbot = gr.Chatbot(height=600)
            msg = gr.Textbox(label="Ask a question about the PDF")
            clear = gr.Button("Clear")
        
        with gr.Column(scale=1):
            pdf_display = gr.HTML('<div id="pdf-viewer" style="height: 600px;"></div>')
    
    file_output = gr.Textbox(label="PDF Processing Status")
    upload_button = gr.UploadButton("Upload PDF", file_types=[".pdf"])
    
    highlights = gr.JSON(visible=False)
    
    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, chatbot, [chatbot, highlights]
    )
    clear.click(lambda: None, None, chatbot, queue=False)
    upload_button.upload(chatbot.process_pdf, upload_button, file_output)

    gr.HTML("""
    <script src="https://cdnjs.cloudflare.com/ajax/libs/pdf.js/2.9.359/pdf.min.js"></script>
    <script>
    let pdfDoc = null;
    let pageNum = 1;
    let pageRendering = false;
    let pageNumPending = null;
    let scale = 1.5;
    let canvas = null;
    let ctx = null;
    function renderPage(num) {
        pageRendering = true;
        pdfDoc.getPage(num).then(function(page) {
            let viewport = page.getViewport({scale: scale});
            canvas.height = viewport.height;
            canvas.width = viewport.width;
            let renderContext = {
                canvasContext: ctx,
                viewport: viewport
            };
            let renderTask = page.render(renderContext);
            renderTask.promise.then(function() {
                pageRendering = false;
                if (pageNumPending !== null) {
                    renderPage(pageNumPending);
                    pageNumPending = null;
                }
            });
        });
    }
    function queueRenderPage(num) {
        if (pageRendering) {
            pageNumPending = num;
        } else {
            renderPage(num);
        }
    }
    function highlightText(page, text) {
        pdfDoc.getPage(page).then(function(pdfPage) {
            pdfPage.getTextContent().then(function(textContent) {
                let textItems = textContent.items;
                let viewport = pdfPage.getViewport({scale: scale});
                ctx.save();
                ctx.fillStyle = "yellow";
                for (let i = 0; i < textItems.length; i++) {
                    if (textItems[i].str.includes(text)) {
                        let tx = textItems[i].transform;
                        let pos = viewport.convertToViewportPoint(tx[4], tx[5]);
                        ctx.fillRect(pos[0], pos[1], textItems[i].width * scale, textItems[i].height * scale);
                    }
                }
                ctx.restore();
                renderPage(page);
            });
        });
    }
    gradio_config.components.forEach(component => {
        if (component.props.label === "Upload PDF") {
            component.props.onChange = (files) => {
                if (files.length > 0) {
                    let fileReader = new FileReader();
                    fileReader.onload = function() {
                        let typedarray = new Uint8Array(this.result);
                        pdfjsLib.getDocument({data: typedarray}).promise.then(function(pdf) {
                            pdfDoc = pdf;
                            canvas = document.createElement('canvas');
                            ctx = canvas.getContext('2d');
                            document.getElementById('pdf-viewer').innerHTML = '';
                            document.getElementById('pdf-viewer').appendChild(canvas);
                            renderPage(pageNum);
                        });
                    };
                    fileReader.readAsArrayBuffer(files[0]);
                }
            };
        } else if (component.props.elem_id === "highlights") {
            component.props.onChange = (highlights) => {
                if (highlights && highlights.length > 0) {
                    highlights.forEach(highlight => {
                        highlightText(highlight.page, highlight.content);
                    });
                }
            };
        }
    });
    </script>
    """)

demo.launch()

	import os
	import gradio as gr
	import tempfile
	from langchain.document_loaders import PyPDFLoader
	from langchain.text_splitter import RecursiveCharacterTextSplitter
	from langchain.embeddings import HuggingFaceEmbeddings
	from langchain.vectorstores import FAISS
	from langchain.chat_models import ChatOpenAI
	from langchain.chains import ConversationalRetrievalChain

	# Set your OpenAI API key
	os.environ["OPENAI_API_KEY"] = "your-api-key-here"

	class RAGChatbot:
	def __init__(self):
	self.vector_store = None
	self.qa_chain = None

	def process_pdf(self, pdf_file):
	# Save the uploaded file temporarily
	with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
	temp_file.write(pdf_file.read())
	temp_path = temp_file.name

	# Load the PDF
	loader = PyPDFLoader(temp_path)
	documents = loader.load()

	# Split the documents into chunks
	text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
	splits = text_splitter.split_documents(documents)

	# Create embeddings and vector store
	embeddings = HuggingFaceEmbeddings()
	self.vector_store = FAISS.from_documents(splits, embeddings)

	# Create the conversational chain
	llm = ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo")
	self.qa_chain = ConversationalRetrievalChain.from_llm(
	llm, retriever=self.vector_store.as_retriever(), return_source_documents=True
	)

	# Clean up the temporary file
	os.unlink(temp_path)

	return "PDF processed successfully. You can now ask questions about the document."

	def ask_question(self, question, history):
	if not self.qa_chain:
	return "Please upload a PDF document first.", []

	# Prepare the chat history
	chat_history = [(h[0], h[1]) for h in history]

	# Get the answer from the QA chain
	result = self.qa_chain({"question": question, "chat_history": chat_history})
	answer = result["answer"]
	source_docs = result["source_documents"]

	# Extract page numbers and content for highlighting
	highlights = []
	for doc in source_docs:
	page_number = doc.metadata.get("page", 0) + 1 # PyPDFLoader uses 0-based indexing
	content = doc.page_content
	highlights.append({"page": page_number, "content": content})

	return answer, highlights

	chatbot = RAGChatbot()

	def user(message, history):
	return "", history + [[message, None]]

	def bot(history):
	question = history[-1][0]
	answer, highlights = chatbot.ask_question(question, history[:-1])
	history[-1][1] = answer
	return history, highlights

	with gr.Blocks() as demo:
	gr.Markdown("# RAG Chatbot for PDF Documents")

	with gr.Row():
	with gr.Column(scale=1):
	chatbot = gr.Chatbot(height=600)
	msg = gr.Textbox(label="Ask a question about the PDF")
	clear = gr.Button("Clear")

	with gr.Column(scale=1):
	pdf_display = gr.HTML('<div id="pdf-viewer" style="height: 600px;"></div>')

	file_output = gr.Textbox(label="PDF Processing Status")
	upload_button = gr.UploadButton("Upload PDF", file_types=[".pdf"])

	highlights = gr.JSON(visible=False)

	msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
	bot, chatbot, [chatbot, highlights]
	)
	clear.click(lambda: None, None, chatbot, queue=False)
	upload_button.upload(chatbot.process_pdf, upload_button, file_output)

	gr.HTML("""
	<script src="https://cdnjs.cloudflare.com/ajax/libs/pdf.js/2.9.359/pdf.min.js"></script>
	<script>
	let pdfDoc = null;
	let pageNum = 1;
	let pageRendering = false;
	let pageNumPending = null;
	let scale = 1.5;
	let canvas = null;
	let ctx = null;

	function renderPage(num) {
	pageRendering = true;
	pdfDoc.getPage(num).then(function(page) {
	let viewport = page.getViewport({scale: scale});
	canvas.height = viewport.height;
	canvas.width = viewport.width;
	let renderContext = {
	canvasContext: ctx,
	viewport: viewport
	};
	let renderTask = page.render(renderContext);
	renderTask.promise.then(function() {
	pageRendering = false;
	if (pageNumPending !== null) {
	renderPage(pageNumPending);
	pageNumPending = null;
	}
	});
	});
	}

	function queueRenderPage(num) {
	if (pageRendering) {
	pageNumPending = num;
	} else {
	renderPage(num);
	}
	}

	function highlightText(page, text) {
	pdfDoc.getPage(page).then(function(pdfPage) {
	pdfPage.getTextContent().then(function(textContent) {
	let textItems = textContent.items;
	let viewport = pdfPage.getViewport({scale: scale});
	ctx.save();
	ctx.fillStyle = "yellow";
	for (let i = 0; i < textItems.length; i++) {
	if (textItems[i].str.includes(text)) {
	let tx = textItems[i].transform;
	let pos = viewport.convertToViewportPoint(tx[4], tx[5]);
	ctx.fillRect(pos[0], pos[1], textItems[i].width * scale, textItems[i].height * scale);
	}
	}
	ctx.restore();
	renderPage(page);
	});
	});
	}

	gradio_config.components.forEach(component => {
	if (component.props.label === "Upload PDF") {
	component.props.onChange = (files) => {
	if (files.length > 0) {
	let fileReader = new FileReader();
	fileReader.onload = function() {
	let typedarray = new Uint8Array(this.result);
	pdfjsLib.getDocument({data: typedarray}).promise.then(function(pdf) {
	pdfDoc = pdf;
	canvas = document.createElement('canvas');
	ctx = canvas.getContext('2d');
	document.getElementById('pdf-viewer').innerHTML = '';
	document.getElementById('pdf-viewer').appendChild(canvas);
	renderPage(pageNum);
	});
	};
	fileReader.readAsArrayBuffer(files[0]);
	}
	};
	} else if (component.props.elem_id === "highlights") {
	component.props.onChange = (highlights) => {
	if (highlights && highlights.length > 0) {
	highlights.forEach(highlight => {
	highlightText(highlight.page, highlight.content);
	});
	}
	};
	}
	});
	</script>
	""")

	demo.launch()