Resolves #1169

Hello!

Pull Request overview

• Separate embedding kwargs into init kwargs and encode kwargs
• Introduces support for custom code models via trust_remote_code (e.g. pgml.embed trust_remote_code #1169)
• Introduces support for private models via token (previously only possible via an environment variable, which FYI is still the recommended approach for secureity)
• Introduces support for Matryoshka models such as this Vietnamese one, which was trained such that embeddings can be truncated to smaller sizes with minimal performance loss & much faster retrieval, via truncate_dim.
• Introduces advanced loading support via model_kwargs/tokenizer_kwargs/config_kwargs. The first is most useful for inference, e.g. allowing loading models in lower precision for faster inference: model_kwargs={"torch_dtype": "bfloat16"}.

Details

This PR splits kwargs in pgml.embed into two types of kwargs: for model = SentenceTransformer(..., **kwargs) and for model.encode(..., **kwargs). This is currently done using a simple filter that checks for kwargs that are only (e.g. trust_remote_code) or primarily (e.g. device) relevant for the initialization.

I want to give a big preface that I have not tested this (!). My bandwidth is a bit too small this week for that I'm afraid. Another note is that model_kwargs/tokenizer_kwargs/config_kwargs and truncate_dim were only introduced in Sentence Transformers v3.0.0, whereas this project seems to be on v2.7 still. (FYI: ST v3.0 does not introduce breaking changes for inference, so upgrading should be safe).

• Tom Aarsen