simonw / llm-mistral

LLM plugin providing access to Mistral models using the Mistral API
Apache License 2.0
144 stars 14 forks source link

Embeddings model #3

Closed simonw closed 9 months ago

simonw commented 9 months ago

https://docs.mistral.ai/api/#operation/createEmbedding

simonw commented 9 months ago

Prototype:

class MistralEmbed(llm.EmbeddingModel):
    model_id = "mistral-embed"

    def embed_batch(self, texts):
        key = llm.get_key("", "mistral", "LLM_MISTRAL_KEY")
        with httpx.Client() as client:
            api_response = client.post(
                "https://api.mistral.ai/v1/embeddings",
                headers={
                    "Content-Type": "application/json",
                    "Accept": "application/json",
                    "Authorization": f"Bearer {key}",
                },
                json={
                    "model": "mistral-embed",
                    "input": texts,
                    "encoding_format": "float",
                },
                timeout=None,
            )
            api_response.raise_for_status()
            return [item["embedding"] for item in api_response.json()["data"]]
simonw commented 9 months ago
$ llm embed -m mistral-embed -c hello | jq length
1024
simonw commented 9 months ago

To test this I'm going to embed all these:

$ ls ../*/llm_*.py
../llm-anyscale-endpoints/llm_anyscale_endpoints.py
../llm-clip/llm_clip.py
../llm-cluster/llm_cluster.py
../llm-embed-cohere/llm_embed_cohere.py
../llm-embed-jina/llm_embed_jina.py
../llm-fireworks/llm_fireworks.py
../llm-functions/llm_functions.py
../llm-gemini/llm_gemini.py
../llm-gpt4all/llm_gpt4all.py
../llm-imagebind/llm_imagebind.py
../llm-llama-cpp/llm_llama_cpp.py
../llm-mistral/llm_mistral.py
../llm-mlc/llm_mlc.py
../llm-mlx-llama/llm_mlx_llama.py
../llm-mpt30b/llm_mpt30b.py
../llm-openai/llm_openai.py
../llm-openrouter/llm_openrouter.py
../llm-python/llm_python.py
../llm-refact/llm_refact.py
../llm-sentence-transformers/llm_sentence_transformers.py
../llm-vllm/llm_vllm.py
simonw commented 9 months ago

Ran this:

llm embed-multi llm_py \
  --model mistral-embed \
  --files .. '**/llm_*.py' \
  --database embed.db --store

Got an error like this: '{"object":"error","message":"Too many tokens in batch. Max is 16384 got 29823","type":"internal_error_proxy","param":null,"code":"1000"}'

So dropped batch size to 10.

That worked, and now:

llm similar -d embed.db llm_py -c 'llama' | jq .id
"llm-llama-cpp/llm_llama_cpp.py"
"llm-mlx-llama/llm_mlx_llama.py"
"llm-vllm/llm_vllm.py"
"llm-mlc/llm_mlc.py"
"llm-embed-cohere/llm_embed_cohere.py"
"llm-plugin/llm-{{cookiecutter.hyphenated}}/llm_{{cookiecutter.underscored}}.py"
"llm_mistral.py"
"llm-mistral/llm_mistral.py"
"llm-anyscale-endpoints/llm_anyscale_endpoints.py"
"llm-fireworks/llm_fireworks.py"
llm similar -d embed.db llm_py -c 'mistral' | jq .id
"llm-mistral/llm_mistral.py"
"llm_mistral.py"
"llm-vllm/llm_vllm.py"
"llm-mlc/llm_mlc.py"
"llm-anyscale-endpoints/llm_anyscale_endpoints.py"
"llm-gemini/llm_gemini.py"
"llm-mpt30b/llm_mpt30b.py"
"llm-mlx-llama/llm_mlx_llama.py"
"llm-openrouter/llm_openrouter.py"
"llm/docs/plugins/llm-markov/llm_markov.py"