lancedb/docs/web/snippets/embedding.mdx

{/* Auto-generated by scripts/mdx_snippets_gen.py. Do not edit manually. */}

export const PyAsyncOpenaiEmbeddings = "db = await lancedb.connect_async(uri)\nfunc = get_registry().get(\"openai\").create(name=\"text-embedding-ada-002\")\n\nclass Words(LanceModel):\n    text: str = func.SourceField()\n    vector: Vector(func.ndims()) = func.VectorField()\n\ntable = await db.create_table(\"words\", schema=Words, mode=\"overwrite\")\nawait table.add([{\"text\": \"hello world\"}, {\"text\": \"goodbye world\"}])\n\nquery = \"greetings\"\nactual = await (await table.search(query)).limit(1).to_pydantic(Words)[0]\nprint(actual.text)\n";

export const PyCreateEmbeddingFunction = "func = get_registry().get(\"openai\").create(\n    name=\"text-embedding-3-small\",\n    max_retries=7,\n)\n";

export const PyEmbeddingFunction = "from functools import cached_property\n\nfrom lancedb.embeddings import TextEmbeddingFunction, register\n\nclass MyEmbeddingModel:\n    def __init__(self, model_name: str):\n        self.model_name = model_name\n\n    def encode(self, texts: list[str]) -> list[list[float]]:\n        return [[1.0, 2.0, 3.0] for _ in texts]\n\n@register(\"my-embedder\")\nclass MyTextEmbedder(TextEmbeddingFunction):\n    model_name: str = \"my-model\"\n\n    def generate_embeddings(self, texts: list[str]) -> list[list[float]]:\n        # Your embedding logic here\n        return self._model.encode(texts)\n\n    def ndims(self) -> int:\n        # Return the dimensionality of the embeddings\n        return len(self.generate_embeddings([\"test\"])[0])\n\n    @cached_property\n    def _model(self) -> MyEmbeddingModel:\n        # Initialize your model once\n        return MyEmbeddingModel(self.model_name)\n";

export const PyImports = "from lancedb.pydantic import LanceModel, Vector\nfrom lancedb.embeddings import get_registry\n";

export const PyManualQueryEmbeddings = "db = lancedb.connect(\"/tmp/db\")\nfunc = get_registry().get(\"openai\").create(name=\"text-embedding-ada-002\")\n\nclass Words(LanceModel):\n    text: str = func.SourceField()\n    vector: Vector(func.ndims()) = func.VectorField()\n\ntable = db.create_table(\"words\", schema=Words, mode=\"overwrite\")\ntable.add([{\"text\": \"hello world\"}, {\"text\": \"goodbye world\"}])\n\nquery_vector = func.generate_embeddings([\"greetings\"])[0]\n# --8<-- [start:manual_query_search]\n# query_vector is assumed to already be generated by your embedding function\nactual = table.search(query_vector).limit(1).to_pydantic(Words)[0]\nprint(actual.text)\n# --8<-- [end:manual_query_search]\n";

export const PyManualQuerySearch = "# query_vector is assumed to already be generated by your embedding function\nactual = table.search(query_vector).limit(1).to_pydantic(Words)[0]\nprint(actual.text)\n";

export const PyOpenaiEmbeddings = "db = lancedb.connect(\"/tmp/db\")\nfunc = get_registry().get(\"openai\").create(name=\"text-embedding-ada-002\")\n\nclass Words(LanceModel):\n    text: str = func.SourceField()\n    vector: Vector(func.ndims()) = func.VectorField()\n\ntable = db.create_table(\"words\", schema=Words, mode=\"overwrite\")\ntable.add([{\"text\": \"hello world\"}, {\"text\": \"goodbye world\"}])\n\nquery = \"greetings\"\nactual = table.search(query).limit(1).to_pydantic(Words)[0]\nprint(actual.text)\n";

export const PyRegisterDevice = "import torch\n\nregistry = get_registry()\nif torch.cuda.is_available():\n    registry.set_var(\"device\", \"cuda\")\n\nfunc = registry.get(\"huggingface\").create(device=\"$var:device:cpu\")\n";

export const PyRegisterSecret = "registry = get_registry()\nregistry.set_var(\"api_key\", \"sk-...\")\n\nfunc = registry.get(\"openai\").create(api_key=\"$var:api_key\")\n";

export const TsCreateEmbeddingFunction = "const func = getRegistry().get(\"openai\")!.create({\n  model: \"text-embedding-3-small\",\n});\n";

export const TsEmbeddingFunction = "const db = await lancedb.connect(databaseDir);\n\n@register(\"my_embedding\")\nclass MyEmbeddingFunction extends EmbeddingFunction<string> {\n  constructor(optionsRaw = {}) {\n    super();\n    const options = this.resolveVariables(optionsRaw);\n    // Initialize using options\n  }\n  ndims() {\n    return 3;\n  }\n  protected getSensitiveKeys(): string[] {\n    return [];\n  }\n  embeddingDataType(): Float {\n    return new Float32();\n  }\n  async computeQueryEmbeddings(_data: string) {\n    // This is a placeholder for a real embedding function\n    return [1, 2, 3];\n  }\n  async computeSourceEmbeddings(data: string[]) {\n    // This is a placeholder for a real embedding function\n    return Array.from({ length: data.length }).fill([\n      1, 2, 3,\n    ]) as number[][];\n  }\n}\n\nconst func = new MyEmbeddingFunction();\n\nconst data = [{ text: \"pepperoni\" }, { text: \"pineapple\" }];\n\n// Option 1: manually specify the embedding function\nconst table = await db.createTable(\"vectors\", data, {\n  embeddingFunction: {\n    function: func,\n    sourceColumn: \"text\",\n    vectorColumn: \"vector\",\n  },\n  mode: \"overwrite\",\n});\n\n// Option 2: provide the embedding function through a schema\n\nconst schema = LanceSchema({\n  text: func.sourceField(new Utf8()),\n  vector: func.vectorField(),\n});\n\nconst table2 = await db.createTable(\"vectors2\", data, {\n  schema,\n  mode: \"overwrite\",\n});\n";

export const TsImports = "import * as lancedb from \"@lancedb/lancedb\";\nimport \"@lancedb/lancedb/embedding/openai\";\nimport { LanceSchema, getRegistry, register } from \"@lancedb/lancedb/embedding\";\nimport { EmbeddingFunction } from \"@lancedb/lancedb/embedding\";\nimport { type Float, Float32, Utf8 } from \"apache-arrow\";\n";

export const TsManualQueryEmbeddings = "const db = await lancedb.connect(databaseDir);\nconst func = getRegistry()\n  .get(\"openai\")\n  ?.create({ model: \"text-embedding-ada-002\" }) as EmbeddingFunction;\n\nconst wordsSchema = LanceSchema({\n  text: func.sourceField(new Utf8()),\n  vector: func.vectorField(),\n});\nconst tbl = await db.createEmptyTable(\"words\", wordsSchema, {\n  mode: \"overwrite\",\n});\nawait tbl.add([{ text: \"hello world\" }, { text: \"goodbye world\" }]);\n\nconst queryVector = await func.computeQueryEmbeddings(\"greetings\");\n// --8<-- [start:manual_query_search]\n// queryVector is assumed to already be generated by your embedding function\nconst actual = (await tbl.search(queryVector).limit(1).toArray())[0];\n// --8<-- [end:manual_query_search]\n";

export const TsManualQuerySearch = "// queryVector is assumed to already be generated by your embedding function\nconst actual = (await tbl.search(queryVector).limit(1).toArray())[0];\n";

export const TsOpenaiEmbeddings = "const db = await lancedb.connect(databaseDir);\nconst func = getRegistry()\n  .get(\"openai\")\n  ?.create({ model: \"text-embedding-ada-002\" }) as EmbeddingFunction;\n\nconst wordsSchema = LanceSchema({\n  text: func.sourceField(new Utf8()),\n  vector: func.vectorField(),\n});\nconst tbl = await db.createEmptyTable(\"words\", wordsSchema, {\n  mode: \"overwrite\",\n});\nawait tbl.add([{ text: \"hello world\" }, { text: \"goodbye world\" }]);\n\nconst query = \"greetings\";\nconst actual = (await tbl.search(query).limit(1).toArray())[0];\n";

export const TsRegisterModelFallback = "const registry = getRegistry();\nregistry.setVar(\"openai_model\", \"text-embedding-3-large\");\n\nconst func = registry.get(\"openai\")!.create({\n  model: \"$var:openai_model:text-embedding-3-small\",\n});\n";

export const TsRegisterSecret = "const registry = getRegistry();\nregistry.setVar(\"api_key\", \"sk-...\");\n\nconst func = registry.get(\"openai\")!.create({\n  apiKey: \"$var:api_key\",\n});\n";

export const RsCreateEmbeddingFunction = "use std::sync::Arc;\n\nuse lancedb::embeddings::openai::OpenAIEmbeddingFunction;\n\nlet api_key = std::env::var(\"OPENAI_API_KEY\").expect(\"OPENAI_API_KEY is not set\");\nlet embedding = Arc::new(\n    OpenAIEmbeddingFunction::new_with_model(api_key, \"text-embedding-3-small\")\n        .expect(\"failed to create OpenAI embedding function\"),\n);\n";

export const RsEmbeddingFunction = "use std::{borrow::Cow, sync::Arc};\n\nuse arrow_array::{Array, FixedSizeListArray, Float32Array};\nuse arrow_schema::{DataType, Field, Schema};\nuse lancedb::{\n    connect,\n    embeddings::{EmbeddingDefinition, EmbeddingFunction},\n    Result,\n};\n\n#[derive(Debug, Clone)]\nstruct MyTextEmbedder {\n    dim: usize,\n}\n\nimpl EmbeddingFunction for MyTextEmbedder {\n    fn name(&self) -> &str {\n        \"my-embedder\"\n    }\n\n    fn source_type(&self) -> Result<Cow<'_, DataType>> {\n        Ok(Cow::Owned(DataType::Utf8))\n    }\n\n    fn dest_type(&self) -> Result<Cow<'_, DataType>> {\n        Ok(Cow::Owned(DataType::new_fixed_size_list(\n            DataType::Float32,\n            self.dim as i32,\n            true,\n        )))\n    }\n\n    fn compute_source_embeddings(&self, source: Arc<dyn Array>) -> Result<Arc<dyn Array>> {\n        let values = Arc::new(Float32Array::from(vec![1.0f32; source.len() * self.dim]));\n        let field = Arc::new(Field::new(\"item\", DataType::Float32, true));\n        Ok(Arc::new(FixedSizeListArray::new(\n            field,\n            self.dim as i32,\n            values,\n            None,\n        )))\n    }\n\n    fn compute_query_embeddings(&self, _input: Arc<dyn Array>) -> Result<Arc<dyn Array>> {\n        unimplemented!()\n    }\n}\n\n#[tokio::main]\nasync fn main() -> Result<()> {\n    let db = connect(\"./mydb\").execute().await?;\n    db.embedding_registry()\n        .register(\"my-embedder\", Arc::new(MyTextEmbedder { dim: 3 }))?;\n\n    let schema = Arc::new(Schema::new(vec![Field::new(\"text\", DataType::Utf8, false)]));\n    db.create_empty_table(\"mytable\", schema)\n        .add_embedding(EmbeddingDefinition::new(\n            \"text\",\n            \"my-embedder\",\n            Some(\"vector\"),\n        ))?\n        .execute()\n        .await?;\n\n    Ok(())\n}\n";

export const RsManualQueryEmbeddings = "use std::{iter::once, sync::Arc};\n\nuse arrow_array::{record_batch, StringArray};\nuse arrow_schema::{DataType, Field, Schema};\nuse futures::StreamExt;\nuse lancedb::{\n    connect,\n    embeddings::{openai::OpenAIEmbeddingFunction, EmbeddingDefinition, EmbeddingFunction},\n    query::{ExecutableQuery, QueryBase},\n    Result,\n};\n\n#[tokio::main]\nasync fn main() -> Result<()> {\n    let db = connect(\"./mydb\").execute().await?;\n    let api_key = std::env::var(\"OPENAI_API_KEY\").expect(\"OPENAI_API_KEY is not set\");\n    let embedding = Arc::new(OpenAIEmbeddingFunction::new_with_model(\n        api_key,\n        \"text-embedding-3-large\",\n    )?);\n    db.embedding_registry().register(\"openai\", embedding.clone())?;\n\n    let schema = Arc::new(Schema::new(vec![Field::new(\"text\", DataType::Utf8, false)]));\n    let table = db\n        .create_empty_table(\"mytable\", schema)\n        .add_embedding(EmbeddingDefinition::new(\"text\", \"openai\", Some(\"vector\")))?\n        .execute()\n        .await?;\n\n    table\n        .add(record_batch!((\"text\", Utf8, [\"This is a test.\", \"Another example.\"]))?)\n        .execute()\n        .await?;\n\n    // Manually generate embeddings for the query (Enterprise path)\n    let query = Arc::new(StringArray::from_iter_values(once(\"test example\")));\n    let query_vector = embedding.compute_query_embeddings(query)?;\n    // --8<-- [start:manual_query_search]\n    // query_vector is assumed to already be generated by your embedding function\n    let mut results = table.vector_search(query_vector)?.limit(5).execute().await?;\n\n    while let Some(batch) = results.next().await {\n        println!(\"{:?}\", batch?);\n    }\n    // --8<-- [end:manual_query_search]\n\n    Ok(())\n}\n";

export const RsManualQuerySearch = "// query_vector is assumed to already be generated by your embedding function\nlet mut results = table.vector_search(query_vector)?.limit(5).execute().await?;\n\nwhile let Some(batch) = results.next().await {\n    println!(\"{:?}\", batch?);\n}\n";

export const RsOpenaiEmbeddings = "use std::{iter::once, sync::Arc};\n\nuse arrow_array::{record_batch, StringArray};\nuse arrow_schema::{DataType, Field, Schema};\nuse futures::StreamExt;\nuse lancedb::{\n    connect,\n    embeddings::{openai::OpenAIEmbeddingFunction, EmbeddingDefinition, EmbeddingFunction},\n    query::{ExecutableQuery, QueryBase},\n    Result,\n};\n\n#[tokio::main]\nasync fn main() -> Result<()> {\n    let db = connect(\"./mydb\").execute().await?;\n    let api_key = std::env::var(\"OPENAI_API_KEY\").expect(\"OPENAI_API_KEY is not set\");\n    let embedding = Arc::new(OpenAIEmbeddingFunction::new_with_model(\n        api_key,\n        \"text-embedding-3-large\",\n    )?);\n\n    db.embedding_registry().register(\"openai\", embedding.clone())?;\n\n    let schema = Arc::new(Schema::new(vec![Field::new(\"text\", DataType::Utf8, false)]));\n    let table = db\n        .create_empty_table(\"mytable\", schema)\n        .add_embedding(EmbeddingDefinition::new(\"text\", \"openai\", Some(\"vector\")))?\n        .execute()\n        .await?;\n\n    table\n        .add(record_batch!((\"text\", Utf8, [\"This is a test.\", \"Another example.\"]))?)\n        .execute()\n        .await?;\n\n    let query = Arc::new(StringArray::from_iter_values(once(\"test example\")));\n    let query_vector = embedding.compute_query_embeddings(query)?;\n    let mut results = table.vector_search(query_vector)?.limit(5).execute().await?;\n\n    while let Some(batch) = results.next().await {\n        println!(\"{:?}\", batch?);\n    }\n\n    Ok(())\n}\n";