AdrienB134 commited on
Commit
ceaa8f3
·
verified ·
1 Parent(s): 0677e27

Update rag_demo/preprocessing/base/embeddings.py

Browse files
rag_demo/preprocessing/base/embeddings.py CHANGED
@@ -1,15 +1,9 @@
1
  from functools import cached_property
2
  from pathlib import Path
3
  from typing import Optional, ClassVar
4
- from threading import Lock
5
 
6
- import numpy as np
7
- from loguru import logger
8
- from numpy.typing import NDArray
9
- from sentence_transformers.SentenceTransformer import SentenceTransformer
10
- from transformers import AutoTokenizer
11
 
12
- from rag_demo.settings import settings
13
 
14
 
15
  class SingletonMeta(type):
@@ -49,97 +43,3 @@ class SingletonMeta(type):
49
  return cls._instances[cls]
50
 
51
 
52
- class EmbeddingModelSingleton(metaclass=SingletonMeta):
53
- """
54
- A singleton class that provides a pre-trained transformer model for generating embeddings of input text.
55
- """
56
-
57
- def __init__(
58
- self,
59
- model_id: str = settings.TEXT_EMBEDDING_MODEL_ID,
60
- device: str = settings.RAG_MODEL_DEVICE,
61
- cache_dir: Optional[Path] = None,
62
- ) -> None:
63
- self._model_id = model_id
64
- self._device = device
65
-
66
- self._model = SentenceTransformer(
67
- self._model_id,
68
- device=self._device,
69
- cache_folder=str(cache_dir) if cache_dir else None,
70
- )
71
- self._model.eval()
72
-
73
- @property
74
- def model_id(self) -> str:
75
- """
76
- Returns the identifier of the pre-trained transformer model to use.
77
-
78
- Returns:
79
- str: The identifier of the pre-trained transformer model to use.
80
- """
81
-
82
- return self._model_id
83
-
84
- @cached_property
85
- def embedding_size(self) -> int:
86
- """
87
- Returns the size of the embeddings generated by the pre-trained transformer model.
88
-
89
- Returns:
90
- int: The size of the embeddings generated by the pre-trained transformer model.
91
- """
92
-
93
- dummy_embedding = self._model.encode("")
94
-
95
- return dummy_embedding.shape[0]
96
-
97
- @property
98
- def max_input_length(self) -> int:
99
- """
100
- Returns the maximum length of input text to tokenize.
101
-
102
- Returns:
103
- int: The maximum length of input text to tokenize.
104
- """
105
-
106
- return self._model.max_seq_length
107
-
108
- @property
109
- def tokenizer(self) -> AutoTokenizer:
110
- """
111
- Returns the tokenizer used to tokenize input text.
112
-
113
- Returns:
114
- AutoTokenizer: The tokenizer used to tokenize input text.
115
- """
116
-
117
- return self._model.tokenizer
118
-
119
- def __call__(
120
- self, input_text: str | list[str], to_list: bool = True
121
- ) -> NDArray[np.float32] | list[float] | list[list[float]]:
122
- """
123
- Generates embeddings for the input text using the pre-trained transformer model.
124
-
125
- Args:
126
- input_text (str): The input text to generate embeddings for.
127
- to_list (bool): Whether to return the embeddings as a list or numpy array. Defaults to True.
128
-
129
- Returns:
130
- Union[np.ndarray, list]: The embeddings generated for the input text.
131
- """
132
-
133
- try:
134
- embeddings = self._model.encode(input_text)
135
- except Exception:
136
- logger.error(
137
- f"Error generating embeddings for {self._model_id=} and {input_text=}"
138
- )
139
-
140
- return [] if to_list else np.array([])
141
-
142
- if to_list:
143
- embeddings = embeddings.tolist()
144
-
145
- return embeddings
 
1
  from functools import cached_property
2
  from pathlib import Path
3
  from typing import Optional, ClassVar
 
4
 
 
 
 
 
 
5
 
6
+
7
 
8
 
9
  class SingletonMeta(type):
 
43
  return cls._instances[cls]
44
 
45