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 pixeltable.functions.fireworks import chat_completions as f_chat_completions from pixeltable.functions.mistralai import chat_completions from gradio.themes import Monochrome import os import getpass """## Store OpenAI API Key""" if 'OPENAI_API_KEY' not in os.environ: os.environ['OPENAI_API_KEY'] = getpass.getpass('OpenAI API key:') if 'FIREWORKS_API_KEY' not in os.environ: os.environ['FIREWORKS_API_KEY'] = getpass.getpass('Fireworks API Key:') if 'MISTRAL_API_KEY' not in os.environ: os.environ['MISTRAL_API_KEY'] = getpass.getpass('Mistral AI API Key:') """## Creating UDFs: Embedding and Prompt Functions""" # 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}''' """Gradio Application""" def process_files(ground_truth_file, pdf_files, chunk_limit, chunk_separator, show_question, show_correct_answer, show_gpt4omini, show_llamav3p23b, show_mistralsmall, progress=gr.Progress()): # Ensure a clean slate for the demo by removing and recreating the 'rag_demo' directory progress(0, desc="Initializing...") pxt.drop_dir('rag_demo', force=True) pxt.create_dir('rag_demo') # Process the ground truth file, which contains questions and correct answers # Import as CSV or Excel depending on the file extension if ground_truth_file.name.endswith('.csv'): queries_t = pxt.io.import_csv('rag_demo.queries', ground_truth_file.name) else: queries_t = pxt.io.import_excel('rag_demo.queries', ground_truth_file.name) progress(0.2, desc="Processing documents...") # Create a table to store the uploaded PDF documents documents_t = pxt.create_table( 'rag_demo.documents', {'document': pxt.DocumentType()} ) # Insert the PDF files into the documents table documents_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( 'rag_demo.chunks', documents_t, iterator=DocumentSplitter.create( document=documents_t.document, separators=chunk_separator, limit=chunk_limit if chunk_separator in ["token_limit", "char_limit"] else None ) ) progress(0.4, desc="Generating embeddings...") # Add an embedding index to the chunks for similarity search chunks_t.add_embedding_index('text', string_embed=e5_embed) # Define a query function to retrieve the top-k most similar chunks for a given question @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 queries table for context retrieval and prompt creation 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 the OpenAI API, including system instructions and user prompt msgs = [ { 'role': 'system', 'content': 'Read the following passages and answer the question based on their contents.' }, { 'role': 'user', 'content': queries_t.prompt } ] progress(0.6, desc="Querying models...") # Add OpenAI response column queries_t['response'] = openai.chat_completions( model='gpt-4o-mini-2024-07-18', messages=msgs, max_tokens=300, top_p=0.9, temperature=0.7 ) # Create a table in Pixeltable and pick a model hosted on Anthropic with some parameters queries_t['response_2'] = f_chat_completions( messages=msgs, model='accounts/fireworks/models/llama-v3p2-3b-instruct', # These parameters are optional and can be used to tune model behavior: max_tokens=300, top_p=0.9, temperature=0.7 ) queries_t['response_3'] = chat_completions( messages=msgs, model='mistral-small-latest', # These parameters are optional and can be used to tune model behavior: max_tokens=300, top_p=0.9, temperature=0.7 ) # Extract the answer text from the API response queries_t['gpt4omini'] = queries_t.response.choices[0].message.content queries_t['llamav3p23b'] = queries_t.response_2.choices[0].message.content queries_t['mistralsmall'] = queries_t.response_3.choices[0].message.content # Prepare the output dataframe with selected columns columns_to_show = [] if show_question: columns_to_show.append(queries_t.question) if show_correct_answer: columns_to_show.append(queries_t.correct_answer) if show_gpt4omini: columns_to_show.append(queries_t.gpt4omini) if show_llamav3p23b: columns_to_show.append(queries_t.llamav3p23b) if show_mistralsmall: columns_to_show.append(queries_t.mistralsmall) df_output = queries_t.select(*columns_to_show).collect().to_pandas() try: # Return the output dataframe for display return df_output except Exception as e: return f"An error occurred: {str(e)}", None def save_dataframe_as_csv(data): print(f"Type of data: {type(data)}") if isinstance(data, pd.DataFrame): print(f"Shape of DataFrame: {data.shape}") if isinstance(data, pd.DataFrame) and not data.empty: filename = f"results_{uuid.uuid4().hex[:8]}.csv" filepath = os.path.join('tmp', filename) os.makedirs('tmp', exist_ok=True) data.to_csv(filepath, index=False) return filepath return None # Gradio interface with gr.Blocks(theme=Monochrome) as demo: gr.Markdown( """
Pixeltable is a declarative interface for working with text, images, embeddings, and even video, enabling you to store, transform, index, and iterate on data.
""" ) # Add the disclaimer gr.HTML( """