Spaces:
Paused
Paused
Commit ·
2a8bb45
1
Parent(s): 62a10a5
feat: Optimize embedder runtime memory allocation
Browse files- 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 |
|