Add inference script

hindi_embeddings.py

@@ -0,0 +1,724 @@
+import os
+import torch
+import json
+import numpy as np
+from torch import nn
+from torch.nn import functional as F
+import sentencepiece as spm
+from sklearn.metrics.pairwise import cosine_similarity
+from tqdm import tqdm
+import matplotlib.pyplot as plt
+from sklearn.manifold import TSNE
+
+# Tokenizer wrapper class
+class SentencePieceTokenizerWrapper:
+    def __init__(self, sp_model_path):
+        self.sp_model = spm.SentencePieceProcessor()
+        self.sp_model.Load(sp_model_path)
+        self.vocab_size = self.sp_model.GetPieceSize()
+        
+        # Special token IDs from tokenizer training
+        self.pad_token_id = 0
+        self.bos_token_id = 1
+        self.eos_token_id = 2
+        self.unk_token_id = 3
+        
+        # Set special tokens
+        self.pad_token = "<pad>"
+        self.bos_token = "<s>"
+        self.eos_token = "</s>"
+        self.unk_token = "<unk>"
+        self.mask_token = "<mask>"
+    
+    def __call__(self, text, padding=False, truncation=False, max_length=None, return_tensors=None):
+        # Handle both string and list inputs
+        if isinstance(text, str):
+            # Encode a single string
+            ids = self.sp_model.EncodeAsIds(text)
+            
+            # Handle truncation
+            if truncation and max_length and len(ids) > max_length:
+                ids = ids[:max_length]
+                
+            attention_mask = [1] * len(ids)
+            
+            # Handle padding
+            if padding and max_length:
+                padding_length = max(0, max_length - len(ids))
+                ids = ids + [self.pad_token_id] * padding_length
+                attention_mask = attention_mask + [0] * padding_length
+            
+            result = {
+                'input_ids': ids,
+                'attention_mask': attention_mask
+            }
+            
+            # Convert to tensors if requested
+            if return_tensors == 'pt':
+                import torch
+                result = {k: torch.tensor([v]) for k, v in result.items()}
+            
+            return result
+        
+        # Process a batch of texts
+        batch_encoded = [self.sp_model.EncodeAsIds(t) for t in text]
+        
+        # Apply truncation if needed
+        if truncation and max_length:
+            batch_encoded = [ids[:max_length] for ids in batch_encoded]
+        
+        # Create attention masks
+        batch_attention_mask = [[1] * len(ids) for ids in batch_encoded]
+        
+        # Apply padding if needed
+        if padding:
+            if max_length:
+                max_len = max_length
+            else:
+                max_len = max(len(ids) for ids in batch_encoded)
+            
+            # Pad sequences to max_len
+            batch_encoded = [ids + [self.pad_token_id] * (max_len - len(ids)) for ids in batch_encoded]
+            batch_attention_mask = [mask + [0] * (max_len - len(mask)) for mask in batch_attention_mask]
+        
+        result = {
+            'input_ids': batch_encoded,
+            'attention_mask': batch_attention_mask
+        }
+        
+        # Convert to tensors if requested
+        if return_tensors == 'pt':
+            import torch
+            result = {k: torch.tensor(v) for k, v in result.items()}
+        
+        return result
+
+# Model architecture components
+class MultiHeadAttention(nn.Module):
+    """Multi-headed attention mechanism"""
+    def __init__(self, config):
+        super().__init__()
+        self.num_attention_heads = config["num_attention_heads"]
+        self.attention_head_size = config["hidden_size"] // config["num_attention_heads"]
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+        
+        # Query, Key, Value projections
+        self.query = nn.Linear(config["hidden_size"], self.all_head_size)
+        self.key = nn.Linear(config["hidden_size"], self.all_head_size)
+        self.value = nn.Linear(config["hidden_size"], self.all_head_size)
+        
+        # Output projection
+        self.output = nn.Sequential(
+            nn.Linear(self.all_head_size, config["hidden_size"]),
+            nn.Dropout(config["attention_probs_dropout_prob"])
+        )
+        
+        # Simplified relative position bias
+        self.max_position_embeddings = config["max_position_embeddings"]
+        self.relative_attention_bias = nn.Embedding(
+            2 * config["max_position_embeddings"] - 1, 
+            config["num_attention_heads"]
+        )
+        
+    def transpose_for_scores(self, x):
+        new_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
+        x = x.view(*new_shape)
+        return x.permute(0, 2, 1, 3)
+    
+    def forward(self, hidden_states, attention_mask=None):
+        batch_size, seq_length = hidden_states.size()[:2]
+        
+        # Project inputs to queries, keys, and values
+        query_layer = self.transpose_for_scores(self.query(hidden_states))
+        key_layer = self.transpose_for_scores(self.key(hidden_states))
+        value_layer = self.transpose_for_scores(self.value(hidden_states))
+        
+        # Take the dot product between query and key to get the raw attention scores
+        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+        
+        # Generate relative position matrix
+        position_ids = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device)
+        relative_position = position_ids.unsqueeze(1) - position_ids.unsqueeze(0)  # [seq_len, seq_len]
+        # Shift values to be >= 0
+        relative_position = relative_position + self.max_position_embeddings - 1  
+        # Ensure indices are within bounds
+        relative_position = torch.clamp(relative_position, 0, 2 * self.max_position_embeddings - 2)  
+        
+        # Get relative position embeddings [seq_len, seq_len, num_heads]
+        rel_attn_bias = self.relative_attention_bias(relative_position)  # [seq_len, seq_len, num_heads]
+        
+        # Reshape to add to attention heads [1, num_heads, seq_len, seq_len]
+        rel_attn_bias = rel_attn_bias.permute(2, 0, 1).unsqueeze(0)
+        
+        # Add to attention scores - now dimensions will match
+        attention_scores = attention_scores + rel_attn_bias
+        
+        # Scale attention scores
+        attention_scores = attention_scores / (self.attention_head_size ** 0.5)
+        
+        # Apply attention mask
+        if attention_mask is not None:
+            attention_scores = attention_scores + attention_mask
+        
+        # Normalize the attention scores to probabilities
+        attention_probs = F.softmax(attention_scores, dim=-1)
+        
+        # Apply dropout
+        attention_probs = F.dropout(attention_probs, p=0.1, training=self.training)
+        
+        # Apply attention to values
+        context_layer = torch.matmul(attention_probs, value_layer)
+        
+        # Reshape back to [batch_size, seq_length, hidden_size]
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(*new_shape)
+        
+        # Final output projection
+        output = self.output(context_layer)
+        
+        return output
+
+class EnhancedTransformerLayer(nn.Module):
+    """Advanced transformer layer with pre-layer norm and enhanced attention"""
+    def __init__(self, config):
+        super().__init__()
+        self.attention_pre_norm = nn.LayerNorm(config["hidden_size"], eps=config["layer_norm_eps"])
+        self.attention = MultiHeadAttention(config)
+        
+        self.ffn_pre_norm = nn.LayerNorm(config["hidden_size"], eps=config["layer_norm_eps"])
+        
+        # Feed-forward network
+        self.ffn = nn.Sequential(
+            nn.Linear(config["hidden_size"], config["intermediate_size"]),
+            nn.GELU(),
+            nn.Dropout(config["hidden_dropout_prob"]),
+            nn.Linear(config["intermediate_size"], config["hidden_size"]),
+            nn.Dropout(config["hidden_dropout_prob"])
+        )
+        
+    def forward(self, hidden_states, attention_mask=None):
+        # Pre-layer norm for attention
+        attn_norm_hidden = self.attention_pre_norm(hidden_states)
+        
+        # Self-attention
+        attention_output = self.attention(attn_norm_hidden, attention_mask)
+        
+        # Residual connection
+        hidden_states = hidden_states + attention_output
+        
+        # Pre-layer norm for feed-forward
+        ffn_norm_hidden = self.ffn_pre_norm(hidden_states)
+        
+        # Feed-forward
+        ffn_output = self.ffn(ffn_norm_hidden)
+        
+        # Residual connection
+        hidden_states = hidden_states + ffn_output
+        
+        return hidden_states
+
+class AdvancedTransformerModel(nn.Module):
+    """Advanced Transformer model for inference"""
+    
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        
+        # Embeddings
+        self.word_embeddings = nn.Embedding(
+            config["vocab_size"], 
+            config["hidden_size"], 
+            padding_idx=config["pad_token_id"]
+        )
+        
+        # Position embeddings
+        self.position_embeddings = nn.Embedding(config["max_position_embeddings"], config["hidden_size"])
+        
+        # Embedding dropout
+        self.embedding_dropout = nn.Dropout(config["hidden_dropout_prob"])
+        
+        # Transformer layers
+        self.layers = nn.ModuleList([
+            EnhancedTransformerLayer(config) for _ in range(config["num_hidden_layers"])
+        ])
+        
+        # Final layer norm
+        self.final_layer_norm = nn.LayerNorm(config["hidden_size"], eps=config["layer_norm_eps"])
+        
+    def forward(self, input_ids, attention_mask=None):
+        input_shape = input_ids.size()
+        batch_size, seq_length = input_shape
+        
+        # Get position ids
+        position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device)
+        position_ids = position_ids.unsqueeze(0).expand(batch_size, -1)
+        
+        # Get embeddings
+        word_embeds = self.word_embeddings(input_ids)
+        position_embeds = self.position_embeddings(position_ids)
+        
+        # Sum embeddings
+        embeddings = word_embeds + position_embeds
+        
+        # Apply dropout
+        embeddings = self.embedding_dropout(embeddings)
+        
+        # Default attention mask
+        if attention_mask is None:
+            attention_mask = torch.ones(input_shape, device=input_ids.device)
+        
+        # Extended attention mask for transformer layers (1 for tokens to attend to, 0 for masked tokens)
+        extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
+        
+        # Apply transformer layers
+        hidden_states = embeddings
+        for layer in self.layers:
+            hidden_states = layer(hidden_states, extended_attention_mask)
+        
+        # Final layer norm
+        hidden_states = self.final_layer_norm(hidden_states)
+        
+        return hidden_states
+
+class AdvancedPooling(nn.Module):
+    """Advanced pooling module supporting multiple pooling strategies"""
+    def __init__(self, config):
+        super().__init__()
+        self.pooling_mode = config["pooling_mode"]  # 'mean', 'max', 'cls', 'attention'
+        self.hidden_size = config["hidden_size"]
+        
+        # For attention pooling
+        if self.pooling_mode == 'attention':
+            self.attention_weights = nn.Linear(config["hidden_size"], 1)
+            
+        # For weighted pooling
+        elif self.pooling_mode == 'weighted':
+            self.weight_layer = nn.Linear(config["hidden_size"], 1)
+            
+    def forward(self, token_embeddings, attention_mask=None):
+        if attention_mask is None:
+            attention_mask = torch.ones_like(token_embeddings[:, :, 0])
+            
+        mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+        
+        if self.pooling_mode == 'cls':
+            # Use [CLS] token (first token)
+            pooled = token_embeddings[:, 0]
+            
+        elif self.pooling_mode == 'max':
+            # Max pooling
+            token_embeddings = token_embeddings.clone()
+            # Set padding tokens to large negative value to exclude them from max
+            token_embeddings[mask_expanded == 0] = -1e9
+            pooled = torch.max(token_embeddings, dim=1)[0]
+            
+        elif self.pooling_mode == 'attention':
+            # Attention pooling
+            weights = self.attention_weights(token_embeddings).squeeze(-1)
+            # Mask out padding tokens
+            weights = weights.masked_fill(attention_mask == 0, -1e9)
+            weights = F.softmax(weights, dim=1).unsqueeze(-1)
+            pooled = torch.sum(token_embeddings * weights, dim=1)
+            
+        elif self.pooling_mode == 'weighted':
+            # Weighted average pooling
+            weights = torch.sigmoid(self.weight_layer(token_embeddings)).squeeze(-1)
+            # Apply mask
+            weights = weights * attention_mask
+            # Normalize weights
+            sum_weights = torch.sum(weights, dim=1, keepdim=True)
+            sum_weights = torch.clamp(sum_weights, min=1e-9)
+            weights = weights / sum_weights
+            # Apply weights
+            pooled = torch.sum(token_embeddings * weights.unsqueeze(-1), dim=1)
+            
+        else:  # Default to mean pooling
+            # Mean pooling
+            sum_embeddings = torch.sum(token_embeddings * mask_expanded, dim=1)
+            sum_mask = torch.clamp(mask_expanded.sum(1), min=1e-9)
+            pooled = sum_embeddings / sum_mask
+            
+        # L2 normalize
+        pooled = F.normalize(pooled, p=2, dim=1)
+        
+        return pooled
+
+class SentenceEmbeddingModel(nn.Module):
+    """Complete sentence embedding model for inference"""
+    def __init__(self, config):
+        super(SentenceEmbeddingModel, self).__init__()
+        self.config = config
+        
+        # Create transformer model
+        self.transformer = AdvancedTransformerModel(config)
+        
+        # Create pooling module
+        self.pooling = AdvancedPooling(config)
+        
+        # Build projection module if needed
+        if "projection_dim" in config and config["projection_dim"] > 0:
+            self.use_projection = True
+            self.projection = nn.Sequential(
+                nn.Linear(config["hidden_size"], config["hidden_size"]),
+                nn.GELU(),
+                nn.Linear(config["hidden_size"], config["projection_dim"]),
+                nn.LayerNorm(config["projection_dim"], eps=config["layer_norm_eps"])
+            )
+        else:
+            self.use_projection = False
+            
+    def forward(self, input_ids, attention_mask=None):
+        # Get token embeddings from transformer
+        token_embeddings = self.transformer(input_ids, attention_mask)
+        
+        # Pool token embeddings
+        pooled_output = self.pooling(token_embeddings, attention_mask)
+        
+        # Apply projection if enabled
+        if self.use_projection:
+            pooled_output = self.projection(pooled_output)
+            pooled_output = F.normalize(pooled_output, p=2, dim=1)
+        
+        return pooled_output
+
+class HindiEmbedder:
+    def __init__(self, model_path="/home/ubuntu/output/hindi-embeddings-custom-tokenizer/final", tokenizer_path=None):
+        """
+        Initialize the Hindi sentence embedder.
+        
+        Args:
+            model_path: Path to the model directory
+            tokenizer_path: Optional path to tokenizer. If None, will look in the model directory.
+        """
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        print(f"Using device: {self.device}")
+        
+        # Load tokenizer
+        if tokenizer_path is None:
+            # Try standard location in model directory
+            tokenizer_path = os.path.join(model_path, "tokenizer.model")
+            if not os.path.exists(tokenizer_path):
+                # Try original location
+                tokenizer_path = "/home/ubuntu/hindi_tokenizer/tokenizer.model"
+        
+        if not os.path.exists(tokenizer_path):
+            raise FileNotFoundError(f"Could not find tokenizer at {tokenizer_path}")
+            
+        self.tokenizer = SentencePieceTokenizerWrapper(tokenizer_path)
+        print(f"Loaded tokenizer from {tokenizer_path} with vocabulary size: {self.tokenizer.vocab_size}")
+        
+        # Load model config
+        config_path = os.path.join(model_path, "config.json")
+        with open(config_path, "r") as f:
+            self.config = json.load(f)
+        print(f"Loaded model config with hidden_size={self.config['hidden_size']}")
+        
+        # Load model
+        model_pt_path = os.path.join(model_path, "embedding_model.pt")
+        
+        try:
+            # Support both PyTorch 2.6+ and older versions
+            try:
+                checkpoint = torch.load(model_pt_path, map_location=self.device, weights_only=False)
+                print("Loaded model using PyTorch 2.6+ style loading")
+            except TypeError:
+                checkpoint = torch.load(model_pt_path, map_location=self.device)
+                print("Loaded model using older PyTorch style loading")
+                
+            # Create model
+            self.model = SentenceEmbeddingModel(self.config)
+            
+            # Load state dict
+            if "model_state_dict" in checkpoint:
+                state_dict = checkpoint["model_state_dict"]
+            else:
+                state_dict = checkpoint
+                
+            missing_keys, unexpected_keys = self.model.load_state_dict(state_dict, strict=False)
+            print(f"Loaded model with {len(missing_keys)} missing keys and {len(unexpected_keys)} unexpected keys")
+            
+            # Move to device
+            self.model.to(self.device)
+            self.model.eval()
+            print("Model loaded successfully and placed in evaluation mode")
+            
+        except Exception as e:
+            print(f"Error loading model: {e}")
+            raise RuntimeError(f"Failed to load the model: {e}")
+    
+    def encode(self, sentences, batch_size=32, normalize=True):
+        """
+        Encode sentences to embeddings.
+        
+        Args:
+            sentences: A string or list of strings to encode
+            batch_size: Batch size for encoding
+            normalize: Whether to normalize the embeddings
+            
+        Returns:
+            Numpy array of embeddings
+        """
+        # Handle single sentence
+        if isinstance(sentences, str):
+            sentences = [sentences]
+            
+        all_embeddings = []
+        
+        # Process in batches
+        with torch.no_grad():
+            for i in range(0, len(sentences), batch_size):
+                batch = sentences[i:i+batch_size]
+                
+                # Tokenize
+                inputs = self.tokenizer(
+                    batch,
+                    padding=True,
+                    truncation=True,
+                    max_length=self.config.get("max_position_embeddings", 128),
+                    return_tensors="pt"
+                )
+                
+                # Move to device
+                input_ids = inputs["input_ids"].to(self.device)
+                attention_mask = inputs["attention_mask"].to(self.device)
+                
+                # Get embeddings
+                embeddings = self.model(input_ids, attention_mask)
+                
+                # Move to CPU and convert to numpy
+                all_embeddings.append(embeddings.cpu().numpy())
+                
+        # Concatenate all embeddings
+        all_embeddings = np.vstack(all_embeddings)
+        
+        # Normalize if requested
+        if normalize:
+            all_embeddings = all_embeddings / np.linalg.norm(all_embeddings, axis=1, keepdims=True)
+            
+        return all_embeddings
+    
+    def compute_similarity(self, texts1, texts2=None):
+        """
+        Compute cosine similarity between texts.
+        
+        Args:
+            texts1: First set of texts
+            texts2: Second set of texts. If None, compute similarity matrix within texts1.
+            
+        Returns:
+            Similarity scores
+        """
+        embeddings1 = self.encode(texts1)
+        
+        if texts2 is None:
+            # Compute similarity matrix within texts1
+            similarities = cosine_similarity(embeddings1)
+            return similarities
+        else:
+            # Compute similarity between texts1 and texts2
+            embeddings2 = self.encode(texts2)
+            
+            if len(texts1) == len(texts2):
+                # Compute pairwise similarity when the number of texts match
+                return np.array([
+                    cosine_similarity([e1], [e2])[0][0]
+                    for e1, e2 in zip(embeddings1, embeddings2)
+                ])
+            else:
+                # Return full similarity matrix
+                return cosine_similarity(embeddings1, embeddings2)
+    
+    def search(self, query, documents, top_k=5):
+        """
+        Search for similar documents to a query.
+        
+        Args:
+            query: The query text
+            documents: List of documents to search
+            top_k: Number of top results to return
+            
+        Returns:
+            List of dictionaries with document and score
+        """
+        # Get embeddings
+        query_embedding = self.encode([query])[0]
+        document_embeddings = self.encode(documents)
+        
+        # Compute similarities
+        similarities = np.dot(document_embeddings, query_embedding)
+        
+        # Get top indices
+        top_indices = np.argsort(similarities)[-top_k:][::-1]
+        
+        # Return results
+        results = []
+        for idx in top_indices:
+            results.append({
+                "document": documents[idx],
+                "score": float(similarities[idx])
+            })
+            
+        return results
+        
+    def evaluate_similarity_samples(self):
+        """Evaluate model on some standard similarity examples for Hindi"""
+        test_pairs = [
+            (
+                "मुझे हिंदी में पढ़ना बहुत पसंद है।",
+                "मैं हिंदी किताबें बहुत पसंद करता हूँ।"
+            ),
+            (
+                "आज मौसम बहुत अच्छा है।",
+                "आज बारिश हो रही है।"
+            ),
+            (
+                "भारत एक विशाल देश है।",
+                "भारत में कई भाषाएँ बोली जाती हैं।"
+            ),
+            (
+                "कंप्यूटर विज्ञान एक रोचक विषय है।",
+                "मैं कंप्यूटर साइंस का छात्र हूँ।"
+            ),
+            (
+                "मैं रोज सुबह योग करता हूँ।",
+                "स्वस्थ रहने के लिए व्यायाम जरूरी है।"
+            ),
+            # Add contrasting pairs to test discrimination
+            (
+                "मुझे हिंदी में पढ़ना बहुत पसंद है।", 
+                "क्रिकेट भारत में सबसे लोकप्रिय खेल है।"
+            ),
+            (
+                "आज मौसम बहुत अच्छा है।",
+                "भारतीय व्यंजन दुनिया भर में मशहूर हैं।"
+            ),
+            (
+                "कंप्यूटर विज्ञान एक रोचक विषय है।",
+                "हिमालय दुनिया का सबसे ऊंचा पर्वत है।"
+            )
+        ]
+        
+        print("Evaluating model on standard similarity samples:")
+        for i, (text1, text2) in enumerate(test_pairs):
+            similarity = self.compute_similarity([text1], [text2])[0]
+            print(f"\nPair {i+1}:")
+            print(f"  Sentence 1: {text1}")
+            print(f"  Sentence 2: {text2}")
+            print(f"  Similarity: {similarity:.4f}")
+            
+        return
+        
+    def visualize_embeddings(self, sentences, labels=None, output_path="hindi_embeddings_visualization.png"):
+        """
+        Create a t-SNE visualization of the embeddings.
+        
+        Args:
+            sentences: List of sentences to visualize
+            labels: Optional list of labels for the points
+            output_path: Path to save the visualization
+            
+        Returns:
+            Path to the saved visualization
+        """
+        # Encode sentences
+        embeddings = self.encode(sentences)
+        
+        # Apply t-SNE
+        tsne = TSNE(n_components=2, random_state=42, perplexity=min(30, len(embeddings)-1))
+        reduced_embeddings = tsne.fit_transform(embeddings)
+        
+        # Create plot
+        plt.figure(figsize=(12, 10))
+        
+        # Plot points
+        scatter = plt.scatter(
+            reduced_embeddings[:, 0], 
+            reduced_embeddings[:, 1],
+            c=range(len(reduced_embeddings)),
+            cmap='viridis',
+            alpha=0.8,
+            s=100
+        )
+        
+        # Add labels if provided
+        if labels:
+            for i, label in enumerate(labels):
+                plt.annotate(
+                    label,
+                    (reduced_embeddings[i, 0], reduced_embeddings[i, 1]),
+                    fontsize=10,
+                    alpha=0.7
+                )
+        
+        plt.title("t-SNE Visualization of Hindi Sentence Embeddings", fontsize=16)
+        plt.xlabel("Dimension 1", fontsize=12)
+        plt.ylabel("Dimension 2", fontsize=12)
+        plt.colorbar(scatter, label="Sentence Index")
+        plt.grid(alpha=0.3)
+        
+        # Save the figure
+        plt.tight_layout()
+        plt.savefig(output_path, dpi=300, bbox_inches='tight')
+        plt.close()
+        
+        print(f"Visualization saved to {output_path}")
+        return output_path
+
+def main():
+    # Create embedder
+    embedder = HindiEmbedder()
+    
+    # Run sample evaluation
+    embedder.evaluate_similarity_samples()
+    
+    # Example of semantic search
+    print("\nSemantic Search Example:")
+    query = "भारत की संस्कृति"
+    documents = [
+        "भारतीय संस्कृति दुनिया की सबसे प्राचीन संस्कृतियों में से एक है।",
+        "भारत की आबादी 1.3 अरब से अधिक है।",
+        "हिमालय पर्वत श्रृंखला भारत के उत्तर में स्थित है।",
+        "भारतीय व्यंजन में मसालों का प्रयोग किया जाता है।",
+        "भारत में 22 आधिकारिक भाषाएँ हैं।",
+        "संस्कृति लोगों के रहन-सहन का तरीका है।",
+        "भारत के विभिन्न राज्यों की अपनी अलग संस्कृति है।",
+        "रामायण और महाभारत भारतीय संस्कृति के महत्वपूर्ण हिस्से हैं।",
+    ]
+    
+    results = embedder.search(query, documents)
+    
+    print(f"Query: {query}")
+    print("Top results:")
+    for i, result in enumerate(results):
+        print(f"{i+1}. Score: {result['score']:.4f}")
+        print(f"   {result['document']}")
+    
+    # Create visualization example
+    print("\nCreating embedding visualization...")
+    visualization_sentences = [
+        "मुझे हिंदी में पढ़ना बहुत पसंद है।",
+        "मैं हिंदी किताबें बहुत पसंद करता हूँ।",
+        "आज मौसम बहुत अच्छा है।",
+        "आज बारिश हो रही है।",
+        "भारत एक विशाल देश है।",
+        "भारत में कई भाषाएँ बोली जाती हैं।",
+        "कंप्यूटर विज्ञान एक रोचक विषय है।",
+        "मैं कंप्यूटर साइंस का छात्र हूँ।",
+        "क्रिकेट भारत में सबसे लोकप्रिय खेल है।",
+        "भारतीय व्यंजन दुनिया भर में मशहूर हैं।",
+        "हिमालय दुनिया का सबसे ऊंचा पर्वत है।",
+        "गंगा भारत की सबसे पवित्र नदी है।",
+        "दिल्ली भारत की राजधानी है।",
+        "मुंबई भारत का आर्थिक केंद्र है।",
+        "तमिल, तेलुगु, कन्नड़ और मलयालम दक्षिण भारत की प्रमुख भाषाएँ हैं।"
+    ]
+    
+    labels = ["पढ़ना", "किताबें", "मौसम", "बारिश", "भारत", "भाषाएँ", "कंप्यूटर", 
+              "छात्र", "क्रिकेट", "व्यंजन", "हिमालय", "गंगा", "दिल्ली", "मुंबई", "भाषाएँ"]
+    
+    embedder.visualize_embeddings(visualization_sentences, labels)
+
+if __name__ == "__main__":
+    main()