Repository: meilisearch/meilisearch
File: crates/index-scheduler/src/scheduler/test_embedders.rs (Rust)

Code before:
    let field_ids_map = index.fields_ids_map(&rtxn).unwrap();
    let field_ids = field_ids_map.ids().collect::<Vec<_>>();
    let documents = index
        .all_documents(&rtxn)
        .unwrap()
        .map(|ret| obkv_to_json(&field_ids, &field_ids_map, ret.unwrap().1).unwrap())
        .collect::<Vec<_>>();
    // the all the vectors linked to the new specified embedder have been removed
    // Only the unknown embedders stays in the document DB
    snapshot!(serde_json::to_string(&documents).unwrap(), @r###"[{"id":0,"doggo":"kefir"},{"id":1,"doggo":"intel","_vectors":{"unknown embedder":[1.0,2.0,3.0]}},{"id":2,"doggo":"max","_vectors":{"unknown embedder":[4.0,5.0]}},{"id":3,"doggo":"marcel"},{"id":4,"doggo":"sora"}]"###);
    let conf = index.embedding_configs(&rtxn).unwrap();
    // even though we specified the vector for the ID 3, it shouldn't be marked
    // as user provided since we explicitely marked it as NOT user provided.
    snapshot!(format!("{conf:#?}"), @r###"
        [
            IndexEmbeddingConfig {
                name: "my_doggo_embedder",
                config: EmbeddingConfig {
                    embedder_options: HuggingFace(
                        EmbedderOptions {
                            model: "sentence-transformers/all-MiniLM-L6-v2",
                            revision: Some(
                                "e4ce9877abf3edfe10b0d82785e83bdcb973e22e",
                            ),
                            distribution: None,
                        },
                    ),
                    prompt: PromptData {
                        template: "{{doc.doggo}}",
                        max_bytes: Some(
                            400,
                        ),
                    },
                    quantized: None,
                },
                user_provided: RoaringBitmap<[1, 2]>,
            },
        ]
        "###);
    let docid = index.external_documents_ids.get(&rtxn, "0").unwrap().unwrap();
    let embeddings = index.embeddings(&rtxn, docid).unwrap();
    let embedding = &embeddings["my_doggo_embedder"];
    assert!(!embedding.is_empty(), "{embedding:?}");

    // the document with the id 3 should keep its original embedding
    let docid = index.external_documents_ids.get(&rtxn, "3").unwrap().unwrap();
    let embeddings = index.embeddings(&rtxn, docid).unwrap();
    let embeddings = &embeddings["my_doggo_embedder"];

    snapshot!(embeddings.len(), @"1");
    assert!(embeddings[0].iter().all(|i| *i == 3.0), "{:?}", embeddings[0]);

Review comment (suggestion):
newly created embedder has `pooling` set to `UseModel`

Code after:
    let field_ids_map = index.fields_ids_map(&rtxn).unwrap();
    let field_ids = field_ids_map.ids().collect::<Vec<_>>();
    let documents = index
        .all_documents(&rtxn)
        .unwrap()
        .map(|ret| obkv_to_json(&field_ids, &field_ids_map, ret.unwrap().1).unwrap())
        .collect::<Vec<_>>();
    // the all the vectors linked to the new specified embedder have been removed
    // Only the unknown embedders stays in the document DB
    snapshot!(serde_json::to_string(&documents).unwrap(), @r###"[{"id":0,"doggo":"kefir"},{"id":1,"doggo":"intel","_vectors":{"unknown embedder":[1.0,2.0,3.0]}},{"id":2,"doggo":"max","_vectors":{"unknown embedder":[4.0,5.0]}},{"id":3,"doggo":"marcel"},{"id":4,"doggo":"sora"}]"###);
    let conf = index.embedding_configs(&rtxn).unwrap();
    // even though we specified the vector for the ID 3, it shouldn't be marked
    // as user provided since we explicitely marked it as NOT user provided.
    snapshot!(format!("{conf:#?}"), @r###"
    [
        IndexEmbeddingConfig {
            name: "my_doggo_embedder",
            config: EmbeddingConfig {
                embedder_options: HuggingFace(
                    EmbedderOptions {
                        model: "sentence-transformers/all-MiniLM-L6-v2",
                        revision: Some(
                            "e4ce9877abf3edfe10b0d82785e83bdcb973e22e",
                        ),
                        distribution: None,
                        pooling: UseModel,
                    },
                ),
                prompt: PromptData {
                    template: "{{doc.doggo}}",
                    max_bytes: Some(
                        400,
                    ),
                },
                quantized: None,
            },
            user_provided: RoaringBitmap<[1, 2]>,
        },
    ]
    "###);
    let docid = index.external_documents_ids.get(&rtxn, "0").unwrap().unwrap();
    let embeddings = index.embeddings(&rtxn, docid).unwrap();
    let embedding = &embeddings["my_doggo_embedder"];
    assert!(!embedding.is_empty(), "{embedding:?}");

    // the document with the id 3 should keep its original embedding
    let docid = index.external_documents_ids.get(&rtxn, "3").unwrap().unwrap();
    let embeddings = index.embeddings(&rtxn, docid).unwrap();
    let embeddings = &embeddings["my_doggo_embedder"];

    snapshot!(embeddings.len(), @"1");