Spaces:
Sleeping
Sleeping
File size: 2,045 Bytes
f5b0cd7 | 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 | """
Module for generating document embeddings using FastEmbed.
"""
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from typing import List
import logging
from fastembed import TextEmbedding
class EmbeddingService:
def __init__(self, model_name: str = "BAAI/bge-small-en-v1.5"):
"""
Initialize the embedding service with the specified model.
By default, uses the BGE small model which generates embeddings.
"""
try:
self.model = TextEmbedding(model_name=model_name)
# Verify that the model produces 384-dimensional embeddings
sample_embedding = list(self.model.embed(["test"]).__next__())
if len(sample_embedding) != 384:
raise ValueError(f"Model {model_name} does not produce 384-dimensional embeddings")
except Exception as e:
logging.error(f"Failed to initialize embedding model: {e}")
raise
def embed_text(self, text: str) -> List[float]:
"""
Generate embedding for a single text string.
Args:
text: Input text to embed
Returns:
384-dimensional embedding vector as a list of floats
"""
try:
embeddings = list(self.model.embed([text]))
return embeddings[0]
except Exception as e:
logging.error(f"Failed to generate embedding for text: {e}")
raise
def embed_texts(self, texts: List[str]) -> List[List[float]]:
"""
Generate embeddings for multiple text strings.
Args:
texts: List of input texts to embed
Returns:
List of 384-dimensional embedding vectors
"""
try:
embeddings = list(self.model.embed(texts))
return [emb.tolist() for emb in embeddings]
except Exception as e:
logging.error(f"Failed to generate embeddings for texts: {e}")
raise |