rishik1111 commited on
Commit
2a8bb45
·
1 Parent(s): 62a10a5

feat: Optimize embedder runtime memory allocation

Browse files
Files changed (1) hide show
  1. vector_search/embed.py +4 -0
vector_search/embed.py CHANGED
@@ -190,6 +190,9 @@ class PyTorchMultilingualE5Embedder:
190
  PyTorch fallback wrapper for sentence-transformers multilingual-e5-small.
191
  """
192
  def __init__(self, model_name: str = config.EMBEDDING_MODEL_NAME):
 
 
 
193
  from sentence_transformers import SentenceTransformer
194
  logger.info(f"Loading PyTorch fallback embedding model: '{model_name}'...")
195
  self.model_name = model_name
@@ -197,6 +200,7 @@ class PyTorchMultilingualE5Embedder:
197
  self.model = SentenceTransformer(model_name, local_files_only=True)
198
  except Exception:
199
  self.model = SentenceTransformer(model_name)
 
200
  self.dim = config.EMBEDDING_DIM
201
  logger.info(f"PyTorch embedding model loaded (dim={self.dim}).")
202
 
 
190
  PyTorch fallback wrapper for sentence-transformers multilingual-e5-small.
191
  """
192
  def __init__(self, model_name: str = config.EMBEDDING_MODEL_NAME):
193
+ import torch
194
+ torch.set_grad_enabled(False)
195
+ torch.set_num_threads(1)
196
  from sentence_transformers import SentenceTransformer
197
  logger.info(f"Loading PyTorch fallback embedding model: '{model_name}'...")
198
  self.model_name = model_name
 
200
  self.model = SentenceTransformer(model_name, local_files_only=True)
201
  except Exception:
202
  self.model = SentenceTransformer(model_name)
203
+ self.model.eval()
204
  self.dim = config.EMBEDDING_DIM
205
  logger.info(f"PyTorch embedding model loaded (dim={self.dim}).")
206