documind / embedder.py
Aaravkumar's picture
Update embedder.py
3020061 verified
Raw
History Blame Contribute Delete
985 Bytes
from sentence_transformers import SentenceTransformer
import numpy as np
import time
class Embedder:
def __init__(self):
print("Loading embedding model...")
self.model = SentenceTransformer('sentence-transformers/paraphrase-MiniLM-L6-v2')
print("Embedder loaded: BAAI/bge-large-en-v1.5")
def embed(self, chunks, batch_size=32):
try:
vectors = self.model.encode(
chunks,
batch_size=batch_size,
show_progress_bar=False,
convert_to_numpy=True
)
return vectors.tolist()
except Exception as e:
print(f"Embedding failed: {e}")
raise e
def embed_q(self, query):
try:
v = self.model.encode(query, convert_to_numpy=True)
return v.tolist()
except Exception as e:
print(f"Query embedding failed: {e}")
raise e