File size: 2,004 Bytes
b5acc22 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | import fitz
import gradio as gr
from dotenv import load_dotenv
load_dotenv(override=True)
from rag import RAG
from utils import format_page
def add_message(history, message):
if message["files"] is not None and len(message["files"]) > 0:
file_path = message["files"][0]
history.append([file_path, ""])
else:
if message["text"] is not None:
history.append((message["text"], None))
return history
def bot(history, message, pdf_file_content):
rag = RAG()
if message["files"] is not None and len(message["files"]) > 0:
file_path = message["files"][0]
if file_path.endswith(".pdf"):
with fitz.open(file_path) as doc:
history[-1][1] = ""
for page in doc:
cont, page_content = format_page(page)
if cont:
continue
pdf_file_content += "\n\n" + page_content
history[-1][1] += "\n\n" + page_content
yield history, gr.MultimodalTextbox(value=None, interactive=False), pdf_file_content
rag.ingest_text(pdf_file_content)
else:
for answer in rag(message, stream=True):
history[-1][1] = answer
yield history, gr.MultimodalTextbox(value=None, interactive=False), pdf_file_content
with gr.Blocks() as demo:
pdf_file_content = gr.State("")
chatbot = gr.Chatbot(
[],
elem_id="chatbot",
bubble_full_width=False
)
chat_input = gr.MultimodalTextbox(interactive=True, file_types=[".pdf"], placeholder="Enter message or upload file...", show_label=False)
chat_msg = chat_input.submit(add_message, [chatbot, chat_input], chatbot)
bot_msg = chat_msg.then(bot, [chatbot, chat_input, pdf_file_content], [chatbot, chat_input, pdf_file_content], api_name="bot_response")
bot_msg.then(lambda: gr.MultimodalTextbox(interactive=True), None, [chat_input])
demo.queue()
demo.launch()
|