fix: kwargs in custom Sentence Transformer

custom_st.py

@@ -2,7 +2,7 @@ import base64
 import json
 import os
 from io import BytesIO
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Literal, Optional, Union
 
 import requests
 import torch
@@ -14,23 +14,91 @@ from transformers import AutoConfig, AutoImageProcessor, AutoModel, AutoTokenize
 class Transformer(nn.Module):
     def __init__(
         self,
-        model_name_or_path: str,
+        model_name_or_path: str = 'jinaai/jina-clip-v2',
         tokenizer_name_or_path: Optional[str] = None,
         image_processor_name_or_path: Optional[str] = None,
         max_seq_length: Optional[int] = None,
-        config_kwargs: Optional[Dict[str, Any]] = None,
-        model_kwargs: Optional[Dict[str, Any]] = None,
-        tokenizer_kwargs: Optional[Dict[str, Any]] = None,
-        image_processor_kwargs: Optional[Dict[str, Any]] = None,
-        cache_dir: str = None,
+        config_args: Optional[Dict[str, Any]] = None,
+        model_args: Optional[Dict[str, Any]] = None,
+        tokenizer_args: Optional[Dict[str, Any]] = None,
+        image_processor_args: Optional[Dict[str, Any]] = None,
+        cache_dir: Optional[str] = None,
+        backend: Literal['torch', 'onnx', 'openvino'] = 'torch',
         **_,
     ) -> None:
+        """
+        Creates a custom SentenceTransformer module that uses `jinai/jina-clip-v2` to
+        map sentences/images to embeddings
+
+        Args:
+            model_name_or_path (str, optional): If it is a filepath on disc, it loads
+                the model from that path. If it is not a path, tries to construct a
+                model from the Hugging Face Hub with that name. Defaults to
+                'jinaai/jina-clip-v2'
+            tokenizer_name_or_path (str, optional): If it is a filepath on disc, it
+                loads the tokenizer from that path. If it is not a path, tries to
+                construct a tokenizer from the Hugging Face Hub with that name.
+                If `None` it is automatically set to the value of `model_name_or_path`
+            image_processor_name_or_path (str, optional): If it is a filepath on disc,
+                it loads the image processor from that path. If it is not a path, tries
+                to construct an image processor from the Hugging Face Hub with that
+                name. If `None` it is automatically set to the value of
+                `model_name_or_path`
+            max_seq_length (int, optional): The maximum sequence length of the model.
+                If not provided, will be inferred from model or tokenizer
+            config_args (Dict[str, Any], optional): Additional model configuration
+                parameters to be passed to the Hugging Face Transformers config
+            model_args (Dict[str, Any], optional): Additional model configuration
+                parameters to be passed to the Hugging Face Transformers model
+            tokenizer_args (Dict[str, Any], optional): Additional tokenizer
+                configuration parameters to be passed to the Hugging Face Transformers
+                tokenizer
+            image_processor_args (Dict[str, Any], optional): Additional image processor
+                configuration parameters to be passed to the Hugging Face Transformers
+                image processor
+            cache_dir (str, optional): The Hugging Face Hub cache directory
+            backend (str, optional): Computational backend, only 'torch' is supported
+
+        Example:
+            ::
+
+                from sentence_transformers import SentenceTransformer
+
+                model = SentenceTransformer(
+                    'jinaai/jina-clip-v2', trust_remote_code=True
+                )
+                sentences_or_images = [
+                    "The weather is lovely today.",
+                    "It's so sunny outside!",
+                    "/path/to/stadium.jpg",
+                ]
+                embeddings = model.encode(sentences_or_images)
+                print(embeddings.shape)
+                # (3, 1024)
+
+                # Get the similarity scores between all inputs
+                similarities = model.similarity(embeddings, embeddings)
+                print(similarities)
+                # tensor([[1.0000, 0.6817, 0.0492],
+                #         [0.6817, 1.0000, 0.0421],
+                #         [0.0492, 0.0421, 1.0000]])
+        """
         super(Transformer, self).__init__()
-
-        config_kwargs = config_kwargs or {}
-        model_kwargs = model_kwargs or {}
-        tokenizer_kwargs = tokenizer_kwargs or {}
-        image_processor_kwargs = image_processor_kwargs or {}
+        if backend != 'torch':
+            raise ValueError(
+                f'Backend \'{backend}\' is not supported, please use \'torch\' instead'
+            )
+
+        config_kwargs = config_args or {}
+        model_kwargs = model_args or {}
+        tokenizer_kwargs = tokenizer_args or {}
+        image_processor_kwargs = {
+            'token': model_kwargs.get('token', None),
+            'trust_remote_code': model_kwargs.get('trust_remote_code', False),
+            'revision': model_kwargs.get('revision', None),
+            'local_files_only': model_kwargs.get('local_files_only', None),
+        }
+        image_processor_kwargs.update(image_processor_args)
 
         config = AutoConfig.from_pretrained(
             model_name_or_path, cache_dir=cache_dir, **config_kwargs