diff --git a/semantic_search/__init__.py b/semantic_search/__init__.py new file mode 100644 index 0000000..c1e9e74 --- /dev/null +++ b/semantic_search/__init__.py @@ -0,0 +1,7 @@ +# semantic_search/__init__.py + +from .embedder import EmbeddingGemma +from .vector_store import VectorStore +from .processor import DocumentProcessor + +__all__ = ['EmbeddingGemma', 'VectorStore', 'DocumentProcessor'] \ No newline at end of file diff --git a/semantic_search/embedder.py b/semantic_search/embedder.py new file mode 100644 index 0000000..29f9452 --- /dev/null +++ b/semantic_search/embedder.py @@ -0,0 +1,173 @@ +# semantic_search/embedder.py + +import logging +from typing import List, Optional +import numpy as np +from sentence_transformers import SentenceTransformer +import torch + +logger = logging.getLogger(__name__) + +class EmbeddingGemma: + """ + Wrapper for Google's EmbeddingGemma model. + Handles query and document encoding with proper prompt templates. + """ + + def __init__(self, model_name: str = "google/embeddinggemma-300m", device: Optional[str] = None): + """ + Initialize EmbeddingGemma model. + + Args: + model_name: HuggingFace model name + device: Device to run model on ('cuda', 'cpu', or None for auto) + """ + 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}") + + 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 + + def encode_query(self, query: str, task: str = "search result") -> np.ndarray: + """ + Encode a search query with appropriate prompt template. + + Args: + query: The search query text + task: Task type for prompt template (search result, question answering, etc.) + + Returns: + Numpy array of embeddings (768 dimensions) + """ + # Apply query prompt template + prompted_query = 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 + + 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. + + Args: + documents: List of document texts + titles: Optional list of document titles + + Returns: + Numpy array of embeddings (N x 768 dimensions) + """ + if not documents: + return np.array([]) + + # Apply document prompt template + prompted_docs = [] + 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) + + 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 + ) + + 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) + + 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 \ No newline at end of file diff --git a/semantic_search/processor.py b/semantic_search/processor.py new file mode 100644 index 0000000..1539a6d --- /dev/null +++ b/semantic_search/processor.py @@ -0,0 +1,305 @@ +# semantic_search/processor.py + +import logging +import re +from typing import List, Dict, Any, Optional +from dataclasses import dataclass +import hashlib + +logger = logging.getLogger(__name__) + +@dataclass +class DocumentChunk: + """Represents a chunk of a document.""" + chunk_id: str + document_id: str + text: str + metadata: Dict[str, Any] + chunk_index: int + total_chunks: int + +class DocumentProcessor: + """ + Processes legal documents for semantic search. + Handles chunking, cleaning, and metadata extraction. + """ + + def __init__(self, + chunk_size: int = 1000, + chunk_overlap: int = 200, + min_chunk_size: int = 100): + """ + Initialize document processor. + + Args: + chunk_size: Target size for each chunk in characters + chunk_overlap: Number of overlapping characters between chunks + min_chunk_size: Minimum chunk size to keep + """ + self.chunk_size = chunk_size + self.chunk_overlap = chunk_overlap + self.min_chunk_size = min_chunk_size + + logger.info(f"Initialized DocumentProcessor (chunk_size={chunk_size}, overlap={chunk_overlap})") + + def process_document(self, + document_id: str, + text: str, + metadata: Optional[Dict[str, Any]] = None) -> List[DocumentChunk]: + """ + Process a single document into chunks. + + Args: + document_id: Unique document identifier + text: Document text content + metadata: Optional document metadata + + Returns: + List of document chunks + """ + if not text or len(text.strip()) < self.min_chunk_size: + logger.warning(f"Document {document_id} too short to process") + return [] + + # Clean text + cleaned_text = self._clean_text(text) + + # Extract metadata from text if not provided + if metadata is None: + metadata = {} + + # Add extracted metadata + extracted_metadata = self._extract_metadata(cleaned_text) + metadata.update(extracted_metadata) + + # Create chunks + chunks = self._create_chunks(cleaned_text) + + # Create DocumentChunk objects + document_chunks = [] + for i, chunk_text in enumerate(chunks): + chunk_id = self._generate_chunk_id(document_id, i) + + chunk = DocumentChunk( + chunk_id=chunk_id, + document_id=document_id, + text=chunk_text, + metadata={ + **metadata, + 'chunk_index': i, + 'total_chunks': len(chunks) + }, + chunk_index=i, + total_chunks=len(chunks) + ) + document_chunks.append(chunk) + + logger.info(f"Processed document {document_id} into {len(chunks)} chunks") + return document_chunks + + def _clean_text(self, text: str) -> str: + """ + Clean and normalize text for processing. + + Args: + text: Raw text + + Returns: + Cleaned text + """ + # Remove excessive whitespace + text = re.sub(r'\s+', ' ', text) + + # Remove special characters but keep Turkish characters + # Keep: letters, numbers, spaces, and common punctuation + text = re.sub(r'[^\w\s\.\,\;\:\!\?\-\(\)\"\'ÇĞIİÖŞÜçğıiöşü]', ' ', text) + + # Remove multiple spaces + text = re.sub(r' +', ' ', text) + + # Trim + text = text.strip() + + return text + + def _extract_metadata(self, text: str) -> Dict[str, Any]: + """ + Extract metadata from legal document text. + + Args: + text: Document text + + Returns: + Extracted metadata + """ + metadata = {} + + # Extract case numbers (Esas/Karar) + esas_pattern = r'E(?:sas)?[\s\.\:]*(\d{4})[\/\-](\d+)' + karar_pattern = r'K(?:arar)?[\s\.\:]*(\d{4})[\/\-](\d+)' + + esas_match = re.search(esas_pattern, text[:500]) # Look in first 500 chars + if esas_match: + metadata['esas_no'] = f"E.{esas_match.group(1)}/{esas_match.group(2)}" + + karar_match = re.search(karar_pattern, text[:500]) + if karar_match: + metadata['karar_no'] = f"K.{karar_match.group(1)}/{karar_match.group(2)}" + + # Extract dates (DD.MM.YYYY or DD/MM/YYYY format) + date_pattern = r'(\d{1,2})[\.\/](\d{1,2})[\.\/](\d{4})' + dates = re.findall(date_pattern, text[:1000]) # Look in first 1000 chars + if dates: + # Take the first date as decision date + day, month, year = dates[0] + metadata['karar_tarihi'] = f"{year}-{month.zfill(2)}-{day.zfill(2)}" + + # Extract court/chamber name + chamber_patterns = [ + r'(\d+)\.\s*Hukuk\s+Dairesi', + r'(\d+)\.\s*Ceza\s+Dairesi', + r'Hukuk\s+Genel\s+Kurulu', + r'Ceza\s+Genel\s+Kurulu', + r'(\d+)\.\s*Daire' + ] + + for pattern in chamber_patterns: + match = re.search(pattern, text[:500], re.IGNORECASE) + if match: + metadata['chamber'] = match.group(0) + break + + return metadata + + def _create_chunks(self, text: str) -> List[str]: + """ + Create overlapping chunks from text. + + Args: + text: Cleaned document text + + Returns: + List of text chunks + """ + chunks = [] + + # Split by sentences for better semantic coherence + sentences = self._split_sentences(text) + + current_chunk = [] + current_size = 0 + + for sentence in sentences: + sentence_size = len(sentence) + + # If adding this sentence exceeds chunk size + if current_size + sentence_size > self.chunk_size and current_chunk: + # Save current chunk + chunk_text = ' '.join(current_chunk) + chunks.append(chunk_text) + + # Create overlap for next chunk + overlap_size = 0 + overlap_sentences = [] + + # Add sentences from the end until we reach overlap size + for sent in reversed(current_chunk): + overlap_size += len(sent) + overlap_sentences.insert(0, sent) + if overlap_size >= self.chunk_overlap: + break + + # Start new chunk with overlap + current_chunk = overlap_sentences + current_size = sum(len(s) for s in current_chunk) + + # Add sentence to current chunk + current_chunk.append(sentence) + current_size += sentence_size + + # Add final chunk if not empty + if current_chunk: + chunk_text = ' '.join(current_chunk) + if len(chunk_text) >= self.min_chunk_size: + chunks.append(chunk_text) + + return chunks + + def _split_sentences(self, text: str) -> List[str]: + """ + Split text into sentences. + + Args: + text: Text to split + + Returns: + List of sentences + """ + # Simple sentence splitting for Turkish text + # Split on period, question mark, exclamation, but not on abbreviations + + # Common Turkish abbreviations to preserve + abbreviations = ['Dr', 'Prof', 'Av', 'Md', 'Yrd', 'Doç', 'No', 'S', 'vs', 'vb', 'bkz'] + + # Replace abbreviations temporarily + temp_text = text + replacements = {} + for i, abbr in enumerate(abbreviations): + placeholder = f"__ABBR{i}__" + temp_text = temp_text.replace(f"{abbr}.", placeholder) + replacements[placeholder] = f"{abbr}." + + # Split sentences + sentence_endings = re.compile(r'[.!?]+') + sentences = sentence_endings.split(temp_text) + + # Restore abbreviations and clean + cleaned_sentences = [] + for sentence in sentences: + # Restore abbreviations + for placeholder, original in replacements.items(): + sentence = sentence.replace(placeholder, original) + + # Clean and add if not empty + sentence = sentence.strip() + if sentence and len(sentence) > 10: # Minimum sentence length + cleaned_sentences.append(sentence) + + return cleaned_sentences + + def _generate_chunk_id(self, document_id: str, chunk_index: int) -> str: + """ + Generate unique chunk ID. + + Args: + document_id: Parent document ID + chunk_index: Index of chunk in document + + Returns: + Unique chunk ID + """ + chunk_string = f"{document_id}_chunk_{chunk_index}" + chunk_hash = hashlib.md5(chunk_string.encode()).hexdigest()[:8] + return f"{document_id}_c{chunk_index}_{chunk_hash}" + + def combine_chunks(self, chunks: List[DocumentChunk]) -> str: + """ + Combine chunks back into full document text. + + Args: + chunks: List of document chunks + + Returns: + Combined text + """ + if not chunks: + return "" + + # Sort by chunk index + sorted_chunks = sorted(chunks, key=lambda x: x.chunk_index) + + # For overlapping chunks, we need to be careful about duplication + # Simple approach: just concatenate with space + combined = " ".join([chunk.text for chunk in sorted_chunks]) + + return combined \ No newline at end of file diff --git a/semantic_search/vector_store.py b/semantic_search/vector_store.py new file mode 100644 index 0000000..a1b0174 --- /dev/null +++ b/semantic_search/vector_store.py @@ -0,0 +1,235 @@ +# semantic_search/vector_store.py + +import logging +import numpy as np +from typing import List, Dict, Any, Tuple, Optional +from dataclasses import dataclass +import json + +logger = logging.getLogger(__name__) + +@dataclass +class Document: + """Represents a document with its embedding and metadata.""" + id: str + text: str + embedding: np.ndarray + metadata: Dict[str, Any] + + def to_dict(self) -> Dict[str, Any]: + """Convert to dictionary (excluding embedding for serialization).""" + return { + 'id': self.id, + 'text': self.text, + 'metadata': self.metadata + } + +class VectorStore: + """ + In-memory vector storage with similarity search capabilities. + Future versions can use Faiss, ChromaDB, or other vector databases. + """ + + def __init__(self, dimension: int = 768): + """ + Initialize vector store. + + Args: + dimension: Embedding dimension size + """ + self.dimension = dimension + self.documents: List[Document] = [] + self.embeddings: Optional[np.ndarray] = None + self.index_built = False + + logger.info(f"Initialized VectorStore with dimension: {dimension}") + + def add_documents(self, + ids: List[str], + texts: List[str], + embeddings: np.ndarray, + metadata: Optional[List[Dict[str, Any]]] = None) -> int: + """ + Add documents to the vector store. + + Args: + ids: Document IDs + texts: Document texts + embeddings: Document embeddings (N x dimension) + metadata: Optional metadata for each document + + Returns: + Number of documents added + """ + if len(ids) != len(texts) or len(ids) != embeddings.shape[0]: + raise ValueError("Mismatched lengths for ids, texts, and embeddings") + + if metadata and len(metadata) != len(ids): + raise ValueError("Metadata length doesn't match document count") + + # Add documents + for i in range(len(ids)): + doc = Document( + id=ids[i], + text=texts[i], + embedding=embeddings[i], + metadata=metadata[i] if metadata else {} + ) + self.documents.append(doc) + + # Rebuild index + self._build_index() + + logger.info(f"Added {len(ids)} documents to vector store. Total: {len(self.documents)}") + return len(ids) + + def _build_index(self): + """Build or rebuild the embedding index.""" + if not self.documents: + self.embeddings = None + self.index_built = False + return + + # Stack all embeddings into a single array + self.embeddings = np.vstack([doc.embedding for doc in self.documents]) + self.index_built = True + + logger.debug(f"Built index with shape: {self.embeddings.shape}") + + def search(self, + query_embedding: np.ndarray, + top_k: int = 10, + threshold: Optional[float] = None) -> List[Tuple[Document, float]]: + """ + Search for similar documents using cosine similarity. + + Args: + query_embedding: Query embedding vector + top_k: Number of results to return + threshold: Optional similarity threshold (0-1) + + Returns: + List of (Document, similarity_score) tuples + """ + if not self.index_built or self.embeddings is None: + logger.warning("No documents in vector store") + return [] + + # Ensure query is 2D + if len(query_embedding.shape) == 1: + query_embedding = query_embedding.reshape(1, -1) + + # Compute cosine similarities (assuming normalized embeddings) + similarities = np.dot(self.embeddings, query_embedding.T).squeeze() + + # Apply threshold if specified + if threshold is not None: + valid_indices = np.where(similarities >= threshold)[0] + if len(valid_indices) == 0: + logger.info(f"No documents above threshold {threshold}") + return [] + similarities = similarities[valid_indices] + valid_docs = [self.documents[i] for i in valid_indices] + else: + valid_docs = self.documents + + # Get top-k indices + top_k = min(top_k, len(valid_docs)) + if top_k == 0: + return [] + + # Use argpartition for efficiency with large arrays + if len(similarities) > top_k: + top_indices = np.argpartition(similarities, -top_k)[-top_k:] + top_indices = top_indices[np.argsort(similarities[top_indices])[::-1]] + else: + top_indices = np.argsort(similarities)[::-1] + + # Create results + results = [] + for idx in top_indices: + doc = valid_docs[idx] if threshold else self.documents[idx] + score = float(similarities[idx]) + results.append((doc, score)) + + logger.info(f"Search returned {len(results)} results (top_k={top_k})") + return results + + def hybrid_search(self, + query_embedding: np.ndarray, + keyword_scores: Dict[str, float], + top_k: int = 10, + alpha: float = 0.5) -> List[Tuple[Document, float]]: + """ + Hybrid search combining vector similarity and keyword scores. + + Args: + query_embedding: Query embedding vector + keyword_scores: Document ID to keyword relevance score mapping + top_k: Number of results to return + alpha: Weight for vector similarity (1-alpha for keyword score) + + Returns: + List of (Document, combined_score) tuples + """ + if not self.index_built: + logger.warning("No documents in vector store") + return [] + + # Get vector similarities + vector_results = self.search(query_embedding, top_k=len(self.documents)) + + # Combine scores + combined_scores = [] + for doc, vector_score in vector_results: + keyword_score = keyword_scores.get(doc.id, 0.0) + # Normalize keyword score to 0-1 range if needed + if keyword_score > 1.0: + keyword_score = keyword_score / max(keyword_scores.values()) + + combined_score = alpha * vector_score + (1 - alpha) * keyword_score + combined_scores.append((doc, combined_score)) + + # Sort by combined score and return top-k + combined_scores.sort(key=lambda x: x[1], reverse=True) + results = combined_scores[:top_k] + + logger.info(f"Hybrid search returned {len(results)} results") + return results + + def clear(self): + """Clear all documents from the store.""" + self.documents = [] + self.embeddings = None + self.index_built = False + logger.info("Cleared vector store") + + def size(self) -> int: + """Get number of documents in store.""" + return len(self.documents) + + def get_by_id(self, doc_id: str) -> Optional[Document]: + """Get document by ID.""" + for doc in self.documents: + if doc.id == doc_id: + return doc + return None + + def get_stats(self) -> Dict[str, Any]: + """Get statistics about the vector store.""" + stats = { + 'num_documents': len(self.documents), + 'dimension': self.dimension, + 'index_built': self.index_built, + 'memory_usage_mb': 0 + } + + if self.embeddings is not None: + # Estimate memory usage + memory_bytes = self.embeddings.nbytes + for doc in self.documents: + memory_bytes += len(doc.text.encode('utf-8')) + memory_bytes += len(json.dumps(doc.metadata).encode('utf-8')) + stats['memory_usage_mb'] = memory_bytes / (1024 * 1024) + + return stats \ No newline at end of file