bwang0911 commited on
Commit
4be32c2
1 Parent(s): 837f25a

Add compatibility with Sentence Transformers v3.2.0 (#45)

Browse files

- Add **kwargs to Transformer to avoid issues when ST adds new arguments (40df588888e65ac3fa8ae1262235425a79ebbd22)

Files changed (1) hide show
  1. custom_st.py +9 -0
custom_st.py CHANGED
@@ -1,4 +1,5 @@
1
  import json
 
2
  import os
3
  from io import BytesIO
4
  from typing import Any, Dict, List, Optional, Tuple, Union
@@ -7,6 +8,8 @@ import torch
7
  from torch import nn
8
  from transformers import AutoConfig, AutoModel, AutoTokenizer
9
 
 
 
10
 
11
  class Transformer(nn.Module):
12
  """Huggingface AutoModel to generate token embeddings.
@@ -40,6 +43,7 @@ class Transformer(nn.Module):
40
  cache_dir: str = None,
41
  do_lower_case: bool = False,
42
  tokenizer_name_or_path: str = None,
 
43
  ) -> None:
44
  super().__init__()
45
  self.config_keys = ["max_seq_length", "do_lower_case"]
@@ -51,6 +55,11 @@ class Transformer(nn.Module):
51
  if config_args is None:
52
  config_args = {}
53
 
 
 
 
 
 
54
 
55
  self.config = AutoConfig.from_pretrained(model_name_or_path, **config_args, cache_dir=cache_dir)
56
 
 
1
  import json
2
+ import logging
3
  import os
4
  from io import BytesIO
5
  from typing import Any, Dict, List, Optional, Tuple, Union
 
8
  from torch import nn
9
  from transformers import AutoConfig, AutoModel, AutoTokenizer
10
 
11
+ logger = logging.getLogger(__name__)
12
+
13
 
14
  class Transformer(nn.Module):
15
  """Huggingface AutoModel to generate token embeddings.
 
43
  cache_dir: str = None,
44
  do_lower_case: bool = False,
45
  tokenizer_name_or_path: str = None,
46
+ **kwargs,
47
  ) -> None:
48
  super().__init__()
49
  self.config_keys = ["max_seq_length", "do_lower_case"]
 
55
  if config_args is None:
56
  config_args = {}
57
 
58
+ if kwargs.get("backend", "torch") != "torch":
59
+ logger.warning(
60
+ f'"jinaai/jina-embeddings-v3" is currently not compatible with the {kwargs["backend"]} backend. '
61
+ 'Continuing with the "torch" backend.'
62
+ )
63
 
64
  self.config = AutoConfig.from_pretrained(model_name_or_path, **config_args, cache_dir=cache_dir)
65