Spaces:
Sleeping
Sleeping
File size: 4,676 Bytes
edac567 |
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 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 |
"""Generate, store, and query embeddings via Gemini API."""
import numpy as np
import google.generativeai as genai
from typing import List, Dict, Optional
from config.settings import Config
from src.utils.vector_store import VectorStore
class EmbeddingService:
"""Handles embedding generation, storage, and similarity search."""
def __init__(self) -> None:
"""Configure Gemini and initialize vector store."""
Config.validate()
genai.configure(api_key=Config.GEMINI_API_KEY)
self.embedding_model = Config.EMBEDDING_MODEL
self.vector_store = VectorStore(storage_path=Config.EMBEDDING_STORAGE_PATH)
def generate_embeddings(self, texts: List[str]) -> List[np.ndarray]:
"""
Embed a list of document texts.
Args:
texts: List of strings.
Returns:
List of embedding vectors (np.ndarray).
"""
embeddings: List[np.ndarray] = []
for i, text in enumerate(texts):
try:
result = genai.embed_content(
model=self.embedding_model,
content=text,
task_type="retrieval_document",
)
embeddings.append(np.array(result["embedding"]))
except Exception as e:
print(f"[EmbeddingService] Doc embed error idx {i}: {e}")
embeddings.append(np.zeros(768))
return embeddings
def generate_query_embedding(self, query: str) -> np.ndarray:
"""
Create an embedding for a query.
Args:
query: User query text.
Returns:
Query embedding vector.
"""
try:
result = genai.embed_content(
model=self.embedding_model,
content=query,
task_type="retrieval_query",
)
return np.array(result["embedding"])
except Exception as e:
print(f"[EmbeddingService] Query embed error: {e}")
return np.zeros(768)
def store_pdf_embeddings(self, pdf_id: str, chunks: List[str]) -> None:
"""
Embed and store all chunks for a PDF (replacing previous).
Args:
pdf_id: Unique PDF identifier.
chunks: List of chunk strings.
"""
self.clear_pdf_embeddings(pdf_id)
for idx, (chunk, vec) in enumerate(zip(chunks, self.generate_embeddings(chunks))):
key = f"{pdf_id}_chunk_{idx}"
self.vector_store.add_embedding(
key=key,
vector=vec.tolist(),
metadata={"pdf_id": pdf_id, "chunk_index": idx, "text": chunk},
)
def find_similar_chunks(self, query: str, pdf_id: Optional[str] = None, top_k: int = 3) -> List[Dict]:
"""
Retrieve top_k most similar stored chunks.
Args:
query: User query string.
pdf_id: Restrict to given PDF id if set.
top_k: Number of results.
Returns:
List of similarity result dicts.
"""
q_vec = self.generate_query_embedding(query)
results = []
for key in self.vector_store.get_all_embeddings():
if pdf_id and not key.startswith(f"{pdf_id}_"):
continue
data = self.vector_store.get_embedding_data(key)
if not data:
continue
vec = np.array(data["vector"])
sim = self._cosine_similarity(q_vec, vec)
md = data.get("metadata", {})
results.append(
{
"key": key,
"similarity": sim,
"text": md.get("text", ""),
"chunk_index": md.get("chunk_index", 0),
"pdf_id": md.get("pdf_id", ""),
}
)
results.sort(key=lambda r: r["similarity"], reverse=True)
return results[:top_k]
def clear_pdf_embeddings(self, pdf_id: str) -> None:
"""
Remove all embeddings tied to a PDF.
Args:
pdf_id: Identifier.
"""
self.vector_store.remove_embeddings_by_prefix(f"{pdf_id}_")
def _cosine_similarity(self, a: np.ndarray, b: np.ndarray) -> float:
"""
Compute cosine similarity.
Args:
a: Vector A
b: Vector B
Returns:
Cosine similarity or 0.0 on failure.
"""
if not a.any() or not b.any():
return 0.0
denom = (np.linalg.norm(a) * np.linalg.norm(b))
if denom == 0:
return 0.0
return float(np.dot(a, b) / denom) |