Spaces:
Running
Running
File size: 1,939 Bytes
2b523d0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | """
Embedding model for document and query vectorization
"""
import torch
from sentence_transformers import SentenceTransformer
from typing import List
import numpy as np
from config import EMBEDDING_MODEL, DEVICE, MODELS_DIR
class EmbeddingModel:
def __init__(self):
print(f"Loading embedding model: {EMBEDDING_MODEL}")
self.model = SentenceTransformer(
EMBEDDING_MODEL,
cache_folder=str(MODELS_DIR),
device=DEVICE
)
self.dimension = self.model.get_sentence_embedding_dimension()
print(f"✓ Embedding model loaded (dimension: {self.dimension})")
def encode(self, texts: List[str], batch_size: int = 32) -> np.ndarray:
"""
Encode texts into embeddings
Args:
texts: List of text strings
batch_size: Batch size for encoding
Returns:
Numpy array of embeddings
"""
if not texts:
return np.array([])
embeddings = self.model.encode(
texts,
batch_size=batch_size,
show_progress_bar=True,
convert_to_numpy=True,
normalize_embeddings=True # L2 normalization for cosine similarity
)
return embeddings
def encode_query(self, query: str) -> np.ndarray:
"""
Encode a single query
Args:
query: Query string
Returns:
Numpy array of embedding
"""
return self.model.encode(
query,
convert_to_numpy=True,
normalize_embeddings=True
)
# Singleton instance
_embedding_model = None
def get_embedding_model() -> EmbeddingModel:
"""Get or create embedding model instance"""
global _embedding_model
if _embedding_model is None:
_embedding_model = EmbeddingModel()
return _embedding_model
|