4f59ed38b0
* Adds support for stella_en_v5 embedding model -400M variant * Unified stella * WIP: Unified Stella * Combined stella for both 1.5B and 400M variants * Cargo fmt for the CI * removed redundant stella-400m model and example after merge into stella-en-v5 * cargo fmt --all --------- Co-authored-by: Anubhab Bandyopadhyay <4890833+AnubhabB@users.noreply.github.com> Co-authored-by: laurent <laurent.mazare@gmail.com>
3.9 KiB
candle-stella-en-v5: Implementation of stella_en_1.5B_v5 embedding model
As of 7th Oct 2024, Stella_en_1.5B_v5 is one of the top ranking model on retrieval and reranking tasks in MTEB leaderboard.
Model card on the HuggingFace Hub.
Running the example
Stella_en_1.5B_v5 is used to generate text embeddings embeddings for a prompt. The model weights are downloaded from the hub on the first run.
$ cargo run --example stella-en-v5 --release -- --query "What are safetensors?"
> [[ 0.3905, -0.0130, 0.2072, ..., -0.1100, -0.0086, 0.6002]]
> Tensor[[1, 1024], f32]
Stella_en_1.5B_v5 is trained by MRL enabling multiple embedding dimensions.
The following reproduces the example in the model card for a retrieval task (s2p). The sample queries and docs are hardcoded in the example.
$ cargo run --example stella-en-v5 --release --features <metal | cuda> -- --which 1.5b
>
> Score: 0.8178786
> Query: What are some ways to reduce stress?
> Answer: There are many effective ways to reduce stress. Some common techniques include deep breathing, meditation, and physical activity. Engaging in hobbies, spending
> time in nature, and connecting with loved ones can also help alleviate stress. Additionally, setting boundaries, practicing self-care, and learning to say no can prevent
> stress from building up.
>
>
> Score: 0.7853528
> Query: What are the benefits of drinking green tea?
> Answer: Green tea has been consumed for centuries and is known for its potential health benefits. It contains antioxidants that may help protect the body against damage
> caused by free radicals. Regular consumption of green tea has been associated with improved heart health, enhanced cognitive function, and a reduced risk of certain types >
> of cancer. The polyphenols in green tea may also have anti-inflammatory and weight loss properties.
>
$ cargo run --example stella-en-v5 --release --features <metal | cuda> -- --which 400m
>
> Score: 0.8397539
> Query: What are some ways to reduce stress?
> Answer: There are many effective ways to reduce stress. Some common techniques include deep breathing, meditation, and physical activity. Engaging in hobbies, spending
> time in nature, and connecting with loved ones can also help alleviate stress. Additionally, setting boundaries, practicing self-care, and learning to say no can prevent
> stress from building up.
>
>
>
> Score: 0.809545
> Query: What are the benefits of drinking green tea?
> Answer: Green tea has been consumed for centuries and is known for its potential health benefits. It contains antioxidants that may help protect the body against damage
> caused by free radicals. Regular consumption of green tea has been associated with improved heart health, enhanced cognitive function, and a reduced risk of certain types
> of cancer. The polyphenols in green tea may also have anti-inflammatory and weight loss properties.
>
Supported options:
-
Stella_en_v5has 2 model variants published - a 1.5B variant and 400M variant. This is enabled through the flag--which. E.g.--which 400mor--which 1.5b. -
Stella_en_v5supports 256, 768, 1024, 2048, 4096, 6144 and 8192 embedding dimensions (though the model card mentions 512, I couldn't find weights for the same). In the example run this is supported with--embed-dimoption. E.g.... --embed-dim 4096. Defaults to1024. -
As per the model card, the model has been primarily trained on
s2s(similarity) ands2p(retrieval) tasks. These require a slightly differentquerypreprocessing (a different prompt template for each). In this example this is enabled though--taskoption.