feat(semantic-search): Replace local embedding model with OpenRouter API
- Replace EmbeddingGemma local model with OpenRouter API integration - Use google/gemini-embedding-001 model via OpenRouter (3072 dimensions) - Add conditional tool registration: auto-disable if OPENROUTER_API_KEY not set - Add openai and numpy dependencies to pyproject.toml - Update .env.example with OPENROUTER_API_KEY configuration - Fix ruff lint issues in semantic_search module
This commit is contained in:
saidsurucu
+100
-119
@@ -1,173 +1,154 @@
|
||||
# semantic_search/embedder.py
|
||||
|
||||
import logging
|
||||
import os
|
||||
from typing import List, Optional
|
||||
import numpy as np
|
||||
from sentence_transformers import SentenceTransformer
|
||||
import torch
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class EmbeddingGemma:
|
||||
|
||||
def is_openrouter_available() -> bool:
|
||||
"""Check if OpenRouter API key is available."""
|
||||
return bool(os.getenv("OPENROUTER_API_KEY"))
|
||||
|
||||
|
||||
class OpenRouterEmbedder:
|
||||
"""
|
||||
Wrapper for Google's EmbeddingGemma model.
|
||||
Handles query and document encoding with proper prompt templates.
|
||||
Embedder using OpenRouter API with Google's Gemini Embedding model.
|
||||
Requires OPENROUTER_API_KEY environment variable.
|
||||
"""
|
||||
|
||||
def __init__(self, model_name: str = "google/embeddinggemma-300m", device: Optional[str] = None):
|
||||
|
||||
def __init__(self):
|
||||
"""
|
||||
Initialize EmbeddingGemma model.
|
||||
|
||||
Args:
|
||||
model_name: HuggingFace model name
|
||||
device: Device to run model on ('cuda', 'cpu', or None for auto)
|
||||
Initialize OpenRouter Embedder.
|
||||
|
||||
Raises:
|
||||
ValueError: If OPENROUTER_API_KEY is not set
|
||||
ImportError: If openai package is not installed
|
||||
"""
|
||||
self.model_name = model_name
|
||||
|
||||
# Auto-detect device if not specified
|
||||
if device is None:
|
||||
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||||
else:
|
||||
self.device = device
|
||||
|
||||
logger.info(f"Initializing EmbeddingGemma on device: {self.device}")
|
||||
|
||||
api_key = os.getenv("OPENROUTER_API_KEY")
|
||||
if not api_key:
|
||||
raise ValueError("OPENROUTER_API_KEY environment variable is not set")
|
||||
|
||||
try:
|
||||
# Load model with float32 precision (EmbeddingGemma doesn't support float16)
|
||||
self.model = SentenceTransformer(model_name, device=self.device)
|
||||
self.model.eval() # Set to evaluation mode
|
||||
|
||||
# Set precision to float32 or bfloat16
|
||||
if self.device == 'cuda' and torch.cuda.is_bf16_supported():
|
||||
logger.info("Using bfloat16 precision for CUDA")
|
||||
self.dtype = torch.bfloat16
|
||||
else:
|
||||
logger.info("Using float32 precision")
|
||||
self.dtype = torch.float32
|
||||
|
||||
logger.info(f"Successfully loaded model: {model_name}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load EmbeddingGemma model: {e}")
|
||||
raise
|
||||
|
||||
from openai import OpenAI
|
||||
except ImportError:
|
||||
raise ImportError("openai package is required. Install with: pip install openai")
|
||||
|
||||
self.client = OpenAI(
|
||||
base_url="https://openrouter.ai/api/v1",
|
||||
api_key=api_key,
|
||||
)
|
||||
self.model = "google/gemini-embedding-001"
|
||||
self.dimension = 3072
|
||||
|
||||
logger.info(f"OpenRouter Embedder initialized with model: {self.model}")
|
||||
|
||||
def encode_query(self, query: str, task: str = "search result") -> np.ndarray:
|
||||
"""
|
||||
Encode a search query with appropriate prompt template.
|
||||
|
||||
Encode a search query.
|
||||
|
||||
Args:
|
||||
query: The search query text
|
||||
task: Task type for prompt template (search result, question answering, etc.)
|
||||
|
||||
task: Task type for prompt template
|
||||
|
||||
Returns:
|
||||
Numpy array of embeddings (768 dimensions)
|
||||
Numpy array of embeddings (3072 dimensions)
|
||||
"""
|
||||
# Apply query prompt template
|
||||
prompted_query = f"task: {task} | query: {query}"
|
||||
|
||||
text = f"task: {task} | query: {query}"
|
||||
|
||||
try:
|
||||
with torch.no_grad():
|
||||
# Encode with model
|
||||
embeddings = self.model.encode(
|
||||
prompted_query,
|
||||
convert_to_numpy=True,
|
||||
normalize_embeddings=True, # L2 normalization for cosine similarity
|
||||
show_progress_bar=False
|
||||
)
|
||||
|
||||
logger.debug(f"Encoded query: {query[:50]}... -> shape: {embeddings.shape}")
|
||||
return embeddings
|
||||
|
||||
response = self.client.embeddings.create(
|
||||
model=self.model,
|
||||
input=text,
|
||||
encoding_format="float",
|
||||
extra_headers={
|
||||
"HTTP-Referer": "https://yargimcp.com",
|
||||
"X-Title": "Yargi MCP Server",
|
||||
}
|
||||
)
|
||||
|
||||
embedding = np.array(response.data[0].embedding, dtype=np.float32)
|
||||
|
||||
# L2 normalize for cosine similarity
|
||||
norm = np.linalg.norm(embedding)
|
||||
if norm > 0:
|
||||
embedding = embedding / norm
|
||||
|
||||
logger.debug(f"Encoded query: {query[:50]}... -> shape: {embedding.shape}")
|
||||
return embedding
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to encode query: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def encode_documents(self, documents: List[str], titles: Optional[List[str]] = None) -> np.ndarray:
|
||||
"""
|
||||
Encode multiple documents with appropriate prompt template.
|
||||
|
||||
Encode multiple documents with batch API call.
|
||||
|
||||
Args:
|
||||
documents: List of document texts
|
||||
titles: Optional list of document titles
|
||||
|
||||
|
||||
Returns:
|
||||
Numpy array of embeddings (N x 768 dimensions)
|
||||
Numpy array of embeddings (N x 3072 dimensions)
|
||||
"""
|
||||
if not documents:
|
||||
return np.array([])
|
||||
|
||||
|
||||
# Apply document prompt template
|
||||
prompted_docs = []
|
||||
texts = []
|
||||
for i, doc in enumerate(documents):
|
||||
title = titles[i] if titles and i < len(titles) else "none"
|
||||
prompted_doc = f"title: {title} | text: {doc}"
|
||||
prompted_docs.append(prompted_doc)
|
||||
|
||||
text = f"title: {title} | text: {doc}"
|
||||
texts.append(text)
|
||||
|
||||
try:
|
||||
with torch.no_grad():
|
||||
# Batch encode documents
|
||||
embeddings = self.model.encode(
|
||||
prompted_docs,
|
||||
convert_to_numpy=True,
|
||||
normalize_embeddings=True,
|
||||
show_progress_bar=len(documents) > 10,
|
||||
batch_size=8 # Adjust based on memory
|
||||
)
|
||||
|
||||
response = self.client.embeddings.create(
|
||||
model=self.model,
|
||||
input=texts,
|
||||
encoding_format="float",
|
||||
extra_headers={
|
||||
"HTTP-Referer": "https://yargimcp.com",
|
||||
"X-Title": "Yargi MCP Server",
|
||||
}
|
||||
)
|
||||
|
||||
# Extract embeddings in order
|
||||
embeddings = np.array(
|
||||
[d.embedding for d in sorted(response.data, key=lambda x: x.index)],
|
||||
dtype=np.float32
|
||||
)
|
||||
|
||||
# L2 normalize each embedding for cosine similarity
|
||||
norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
|
||||
embeddings = embeddings / (norms + 1e-8)
|
||||
|
||||
logger.info(f"Encoded {len(documents)} documents -> shape: {embeddings.shape}")
|
||||
return embeddings
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to encode documents: {e}")
|
||||
raise
|
||||
|
||||
def reduce_dimensions(self, embeddings: np.ndarray, target_dim: int = 512) -> np.ndarray:
|
||||
"""
|
||||
Reduce embedding dimensions using Matryoshka Representation Learning.
|
||||
|
||||
Args:
|
||||
embeddings: Original embeddings (N x 768)
|
||||
target_dim: Target dimension (512, 256, or 128)
|
||||
|
||||
Returns:
|
||||
Reduced embeddings (N x target_dim)
|
||||
"""
|
||||
if target_dim not in [512, 256, 128]:
|
||||
raise ValueError(f"Target dimension must be 512, 256, or 128, got {target_dim}")
|
||||
|
||||
if len(embeddings.shape) == 1:
|
||||
# Single embedding
|
||||
reduced = embeddings[:target_dim]
|
||||
# Re-normalize after truncation
|
||||
norm = np.linalg.norm(reduced)
|
||||
if norm > 0:
|
||||
reduced = reduced / norm
|
||||
else:
|
||||
# Multiple embeddings
|
||||
reduced = embeddings[:, :target_dim]
|
||||
# Re-normalize each embedding
|
||||
norms = np.linalg.norm(reduced, axis=1, keepdims=True)
|
||||
reduced = reduced / (norms + 1e-8) # Avoid division by zero
|
||||
|
||||
logger.debug(f"Reduced dimensions: {embeddings.shape} -> {reduced.shape}")
|
||||
return reduced
|
||||
|
||||
|
||||
def compute_similarity(self, query_embedding: np.ndarray, document_embeddings: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Compute cosine similarity between query and documents.
|
||||
|
||||
|
||||
Args:
|
||||
query_embedding: Query embedding (768,)
|
||||
document_embeddings: Document embeddings (N x 768)
|
||||
|
||||
query_embedding: Query embedding (3072,)
|
||||
document_embeddings: Document embeddings (N x 3072)
|
||||
|
||||
Returns:
|
||||
Similarity scores (N,)
|
||||
"""
|
||||
# Ensure query is 2D for matrix multiplication
|
||||
if len(query_embedding.shape) == 1:
|
||||
query_embedding = query_embedding.reshape(1, -1)
|
||||
|
||||
|
||||
# Compute cosine similarity (embeddings are already normalized)
|
||||
similarities = np.dot(document_embeddings, query_embedding.T).squeeze()
|
||||
|
||||
return similarities
|
||||
|
||||
return similarities
|
||||
|
||||
Reference in New Issue
Block a user