Spaces:
Runtime error
Runtime error
File size: 2,844 Bytes
f3997d4 | 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 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | """
Embedding generation utilities using sentence transformers.
"""
from typing import List, Union
from sentence_transformers import SentenceTransformer
import numpy as np
class EmbeddingGenerator:
"""Handles text embedding generation."""
def __init__(self, model_name: str = 'all-MiniLM-L6-v2'):
"""
Initialize embedding generator.
Args:
model_name: Name of the sentence transformer model
"""
self.model_name = model_name
self.model = SentenceTransformer(model_name)
print(f"[Embeddings] Loaded model: {model_name}")
def generate_embedding(self, text: str) -> List[float]:
"""
Generate embedding for a single text.
Args:
text: Input text
Returns:
Embedding vector as list of floats
"""
embedding = self.model.encode([text])[0]
return embedding.tolist()
def generate_embeddings(self, texts: List[str]) -> List[List[float]]:
"""
Generate embeddings for multiple texts (batch processing).
Args:
texts: List of input texts
Returns:
List of embedding vectors
"""
if not texts:
return []
embeddings = self.model.encode(texts, show_progress_bar=len(texts) > 10)
return embeddings.tolist()
def compute_similarity(
self,
embedding1: Union[List[float], np.ndarray],
embedding2: Union[List[float], np.ndarray]
) -> float:
"""
Compute cosine similarity between two embeddings.
Args:
embedding1: First embedding vector
embedding2: Second embedding vector
Returns:
Cosine similarity score (0-1)
"""
# Convert to numpy arrays if needed
emb1 = np.array(embedding1) if isinstance(embedding1, list) else embedding1
emb2 = np.array(embedding2) if isinstance(embedding2, list) else embedding2
# Compute cosine similarity
dot_product = np.dot(emb1, emb2)
norm1 = np.linalg.norm(emb1)
norm2 = np.linalg.norm(emb2)
if norm1 == 0 or norm2 == 0:
return 0.0
return float(dot_product / (norm1 * norm2))
def get_model_info(self) -> dict:
"""
Get information about the loaded model.
Returns:
Dictionary with model information
"""
return {
"model_name": self.model_name,
"embedding_dimension": self.model.get_sentence_embedding_dimension(),
"max_seq_length": self.model.max_seq_length
}
# Global embedding generator instance
embedding_generator = EmbeddingGenerator()
|