Create custom_st.py

custom_st.py

@@ -0,0 +1,196 @@
+import base64
+import json
+import os
+from io import BytesIO
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import requests
+import torch
+from PIL import Image
+from torch import nn
+from transformers import AutoConfig, AutoImageProcessor, AutoModel, AutoTokenizer
+
+
+class Transformer(nn.Module):
+    """Huggingface AutoModel to generate token embeddings.
+    Loads the correct class, e.g. BERT / RoBERTa etc.
+
+    Args:
+        model_name_or_path: Huggingface models name
+            (https://huggingface.co/models)
+        max_seq_length: Truncate any inputs longer than max_seq_length
+        model_args: Keyword arguments passed to the Huggingface
+            Transformers model
+        tokenizer_args: Keyword arguments passed to the Huggingface
+            Transformers tokenizer
+        config_args: Keyword arguments passed to the Huggingface
+            Transformers config
+        cache_dir: Cache dir for Huggingface Transformers to store/load
+            models
+        do_lower_case: If true, lowercases the input (independent if the
+            model is cased or not)
+        tokenizer_name_or_path: Name or path of the tokenizer. When
+            None, then model_name_or_path is used
+    """
+
+    def __init__(
+        self,
+        model_name_or_path: str,
+        max_seq_length: Optional[int] = None,
+        model_args: Optional[Dict[str, Any]] = None,
+        tokenizer_args: Optional[Dict[str, Any]] = None,
+        config_args: Optional[Dict[str, Any]] = None,
+        cache_dir: Optional[str] = None,
+        do_lower_case: bool = False,
+        tokenizer_name_or_path: str = None,
+    ) -> None:
+        super(Transformer, self).__init__()
+        self.config_keys = ["max_seq_length", "do_lower_case"]
+        self.do_lower_case = do_lower_case
+        if model_args is None:
+            model_args = {}
+        if tokenizer_args is None:
+            tokenizer_args = {}
+        if config_args is None:
+            config_args = {}
+
+        config = AutoConfig.from_pretrained(
+            model_name_or_path, **config_args, cache_dir=cache_dir
+        )
+        self.jina_clip = AutoModel.from_pretrained(
+            model_name_or_path, config=config, cache_dir=cache_dir, **model_args
+        )
+
+        if max_seq_length is not None and "model_max_length" not in tokenizer_args:
+            tokenizer_args["model_max_length"] = max_seq_length
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            (
+                tokenizer_name_or_path
+                if tokenizer_name_or_path is not None
+                else model_name_or_path
+            ),
+            cache_dir=cache_dir,
+            **tokenizer_args,
+        )
+        self.preprocessor = AutoImageProcessor.from_pretrained(
+            (
+                tokenizer_name_or_path
+                if tokenizer_name_or_path is not None
+                else model_name_or_path
+            ),
+            cache_dir=cache_dir,
+            **tokenizer_args,
+        )
+
+        # No max_seq_length set. Try to infer from model
+        if max_seq_length is None:
+            if (
+                hasattr(self.jina_clip, "config")
+                and hasattr(self.jina_clip.config, "max_position_embeddings")
+                and hasattr(self.tokenizer, "model_max_length")
+            ):
+                max_seq_length = min(
+                    self.jina_clip.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.jina_clip.config.tokenizer_class = self.tokenizer.__class__.__name__
+
+    def forward(
+        self, features: Dict[str, torch.Tensor]
+    ) -> Dict[str, torch.Tensor]:
+        """Returns token_embeddings, cls_token"""
+        if "input_ids" in features:
+            embedding = self.jina_clip.get_text_features(
+                input_ids=features["input_ids"]
+            )
+        else:
+            embedding = self.jina_clip.get_image_features(
+                pixel_values=features["pixel_values"]
+            )
+        return {"sentence_embedding": embedding}
+
+    def get_word_embedding_dimension(self) -> int:
+        return self.config.text_config.embed_dim
+
+    def decode_data_image(data_image_str):
+        header, data = data_image_str.split(',', 1)
+        image_data = base64.b64decode(data)
+        return Image.open(BytesIO(image_data))
+
+    def tokenize(
+        self, batch: Union[List[str]], padding: Union[str, bool] = True
+    ) -> Dict[str, torch.Tensor]:
+        """Tokenizes a text and maps tokens to token-ids"""
+        images = []
+        texts = []
+        for sample in batch:
+            if isinstance(sample, str):
+                if sample.startswith('http'):
+                    response = requests.get(sample)
+                    images.append(Image.open(BytesIO(response.content)).convert('RGB'))
+                elif sample.startswith('data:image/'):
+                    images.append(self.decode_data_image(sample).convert('RGB'))
+                else:
+                    # TODO: Make sure that Image.open fails for non-image files
+                    try:
+                        images.append(Image.open(sample).convert('RGB'))
+                    except:
+                        texts.append(sample)
+            elif isinstance(sample, Image.Image):
+                images.append(sample.convert('RGB'))
+
+        if images and texts:
+            raise ValueError('Batch must contain either images or texts, not both')
+
+        if texts:
+            return self.tokenizer(
+                texts,
+                padding=padding,
+                truncation="longest_first",
+                return_tensors="pt",
+                max_length=self.max_seq_length,
+            )
+        elif images:
+            return self.preprocessor(images)
+        return {}
+
+    def save(self, output_path: str, safe_serialization: bool = True) -> None:
+        self.jina_clip.save_pretrained(
+            output_path, safe_serialization=safe_serialization
+        )
+        self.tokenizer.save_pretrained(output_path)
+        self.preprocessor.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 "model_args" in config and "trust_remote_code" in config["model_args"]:
+            config["model_args"].pop("trust_remote_code")
+        if (
+            "tokenizer_args" in config
+            and "trust_remote_code" in config["tokenizer_args"]
+        ):
+            config["tokenizer_args"].pop("trust_remote_code")
+        if "config_args" in config and "trust_remote_code" in config["config_args"]:
+            config["config_args"].pop("trust_remote_code")
+        return Transformer(model_name_or_path=input_path, **config)