|
import gradio as gr |
|
from pipeline import preprocessing_pipeline, conversational_rag |
|
from pipeline import system_message, user_message |
|
from haystack.dataclasses import ChatMessage |
|
import time |
|
import os |
|
|
|
def process_files_into_docs(files,progress=gr.Progress()): |
|
if isinstance(files, dict): |
|
files = [files] |
|
if not files: |
|
return 'No file uploaded!' |
|
|
|
preprocessing_pipeline.run({'file_type_router': {'sources': files}}) |
|
|
|
return "Database created🤗🤗" |
|
|
|
|
|
def rag(history,question): |
|
|
|
if history is None: |
|
history=[] |
|
messages = [system_message, user_message] |
|
res = conversational_rag.run( |
|
data = {'query_rephrase_prompt_builder' : {'query': question}, |
|
'prompt_builder': {'template': messages, 'query': question}, |
|
'memory_joiner': {'values': [ChatMessage.from_user(question)]}}, |
|
include_outputs_from=['llm','query_rephrase_llm']) |
|
|
|
bot_message = res['llm']['replies'][0].content |
|
|
|
streamed_message = "" |
|
for token in bot_message.split(): |
|
streamed_message += f"{token} " |
|
yield history + [(question, streamed_message.strip())], " " |
|
time.sleep(0.05) |
|
|
|
history.append((question,bot_message)) |
|
|
|
yield history, " " |
|
|
|
EXAMPLE_FILE = "RAG Survey.pdf" |
|
|
|
with gr.Blocks(theme=gr.themes.Soft())as demo: |
|
|
|
gr.HTML("<center><h1>TalkToFiles - Query your documents! 📂📄</h1><center>") |
|
gr.Markdown("""##### This AI chatbot🤖 can help you chat with your documents. Can upload <b>Text(.txt), PDF(.pdf) and Markdown(.md)</b> files.\ |
|
<b>Please do not upload confidential documents.</b>""") |
|
with gr.Row(): |
|
with gr.Column(scale=86): |
|
gr.Markdown("""#### ***Step 1 - Upload Documents and Initialize RAG pipeline***</br> |
|
Can upload Multiple documents""") |
|
with gr.Row(): |
|
file_input = gr.File(label='Upload Files', file_count='multiple',file_types=['.pdf', '.txt', '.md'],interactive=True) |
|
with gr.Row(): |
|
process_files = gr.Button('Create Document store') |
|
with gr.Row(): |
|
result = gr.Textbox(label="Document store", value='Document store not initialized') |
|
|
|
process_files.click(fn=process_files_into_docs, inputs=file_input, outputs=result ,show_progress=True) |
|
|
|
def load_example(): |
|
return [EXAMPLE_FILE] |
|
|
|
with gr.Row(): |
|
gr.Examples( |
|
examples=[[EXAMPLE_FILE]], |
|
inputs=file_input, |
|
examples_per_page=1, |
|
label="Click to upload an example" |
|
).dataset.click(fn=load_example, inputs=[], outputs=file_input) |
|
|
|
|
|
|
|
with gr.Column(scale=200): |
|
gr.Markdown("""#### ***Step 2 - Chat with your docs*** """) |
|
chatbot = gr.Chatbot(label='ChatBot') |
|
user_input = gr.Textbox(label='Enter your query', placeholder='Type here...') |
|
|
|
with gr.Row(): |
|
submit_button = gr.Button("Submit") |
|
clear_btn = gr.ClearButton([user_input, chatbot], value='Clear') |
|
submit_button.click(rag, inputs=[chatbot, user_input], outputs=[chatbot, user_input]) |
|
|
|
|
|
demo.launch() |
|
|
|
|
