Spaces:
Running
Running
| # src/embeddings/huggingface_embedding.py | |
| from typing import List | |
| from sentence_transformers import SentenceTransformer | |
| from .base_embedding import BaseEmbedding | |
| class HuggingFaceEmbedding(BaseEmbedding): | |
| def __init__(self, model_name: str = 'all-MiniLM-L6-v2'): | |
| """ | |
| Initialize HuggingFace embedding model | |
| Args: | |
| model_name (str): Name of the embedding model | |
| """ | |
| self.model = SentenceTransformer(model_name) | |
| def embed_documents(self, texts: List[str]) -> List[List[float]]: | |
| """ | |
| Embed a list of documents | |
| Args: | |
| texts (List[str]): List of texts to embed | |
| Returns: | |
| List[List[float]]: List of embeddings | |
| """ | |
| return self.model.encode(texts).tolist() | |
| def embed_query(self, text: str) -> List[float]: | |
| """ | |
| Embed a single query | |
| Args: | |
| text (str): Text to embed | |
| Returns: | |
| List[float]: Embedding vector | |
| """ | |
| return self.model.encode(text).tolist() |