refactor: refine encode_text

modeling_clip.py

@@ -18,6 +18,12 @@ from transformers.models.clip.modeling_clip import (
     CLIPVisionModelOutput,
     clip_loss,
 )
+try:
+    from tqdm.autonotebook import trange
+
+    has_tqdm = True
+except ImportError:
+    has_tqdm = False
 
 from .configuration_clip import JinaCLIPConfig, JinaCLIPTextConfig, JinaCLIPVisionConfig
 from .eva_model import EVAVisionTransformer
@@ -215,6 +221,7 @@ class JinaCLIPModel(JinaCLIPPreTrainedModel):
             self.visual_projection = nn.Identity()
             self.text_projection = nn.Identity()
 
+        self.tokenizer = AutoTokenizer.from_pretrained(config._name_or_path)
         self.post_init()
 
     def get_text_features(
@@ -239,19 +246,92 @@ class JinaCLIPModel(JinaCLIPPreTrainedModel):
         )
         return self.visual_projection(self.vision_model(x=x))
 
+    @torch.inference_mode()
     def encode_text(
         self,
-        input_ids: Union[None, torch.Tensor, BatchEncoding] = None,
-        return_dict: Optional[bool] = None,
-        *_,
-        **__,
+        sentences: Union[str, List[str]],
+        batch_size: int = 32,
+        show_progress_bar: Optional[bool] = None,
+        output_value: str = 'sentence_embedding',
+        convert_to_numpy: bool = True,
+        convert_to_tensor: bool = False,
+        device: Optional[torch.device] = None,
+        normalize_embeddings: bool = False,
+        **tokenizer_kwargs,
     ) -> Union[Tuple[Optional[torch.FloatTensor], ...], CLIPTextModelOutput]:
-        return_dict = (
-            return_dict if return_dict is not None else self.config.use_return_dict
-        )
-        feats = self.get_text_features(input_ids=input_ids)
-        out = CLIPTextModelOutput(text_embeds=feats)
-        return out if return_dict else out.to_tuple()
+        
+        self.eval()
+        
+        if show_progress_bar is None:
+            show_progress_bar = (
+                logger.getEffectiveLevel() == logging.INFO
+                or logger.getEffectiveLevel() == logging.DEBUG
+            )
+
+        if convert_to_tensor:
+            convert_to_numpy = False
+
+        if output_value != 'sentence_embedding':
+            convert_to_tensor = False
+            convert_to_numpy = False
+
+        input_was_string = False
+        if isinstance(sentences, str) or not hasattr(sentences, '__len__'):
+            sentences = [sentences]
+            input_was_string = True
+
+        if device is not None:
+            self.to(device)
+
+        permutation = np.argsort([-len(i) for i in sentences])
+        inverse_permutation = np.argsort(permutation)
+        sentences = [sentences[idx] for idx in permutation]
+
+        tokenizer_kwargs['padding'] = tokenizer_kwargs.get('padding', True)
+        tokenizer_kwargs['max_length'] = tokenizer_kwargs.get('max_length', 512)
+        tokenizer_kwargs['truncation'] = tokenizer_kwargs.get('truncation', True)
+
+        if has_tqdm:
+            range_iter = trange(
+                0,
+                len(sentences),
+                batch_size,
+                desc="Encoding",
+                disable=not show_progress_bar,
+            )
+        else:
+            range_iter = range(0, len(sentences), batch_size)
+
+        for i in range_iter:
+            encoded_input = self.tokenizer(
+                sentences[i : i + batch_size],
+                return_tensors='pt',
+                **tokenizer_kwargs,
+            ).to(self.device)
+
+            if output_value == 'token_embeddings':
+                raise NotImplementedError
+            elif output_value is None:
+                raise NotImplementedError
+            else:
+                embeddings = self.get_text_features(input_ids=encoded_input)
+                if normalize_embeddings:
+                    embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
+                if convert_to_numpy:
+                    embeddings = embeddings.cpu()
+            all_embeddings.extend(embeddings)
+
+        all_embeddings = [all_embeddings[idx] for idx in inverse_permutation]
+
+        if convert_to_tensor:
+            all_embeddings = torch.stack(all_embeddings)
+        elif convert_to_numpy:
+            all_embeddings = np.asarray([emb.numpy() for emb in all_embeddings])
+
+        if input_was_string:
+            all_embeddings = all_embeddings[0]
+
+        return all_embeddings
 
     def encode_image(
         self,