Spaces:
Running on Zero
Running on Zero
| import os | |
| import json | |
| import re | |
| import numpy as np | |
| from sentence_transformers import SentenceTransformer | |
| from .search_utils import load_movies, PROJECT_ROOT | |
| def cosine_similarity(vec1, vec2): | |
| """Calculate cosine similarity between two vectors""" | |
| dot_product = np.dot(vec1, vec2) | |
| norm1 = np.linalg.norm(vec1) | |
| norm2 = np.linalg.norm(vec2) | |
| if norm1 == 0 or norm2 == 0: | |
| return 0.0 | |
| return dot_product / (norm1 * norm2) | |
| def chunk_text(text: str, chunk_size: int = 200, overlap: int = 0) -> list[str]: | |
| """Split text into fixed-size chunks by word count with optional overlap""" | |
| # Split text into words | |
| words = text.split() | |
| # Handle edge cases | |
| if not words: | |
| return [] | |
| chunks = [] | |
| i = 0 | |
| while i < len(words): | |
| # Get chunk_size words starting from position i | |
| chunk_words = words[i:i + chunk_size] | |
| chunk = " ".join(chunk_words) | |
| chunks.append(chunk) | |
| # Move forward by (chunk_size - overlap) words | |
| # If this is the last chunk, break to avoid infinite loop | |
| if i + chunk_size >= len(words): | |
| break | |
| i += chunk_size - overlap | |
| return chunks | |
| def semantic_chunk_text(text: str, max_chunk_size: int = 4, overlap: int = 0) -> list[str]: | |
| """Split text into chunks by sentence boundaries with optional overlap""" | |
| # Strip leading and trailing whitespace from input | |
| text = text.strip() | |
| # If nothing left after stripping, return empty list | |
| if not text: | |
| return [] | |
| # Split text into sentences using regex | |
| # (?<=[.!?]) is a positive lookbehind: split after punctuation | |
| # \s+ matches one or more whitespace characters | |
| sentences = re.split(r"(?<=[.!?])\s+", text) | |
| # Remove empty strings and strip whitespace from each sentence | |
| sentences = [s.strip() for s in sentences if s.strip()] | |
| # Handle edge cases | |
| if not sentences: | |
| return [] | |
| # If only one sentence and it doesn't end with punctuation, treat whole text as one sentence | |
| if len(sentences) == 1 and not text.endswith(('.', '!', '?')): | |
| return [text] | |
| chunks = [] | |
| i = 0 | |
| while i < len(sentences): | |
| # Get max_chunk_size sentences starting from position i | |
| chunk_sentences = sentences[i:i + max_chunk_size] | |
| chunk = " ".join(chunk_sentences).strip() | |
| # Only append chunk if it has content after stripping | |
| if chunk: | |
| chunks.append(chunk) | |
| # Break if this is the last chunk | |
| if i + max_chunk_size >= len(sentences): | |
| break | |
| # Move forward by (max_chunk_size - overlap) sentences | |
| i += max_chunk_size - overlap | |
| return chunks | |
| class SemanticSearch: | |
| def __init__(self, model_name="all-MiniLM-L6-v2"): | |
| """Initialize the semantic search model""" | |
| self.model = SentenceTransformer(model_name) | |
| self.embeddings = None | |
| self.documents = None | |
| self.document_map = {} | |
| def generate_embedding(self, text: str): | |
| """Generate an embedding for a single text input""" | |
| if not text or text.strip() == "": | |
| raise ValueError("Input text cannot be empty or contain only whitespace") | |
| embeddings = self.model.encode([text]) | |
| return embeddings[0] | |
| def build_embeddings(self, documents: list[dict]): | |
| """Generate embeddings for all documents and save to cache""" | |
| self.documents = documents | |
| # Build document map | |
| for doc in documents: | |
| self.document_map[doc['id']] = doc | |
| # Create string representations for each movie | |
| movie_strings = [f"{doc['title']}: {doc['description']}" for doc in documents] | |
| # Generate embeddings with progress bar | |
| print("Generating embeddings...") | |
| self.embeddings = self.model.encode(movie_strings, show_progress_bar=True) | |
| # Save to cache | |
| cache_dir = os.path.join(PROJECT_ROOT, "cache") | |
| os.makedirs(cache_dir, exist_ok=True) | |
| embeddings_path = os.path.join(cache_dir, "movie_embeddings.npy") | |
| np.save(embeddings_path, self.embeddings) | |
| return self.embeddings | |
| def load_or_create_embeddings(self, documents: list[dict]): | |
| """Load embeddings from cache or create them if not found""" | |
| self.documents = documents | |
| # Build document map | |
| for doc in documents: | |
| self.document_map[doc['id']] = doc | |
| # Check if cached embeddings exist | |
| cache_dir = os.path.join(PROJECT_ROOT, "cache") | |
| embeddings_path = os.path.join(cache_dir, "movie_embeddings.npy") | |
| if os.path.exists(embeddings_path): | |
| print("Loading cached embeddings...") | |
| self.embeddings = np.load(embeddings_path) | |
| # Verify cache is valid | |
| if len(self.embeddings) == len(documents): | |
| return self.embeddings | |
| else: | |
| print("Cache size mismatch, rebuilding...") | |
| # Build embeddings from scratch | |
| return self.build_embeddings(documents) | |
| def search(self, query: str, limit: int = 5) -> list[dict]: | |
| """Search for documents using semantic similarity""" | |
| # Check if embeddings are loaded | |
| if self.embeddings is None: | |
| raise ValueError("No embeddings loaded. Call `load_or_create_embeddings` first.") | |
| # Generate query embedding | |
| query_embedding = self.generate_embedding(query) | |
| # Calculate similarity scores for all documents | |
| results = [] | |
| for i, doc_embedding in enumerate(self.embeddings): | |
| similarity = cosine_similarity(query_embedding, doc_embedding) | |
| doc = self.documents[i] | |
| results.append({ | |
| 'score': similarity, | |
| 'title': doc['title'], | |
| 'description': doc['description'] | |
| }) | |
| # Sort by similarity score (descending) | |
| results.sort(key=lambda x: x['score'], reverse=True) | |
| # Return top results | |
| return results[:limit] | |
| class ChunkedSemanticSearch(SemanticSearch): | |
| def __init__(self, model_name="all-MiniLM-L6-v2") -> None: | |
| super().__init__(model_name) | |
| self.chunk_embeddings = None | |
| self.chunk_metadata = None | |
| def build_chunk_embeddings(self, documents: list[dict]): | |
| """Generate embeddings for document chunks""" | |
| self.documents = documents | |
| # Build document map | |
| for doc in documents: | |
| self.document_map[doc['id']] = doc | |
| all_chunks = [] # List of chunk strings | |
| chunk_metadata = [] # List of metadata dicts | |
| # Process each document | |
| for movie_idx, doc in enumerate(documents): | |
| # Skip empty descriptions | |
| description = doc.get('description', '').strip() | |
| if not description: | |
| continue | |
| # Chunk the description (4 sentences per chunk, 1 sentence overlap) | |
| chunks = semantic_chunk_text(description, max_chunk_size=4, overlap=1) | |
| # Add chunks and metadata | |
| for chunk_idx, chunk in enumerate(chunks): | |
| all_chunks.append(chunk) | |
| chunk_metadata.append({ | |
| 'movie_idx': movie_idx, | |
| 'chunk_idx': chunk_idx, | |
| 'total_chunks': len(chunks) | |
| }) | |
| # Generate embeddings for all chunks | |
| print(f"Generating embeddings for {len(all_chunks)} chunks...") | |
| self.chunk_embeddings = self.model.encode(all_chunks, show_progress_bar=True) | |
| self.chunk_metadata = chunk_metadata | |
| # Save to cache | |
| cache_dir = os.path.join(PROJECT_ROOT, "cache") | |
| os.makedirs(cache_dir, exist_ok=True) | |
| # Save embeddings | |
| embeddings_path = os.path.join(cache_dir, "chunk_embeddings.npy") | |
| np.save(embeddings_path, self.chunk_embeddings) | |
| # Save metadata as JSON | |
| metadata_path = os.path.join(cache_dir, "chunk_metadata.json") | |
| with open(metadata_path, "w") as f: | |
| json.dump({ | |
| "chunks": chunk_metadata, | |
| "total_chunks": len(all_chunks) | |
| }, f, indent=2) | |
| return self.chunk_embeddings | |
| def load_or_create_chunk_embeddings(self, documents: list[dict]) -> np.ndarray: | |
| """Load chunk embeddings from cache or create them""" | |
| self.documents = documents | |
| # Build document map | |
| for doc in documents: | |
| self.document_map[doc['id']] = doc | |
| # Check if cache exists | |
| cache_dir = os.path.join(PROJECT_ROOT, "cache") | |
| embeddings_path = os.path.join(cache_dir, "chunk_embeddings.npy") | |
| metadata_path = os.path.join(cache_dir, "chunk_metadata.json") | |
| if os.path.exists(embeddings_path) and os.path.exists(metadata_path): | |
| print("Loading cached chunk embeddings...") | |
| # Load embeddings | |
| self.chunk_embeddings = np.load(embeddings_path) | |
| # Load metadata | |
| with open(metadata_path, "r") as f: | |
| metadata = json.load(f) | |
| self.chunk_metadata = metadata['chunks'] | |
| return self.chunk_embeddings | |
| # Build from scratch | |
| return self.build_chunk_embeddings(documents) | |
| def search_chunks(self, query: str, limit: int = 10) -> list[dict]: | |
| """Search using chunk embeddings and aggregate results by movie""" | |
| # Check if chunk embeddings are loaded | |
| if self.chunk_embeddings is None: | |
| raise ValueError("No chunk embeddings loaded. Call `load_or_create_chunk_embeddings` first.") | |
| # Generate query embedding | |
| query_embedding = self.generate_embedding(query) | |
| # Store chunk scores | |
| chunk_scores = [] | |
| # Calculate similarity for each chunk | |
| for chunk_idx, chunk_embedding in enumerate(self.chunk_embeddings): | |
| similarity = cosine_similarity(query_embedding, chunk_embedding) | |
| # Get movie index from metadata | |
| movie_idx = self.chunk_metadata[chunk_idx]['movie_idx'] | |
| chunk_scores.append({ | |
| 'chunk_idx': chunk_idx, | |
| 'movie_idx': movie_idx, | |
| 'score': similarity | |
| }) | |
| # Aggregate scores by movie (keep highest score per movie) | |
| movie_scores = {} | |
| for chunk_score in chunk_scores: | |
| movie_idx = chunk_score['movie_idx'] | |
| score = chunk_score['score'] | |
| # Update if this is the first score or a higher score | |
| if movie_idx not in movie_scores or score > movie_scores[movie_idx]: | |
| movie_scores[movie_idx] = score | |
| # Convert to list and sort by score (descending) | |
| sorted_movies = sorted(movie_scores.items(), key=lambda x: x[1], reverse=True) | |
| # Take top limit movies | |
| top_movies = sorted_movies[:limit] | |
| # Format results | |
| SCORE_PRECISION = 4 | |
| results = [] | |
| for movie_idx, score in top_movies: | |
| doc = self.documents[movie_idx] | |
| result = { | |
| 'id': doc['id'], | |
| 'title': doc['title'], | |
| 'document': doc['description'][:100], | |
| 'score': round(score, SCORE_PRECISION), | |
| 'metadata': doc.get('metadata', {}) | |
| } | |
| results.append(result) | |
| return results | |
| def verify_model(): | |
| """Verify the embedding model is loaded correctly""" | |
| semantic_search = SemanticSearch() | |
| print(f"Model loaded: {semantic_search.model}") | |
| print(f"Max sequence length: {semantic_search.model.max_seq_length}") | |
| def embed_text(text: str): | |
| """Generate and display embedding for input text""" | |
| semantic_search = SemanticSearch() | |
| embedding = semantic_search.generate_embedding(text) | |
| print(f"Text: {text}") | |
| print(f"First 3 dimensions: {embedding[:3]}") | |
| print(f"Dimensions: {embedding.shape[0]}") | |
| def verify_embeddings(): | |
| """Verify embeddings are generated correctly for all movies""" | |
| semantic_search = SemanticSearch() | |
| # Load movie documents | |
| documents = load_movies() | |
| # Load or create embeddings | |
| embeddings = semantic_search.load_or_create_embeddings(documents) | |
| # Print verification info | |
| print(f"Number of docs: {len(documents)}") | |
| print(f"Embeddings shape: {embeddings.shape[0]} vectors in {embeddings.shape[1]} dimensions") | |
| def embed_query_text(query: str): | |
| """Generate and display embedding for a search query""" | |
| semantic_search = SemanticSearch() | |
| # Generate embedding using existing method | |
| embedding = semantic_search.generate_embedding(query) | |
| # Print query and embedding info | |
| print(f"Query: {query}") | |
| print(f"First 5 dimensions: {embedding[:5]}") | |
| print(f"Shape: {embedding.shape}") | |