AI on Your Laptop: How to Integrate AlloyDB Omni with Ollama Models.

After my session at the Google Cloud Community Day in Mumbai, one of the most frequently asked questions was how I managed to call my local model — specifically, one running on my laptop using Ollama — to generate embeddings based on user input for augmentation, similarity, or hybrid searches.

For a comprehensive guide on setting up AlloyDB OMNI with AI, check out my previous blog post. There, I cover all the necessary steps to get started.

Instant IPL Facts with AlloyDB Omni: Bringing Data and Insights Closer with Natural Language to…

Leveraging local models can help you quickly grasp the core concepts and building blocks of GENAI-enabled applications. We’ll use Ollama, an open-source tool designed to facilitate the local execution of large language models (LLMs) on personal computers with AlloyDB Omni.

AlloyDB OMNI Integration with Ollama models

Step 1: Setting Up Ollama on Your Local Laptop with an Embedding Model

First, we’ll set up a local embedding model using Ollama. We’ll use the nomic-embed-text model to generate embeddings from user inputs stored in databases, utilizing the embedding function provided by the google_ml_integration extension.

nomic-embed-text

ollama show nomic-embed-text:latest
  Model
   parameters       137M
   quantization     F16
   arch             nomic-bert
   context length   2048
   embedding length 768

  Parameters
   num_ctx 8192

  License
   Apache License
   Version 2.0, January 2004

Step 2: Integrating AlloyDB Omni with Vertex AI

We’ll follow the same steps outlined in my previous blog post to configure PostgreSQL, mount a local directory, and upload the private key JSON for Vertex AI. Here’s a snippet to get you started:

docker run -i --rm google/alloydbomni cat /usr/share/postgresql/postgresql.conf.sample > alloydb-omni-postgres.conf

echo "omni_enable_ml_agent_process = 'on'" >> alloydb-omni-postgres.conf
echo "google_ml_integration.enable_model_support= 'on'" >> alloydb-omni-postgres.conf
echo "omni_google_cloud_private_key_file_path = '/etc/postgresql/private-key.json'" >> alloydb-omni-postgres.conf


docker run --name alloydbomni_local -e POSTGRES_PASSWORD=******** -p 5434:5432 \
-v /Users/deepakmahto/Desktop/deepak/gcp/IPL_Trivia/vertex-ai-connect-omni.json:/etc/postgresql/private-key.json \
-v "/Users/deepakmahto/Desktop/deepak/gcp/alloydb-omni-postgres.conf":/etc/postgresql/postgresql.conf  \
-v /Users/deepakmahto/Desktop/deepak/gcp/alloydbomnidata1:/var/lib/postgresql/data -d google/alloydbomni \
-c 'config_file=/etc/postgresql/postgresql.conf'

Step 3: Setting Up Input and Output Transformation and Loading the Local Model into AlloyDB Omni

For third-party models, we’ll need to define the necessary transformation functions to handle requests to the embedding API and retrieve embeddings from the JSON response within the database.

Input Transform Function for the Custom Model:

-- Input Transform Function corresponding to the custom model
CREATE OR REPLACE FUNCTION local_embedding_input_transform(model_id VARCHAR(100), input_text TEXT)
RETURNS JSON
LANGUAGE plpgsql
AS $$
DECLARE
  transformed_input JSON;
  model_qualified_name TEXT;
BEGIN
  SELECT json_build_object('prompt', input_text,'model','nomic-embed-text')::JSON INTO transformed_input;
  RETURN transformed_input;
END;
$$;

Output Transform Function for the Custom Model:

-- Output Transform Function corresponding to the custom model
CREATE OR REPLACE FUNCTION local_embedding_output_transform(model_id VARCHAR(100), response_json JSON)
RETURNS REAL[]
LANGUAGE plpgsql
AS $$
DECLARE
  transformed_output REAL[];
BEGIN
  select ARRAY(select json_array_elements_text(response_json->'embedding'))::real[] into transformed_output;
  RETURN transformed_output;
END;
$$;

Next, we can load the model using the google_ml.create_model function:

CALL google_ml.create_model(
    model_id => 'local_embedding',
    model_request_url => 'http://host.docker.internal:11434/api/embeddings',
    model_provider => 'custom',
    model_type => 'text_embedding',
    model_qualified_name => 'nomic-embed-text',
    model_in_transform_fn => 'local_embedding_input_transform',
    model_out_transform_fn => 'local_embedding_output_transform'
);

Note: Since AlloyDB Omni is running within a Docker container, we have modified the model_request_url accordingly.

Step 4 : Generate Embedding using locally loaded model.

Using the embedding wrapper, we can invoke our local model to generate necessary inputs in near real-time as part of a database transaction or workflow:

SELECT array_dims(embedding( 'local_embedding', 'AlloyDB Omni AI - Local Model')::real[]);

Overall Steps for reference

Enhancing AlloyDB Omni with Local AI Models Using Ollama

Conclusion

Integrating local models with AlloyDB Omni enables quick and efficient experimentation with GENAI applications on personal devices. This approach not only accelerates initial development but also provides a deeper understanding of the underlying mechanisms.