NCEE-Build-Lab/watsonx.ai_Vector_Embedding_Visualizer_MNB
0
1# /// script2# [tool.marimo.display]3# custom_css = ["./custom_header_font.css"]4# ///5 6import marimo7 8__generated_with = "0.13.0"9app = marimo.App(width="full")10 11with app.setup:12 # Initialization code that runs before all other cells13 import marimo as mo14 from typing import Dict, Optional, List, Union, Any15 from ibm_watsonx_ai import APIClient, Credentials16 from pathlib import Path17 import pandas as pd18 import mimetypes19 import requests20 import zipfile21 import tempfile22 import certifi23 import base6424 import polars25 import nltk26 import time27 import json28 import ast 29 import os30 import io31 import re32 33 def get_iam_token(api_key):34 return requests.post(35 'https://iam.cloud.ibm.com/identity/token',36 headers={'Content-Type': 'application/x-www-form-urlencoded'},37 data={'grant_type': 'urn:ibm:params:oauth:grant-type:apikey', 'apikey': api_key},38 verify=certifi.where()39 ).json()['access_token']40 41 def setup_task_credentials(client):42 # Get existing task credentials43 existing_credentials = client.task_credentials.get_details()44 45 # Delete existing credentials if any46 if "resources" in existing_credentials and existing_credentials["resources"]:47 for cred in existing_credentials["resources"]:48 cred_id = client.task_credentials.get_id(cred)49 client.task_credentials.delete(cred_id)50 51 # Store new credentials52 return client.task_credentials.store()53 54 def get_cred_value(key, creds_var_name="baked_in_creds", default=""): ### Helper for working with preset credentials55 """56 Helper function to safely get a value from a credentials dictionary.57 58 Args:59 key: The key to look up in the credentials dictionary.60 creds_var_name: The variable name of the credentials dictionary.61 default: The default value to return if the key is not found.62 63 Returns:64 The value from the credentials dictionary if it exists and contains the key,65 otherwise returns the default value.66 """67 # Check if the credentials variable exists in globals68 if creds_var_name in globals():69 creds_dict = globals()[creds_var_name]70 if isinstance(creds_dict, dict) and key in creds_dict:71 # Return empty string if the value is None72 value = creds_dict[key]73 return "" if value is None else value74 return default75 76@app.cell77def client_variables(client_instantiation_form):78 client_setup = client_instantiation_form.value or None79 80 ### Extract Credential Variables:81 if client_setup:82 wx_url = client_setup["wx_region"] if client_setup["wx_region"] else "EU"83 wx_api_key = client_setup["wx_api_key"].strip() if client_setup["wx_api_key"] else None84 os.environ["WATSONX_APIKEY"] = wx_api_key or ""85 86 project_id = client_setup["project_id"].strip() if client_setup["project_id"] else None87 space_id = client_setup["space_id"].strip() if client_setup["space_id"] else None88 else:89 os.environ["WATSONX_APIKEY"] = ""90 project_id = space_id = wx_api_key = wx_url = None91 return client_setup, project_id, space_id, wx_api_key, wx_url92 93 94# @app.cell95# def _(client_setup, wx_api_key):96# if client_setup:97# token = get_iam_token(wx_api_key)98# else:99# token = None100# return101 102@app.cell103def _():104 from baked_in_credentials.creds import credentials105 from base_variables import wx_regions, wx_platform_url106 from helper_functions.helper_functions import wrap_with_spaces, get_key_by_value, markdown_spacing107 return (108 credentials, 109 wx_regions, 110 wx_platform_url, 111 wrap_with_spaces, 112 get_key_by_value,113 markdown_spacing,114 )115 116@app.cell117def client_instantiation(118 APIClient,119 Credentials,120 client_setup,121 project_id,122 space_id,123 wx_api_key,124 wx_url,125):126 ### Instantiate the watsonx.ai client127 if client_setup:128 try:129 wx_credentials = Credentials(url=wx_url, api_key=wx_api_key)130 project_client = (131 APIClient(credentials=wx_credentials, project_id=project_id)132 if project_id133 else None134 )135 deployment_client = (136 APIClient(credentials=wx_credentials, space_id=space_id)137 if space_id138 else None139 )140 instantiation_success = True141 instantiation_error = None142 except Exception as e:143 instantiation_success = False144 instantiation_error = str(e)145 wx_credentials = project_client = deployment_client = None146 else:147 wx_credentials = project_client = deployment_client = None148 instantiation_success = None149 instantiation_error = None150 151 return (152 deployment_client,153 instantiation_error,154 instantiation_success,155 project_client,156 )157 158 159@app.cell160def _():161 mo.md(162 r"""163 #watsonx.ai Embedding Visualizer - Marimo Notebook164 165 #### This marimo notebook can be used to develop a more intuitive understanding of how vector embeddings work by creating a 3D visualization of vector embeddings based on chunked PDF document pages. 166 167 #### It can also serve as a useful tool for identifying gaps in model choice, chunking strategy or contents used in building collections by showing how far you are from what you want.168 <br>169 170 /// admonition171 Created by ***Milan Mrdenovic*** [milan.mrdenovic@ibm.com] for IBM Ecosystem Client Engineering, NCEE - ***version 5.3** - 20.04.2025*172 ///173 174 175 >Licensed under apache 2.0, users hold full accountability for any use or modification of the code.176 ><br>This asset is part of a set meant to support IBMers, IBM Partners, Clients in developing understanding of how to better utilize various watsonx features and generative AI as a subject matter.177 178 <br>179 """180 )181 return182 183 184@app.cell185def _():186 mo.md("""###Part 1 - Client Setup, File Preparation and Chunking""")187 return188 189 190@app.cell191def accordion_client_setup(client_selector, client_stack):192 ui_accordion_part_1_1 = mo.accordion(193 {194 "Instantiate Client": mo.vstack([client_stack, client_selector], align="center"),195 }196 )197 198 ui_accordion_part_1_1199 return200 201 202@app.cell203def accordion_file_upload(select_stack):204 ui_accordion_part_1_2 = mo.accordion(205 {206 "Select Model & Upload Files": select_stack207 }208 )209 210 ui_accordion_part_1_2211 return212 213 214@app.cell215def loaded_texts(216 create_temp_files_from_uploads,217 file_loader,218 pdf_reader,219 run_upload_button,220 set_text_state,221):222 if file_loader.value is not None and run_upload_button.value:223 filepaths = create_temp_files_from_uploads(file_loader.value)224 loaded_texts = load_pdf_data_with_progress(pdf_reader, filepaths, file_loader.value, show_progress=True)225 226 set_text_state(loaded_texts)227 else:228 filepaths = None229 loaded_texts = None230 return231 232 233@app.cell234def accordion_chunker_setup(chunker_setup):235 ui_accordion_part_1_3 = mo.accordion(236 {237 "Chunker Setup": chunker_setup238 }239 )240 241 ui_accordion_part_1_3242 return243 244 245@app.cell246def chunk_documents_to_nodes(247 get_text_state,248 sentence_splitter,249 sentence_splitter_config,250 set_chunk_state,251):252 if sentence_splitter_config.value and sentence_splitter and get_text_state() is not None:253 chunked_texts = chunk_documents(get_text_state(), sentence_splitter, show_progress=True)254 set_chunk_state(chunked_texts)255 else:256 chunked_texts = None257 return (chunked_texts,)258 259 260@app.cell261def _():262 mo.md(r"""###Part 2 - Query Setup and Visualization""")263 return264 265 266@app.cell267def accordion_chunk_range(chart_range_selection):268 ui_accordion_part_2_1 = mo.accordion(269 {270 "Chunk Range Selection": chart_range_selection271 }272 )273 ui_accordion_part_2_1274 return275 276 277@app.cell278def chunk_embedding(279 chunks_to_process,280 embedding,281 sentence_splitter_config,282 set_embedding_state,283):284 if sentence_splitter_config.value is not None and chunks_to_process is not None:285 with mo.status.spinner(title="Embedding Documents...", remove_on_exit=True) as _spinner:286 output_embeddings = embedding.embed_documents(chunks_to_process)287 _spinner.update("Almost Done")288 time.sleep(1.5)289 set_embedding_state(output_embeddings)290 _spinner.update("Documents Embedded")291 else:292 output_embeddings = None293 return294 295 296@app.cell297def preview_chunks(chunks_dict):298 if chunks_dict is not None:299 stats = create_stats(chunks_dict, 300 bordered=True, 301 object_names=['text','text'], 302 group_by_row=True, 303 items_per_row=5, 304 gap=1,305 label="Chunk")306 ui_chunk_viewer = mo.accordion(307 {308 "View Chunks": stats,309 }310 )311 else:312 ui_chunk_viewer = None313 314 ui_chunk_viewer315 return316 317 318@app.cell319def accordion_query_view(chart_visualization, query_stack):320 ui_accordion_part_2_2 = mo.accordion(321 {322 "Query": mo.vstack([query_stack, mo.hstack([chart_visualization])], align="center", gap=3)323 }324 )325 ui_accordion_part_2_2326 return327 328 329@app.cell330def chunker_setup(sentence_splitter_config):331 chunker_setup = mo.hstack([sentence_splitter_config], justify="space-around", align="center", widths=[0.55])332 return (chunker_setup,)333 334 335@app.cell336def file_and_model_select(337 file_loader,338 get_embedding_model_list,339 run_upload_button,340):341 select_stack = mo.hstack([get_embedding_model_list(), mo.vstack([mo.md("Drag & Drop or Double Click to select PDFs, then press **Load Files**"),file_loader, run_upload_button], align="center")], justify="space-around", align="center", widths=[0.3,0.3])342 return (select_stack,)343 344 345@app.cell346def client_instantiation_form():347 baked_in_creds = credentials348 # Create a form with multiple elements349 client_instantiation_form = (350 mo.md('''351 ###**watsonx.ai credentials:**352 353 {wx_region}354 355 {wx_api_key}356 357 {project_id}358 359 {space_id}360 361 > You can add either a project_id, space_id or both, **only one is required**. 362 > If you provide both you can switch the active one in the dropdown.363 ''')364 .batch(365 wx_region = mo.ui.dropdown(366 wx_regions, 367 label="Select your watsonx.ai region:", 368 value=get_cred_value('region', creds_var_name='baked_in_creds'),369 searchable=True370 ),371 wx_api_key = mo.ui.text(372 placeholder="Add your IBM Cloud api-key...", 373 label="IBM Cloud Api-key:", 374 kind="password", 375 value=get_cred_value('api_key', creds_var_name='baked_in_creds')376 ),377 project_id = mo.ui.text(378 placeholder="Add your watsonx.ai project_id...", 379 label="Project_ID:", 380 kind="text", 381 value=get_cred_value('project_id', creds_var_name='baked_in_creds')382 ),383 space_id = mo.ui.text(384 placeholder="Add your watsonx.ai space_id...", 385 label="Space_ID:", 386 kind="text",387 value=get_cred_value('space_id', creds_var_name='baked_in_creds')388 )389 ,)390 .form(show_clear_button=True, bordered=False)391 )392 return (client_instantiation_form,)393 394 395@app.cell396def instantiation_status(397 client_callout_kind,398 client_instantiation_form,399 client_status,400):401 client_callout = mo.callout(client_status, kind=client_callout_kind)402 client_stack = mo.hstack([client_instantiation_form, client_callout], align="center", justify="space-around", gap=10)403 return (client_stack,)404 405@app.cell406def _(407 client,408 client_key,409 client_options,410 client_selector,411 client_setup,412 get_key_by_value,413 instantiation_error,414 instantiation_success,415 mo,416 wrap_with_spaces,417):418 active_client_name = get_key_by_value(client_options, client_key)419 420 if client_setup:421 if instantiation_success:422 client_status = mo.md(423 f"### ✅ Client Instantiation Successful ✅\n\n"424 f"{client_selector}\n\n"425 f"**Active Client:**{wrap_with_spaces(active_client_name, prefix_spaces=5)}"426 )427 client_callout_kind = "success"428 else:429 client_status = mo.md(430 f"### ❌ Client Instantiation Failed\n**Error:** {instantiation_error}\n\nCheck your region selection and credentials"431 )432 client_callout_kind = "danger"433 else:434 client_status = mo.md(435 f"### Client Instantiation Status will turn Green When Ready\n\n"436 f"{client_selector}\n\n"437 f"**Active Client:**{wrap_with_spaces(active_client_name, prefix_spaces=5)}"438 )439 client_callout_kind = "neutral"440 441 return active_client_name, client_callout_kind, client_status442 443@app.cell444def client_selector(deployment_client, project_client):445 if deployment_client is not None:446 client_options = {"Deployment Client":deployment_client}447 448 elif project_client is not None:449 client_options = {"Project Client":project_client}450 451 elif project_client is not None and deployment_client is not None:452 client_options = {"Project Client":project_client,"Deployment Client":deployment_client}453 454 else:455 client_options = {"No Client": "Instantiate a Client"}456 457 default_client = next(iter(client_options))458 client_selector = mo.ui.dropdown(client_options, value=default_client, label="**Select your active client:**")459 return client_options, client_selector460 461 462@app.cell463def active_client(client_selector):464 client_key = client_selector.value465 if client_key == "Instantiate a Client":466 client = None467 else:468 client = client_key469 return client, client_key470 471@app.cell472def emb_model_selection(client, set_embedding_model_list):473 if client is not None:474 model_specs = client.foundation_models.get_embeddings_model_specs()475 # model_specs = client.foundation_models.get_model_specs()476 resources = model_specs["resources"]477 # Define embedding models reference data478 embedding_models = {479 "ibm/granite-embedding-107m-multilingual": {"max_tokens": 512, "embedding_dimensions": 384},480 "ibm/granite-embedding-278m-multilingual": {"max_tokens": 512, "embedding_dimensions": 768},481 "ibm/slate-125m-english-rtrvr-v2": {"max_tokens": 512, "embedding_dimensions": 768},482 "ibm/slate-125m-english-rtrvr": {"max_tokens": 512, "embedding_dimensions": 768},483 "ibm/slate-30m-english-rtrvr-v2": {"max_tokens": 512, "embedding_dimensions": 384},484 "ibm/slate-30m-english-rtrvr": {"max_tokens": 512, "embedding_dimensions": 384},485 "sentence-transformers/all-minilm-l6-v2": {"max_tokens": 128, "embedding_dimensions": 384},486 "sentence-transformers/all-minilm-l12-v2": {"max_tokens": 128, "embedding_dimensions": 384},487 "intfloat/multilingual-e5-large": {"max_tokens": 512, "embedding_dimensions": 1024}488 }489 490 # Get model IDs from resources491 model_id_list = []492 for resource in resources:493 model_id_list.append(resource["model_id"])494 495 # Create enhanced model data for the table496 embedding_model_data = []497 for model_id in model_id_list:498 model_entry = {"model_id": model_id}499 500 # Add properties if model exists in our reference, otherwise use 0501 if model_id in embedding_models:502 model_entry["max_tokens"] = embedding_models[model_id]["max_tokens"]503 model_entry["embedding_dimensions"] = embedding_models[model_id]["embedding_dimensions"]504 else:505 model_entry["max_tokens"] = 0506 model_entry["embedding_dimensions"] = 0507 508 embedding_model_data.append(model_entry)509 510 embedding_model_selection = mo.ui.table(511 embedding_model_data,512 selection="single", # Only allow selecting one row513 label="Select an embedding model to use.",514 page_size=30,515 initial_selection=[1]516 )517 set_embedding_model_list(embedding_model_selection)518 else:519 default_model_data = [{520 "model_id": "ibm/granite-embedding-107m-multilingual",521 "max_tokens": 512,522 "embedding_dimensions": 384523 }]524 525 set_embedding_model_list(create_emb_model_selection_table(default_model_data, initial_selection=0, selection_type="single", label="Select a model to use."))526 return527 528 529@app.function530def create_emb_model_selection_table(model_data, initial_selection=0, selection_type="single", label="Select a model to use."):531 embedding_model_selection = mo.ui.table(532 model_data,533 selection=selection_type, # Only allow selecting one row534 label=label,535 page_size=30,536 initial_selection=[initial_selection]537 )538 return embedding_model_selection539 540 541@app.cell542def embedding_model():543 get_embedding_model_list, set_embedding_model_list = mo.state(None)544 return get_embedding_model_list, set_embedding_model_list545 546 547@app.cell548def emb_model_parameters(emb_model_max_tk):549 from ibm_watsonx_ai.metanames import EmbedTextParamsMetaNames as EmbedParams550 if embedding_model is not None:551 embed_params = {552 EmbedParams.TRUNCATE_INPUT_TOKENS: emb_model_max_tk,553 EmbedParams.RETURN_OPTIONS: {554 'input_text': True555 }556 }557 else:558 embed_params = {559 EmbedParams.TRUNCATE_INPUT_TOKENS: 128,560 EmbedParams.RETURN_OPTIONS: {561 'input_text': True562 }563 } 564 return embed_params565 566 567@app.cell568def emb_model_state(get_embedding_model_list):569 embedding_model = get_embedding_model_list()570 return (embedding_model,)571 572 573@app.cell574def emb_model_setup(embedding_model):575 if embedding_model is not None:576 emb_model = embedding_model.value[0]['model_id']577 emb_model_max_tk = embedding_model.value[0]['max_tokens']578 emb_model_emb_dim = embedding_model.value[0]['embedding_dimensions']579 else:580 emb_model = None581 emb_model_max_tk = None582 emb_model_emb_dim = None583 return emb_model, emb_model_emb_dim, emb_model_max_tk584 585 586@app.cell587def emb_model_instantiation(client, emb_model, embed_params):588 from ibm_watsonx_ai.foundation_models import Embeddings589 if client is not None:590 embedding = Embeddings(591 model_id=emb_model,592 api_client=client,593 params=embed_params,594 batch_size=1000,595 concurrency_limit=10596 )597 else:598 embedding = None599 return (embedding,)600 601 602@app.cell603def _():604 get_embedding_state, set_embedding_state = mo.state(None)605 return get_embedding_state, set_embedding_state606 607 608@app.cell609def _():610 get_query_state, set_query_state = mo.state(None)611 return get_query_state, set_query_state612 613 614@app.cell615def file_loader_input():616 file_loader = mo.ui.file(617 kind="area",618 filetypes=[".pdf"],619 label=" Load .pdf files ",620 multiple=True621 )622 return (file_loader,)623 624 625@app.cell626def file_loader_run(file_loader):627 if file_loader.value:628 run_upload_button = mo.ui.run_button(label="Load Files")629 else:630 run_upload_button = mo.ui.run_button(disabled=True, label="Load Files")631 return (run_upload_button,)632 633 634@app.cell635def helper_function_tempfiles():636 def create_temp_files_from_uploads(upload_results) -> List[str]:637 """638 Creates temporary files from a tuple of FileUploadResults objects and returns their paths.639 Args:640 upload_results: Object containing a value attribute that is a tuple of FileUploadResults641 Returns:642 List of temporary file paths643 """644 temp_file_paths = []645 646 # Get the number of items in the tuple647 num_items = len(upload_results)648 649 # Process each item by index650 for i in range(num_items):651 result = upload_results[i] # Get item by index652 653 # Create a temporary file with the original filename654 temp_dir = tempfile.gettempdir()655 file_name = result.name656 temp_path = os.path.join(temp_dir, file_name)657 # Write the contents to the temp file658 with open(temp_path, 'wb') as temp_file:659 temp_file.write(result.contents)660 # Add the path to our list661 temp_file_paths.append(temp_path)662 663 return temp_file_paths664 665 def cleanup_temp_files(temp_file_paths: List[str]) -> None:666 """Delete temporary files after use."""667 for path in temp_file_paths:668 if os.path.exists(path):669 os.unlink(path)670 return (create_temp_files_from_uploads,)671 672 673@app.function674def load_pdf_data_with_progress(pdf_reader, filepaths, file_loader_value, show_progress=True):675 """676 Loads PDF data for each file path and organizes results by original filename.677 Args:678 pdf_reader: The PyMuPDFReader instance679 filepaths: List of temporary file paths680 file_loader_value: The original upload results value containing file information681 show_progress: Whether to show a progress bar during loading (default: False)682 Returns:683 Dictionary mapping original filenames to their loaded text content684 """685 results = {}686 687 # Process files with or without progress bar688 if show_progress:689 import marimo as mo690 # Use progress bar with the length of filepaths as total691 with mo.status.progress_bar(692 total=len(filepaths),693 title="Loading PDFs",694 subtitle="Processing documents...",695 completion_title="PDF Loading Complete",696 completion_subtitle=f"{len(filepaths)} documents processed",697 remove_on_exit=True698 ) as bar:699 # Process each file path700 for i, file_path in enumerate(filepaths):701 702 original_file_name = file_loader_value[i].name703 bar.update(subtitle=f"Processing {original_file_name}...")704 loaded_text = pdf_reader.load_data(file_path=file_path, metadata=True)705 706 # Store the result with the original filename as the key707 results[original_file_name] = loaded_text708 # Update progress bar709 bar.update(increment=1)710 else:711 # Original logic without progress bar712 for i, file_path in enumerate(filepaths):713 original_file_name = file_loader_value[i].name714 loaded_text = pdf_reader.load_data(file_path=file_path, metadata=True)715 results[original_file_name] = loaded_text716 717 return results718 719 720@app.cell721def file_readers():722 from llama_index.readers.file import PyMuPDFReader723 from llama_index.readers.file import FlatReader724 from llama_index.core.node_parser import SentenceSplitter725 726 ### File Readers727 pdf_reader = PyMuPDFReader()728 # flat_file_reader = FlatReader()729 return SentenceSplitter, pdf_reader730 731 732@app.cell733def sentence_splitter_setup():734 ### Chunker Setup735 sentence_splitter_config = (736 mo.md('''737 ###**Chunking Setup:**738 739 > Unless you want to do some advanced sentence splitting, it's best to stick to adjusting only the chunk size and overlap. Changing the other settings might result in unexpected results.740 741 Separator value is set to **" "** by default, while the paragraph separator is **"\\n\\n\\n"**.742 743 {chunk_size} 744 745 {chunk_overlap}746 747 {separator} {paragraph_separator}748 749 {secondary_chunking_regex} {include_metadata}750 751 ''')752 .batch(753 chunk_size = mo.ui.slider(start=100, stop=5000, step=1, label="**Chunk Size:**", value=275, show_value=True, full_width=True),754 chunk_overlap = mo.ui.slider(start=0, stop=1000, step=1, label="**Chunk Overlap** *(Must always be smaller than Chunk Size)* **:**", value=0, show_value=True, full_width=True),755 separator = mo.ui.text(placeholder="Define a separator", label="**Separator:**", kind="text", value=" "),756 paragraph_separator = mo.ui.text(placeholder="Define a paragraph separator", 757 label="**Paragraph Separator:**", kind="text", 758 value="\n\n\n"),759 secondary_chunking_regex = mo.ui.text(placeholder="Define a secondary chunking regex", 760 label="**Chunking Regex:**", kind="text", 761 value="[^,.;?!]+[,.;?!]?"),762 include_metadata= mo.ui.checkbox(value=True, label="**Include Metadata**")763 )764 .form(show_clear_button=True, bordered=False, submit_button_label="Chunk Documents")765 )766 return (sentence_splitter_config,)767 768 769@app.cell770def sentence_splitter_instantiation(771 SentenceSplitter,772 sentence_splitter_config,773):774 ### Chunker/Sentence Splitter 775 def simple_whitespace_tokenizer(text):776 return text.split()777 778 if sentence_splitter_config.value is not None:779 sentence_splitter_config_values = sentence_splitter_config.value780 validated_chunk_overlap = min(sentence_splitter_config_values.get("chunk_overlap"), 781 int(sentence_splitter_config_values.get("chunk_size") * 0.3))782 783 sentence_splitter = SentenceSplitter(784 chunk_size=sentence_splitter_config_values.get("chunk_size"),785 chunk_overlap=validated_chunk_overlap,786 separator=sentence_splitter_config_values.get("separator"),787 paragraph_separator=sentence_splitter_config_values.get("paragraph_separator"),788 secondary_chunking_regex=sentence_splitter_config_values.get("secondary_chunking_regex"),789 include_metadata=sentence_splitter_config_values.get("include_metadata"),790 tokenizer=simple_whitespace_tokenizer791 )792 793 else:794 sentence_splitter = SentenceSplitter(795 chunk_size=2048,796 chunk_overlap=204,797 separator=" ",798 paragraph_separator="\n\n\n",799 secondary_chunking_regex="[^,.;?!]+[,.;?!]?",800 include_metadata=True,801 tokenizer=simple_whitespace_tokenizer802 )803 return (sentence_splitter,)804 805 806@app.cell807def text_state():808 get_text_state, set_text_state = mo.state(None)809 return get_text_state, set_text_state810 811 812@app.cell813def chunk_state():814 get_chunk_state, set_chunk_state = mo.state(None)815 return get_chunk_state, set_chunk_state816 817 818@app.function819def chunk_documents(loaded_texts, sentence_splitter, show_progress=True):820 """821 Process each document in the loaded_texts dictionary using the sentence_splitter,822 with an optional marimo progress bar tracking progress at document level.823 824 Args:825 loaded_texts (dict): Dictionary containing lists of Document objects826 sentence_splitter: The sentence splitter object with get_nodes_from_documents method827 show_progress (bool): Whether to show a progress bar during processing828 829 Returns:830 dict: Dictionary with the same structure but containing chunked texts831 """832 chunked_texts_dict = {}833 834 # Get the total number of documents across all keys835 total_docs = sum(len(docs) for docs in loaded_texts.values())836 processed_docs = 0837 838 # Process with or without progress bar839 if show_progress:840 import marimo as mo841 # Use progress bar with the total number of documents as total842 with mo.status.progress_bar(843 total=total_docs,844 title="Processing Documents",845 subtitle="Chunking documents...",846 completion_title="Processing Complete",847 completion_subtitle=f"{total_docs} documents processed",848 remove_on_exit=True849 ) as bar:850 # Process each key-value pair in the loaded_texts dictionary851 for key, documents in loaded_texts.items():852 # Update progress bar subtitle to show current key853 doc_count = len(documents)854 bar.update(subtitle=f"Chunking {key}... ({doc_count} documents)")855 856 # Apply the sentence splitter to each list of documents857 chunked_texts = sentence_splitter.get_nodes_from_documents(858 documents, 859 show_progress=False # Disable internal progress to avoid nested bars860 )861 862 # Store the result with the same key863 chunked_texts_dict[key] = chunked_texts864 time.sleep(0.15)865 866 # Update progress bar with the number of documents in this batch867 bar.update(increment=doc_count)868 processed_docs += doc_count869 else:870 # Process without progress bar871 for key, documents in loaded_texts.items():872 chunked_texts = sentence_splitter.get_nodes_from_documents(873 documents, 874 show_progress=True # Use the internal progress bar if no marimo bar875 )876 chunked_texts_dict[key] = chunked_texts877 878 return chunked_texts_dict879 880 881@app.cell882def chunked_nodes(chunked_texts, get_chunk_state, sentence_splitter):883 if chunked_texts is not None and sentence_splitter:884 chunked_documents = get_chunk_state()885 else:886 chunked_documents = None887 return (chunked_documents,)888 889 890@app.cell891def prep_cumulative_df(chunked_documents, llamaindex_convert_docs_multi):892 if chunked_documents is not None:893 dict_from_nodes = llamaindex_convert_docs_multi(chunked_documents)894 nodes_from_dict = llamaindex_convert_docs_multi(dict_from_nodes)895 else:896 dict_from_nodes = None897 nodes_from_dict = None898 return (dict_from_nodes,)899 900 901@app.cell902def chunks_to_process(903 dict_from_nodes,904 document_range_stack,905 get_data_in_range_triplequote,906):907 if dict_from_nodes is not None and document_range_stack is not None:908 909 chunk_dict_df = create_cumulative_dataframe(dict_from_nodes)910 911 if document_range_stack.value is not None:912 chunk_start_idx = document_range_stack.value[0]913 chunk_end_idx = document_range_stack.value[1]914 else:915 chunk_start_idx = 0916 chunk_end_idx = len(chunk_dict_df)917 918 chunk_range_index = [chunk_start_idx, chunk_end_idx]919 chunks_dict = get_data_in_range_triplequote(chunk_dict_df,920 index_range=chunk_range_index, 921 columns_to_include=["text"])922 923 chunks_to_process = chunks_dict['text'] if 'text' in chunks_dict else []924 else:925 chunk_objects = None926 chunks_dict = None927 chunks_to_process = None928 return chunks_dict, chunks_to_process929 930 931@app.cell932def helper_function_doc_formatting():933 def llamaindex_convert_docs_multi(items):934 """935 Automatically convert between document objects and dictionaries.936 937 This function handles:938 - Converting dictionaries to document objects939 - Converting document objects to dictionaries940 - Processing lists or individual items941 - Supporting dictionary structures where values are lists of documents942 943 Args:944 items: A document object, dictionary, or list of either.945 Can also be a dictionary mapping filenames to lists of documents.946 947 Returns:948 Converted item(s) maintaining the original structure949 """950 # Handle empty or None input951 if not items:952 return []953 954 # Handle dictionary mapping filenames to document lists (from load_pdf_data)955 if isinstance(items, dict) and all(isinstance(v, list) for v in items.values()):956 result = {}957 for filename, doc_list in items.items():958 result[filename] = llamaindex_convert_docs(doc_list)959 return result960 961 # Handle single items (not in a list)962 if not isinstance(items, list):963 # Single dictionary to document964 if isinstance(items, dict):965 # Determine document class966 doc_class = None967 if 'doc_type' in items:968 import importlib969 module_path, class_name = items['doc_type'].rsplit('.', 1)970 module = importlib.import_module(module_path)971 doc_class = getattr(module, class_name)972 if not doc_class:973 from llama_index.core.schema import Document974 doc_class = Document975 return doc_class.from_dict(items)976 # Single document to dictionary977 elif hasattr(items, 'to_dict'):978 return items.to_dict()979 # Return as is if can't convert980 return items981 982 # Handle list input983 result = []984 985 # Handle empty list986 if len(items) == 0:987 return result988 989 # Determine the type of conversion based on the first non-None item990 first_item = next((item for item in items if item is not None), None)991 992 # If we found no non-None items, return empty list993 if first_item is None:994 return result995 996 # Convert dictionaries to documents997 if isinstance(first_item, dict):998 # Get the right document class from the items themselves999 doc_class = None1000 # Try to get doc class from metadata if available1001 if 'doc_type' in first_item:1002 import importlib1003 module_path, class_name = first_item['doc_type'].rsplit('.', 1)1004 module = importlib.import_module(module_path)1005 doc_class = getattr(module, class_name)1006 if not doc_class:1007 # Fallback to default Document class from llama_index1008 from llama_index.core.schema import Document1009 doc_class = Document1010 1011 # Convert each dictionary to document1012 for item in items:1013 if isinstance(item, dict):1014 result.append(doc_class.from_dict(item))1015 elif item is None:1016 result.append(None)1017 elif isinstance(item, list):1018 result.append(llamaindex_convert_docs(item))1019 else:1020 result.append(item)1021 1022 # Convert documents to dictionaries1023 else:1024 for item in items:1025 if hasattr(item, 'to_dict'):1026 result.append(item.to_dict())1027 elif item is None:1028 result.append(None)1029 elif isinstance(item, list):1030 result.append(llamaindex_convert_docs(item))1031 else:1032 result.append(item)1033 1034 return result1035 1036 def llamaindex_convert_docs(items):1037 """1038 Automatically convert between document objects and dictionaries.1039 1040 Args:1041 items: A list of document objects or dictionaries1042 1043 Returns:1044 List of converted items (dictionaries or document objects)1045 """1046 result = []1047 1048 # Handle empty or None input1049 if not items:1050 return result1051 1052 # Determine the type of conversion based on the first item1053 if isinstance(items[0], dict):1054 # Get the right document class from the items themselves1055 # Look for a 'doc_type' or '__class__' field in the dictionary1056 doc_class = None1057 1058 # Try to get doc class from metadata if available1059 if 'doc_type' in items[0]:1060 import importlib1061 module_path, class_name = items[0]['doc_type'].rsplit('.', 1)1062 module = importlib.import_module(module_path)1063 doc_class = getattr(module, class_name)1064 1065 if not doc_class:1066 # Fallback to default Document class from llama_index1067 from llama_index.core.schema import Document1068 doc_class = Document1069 1070 # Convert dictionaries to documents1071 for item in items:1072 if isinstance(item, dict):1073 result.append(doc_class.from_dict(item))1074 else:1075 # Convert documents to dictionaries1076 for item in items:1077 if hasattr(item, 'to_dict'):1078 result.append(item.to_dict())1079 1080 return result1081 return (llamaindex_convert_docs_multi,)1082 1083 1084@app.cell1085def helper_function_create_df():1086 def create_document_dataframes(dict_from_docs):1087 """1088 Creates a pandas DataFrame for each file in the dictionary.1089 1090 Args:1091 dict_from_docs: Dictionary mapping filenames to lists of documents1092 1093 Returns:1094 List of pandas DataFrames, each representing all documents from a single file1095 """1096 dataframes = []1097 1098 for filename, docs in dict_from_docs.items():1099 # Create a list to hold all document records for this file1100 file_records = []1101 1102 for i, doc in enumerate(docs):1103 # Convert the document to a format compatible with DataFrame1104 if hasattr(doc, 'to_dict'):1105 doc_data = doc.to_dict()1106 elif isinstance(doc, dict):1107 doc_data = doc1108 else:1109 doc_data = {'content': str(doc)}1110 1111 # Add document index information1112 doc_data['doc_index'] = i1113 1114 # Add to the list of records for this file1115 file_records.append(doc_data)1116 1117 # Create a single DataFrame for all documents in this file1118 if file_records:1119 df = pd.DataFrame(file_records)1120 df['filename'] = filename # Add filename as a column1121 dataframes.append(df)1122 1123 return dataframes1124 1125 def create_dataframe_previews(dataframe_list, page_size=5):1126 """1127 Creates a list of mo.ui.dataframe components, one for each DataFrame in the input list.1128 1129 Args:1130 dataframe_list: List of pandas DataFrames (output from create_document_dataframes)1131 page_size: Number of rows to show per page for each component1132 1133 Returns:1134 List of mo.ui.dataframe components1135 """1136 # Create a list of mo.ui.dataframe components1137 preview_components = []1138 1139 for df in dataframe_list:1140 # Create a mo.ui.dataframe component for this DataFrame1141 preview = mo.ui.dataframe(df, page_size=page_size)1142 preview_components.append(preview)1143 1144 return preview_components1145 return1146 1147 1148@app.cell1149def helper_function_chart_preparation():1150 import altair as alt1151 import numpy as np1152 import plotly.express as px1153 from sklearn.manifold import TSNE1154 1155 def prepare_embedding_data(embeddings, texts, model_id=None, embedding_dimensions=None):1156 """1157 Prepare embedding data for visualization1158 1159 Args:1160 embeddings: List of embeddings arrays1161 texts: List of text strings1162 model_id: Embedding model ID (optional)1163 embedding_dimensions: Embedding dimensions (optional)1164 1165 Returns:1166 DataFrame with processed data and metadata1167 """1168 # Flatten embeddings (in case they're nested)1169 flattened_embeddings = []1170 for emb in embeddings:1171 if isinstance(emb, list) and len(emb) > 0 and isinstance(emb[0], list):1172 flattened_embeddings.append(emb[0]) # Take first element if nested1173 else:1174 flattened_embeddings.append(emb)1175 1176 # Convert to numpy array1177 embedding_array = np.array(flattened_embeddings)1178 1179 # Apply dimensionality reduction (t-SNE)1180 tsne = TSNE(n_components=2, random_state=42, perplexity=min(30, len(embedding_array)-1))1181 reduced_embeddings = tsne.fit_transform(embedding_array)1182 1183 # Create truncated texts for display1184 truncated_texts = [text[:50] + "..." if len(text) > 50 else text for text in texts]1185 1186 # Create dataframe for visualization1187 df = pd.DataFrame({1188 "x": reduced_embeddings[:, 0],1189 "y": reduced_embeddings[:, 1],1190 "text": truncated_texts,1191 "full_text": texts,1192 "index": range(len(texts))1193 })1194 1195 # Add metadata1196 metadata = {1197 "model_id": model_id,1198 "embedding_dimensions": embedding_dimensions1199 }1200