Update app.py

app.py

@@ -26,13 +26,6 @@ if 'FIREWORKS_API_KEY' not in os.environ:
 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:
@@ -87,8 +80,11 @@ def process_files(ground_truth_file, pdf_files, chunk_limit, chunk_separator, sh
 
     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)
+    chunks_t.add_embedding_index(
+        'text',
+        idx_name='minilm_idx',
+        string_embed=sentence_transformer.using(model_id='sentence-transformers/all-MiniLM-L12-v2')
+    )
 
     # Define a query function to retrieve the top-k most similar chunks for a given question
     @chunks_t.query
@@ -101,10 +97,10 @@ def process_files(ground_truth_file, pdf_files, chunk_limit, chunk_separator, sh
       )
 
     # 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.add_computed_column(question_context=chunks_t.queries.top_k(queries_t.question))
+    queries_t.add_computed_column(prompt=create_prompt(
         queries_t.question_context, queries_t.question
-    )
+    ))
 
     # Prepare messages for the OpenAI API, including system instructions and user prompt
     msgs = [
@@ -121,37 +117,37 @@ def process_files(ground_truth_file, pdf_files, chunk_limit, chunk_separator, sh
     progress(0.6, desc="Querying models...")
 
     # Add OpenAI response column
-    queries_t['response'] = openai.chat_completions(
+    queries_t.add_computed_column(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(
+    queries_t.add_computed_column(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(
+    queries_t.add_computed_column(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
+    queries_t.add_computed_column(gpt4omini=queries_t.response.choices[0].message.content)
+    queries_t.add_computed_column(llamav3p23b=queries_t.response_2.choices[0].message.content)
+    queries_t.add_computed_column(mistralsmall=queries_t.response_3.choices[0].message.content)
 
     # Prepare the output dataframe with selected columns
     columns_to_show = []
@@ -291,4 +287,4 @@ with gr.Blocks(theme=Monochrome) as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(show_api=False)
+    demo.launch()