| import torch |
| from sentence_transformers import SentenceTransformer |
| from collections import OrderedDict |
| import constants |
|
|
| print("AI Model Load हो रहा है...") |
| model = SentenceTransformer(constants.MODEL_NAME, trust_remote_code=True) |
| if torch.cuda.is_available(): |
| model = model.to("cuda") |
|
|
| |
| class BoundedEmbeddingCache: |
| def __init__(self, maxsize=2000): |
| self.cache = OrderedDict() |
| self.maxsize = maxsize |
|
|
| def get(self, text): |
| if text in self.cache: |
| self.cache.move_to_end(text) |
| return self.cache[text] |
| return None |
|
|
| def put(self, text, emb): |
| self.cache[text] = emb |
| self.cache.move_to_end(text) |
| if len(self.cache) > self.maxsize: |
| self.cache.popitem(last=False) |
|
|
| EMBEDDING_CACHE = BoundedEmbeddingCache(maxsize=2000) |
|
|
| def get_batch_embeddings(texts): |
| uncached_texts = [t for t in texts if EMBEDDING_CACHE.get(t) is None] |
| if uncached_texts: |
| with torch.no_grad(): |
| embs = model.encode(uncached_texts, convert_to_tensor=True, normalize_embeddings=True) |
| for t, e in zip(uncached_texts, embs): |
| EMBEDDING_CACHE.put(t, e.cpu()) |
| |
| return torch.stack([EMBEDDING_CACHE.get(t) for t in texts]).to(model.device) |
| |
