"""Embedding utilities for QualiVec.""" import numpy as np import pandas as pd from typing import List, Dict, Any, Optional, Union import torch from tqdm import tqdm from transformers import AutoTokenizer, AutoModel import os import time class EmbeddingModel: """Handles text embedding for QualiVec.""" def __init__(self, model_name: str = "sentence-transformers/all-MiniLM-L6-v2", model_type: str = "huggingface", device: Optional[str] = None, cache_dir: Optional[str] = None, verbose: bool = True): """Initialize the embedding model. Args: model_name: Name of the model to use (HuggingFace model or Gemini model). model_type: Type of model ('huggingface' or 'gemini'). device: Device to use for computation ('cpu' or 'cuda'). Only for HuggingFace models. cache_dir: Directory to cache models. Only for HuggingFace models. verbose: Whether to print status messages. """ self.model_name = model_name self.model_type = model_type.lower() self.verbose = verbose self.cache_dir = cache_dir if self.model_type not in ["huggingface", "gemini"]: raise ValueError(f"model_type must be 'huggingface' or 'gemini', got '{model_type}'") if self.model_type == "huggingface": # Determine device if device is None: self.device = "cuda" if torch.cuda.is_available() else "cpu" else: self.device = device if self.verbose: print(f"Using device: {self.device}") print(f"Loading HuggingFace model: {model_name}") # Load model and tokenizer self.tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir) self.model = AutoModel.from_pretrained(model_name, cache_dir=cache_dir).to(self.device) if self.verbose: print(f"HuggingFace model loaded successfully") elif self.model_type == "gemini": if self.verbose: print(f"Initializing Gemini model: {model_name}") # Import Gemini client try: from google import genai # Get API key from environment variable api_key = os.environ.get("GOOGLE_API_KEY") if not api_key: raise ValueError( "GOOGLE_API_KEY environment variable not set. " "Please set it with your Gemini API key." ) self.genai_client = genai.Client(api_key="API_KEY") if self.verbose: print(f"Gemini client initialized successfully") print(f"⚠️ Free tier limits: 1,500 requests/day, 100 texts per batch") except ImportError: raise ImportError("google-genai library is required for Gemini models. Install with: pip install google-genai") def _mean_pooling(self, model_output, attention_mask): """Mean pooling operation to get sentence embeddings.""" token_embeddings = model_output[0] input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9) def embed_texts(self, texts: List[str], batch_size: int = 32) -> np.ndarray: """Generate embeddings for a list of texts. Args: texts: List of texts to embed. batch_size: Batch size for processing. Returns: Numpy array of embeddings. """ if self.verbose: print(f"Generating embeddings for {len(texts)} texts") if self.model_type == "huggingface": return self._embed_texts_huggingface(texts, batch_size) elif self.model_type == "gemini": return self._embed_texts_gemini(texts, batch_size) else: raise ValueError(f"Unsupported model_type: {self.model_type}") def _embed_texts_huggingface(self, texts: List[str], batch_size: int = 32) -> np.ndarray: """Generate embeddings using HuggingFace model. Args: texts: List of texts to embed. batch_size: Batch size for processing. Returns: Numpy array of embeddings. """ embeddings = [] # Process in batches for i in tqdm(range(0, len(texts), batch_size), disable=not self.verbose): batch_texts = texts[i:i + batch_size] # Tokenize encoded_input = self.tokenizer(batch_texts, padding=True, truncation=True, max_length=512, return_tensors='pt').to(self.device) # Get model output with torch.no_grad(): model_output = self.model(**encoded_input) # Mean pooling batch_embeddings = self._mean_pooling(model_output, encoded_input['attention_mask']) # Normalize embeddings batch_embeddings = torch.nn.functional.normalize(batch_embeddings, p=2, dim=1) # Add to list embeddings.append(batch_embeddings.cpu().numpy()) # Concatenate all batches all_embeddings = np.vstack(embeddings) if self.verbose: print(f"Generated embeddings with shape: {all_embeddings.shape}") return all_embeddings def _embed_texts_gemini(self, texts: List[str], batch_size: int = 100) -> np.ndarray: """Generate embeddings using Gemini model with rate limiting. Args: texts: List of texts to embed. batch_size: Batch size for processing (reduced to 100 to respect rate limits). Returns: Numpy array of embeddings. """ embeddings = [] # Process in batches with rate limiting for i in tqdm(range(0, len(texts), batch_size), disable=not self.verbose): batch_texts = texts[i:i + batch_size] # Retry logic with exponential backoff max_retries = 3 retry_delay = 2 # seconds for attempt in range(max_retries): try: # Get embeddings from Gemini result = self.genai_client.models.embed_content( model=self.model_name, contents=batch_texts # type: ignore ) # Extract embeddings if result.embeddings: batch_embeddings = [emb.values for emb in result.embeddings] embeddings.extend(batch_embeddings) # Add delay between batches to respect rate limits (free tier: 1500 requests/day) # With 100 texts per batch and ~60 second delay, we can process ~1440 texts/day if i + batch_size < len(texts): time.sleep(1) # 1 second delay between batches break # Success, exit retry loop except Exception as e: error_msg = str(e) if "429" in error_msg or "RESOURCE_EXHAUSTED" in error_msg: if attempt < max_retries - 1: if self.verbose: print(f"\nRate limit hit. Waiting {retry_delay} seconds before retry {attempt + 1}/{max_retries}...") time.sleep(retry_delay) retry_delay *= 2 # Exponential backoff else: raise Exception( f"Gemini API quota exceeded. Free tier limits: 1500 requests/day.\n" f"Error: {error_msg}\n\n" f"Solutions:\n" f"1. Wait and try again later (quota resets daily)\n" f"2. Reduce the amount of data being processed\n" f"3. Upgrade to a paid API plan\n" f"4. Use HuggingFace models instead (no API limits)" ) else: raise # Re-raise non-quota errors # Convert to numpy array all_embeddings = np.array(embeddings) if self.verbose: print(f"Generated embeddings with shape: {all_embeddings.shape}") return all_embeddings def embed_dataframe(self, df: pd.DataFrame, text_column: str, batch_size: int = 32) -> np.ndarray: """Generate embeddings for texts in a DataFrame column. Args: df: DataFrame containing texts. text_column: Name of the column containing texts. batch_size: Batch size for processing. Returns: Numpy array of embeddings. """ if text_column not in df.columns: raise ValueError(f"Column '{text_column}' not found in DataFrame.") texts = df[text_column].fillna("").tolist() return self.embed_texts(texts, batch_size) def embed_reference_vectors(self, df: pd.DataFrame, class_column: str = "class", node_column: str = "matching_node", batch_size: int = 32) -> Dict[str, Any]: """Generate embeddings for reference vectors. Args: df: DataFrame containing reference vectors. class_column: Name of the column containing class labels. node_column: Name of the column containing matching nodes. batch_size: Batch size for processing. Returns: Dictionary with class info and embeddings. """ required_columns = [class_column, node_column] missing_columns = [col for col in required_columns if col not in df.columns] if missing_columns: raise ValueError(f"Required columns {missing_columns} not found in DataFrame.") # Get texts and generate embeddings texts = df[node_column].fillna("").tolist() embeddings = self.embed_texts(texts, batch_size) # Create result dictionary result = { "classes": df[class_column].tolist(), "nodes": df[node_column].tolist(), "embeddings": embeddings, "class_to_idx": {cls: i for i, cls in enumerate(df[class_column].unique())} } if self.verbose: print(f"Generated embeddings for {len(result['classes'])} reference vectors") print(f"Unique classes: {len(result['class_to_idx'])}") return result