inference_embeddings.py

"""
Helion-V1-Embeddings Inference Script
Simple interface for generating embeddings and similarity search
"""

import numpy as np
import logging
from typing import List, Union, Optional
from pathlib import Path

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class HelionEmbeddings:
    """Simple interface for Helion-V1-Embeddings model."""
    
    def __init__(self, model_name: str = "DeepXR/Helion-V1-embeddings"):
        """
        Initialize embeddings model.
        
        Args:
            model_name: Model name or path
        """
        try:
            from sentence_transformers import SentenceTransformer
            
            logger.info(f"Loading model: {model_name}")
            self.model = SentenceTransformer(model_name)
            self.dimension = self.model.get_sentence_embedding_dimension()
            logger.info(f"Model loaded. Embedding dimension: {self.dimension}")
            
        except ImportError:
            logger.error("sentence-transformers not installed. Install with: pip install sentence-transformers")
            raise
        except Exception as e:
            logger.error(f"Failed to load model: {e}")
            raise
    
    def encode(
        self,
        texts: Union[str, List[str]],
        batch_size: int = 32,
        show_progress: bool = False,
        normalize: bool = True
    ) -> np.ndarray:
        """
        Generate embeddings for text(s).
        
        Args:
            texts: Single text or list of texts
            batch_size: Batch size for encoding
            show_progress: Show progress bar
            normalize: L2 normalize embeddings
            
        Returns:
            Numpy array of embeddings
        """
        embeddings = self.model.encode(
            texts,
            batch_size=batch_size,
            show_progress_bar=show_progress,
            normalize_embeddings=normalize
        )
        
        return embeddings
    
    def similarity(
        self,
        text1: Union[str, List[str]],
        text2: Union[str, List[str]]
    ) -> Union[float, np.ndarray]:
        """
        Calculate cosine similarity between texts.
        
        Args:
            text1: First text or list of texts
            text2: Second text or list of texts
            
        Returns:
            Similarity score(s)
        """
        from sentence_transformers import util
        
        emb1 = self.encode(text1)
        emb2 = self.encode(text2)
        
        similarity = util.cos_sim(emb1, emb2)
        
        # Return single value if both inputs are single strings
        if isinstance(text1, str) and isinstance(text2, str):
            return float(similarity[0][0])
        
        return similarity.numpy()
    
    def search(
        self,
        query: str,
        documents: List[str],
        top_k: int = 5
    ) -> List[tuple]:
        """
        Search for most similar documents to query.
        
        Args:
            query: Search query
            documents: List of documents to search
            top_k: Number of top results to return
            
        Returns:
            List of (document, score, index) tuples
        """
        from sentence_transformers import util
        
        query_emb = self.encode(query)
        doc_embs = self.encode(documents)
        
        # Calculate similarities
        similarities = util.cos_sim(query_emb, doc_embs)[0]
        
        # Get top-k results
        top_results = np.argsort(-similarities.numpy())[:top_k]
        
        results = []
        for idx in top_results:
            results.append((
                documents[idx],
                float(similarities[idx]),
                int(idx)
            ))
        
        return results
    
    def cluster(
        self,
        texts: List[str],
        num_clusters: int = 5,
        min_cluster_size: int = 2
    ) -> List[List[int]]:
        """
        Cluster texts by similarity.
        
        Args:
            texts: List of texts to cluster
            num_clusters: Number of clusters
            min_cluster_size: Minimum cluster size
            
        Returns:
            List of clusters (each cluster is a list of indices)
        """
        from sentence_transformers import util
        
        embeddings = self.encode(texts)
        
        # Perform clustering
        clusters = util.community_detection(
            embeddings,
            min_community_size=min_cluster_size,
            threshold=0.75
        )
        
        return clusters
    
    def save_embeddings(
        self,
        texts: List[str],
        output_file: str,
        format: str = "npy"
    ):
        """
        Generate and save embeddings to file.
        
        Args:
            texts: Texts to embed
            output_file: Output file path
            format: Format ('npy', 'npz', or 'json')
        """
        embeddings = self.encode(texts, show_progress=True)
        
        if format == "npy":
            np.save(output_file, embeddings)
        elif format == "npz":
            np.savez_compressed(output_file, embeddings=embeddings, texts=texts)
        elif format == "json":
            import json
            data = {
                "embeddings": embeddings.tolist(),
                "texts": texts,
                "dimension": self.dimension
            }
            with open(output_file, 'w') as f:
                json.dump(data, f)
        
        logger.info(f"Saved {len(texts)} embeddings to {output_file}")


def demo_usage():
    """Demonstrate usage examples."""
    
    print("="*60)
    print("Helion-V1-Embeddings Demo")
    print("="*60)
    
    # Initialize
    embedder = HelionEmbeddings("DeepXR/Helion-V1-embeddings")
    
    # Example 1: Basic encoding
    print("\n1. Basic Encoding:")
    text = "Hello, how are you?"
    embedding = embedder.encode(text)
    print(f"Text: {text}")
    print(f"Embedding shape: {embedding.shape}")
    print(f"First 5 values: {embedding[:5]}")
    
    # Example 2: Similarity
    print("\n2. Similarity Calculation:")
    text1 = "How do I reset my password?"
    text2 = "Password reset instructions"
    similarity = embedder.similarity(text1, text2)
    print(f"Text 1: {text1}")
    print(f"Text 2: {text2}")
    print(f"Similarity: {similarity:.4f}")
    
    # Example 3: Search
    print("\n3. Semantic Search:")
    query = "machine learning tutorial"
    documents = [
        "Learn machine learning basics",
        "Cooking recipes for beginners",
        "Introduction to neural networks",
        "Travel guide to Europe",
        "Python programming course"
    ]
    
    results = embedder.search(query, documents, top_k=3)
    print(f"Query: {query}")
    print("\nTop 3 Results:")
    for i, (doc, score, idx) in enumerate(results, 1):
        print(f"{i}. [{score:.4f}] {doc}")
    
    print("\n" + "="*60)


def main():
    """Main CLI interface."""
    import argparse
    
    parser = argparse.ArgumentParser(
        description="Helion-V1-Embeddings Inference"
    )
    parser.add_argument(
        "--model",
        default="DeepXR/Helion-V1-embeddings",
        help="Model name or path"
    )
    
    subparsers = parser.add_subparsers(dest="command", help="Command to run")
    
    # Encode command
    encode_parser = subparsers.add_parser("encode", help="Encode text(s)")
    encode_parser.add_argument("text", nargs="+", help="Text(s) to encode")
    encode_parser.add_argument("--output", help="Save embeddings to file")
    
    # Similarity command
    sim_parser = subparsers.add_parser("similarity", help="Calculate similarity")
    sim_parser.add_argument("text1", help="First text")
    sim_parser.add_argument("text2", help="Second text")
    
    # Search command
    search_parser = subparsers.add_parser("search", help="Search documents")
    search_parser.add_argument("query", help="Search query")
    search_parser.add_argument("--documents", nargs="+", required=True)
    search_parser.add_argument("--top-k", type=int, default=5)
    
    # Demo command
    subparsers.add_parser("demo", help="Run demo")
    
    args = parser.parse_args()
    
    if args.command == "demo":
        demo_usage()
        return
    
    # Initialize model
    embedder = HelionEmbeddings(args.model)
    
    if args.command == "encode":
        embeddings = embedder.encode(args.text, show_progress=True)
        print(f"Generated {len(embeddings)} embeddings")
        print(f"Shape: {embeddings.shape}")
        
        if args.output:
            embedder.save_embeddings(args.text, args.output)
    
    elif args.command == "similarity":
        similarity = embedder.similarity(args.text1, args.text2)
        print(f"Text 1: {args.text1}")
        print(f"Text 2: {args.text2}")
        print(f"Similarity: {similarity:.4f}")
    
    elif args.command == "search":
        results = embedder.search(
            args.query,
            args.documents,
            top_k=args.top_k
        )
        
        print(f"Query: {args.query}")
        print(f"\nTop {args.top_k} results:")
        for i, (doc, score, idx) in enumerate(results, 1):
            print(f"{i}. [{score:.4f}] {doc}")
    
    else:
        parser.print_help()


if __name__ == "__main__":
    main()