Migrate to a new embedding model for a table (Python)

Follow this migration guide if you want to switch embedding models for a column in your table.

This migration only works if you stored the original text in another column in your table. If you did not store the original text in your table, then you must modify the migration script below to get the original text from another source.

If your new embedding model supports a larger context window, then you might also want to re-chunk your data.

Add a new vector column with the desired embedding provider integration to your table. For examples, see Add a vector column and configure an embedding provider integration.

Populate the new vector column with the contents of the column that stores the original text.

The embedding provider integration for your new column will automatically generate vector embeddings based on the text.

For example:

from astrapy import DataAPIClient

client = DataAPIClient("APPLICATION_TOKEN")
database = client.get_database("API_ENDPOINT")

table = database.get_table("TABLE_NAME")

page_state = None
migrated_count = 0

# Use an empty filter to find all rows
filter = {}

# You must include ALL primary key columns for your table
primary_key_columns = [
    "PRIMARY_KEY_1",
    "PRIMARY_KEY_2",
]

original_text_column = "NAME_OF_ORIGINAL_TEXT_COLUMN"

new_vector_column = "NAME_OF_NEW_VECTOR_COLUMN"

# The projection should include ALL primary key columns
# and the column that stores the original text
projection = {
    **{column: True for column in primary_key_columns},
    original_text_column: True,
}

while True:
    if page_state:
        cursor = table.find(
            filter, projection=projection, initial_page_state=page_state
        )
    else:
        cursor = table.find(filter, projection=projection)

    page = cursor.fetch_next_page()
    rows = page.results
    page_state = page.next_page_state

    if not rows:
        print("No more rows. Migration complete.")
        break

    # Build the updates
    updated_rows = []
    for row in rows:
        if text := row.get(original_text_column):
            updated_row = {
                # Include the full primary key
                **{column: row[column] for column in primary_key_columns},
                # Set the new vector column to the original text
                new_vector_column: text,
            }
            updated_rows.append(updated_row)

    # Inserting a row with a primary key that already exists in the table will
    # overwrite the specified column but leave unspecified columns unchanged.
    table.insert_many(updated_rows)
    migrated_count += len(updated_rows)

    print(f"Migrated {migrated_count} rows. Page state: {page_state}")

    if page_state is None:
        print("Reached final page. Migration complete.")
        break

Optionally, delete the column that stores the old vector embeddings.

Migrate to a new embedding model for a table (Python)

Was this helpful?

Give Feedback