Add Dense layer in 2_Dense/ modules

[alvarobartt]

What does this PR do?

This PR adds support for 2_Dense/ modules, since some models as e.g. https://huggingface.co/sentence-transformers/LaBSE require the extra Dense module i.e., an extra Linear layer on top of the pooled embeddings, when generating the embeddings.

So on, this PR introduces the DenseLayer trait, impl for Dense and adds DenseConfig, which are models with basically a single Linear layer, pulling the configuration from 2_Dense/config.json and the model weights from 2_Dense/model.safetensors.

Note

The 2_Dense/ is only required when generating embeddings, meaning that it will only apply to the Embedding model type, whereas the Reranker and Classifier are not affected by this addition, so on, neither the rank or predict methods for the given backend.

This PR solves the issue recently reported at https://discuss.huggingface.co/t/inference-result-not-aligned-with-local-version-of-same-model-and-revision/160514.

Additionally, this PR also fixes a shape mismatch issue produced when performing matrix multiplication of 2D tensors on Metal devices due to the candle Metal kernels expecting the tensors to be contiguous. It seems that the error only arises on Metal for 2D tensors, where as for e.g. 3D tensors it seems to be working just fine without having to use .contiguous() (which is expensive as it needs to clone the tensor).

Reproduce

To ensure that the implementation was working fine and producing successful results i.e., allclose like checks are true, and the cosine similarity is 1.0 (or as close as possible), the following test has been run:

1. Deploy Text Embeddings Inference (TEI) as e.g.:

cargo run --release --features candle,http --no-default-features -- --model-id sentence-transformers/LaBSE --dtype float16

2. Then, once it's running run the following Python script (requires torch, transformers, sentence-transformers, accelerate and numpy):

import numpy as np
import requests
from sentence_transformers import SentenceTransformer

model = SentenceTransformer(
    "sentence-transformers/LaBSE",
    model_kwargs={
        "torch_dtype": "float16",
        "device_map": "mps",
    },
)

out_py = model.encode(
    "What is Deep Learning?",
    normalize_embeddings=True,
    convert_to_numpy=True,
)

response = requests.post(
    "http://localhost:3000/embed",
    json={
        "inputs": "What is Deep Learning?",
        "normalize": True,
    },
)

response.raise_for_status()
out = response.json()[0]
out_http = np.array(out, dtype=np.float16)

print(f"Embeddings are close: {np.allclose(out_py, out_http, atol=1e-3, rtol=1e-4)=}")


def cosine_similarity(x: np.ndarray, y: np.ndarray) -> float:
    dot_product = np.dot(x, y)
    norm_x = np.linalg.norm(x)
    norm_y = np.linalg.norm(y)
    return dot_product / (norm_x * norm_y)


print(f"The similarity score is: {cosine_similarity(out_py, out_http)=}")

It should produce the following on any combination of device (CPU, MPS, CUDA) and dtype (float32, float16):

Embeddings are close: np.allclose(out_py, out_http, atol=1e-3, rtol=1e-4)=True
The similarity score is: cosine_similarity(out_py, out_http)=np.float16(1.0)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline, Pull Request section?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the documentation guidelines, and here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@Narsil

[kozistr]

@alvarobartt Hi! Just for your reference, you may already be aware, Stella v5 model uses Identity layer as its activation function for 2_Dense! 2_Dense/config.json

[Narsil]

Looks good, I think we can simplify a bit the parsing part.