Create app.py

app.py

@@ -0,0 +1,155 @@
+import gradio as gr
+import pandas as pd
+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
+import os
+
+# Ensure a clean slate for the demo
+pxt.drop_dir('rag_demo', force=True)
+pxt.create_dir('rag_demo')
+
+# 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(ground_truth_file, pdf_files):
+    # Process ground truth file
+    if ground_truth_file.name.endswith('.csv'):
+        df = pd.read_csv(ground_truth_file.name)
+    else:
+        df = pd.read_excel(ground_truth_file.name)
+    
+    queries_t = pxt.create_table('rag_demo.queries', df)
+
+    # Process PDF files
+    documents_t = pxt.create_table(
+        'rag_demo.documents',
+        {'document': pxt.DocumentType()}
+    )
+    
+    for pdf_file in pdf_files:
+        documents_t.insert({'document': pdf_file.name})
+
+    # Create chunks view
+    chunks_t = pxt.create_view(
+        'rag_demo.chunks',
+        documents_t,
+        iterator=DocumentSplitter.create(
+            document=documents_t.document,
+            separators='token_limit',
+            limit=300
+        )
+    )
+
+    # Add embedding index
+    chunks_t.add_embedding_index('text', string_embed=e5_embed)
+
+    # Create top_k query
+    @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 queries_t
+    queries_t['question_context'] = chunks_t.top_k(queries_t.Question)
+    queries_t['prompt'] = create_prompt(
+        queries_t.question_context, queries_t.Question
+    )
+
+    # Prepare messages for OpenAI
+    messages = [
+        {
+            'role': 'system',
+            'content': 'Please read the following passages and answer the question based on their contents.'
+        },
+        {
+            'role': 'user',
+            'content': queries_t.prompt
+        }
+    ]
+
+    # Add OpenAI response column
+    queries_t['response'] = openai.chat_completions(
+        model='gpt-4-0125-preview', messages=messages
+    )
+    queries_t['answer'] = queries_t.response.choices[0].message.content
+
+    return "Files processed successfully!"
+
+def query_llm(question):
+    queries_t = pxt.get_table('rag_demo.queries')
+    chunks_t = pxt.get_table('rag_demo.chunks')
+
+    # Perform top-k lookup
+    context = chunks_t.top_k(question).collect()
+
+    # Create prompt
+    prompt = create_prompt(context, question)
+
+    # Prepare messages for OpenAI
+    messages = [
+        {
+            'role': 'system',
+            'content': 'Please read the following passages and answer the question based on their contents.'
+        },
+        {
+            'role': 'user',
+            'content': prompt
+        }
+    ]
+
+    # Get LLM response
+    response = openai.chat_completions(model='gpt-4-0125-preview', messages=messages)
+    answer = response.choices[0].message.content
+
+    # Add new row to queries_t
+    new_row = {'Question': question, 'answer': answer}
+    queries_t.insert([new_row])
+
+    # Return updated dataframe
+    return queries_t.select(queries_t.Question, queries_t.answer).collect()
+
+# Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# RAG Demo App")
+    
+    with gr.Row():
+        ground_truth_file = gr.File(label="Upload Ground Truth (CSV or XLSX)")
+        pdf_files = gr.File(label="Upload PDF Documents", file_count="multiple")
+    
+    process_button = gr.Button("Process Files")
+    process_output = gr.Textbox(label="Processing Output")
+    
+    question_input = gr.Textbox(label="Enter your question")
+    query_button = gr.Button("Query LLM")
+    
+    output_dataframe = gr.Dataframe(label="LLM Outputs")
+
+    process_button.click(process_files, inputs=[ground_truth_file, pdf_files], outputs=process_output)
+    query_button.click(query_llm, inputs=question_input, outputs=output_dataframe)
+
+if __name__ == "__main__":
+    demo.launch()