import gradio as gr
import pandas as pd
import io
import base64
import uuid
import pixeltable as pxt
from pixeltable.iterators import DocumentSplitter
import numpy as np
from pixeltable.functions.huggingface import sentence_transformer
from pixeltable.functions import openai
from gradio.themes import Monochrome
import os
import getpass
# Store API keys
if 'OPENAI_API_KEY' not in os.environ:
os.environ['OPENAI_API_KEY'] = getpass.getpass('OpenAI API key:')
# Set up embedding function
@pxt.expr_udf
def e5_embed(text: str) -> np.ndarray:
return sentence_transformer(text, model_id='intfloat/e5-large-v2')
# Create prompt function
@pxt.udf
def create_prompt(top_k_list: list[dict], question: str) -> str:
concat_top_k = '\n\n'.join(
elt['text'] for elt in reversed(top_k_list)
)
return f'''
PASSAGES:
{concat_top_k}
QUESTION:
{question}'''
def process_files(pdf_files, chunk_limit, chunk_separator):
# Initialize Pixeltable
pxt.drop_dir('chatbot_demo', force=True)
pxt.create_dir('chatbot_demo')
# Create a table to store the uploaded PDF documents
t = pxt.create_table(
'chatbot_demo.documents',
{'document': pxt.DocumentType(nullable=True),
'question': pxt.StringType(nullable=True)}
)
# Insert the PDF files into the documents table
t.insert({'document': file.name} for file in pdf_files if file.name.endswith('.pdf'))
# Create a view that splits the documents into smaller chunks
chunks_t = pxt.create_view(
'chatbot_demo.chunks',
t,
iterator=DocumentSplitter.create(
document=t.document,
separators=chunk_separator,
limit=chunk_limit if chunk_separator in ["token_limit", "char_limit"] else None,
metadata='title,heading,sourceline'
)
)
# Add an embedding index to the chunks for similarity search
chunks_t.add_embedding_index('text', string_embed=e5_embed)
@chunks_t.query
def top_k(query_text: str):
sim = chunks_t.text.similarity(query_text)
return (
chunks_t.order_by(sim, asc=False)
.select(chunks_t.text, sim=sim)
.limit(5)
)
# Add computed columns to the table for context retrieval and prompt creation
t['question_context'] = chunks_t.top_k(t.question)
t['prompt'] = create_prompt(
t.question_context, t.question
)
# Prepare messages for the API
msgs = [
{
'role': 'system',
'content': 'Read the following passages and answer the question based on their contents.'
},
{
'role': 'user',
'content': t.prompt
}
]
# Add OpenAI response column
t['response'] = openai.chat_completions(
model='gpt-4o-mini-2024-07-18',
messages=msgs,
max_tokens=300,
top_p=0.9,
temperature=0.7
)
# Extract the answer text from the API response
t['gpt4omini'] = t.response.choices[0].message.content
return "Files processed successfully!"
def get_answer(msg):
t = pxt.get_table('chatbot_demo.documents')
chunks_t = pxt.get_table('chatbot_demo.chunks')
# Insert the question into the table
t.insert([{'question': msg}])
answer = t.select(t.gpt4omini).where(t.question == msg).collect()['gpt4omini'][0]
return answer
def respond(message, chat_history):
bot_message = get_answer(message)
chat_history.append((message, bot_message))
return "", chat_history
# Gradio interface
with gr.Blocks(theme=Monochrome()) as demo:
gr.Markdown(
"""
AI Chatbot With Retrieval-Augmented Generation (RAG)
"""
)
gr.HTML(
"""
Pixeltable is a declarative interface for working with text, images, embeddings, and even video, enabling you to store, transform, index, and iterate on data.
"""
)
with gr.Row():
with gr.Column():
with gr.Accordion("What This Demo Does", open = True):
gr.Markdown("""
This AI Chatbot application uses Retrieval-Augmented Generation (RAG) to provide intelligent responses based on the content of uploaded PDF documents. It allows users to:
1. Upload multiple PDF documents
2. Process and index the content of these documents
3. Ask questions about the content
4. Receive AI-generated answers that are grounded in the uploaded documents
""")
with gr.Column():
with gr.Accordion("How does it work?", open = True):
gr.Markdown("""
**Question Answering:**
- When a user asks a question, the system searches for the most relevant chunks of text from the uploaded documents.
- It then uses these relevant chunks as context for a large language model (LLM) to generate an answer.
- The LLM (in this case, GPT-4) formulates a response based on the provided context and the user's question.
**Pixeltable Integration:**
- Pixeltable is used to manage the document data, chunks, and embeddings efficiently.
- It provides a declarative interface for complex data operations, making it easier to build and maintain this RAG system.
""")
with gr.Row():
with gr.Column(scale=1):
pdf_files = gr.File(label="Upload PDF Documents", file_count="multiple")
chunk_limit = gr.Slider(minimum=100, maximum=500, value=300, step=5, label="Chunk Size Limit")
chunk_separator = gr.Dropdown(
choices=["token_limit", "char_limit", "sentence", "paragraph", "heading"],
value="token_limit",
label="Chunk Separator"
)
process_button = gr.Button("Process Files")
process_output = gr.Textbox(label="Processing Output")
with gr.Column(scale=2):
chatbot = gr.Chatbot(label="Chat History")
msg = gr.Textbox(label="Your Question", placeholder="Ask a question about the uploaded documents")
submit = gr.Button("Submit")
process_button.click(process_files, inputs=[pdf_files, chunk_limit, chunk_separator], outputs=[process_output])
submit.click(respond, inputs=[msg, chatbot], outputs=[msg, chatbot])
if __name__ == "__main__":
demo.launch()