Upload 6 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

config.json

@@ -1,205 +1,70 @@
 {
-  "_commit_hash": "344d954da76eb8ad47a7aaff42d012e30c15b8fe",
-  "add_projections": false,
-  "architectures": [
-    "JinaCLIPModel"
-  ],
-  "auto_map": {
-    "AutoConfig": "jinaai/jina-clip-implementation--configuration_clip.JinaCLIPConfig",
-    "AutoModel": "jinaai/jina-clip-implementation--modeling_clip.JinaCLIPModel"
-  },
-  "initializer_factor": 1.0,
-  "logit_scale_init_value": 2.6592,
-  "matryoshka_dimensions": [
-    32,
-    64,
-    128,
-    256,
-    512,
-    768,
-    1024
-  ],
-  "model_type": "jina_clip",
-  "projection_dim": 1024,
-  "text_config": {
-    "_attn_implementation_autoset": false,
-    "_name_or_path": "",
-    "add_cross_attention": false,
-    "architectures": null,
-    "bad_words_ids": null,
-    "begin_suppress_tokens": null,
-    "bos_token_id": null,
-    "chunk_size_feed_forward": 0,
-    "cross_attention_hidden_size": null,
-    "decoder_start_token_id": null,
-    "default_instruction_task": null,
-    "default_lora_task": "retrieval.query",
-    "diversity_penalty": 0.0,
-    "do_sample": false,
-    "early_stopping": false,
-    "embed_dim": 1024,
-    "encoder_no_repeat_ngram_size": 0,
-    "eos_token_id": null,
-    "exponential_decay_length_penalty": null,
-    "finetuning_task": null,
-    "forced_bos_token_id": null,
-    "forced_eos_token_id": null,
-    "hf_model_config_kwargs": {
-      "load_trained_adapters": false,
-      "lora_adaptations": [
-        "retrieval.query"
-      ],
-      "lora_alpha": 4,
-      "lora_dropout_p": 0.0,
-      "lora_main_params_trainable": false,
-      "lora_rank": 4,
-      "task_instructions": {
-        "retrieval.query": "Represent the query for retrieving evidence documents: "
-      },
-      "use_flash_attn": false
+    "add_projections": false,
+    "architectures": [
+        "JinaCLIPModel"
+    ],
+    "auto_map": {
+        "AutoConfig": "jinaai/jina-clip-implementation--configuration_clip.JinaCLIPConfig",
+        "AutoModel": "jinaai/jina-clip-implementation--modeling_clip.JinaCLIPModel"
     },
-    "hf_model_name_or_path": "jinaai/jina-embeddings-v3",
-    "id2label": {
-      "0": "LABEL_0",
-      "1": "LABEL_1"
+    "initializer_factor": 1.0,
+    "logit_scale_init_value": 2.6592,
+    "matryoshka_dimensions": [32, 64, 128, 256, 512, 768, 1024],
+    "model_type": "jina_clip",
+    "projection_dim": 1024,
+    "text_config": {
+        "default_instruction_task": null,
+        "default_lora_task": "retrieval.query",
+        "embed_dim": 1024,
+        "hf_model_config_kwargs": {
+            "load_trained_adapters": false,
+            "lora_adaptations": [
+                "retrieval.query"
+            ],
+            "lora_alpha": 4,
+            "lora_dropout_p": 0.0,
+            "lora_main_params_trainable": false,
+            "lora_rank": 4,
+            "task_instructions": {
+                "retrieval.query": "Represent the query for retrieving evidence documents: "
+            },
+            "use_flash_attn": true
+        },
+        "hf_model_name_or_path": "jinaai/jina-embeddings-v3",
+        "model_type": "jina_clip_text",
+        "pooler_type": "mean_pooler",
+        "proj_bias": false,
+        "proj_type": null
     },
-    "is_decoder": false,
-    "is_encoder_decoder": false,
-    "label2id": {
-      "LABEL_0": 0,
-      "LABEL_1": 1
+    "torch_dtype": "bfloat16",
+    "transformers.js_config": {
+        "use_external_data_format": {
+            "model.onnx": true
+        }
     },
-    "length_penalty": 1.0,
-    "max_length": 20,
-    "min_length": 0,
-    "model_type": "jina_clip_text",
-    "no_repeat_ngram_size": 0,
-    "num_beam_groups": 1,
-    "num_beams": 1,
-    "num_return_sequences": 1,
-    "output_attentions": false,
-    "output_hidden_states": false,
-    "output_scores": false,
-    "pad_token_id": null,
-    "pooler_type": "mean_pooler",
-    "prefix": null,
-    "problem_type": null,
-    "proj_bias": false,
-    "proj_type": null,
-    "pruned_heads": {},
-    "remove_invalid_values": false,
-    "repetition_penalty": 1.0,
-    "return_dict": true,
-    "return_dict_in_generate": false,
-    "sep_token_id": null,
-    "suppress_tokens": null,
-    "task_specific_params": null,
-    "temperature": 1.0,
-    "tf_legacy_loss": false,
-    "tie_encoder_decoder": false,
-    "tie_word_embeddings": true,
-    "tokenizer_class": null,
-    "top_k": 50,
-    "top_p": 1.0,
-    "torch_dtype": null,
-    "torchscript": false,
-    "transformers_version": "4.51.3",
-    "typical_p": 1.0,
-    "use_bfloat16": false
-  },
-  "torch_dtype": "bfloat16",
-  "transformers.js_config": {
-    "use_external_data_format": {
-      "model.onnx": true
+    "truncate_dim": null,
+    "use_text_flash_attn": null,
+    "use_vision_xformers": null,
+    "vision_config": {
+        "embed_dim": 1024,
+        "fused_layer_norm": false,
+        "head_width": 64,
+        "image_size": 512,
+        "intp_freq": true,
+        "layers": 24,
+        "ls_init_value": null,
+        "mlp_ratio": 2.6667,
+        "model_type": "jina_clip_vision",
+        "naive_swiglu": true,
+        "patch_dropout": 0.1,
+        "patch_size": 14,
+        "post_norm": false,
+        "proj_type": null,
+        "pt_hw_seq_len": 16,
+        "qkv_bias": true,
+        "rope_embeddings": true,
+        "subln": true,
+        "width": 1024,
+        "x_attention": true
     }
-  },
-  "transformers_version": null,
-  "truncate_dim": null,
-  "use_text_flash_attn": false,
-  "use_vision_xformers": false,
-  "vision_config": {
-    "_attn_implementation_autoset": false,
-    "_name_or_path": "",
-    "add_cross_attention": false,
-    "architectures": null,
-    "bad_words_ids": null,
-    "begin_suppress_tokens": null,
-    "bos_token_id": null,
-    "chunk_size_feed_forward": 0,
-    "cross_attention_hidden_size": null,
-    "decoder_start_token_id": null,
-    "diversity_penalty": 0.0,
-    "do_sample": false,
-    "drop_path_rate": 0.0,
-    "early_stopping": false,
-    "embed_dim": 1024,
-    "encoder_no_repeat_ngram_size": 0,
-    "eos_token_id": null,
-    "exponential_decay_length_penalty": null,
-    "finetuning_task": null,
-    "forced_bos_token_id": null,
-    "forced_eos_token_id": null,
-    "fused_layer_norm": false,
-    "head_width": 64,
-    "id2label": {
-      "0": "LABEL_0",
-      "1": "LABEL_1"
-    },
-    "image_size": 512,
-    "intp_freq": true,
-    "is_decoder": false,
-    "is_encoder_decoder": false,
-    "label2id": {
-      "LABEL_0": 0,
-      "LABEL_1": 1
-    },
-    "layers": 24,
-    "length_penalty": 1.0,
-    "ls_init_value": null,
-    "max_length": 20,
-    "min_length": 0,
-    "mlp_ratio": 2.6667,
-    "model_type": "jina_clip_vision",
-    "naive_swiglu": true,
-    "no_repeat_ngram_size": 0,
-    "num_beam_groups": 1,
-    "num_beams": 1,
-    "num_return_sequences": 1,
-    "output_attentions": false,
-    "output_hidden_states": false,
-    "output_scores": false,
-    "pad_token_id": null,
-    "patch_dropout": 0.1,
-    "patch_size": 14,
-    "post_norm": false,
-    "prefix": null,
-    "problem_type": null,
-    "proj_type": null,
-    "pruned_heads": {},
-    "pt_hw_seq_len": 16,
-    "qkv_bias": true,
-    "remove_invalid_values": false,
-    "repetition_penalty": 1.0,
-    "return_dict": true,
-    "return_dict_in_generate": false,
-    "rope_embeddings": true,
-    "sep_token_id": null,
-    "subln": true,
-    "suppress_tokens": null,
-    "task_specific_params": null,
-    "temperature": 1.0,
-    "tf_legacy_loss": false,
-    "tie_encoder_decoder": false,
-    "tie_word_embeddings": true,
-    "tokenizer_class": null,
-    "top_k": 50,
-    "top_p": 1.0,
-    "torch_dtype": null,
-    "torchscript": false,
-    "transformers_version": "4.51.3",
-    "typical_p": 1.0,
-    "use_bfloat16": false,
-    "width": 1024,
-    "x_attention": false
-  }
-}
+}

config_sentence_transformers.json

@@ -0,0 +1,12 @@
+{
+  "__version__": {
+    "sentence_transformers": "3.3.0",
+    "transformers": "4.46.2",
+    "pytorch": "2.2.2"
+  },
+  "prompts":{
+    "retrieval.query":"Represent the query for retrieving evidence documents: "
+  },
+  "default_prompt_name": null,
+  "similarity_fn_name": "cosine"
+}

custom_st.py

@@ -0,0 +1,275 @@
+import base64
+import json
+import os
+from io import BytesIO
+from typing import Any, Dict, List, Literal, Optional, Union
+
+import requests
+import torch
+from PIL import Image
+from torch import nn
+from transformers import AutoConfig, AutoImageProcessor, AutoModel, AutoTokenizer
+
+
+class Transformer(nn.Module):
+
+    save_in_root: bool = True
+    
+    def __init__(
+        self,
+        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_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,
+        assume_text_inputs: bool = False,
+        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
+            assume_text_inputs (bool, optional): If set to `True`, all inputs are
+                treated as texts. Defaults to `False`
+            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__()
+        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 or {})
+
+        config = AutoConfig.from_pretrained(
+            model_name_or_path, cache_dir=cache_dir, **config_kwargs
+        )
+        self.model = AutoModel.from_pretrained(
+            model_name_or_path, config=config, cache_dir=cache_dir, **model_kwargs
+        )
+        if max_seq_length is not None and 'model_max_length' not in tokenizer_kwargs:
+            tokenizer_kwargs['model_max_length'] = max_seq_length
+
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            tokenizer_name_or_path or model_name_or_path,
+            cache_dir=cache_dir,
+            **tokenizer_kwargs,
+        )
+        self.image_processor = AutoImageProcessor.from_pretrained(
+            image_processor_name_or_path or model_name_or_path,
+            cache_dir=cache_dir,
+            **image_processor_kwargs,
+        )
+        self.assume_text_inputs = assume_text_inputs
+
+        # No max_seq_length set. Try to infer from model
+        if max_seq_length is None:
+            if (
+                hasattr(self.model, 'config')
+                and hasattr(self.model.config, 'max_position_embeddings')
+                and hasattr(self.tokenizer, 'model_max_length')
+            ):
+                max_seq_length = min(
+                    self.model.config.max_position_embeddings,
+                    self.tokenizer.model_max_length,
+                )
+        self.max_seq_length = max_seq_length
+        if tokenizer_name_or_path is not None:
+            self.model.config.tokenizer_class = self.tokenizer.__class__.__name__
+
+    @staticmethod
+    def _decode_data_image(data_image_str: str) -> Image.Image:
+        header, data = data_image_str.split(',', 1)
+        image_data = base64.b64decode(data)
+        return Image.open(BytesIO(image_data))
+
+    def tokenize(
+        self, texts: List[Union[str, Image.Image]], padding: Union[str, bool] = True
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Encodes input samples. Text samples are tokenized. Image URLs, image data
+        buffers and PIL images are passed through the image processor.
+        """
+        _images = []
+        _texts = []
+        _image_or_text_descriptors = []
+
+        if self.assume_text_inputs:
+            for sample in texts:
+                if isinstance(sample, str):
+                    _texts.append(sample)
+                    _image_or_text_descriptors.append(1)
+        else:
+            for sample in texts:
+                if isinstance(sample, str):
+                    if sample.startswith('http'):
+                        try:
+                            response = requests.get(sample)
+                            _images.append(
+                                Image.open(BytesIO(response.content)).convert('RGB')
+                            )
+                            _image_or_text_descriptors.append(0)
+                        except Exception as e:
+                            _ = str(e)
+                            _texts.append(sample)
+                            _image_or_text_descriptors.append(1)
+                    elif sample.startswith('data:image/'):
+                        _images.append(self._decode_data_image(sample).convert('RGB'))
+                        _image_or_text_descriptors.append(0)
+                    else:
+                        try:
+                            _images.append(Image.open(sample).convert('RGB'))
+                            _image_or_text_descriptors.append(0)
+                        except Exception as e:
+                            _ = str(e)
+                            _texts.append(sample)
+                            _image_or_text_descriptors.append(1)
+                elif isinstance(sample, Image.Image):
+                    _images.append(sample.convert('RGB'))
+                    _image_or_text_descriptors.append(0)
+
+        encoding = {}
+        if len(_texts):
+            encoding['input_ids'] = self.tokenizer(
+                _texts,
+                padding=padding,
+                truncation='longest_first',
+                return_tensors='pt',
+                max_length=self.max_seq_length,
+            ).input_ids
+
+        if len(_images):
+            encoding['pixel_values'] = self.image_processor(
+                _images, return_tensors='pt'
+            ).pixel_values
+
+        encoding['image_text_info'] = _image_or_text_descriptors
+        return encoding
+
+    def forward(self, features: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+        image_embeddings = []
+        text_embeddings = []
+
+        if 'pixel_values' in features:
+            image_embeddings = self.model.get_image_features(features['pixel_values'])
+        if 'input_ids' in features:
+            text_embeddings = self.model.get_text_features(features['input_ids'])
+
+        sentence_embedding = []
+        image_features = iter(image_embeddings)
+        text_features = iter(text_embeddings)
+        for _, _input_type in enumerate(features['image_text_info']):
+            if _input_type == 0:
+                sentence_embedding.append(next(image_features))
+            else:
+                sentence_embedding.append(next(text_features))
+
+        features['sentence_embedding'] = torch.stack(sentence_embedding).float()
+        return features
+
+    def save(self, output_path: str, safe_serialization: bool = True) -> None:
+        self.model.save_pretrained(output_path, safe_serialization=safe_serialization)
+        self.tokenizer.save_pretrained(output_path)
+        self.image_processor.save_pretrained(output_path)
+
+    @staticmethod
+    def load(input_path: str) -> 'Transformer':
+        # Old classes used other config names than 'sentence_bert_config.json'
+        for config_name in [
+            'sentence_bert_config.json',
+            'sentence_roberta_config.json',
+            'sentence_distilbert_config.json',
+            'sentence_camembert_config.json',
+            'sentence_albert_config.json',
+            'sentence_xlm-roberta_config.json',
+            'sentence_xlnet_config.json',
+        ]:
+            sbert_config_path = os.path.join(input_path, config_name)
+            if os.path.exists(sbert_config_path):
+                break
+
+        with open(sbert_config_path) as fIn:
+            config = json.load(fIn)
+
+        # Don't allow configs to set trust_remote_code
+        if 'config_kwargs' in config and 'trust_remote_code' in config['config_kwargs']:
+            config['config_kwargs'].pop('trust_remote_code')
+        if 'model_kwargs' in config and 'trust_remote_code' in config['model_kwargs']:
+            config['model_kwargs'].pop('trust_remote_code')
+        if (
+            'tokenizer_kwargs' in config
+            and 'trust_remote_code' in config['tokenizer_kwargs']
+        ):
+            config['tokenizer_kwargs'].pop('trust_remote_code')
+        if (
+            'image_processor_kwargs' in config
+            and 'trust_remote_code' in config['image_processor_kwargs']
+        ):
+            config['image_processor_kwargs'].pop('trust_remote_code')
+
+        return Transformer(model_name_or_path=input_path, **config)

special_tokens_map.json

@@ -0,0 +1,51 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "cls_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "<mask>",
+    "lstrip": true,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6601c4120779a1a3863897ba332fe3481d548e363bec2c91eba10ef8640a5e93
+size 17082997

tokenizer_config.json

@@ -0,0 +1,54 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "250001": {
+      "content": "<mask>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<s>",
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "<s>",
+  "eos_token": "</s>",
+  "mask_token": "<mask>",
+  "model_max_length": 8194,
+  "pad_token": "<pad>",
+  "sep_token": "</s>",
+  "tokenizer_class": "XLMRobertaTokenizer",
+  "unk_token": "<unk>"
+}