Generate embeddings using sentence-transformers.
| 9 | |
| 10 | |
| 11 | class EmbeddingGenerator: |
| 12 | """Generate embeddings using sentence-transformers.""" |
| 13 | |
| 14 | def __init__(self, model_name: str = "all-MiniLM-L6-v2", batch_size: int = 32, cache_dir: Optional[str] = None): |
| 15 | """ |
| 16 | Initialize the embedding generator. |
| 17 | |
| 18 | Args: |
| 19 | model_name: Name of the sentence-transformers model |
| 20 | batch_size: Batch size for encoding |
| 21 | cache_dir: Directory to cache embeddings (default: ~/.cache/beir_embeddings) |
| 22 | """ |
| 23 | print(f"Loading embedding model: {model_name}") |
| 24 | self.model = SentenceTransformer(model_name) |
| 25 | self.model_name = model_name |
| 26 | self.batch_size = batch_size |
| 27 | self.embedding_dim = self.model.get_sentence_embedding_dimension() |
| 28 | |
| 29 | if cache_dir is None: |
| 30 | cache_dir = os.path.join(Path.home(), ".cache", "beir_embeddings") |
| 31 | self.cache_dir = Path(cache_dir) |
| 32 | self.cache_dir.mkdir(parents=True, exist_ok=True) |
| 33 | |
| 34 | print(f"Model loaded. Embedding dimension: {self.embedding_dim}") |
| 35 | print(f"Cache directory: {self.cache_dir}") |
| 36 | |
| 37 | def _get_cache_path(self, dataset_name: str, data_type: str) -> Path: |
| 38 | """Get the cache file path for a dataset.""" |
| 39 | safe_model_name = self.model_name.replace("/", "_") |
| 40 | filename = f"{dataset_name}_{data_type}_{safe_model_name}.npz" |
| 41 | return self.cache_dir / filename |
| 42 | |
| 43 | def _compute_content_hash(self, ids: List[str], texts: List[str]) -> str: |
| 44 | """Compute a hash of the content for cache validation.""" |
| 45 | content = "|".join(f"{id_}:{text[:100]}" for id_, text in zip(ids, texts, strict=True)) |
| 46 | return hashlib.md5(content.encode(), usedforsecurity=False).hexdigest() |
| 47 | |
| 48 | def _load_cached_embeddings(self, cache_path: Path, content_hash: str) -> Optional[Dict[str, np.ndarray]]: |
| 49 | """Load embeddings from cache if valid.""" |
| 50 | if not cache_path.exists(): |
| 51 | return None |
| 52 | |
| 53 | try: |
| 54 | data = np.load(cache_path, allow_pickle=True) |
| 55 | if data.get("content_hash", "") != content_hash: |
| 56 | print("Cache content mismatch, regenerating embeddings...") |
| 57 | return None |
| 58 | |
| 59 | ids = data["ids"] |
| 60 | embeddings = data["embeddings"] |
| 61 | print(f"Loaded {len(ids)} cached embeddings from {cache_path.name}") |
| 62 | return {id_: emb for id_, emb in zip(ids, embeddings, strict=True)} |
| 63 | except Exception as e: |
| 64 | print(f"Failed to load cache: {e}") |
| 65 | return None |
| 66 | |
| 67 | def _save_embeddings_to_cache(self, cache_path: Path, embedding_map: Dict[str, np.ndarray], content_hash: str): |
| 68 | """Save embeddings to cache.""" |