ATS-SCORE-CHECKER / embeddings /model_manager.py
Akash-Dragon's picture
Use lightweight paraphrase-MiniLM-L3-v2 model + CPU-only torch
f9a8f86
Raw
History Blame Contribute Delete
1.47 kB
from sentence_transformers import SentenceTransformer
from typing import List
import numpy as np
class EmbeddingModel:
"""Lightweight sentence-transformer for semantic similarity"""
def __init__(self, model_name: str = "paraphrase-MiniLM-L3-v2"):
# Use smallest model (3 layers, ~60MB, ~100MB RAM)
self.model = SentenceTransformer(model_name)
self.model_name = model_name
self.dimension = self.model.get_sentence_embedding_dimension()
def encode(self, texts: List[str], batch_size: int = 32) -> np.ndarray:
"""Generate embeddings for a list of texts"""
if isinstance(texts, str):
texts = [texts]
cleaned_texts = []
for text in texts:
if text and isinstance(text, str):
cleaned = ' '.join(text.split())
if cleaned.strip():
cleaned_texts.append(cleaned)
if not cleaned_texts:
return np.array([])
embeddings = self.model.encode(cleaned_texts, batch_size=batch_size, convert_to_numpy=True)
return embeddings
def encode_single(self, text: str) -> np.ndarray:
"""Generate embedding for a single text"""
cleaned = ' '.join(text.split()) if text else ""
if not cleaned.strip():
return np.array([])
embedding = self.model.encode([cleaned], convert_to_numpy=True)
return embedding[0]