tritopic 0.1.0__py3-none-any.whl

tritopic/__init__.py

@@ -0,0 +1,46 @@
+"""
+TriTopic: Tri-Modal Graph Topic Modeling with Iterative Refinement
+===================================================================
+
+A state-of-the-art topic modeling library that combines:
+- Semantic embeddings (Sentence-BERT, Instructor, BGE)
+- Lexical similarity (BM25)
+- Metadata context (optional)
+
+With advanced techniques:
+- Leiden clustering with consensus
+- Mutual kNN + SNN graph construction
+- Iterative refinement loop
+- LLM-powered topic labeling
+
+Basic usage:
+-----------
+>>> from tritopic import TriTopic
+>>> model = TriTopic()
+>>> topics = model.fit_transform(documents)
+>>> model.visualize()
+
+Author: Roman Egger
+License: MIT
+"""
+
+__version__ = "0.1.0"
+__author__ = "Roman Egger"
+
+from tritopic.core.model import TriTopic
+from tritopic.core.graph_builder import GraphBuilder
+from tritopic.core.clustering import ConsensusLeiden
+from tritopic.core.embeddings import EmbeddingEngine
+from tritopic.core.keywords import KeywordExtractor
+from tritopic.labeling.llm_labeler import LLMLabeler
+from tritopic.visualization.plotter import TopicVisualizer
+
+__all__ = [
+    "TriTopic",
+    "GraphBuilder", 
+    "ConsensusLeiden",
+    "EmbeddingEngine",
+    "KeywordExtractor",
+    "LLMLabeler",
+    "TopicVisualizer",
+]

tritopic/core/__init__.py

@@ -0,0 +1,17 @@
+"""Core components for TriTopic."""
+
+from tritopic.core.model import TriTopic, TriTopicConfig, TopicInfo
+from tritopic.core.graph_builder import GraphBuilder
+from tritopic.core.clustering import ConsensusLeiden
+from tritopic.core.embeddings import EmbeddingEngine
+from tritopic.core.keywords import KeywordExtractor
+
+__all__ = [
+    "TriTopic",
+    "TriTopicConfig",
+    "TopicInfo",
+    "GraphBuilder",
+    "ConsensusLeiden",
+    "EmbeddingEngine",
+    "KeywordExtractor",
+]

tritopic/core/clustering.py

@@ -0,0 +1,331 @@
+"""
+Consensus Leiden Clustering
+============================
+
+Robust community detection with:
+- Leiden algorithm (better than Louvain)
+- Consensus clustering for stability
+- Resolution parameter tuning
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+from scipy.cluster.hierarchy import linkage, fcluster
+from sklearn.metrics import adjusted_rand_score
+from collections import Counter
+
+
+class ConsensusLeiden:
+    """
+    Leiden clustering with consensus for stability.
+    
+    Runs multiple Leiden clusterings with different seeds and combines
+    results using consensus clustering. This dramatically improves
+    reproducibility and reduces sensitivity to random initialization.
+    
+    Parameters
+    ----------
+    resolution : float
+        Resolution parameter for Leiden. Higher = more clusters. Default: 1.0
+    n_runs : int
+        Number of consensus runs. Default: 10
+    random_state : int
+        Random seed for reproducibility. Default: 42
+    consensus_threshold : float
+        Minimum agreement ratio for consensus. Default: 0.5
+    """
+    
+    def __init__(
+        self,
+        resolution: float = 1.0,
+        n_runs: int = 10,
+        random_state: int = 42,
+        consensus_threshold: float = 0.5,
+    ):
+        self.resolution = resolution
+        self.n_runs = n_runs
+        self.random_state = random_state
+        self.consensus_threshold = consensus_threshold
+        
+        self.labels_: np.ndarray | None = None
+        self.stability_score_: float | None = None
+        self._all_partitions: list[np.ndarray] = []
+    
+    def fit_predict(
+        self,
+        graph: "igraph.Graph",
+        min_cluster_size: int = 5,
+        resolution: float | None = None,
+    ) -> np.ndarray:
+        """
+        Fit Leiden clustering with consensus.
+        
+        Parameters
+        ----------
+        graph : igraph.Graph
+            Input graph with edge weights.
+        min_cluster_size : int
+            Minimum cluster size. Smaller clusters become outliers.
+        resolution : float, optional
+            Override default resolution.
+            
+        Returns
+        -------
+        labels : np.ndarray
+            Cluster assignments. -1 for outliers.
+        """
+        import leidenalg as la
+        
+        res = resolution or self.resolution
+        n_nodes = graph.vcount()
+        
+        # Run multiple Leiden clusterings
+        self._all_partitions = []
+        
+        for run in range(self.n_runs):
+            seed = self.random_state + run
+            
+            # Run Leiden
+            partition = la.find_partition(
+                graph,
+                la.RBConfigurationVertexPartition,
+                weights="weight",
+                resolution_parameter=res,
+                seed=seed,
+            )
+            
+            # Convert to labels
+            labels = np.array(partition.membership)
+            self._all_partitions.append(labels)
+        
+        # Compute consensus
+        self.labels_ = self._compute_consensus(self._all_partitions)
+        
+        # Handle small clusters as outliers
+        self.labels_ = self._handle_small_clusters(self.labels_, min_cluster_size)
+        
+        # Compute stability score
+        self.stability_score_ = self._compute_stability()
+        
+        return self.labels_
+    
+    def _compute_consensus(self, partitions: list[np.ndarray]) -> np.ndarray:
+        """
+        Compute consensus partition from multiple runs.
+        
+        Uses co-occurrence matrix and hierarchical clustering.
+        """
+        n_nodes = len(partitions[0])
+        n_runs = len(partitions)
+        
+        # Build co-occurrence matrix
+        # co_occur[i,j] = fraction of runs where i and j are in same cluster
+        co_occur = np.zeros((n_nodes, n_nodes))
+        
+        for partition in partitions:
+            for cluster_id in np.unique(partition):
+                members = np.where(partition == cluster_id)[0]
+                for i in members:
+                    for j in members:
+                        co_occur[i, j] += 1
+        
+        co_occur /= n_runs
+        
+        # Convert co-occurrence to distance
+        distance = 1 - co_occur
+        
+        # Hierarchical clustering on distance matrix
+        # Use condensed form for linkage
+        condensed = []
+        for i in range(n_nodes):
+            for j in range(i + 1, n_nodes):
+                condensed.append(distance[i, j])
+        condensed = np.array(condensed)
+        
+        # Average linkage tends to work well for consensus
+        Z = linkage(condensed, method="average")
+        
+        # Cut at threshold that matches approximate number of clusters
+        # from the most frequent partition
+        n_clusters_list = [len(np.unique(p)) for p in partitions]
+        median_n_clusters = int(np.median(n_clusters_list))
+        
+        # Find optimal cut
+        best_labels = None
+        best_score = -1
+        
+        for n_clusters in range(max(2, median_n_clusters - 2), median_n_clusters + 3):
+            try:
+                labels = fcluster(Z, n_clusters, criterion="maxclust")
+                labels = labels - 1  # 0-indexed
+                
+                # Score by average ARI with original partitions
+                ari_scores = [adjusted_rand_score(labels, p) for p in partitions]
+                avg_ari = np.mean(ari_scores)
+                
+                if avg_ari > best_score:
+                    best_score = avg_ari
+                    best_labels = labels
+            except Exception:
+                continue
+        
+        if best_labels is None:
+            # Fallback to most common partition
+            best_labels = partitions[0]
+        
+        return best_labels
+    
+    def _handle_small_clusters(
+        self,
+        labels: np.ndarray,
+        min_size: int,
+    ) -> np.ndarray:
+        """Mark small clusters as outliers (-1)."""
+        result = labels.copy()
+        
+        for cluster_id in np.unique(labels):
+            if cluster_id == -1:
+                continue
+            
+            size = np.sum(labels == cluster_id)
+            if size < min_size:
+                result[labels == cluster_id] = -1
+        
+        # Relabel to consecutive integers
+        unique_labels = sorted([l for l in np.unique(result) if l != -1])
+        label_map = {old: new for new, old in enumerate(unique_labels)}
+        label_map[-1] = -1
+        
+        result = np.array([label_map[l] for l in result])
+        
+        return result
+    
+    def _compute_stability(self) -> float:
+        """Compute stability score as average pairwise ARI."""
+        if len(self._all_partitions) < 2:
+            return 1.0
+        
+        ari_scores = []
+        for i in range(len(self._all_partitions)):
+            for j in range(i + 1, len(self._all_partitions)):
+                ari = adjusted_rand_score(
+                    self._all_partitions[i],
+                    self._all_partitions[j]
+                )
+                ari_scores.append(ari)
+        
+        return float(np.mean(ari_scores))
+    
+    def find_optimal_resolution(
+        self,
+        graph: "igraph.Graph",
+        resolution_range: tuple[float, float] = (0.1, 2.0),
+        n_steps: int = 10,
+        target_n_topics: int | None = None,
+    ) -> float:
+        """
+        Find optimal resolution parameter.
+        
+        Parameters
+        ----------
+        graph : igraph.Graph
+            Input graph.
+        resolution_range : tuple
+            Range of resolutions to search.
+        n_steps : int
+            Number of resolutions to try.
+        target_n_topics : int, optional
+            If provided, find resolution closest to this number of topics.
+            
+        Returns
+        -------
+        optimal_resolution : float
+            Best resolution parameter.
+        """
+        import leidenalg as la
+        
+        resolutions = np.linspace(resolution_range[0], resolution_range[1], n_steps)
+        results = []
+        
+        for res in resolutions:
+            partition = la.find_partition(
+                graph,
+                la.RBConfigurationVertexPartition,
+                weights="weight",
+                resolution_parameter=res,
+                seed=self.random_state,
+            )
+            
+            n_clusters = len(set(partition.membership))
+            modularity = partition.modularity
+            
+            results.append({
+                "resolution": res,
+                "n_clusters": n_clusters,
+                "modularity": modularity,
+            })
+        
+        if target_n_topics is not None:
+            # Find closest to target
+            best = min(results, key=lambda x: abs(x["n_clusters"] - target_n_topics))
+        else:
+            # Find highest modularity
+            best = max(results, key=lambda x: x["modularity"])
+        
+        return best["resolution"]
+
+
+class HDBSCANClusterer:
+    """
+    Alternative clustering using HDBSCAN.
+    
+    Useful for datasets with varying density or many outliers.
+    """
+    
+    def __init__(
+        self,
+        min_cluster_size: int = 10,
+        min_samples: int = 5,
+        metric: str = "euclidean",
+    ):
+        self.min_cluster_size = min_cluster_size
+        self.min_samples = min_samples
+        self.metric = metric
+        
+        self.labels_: np.ndarray | None = None
+        self.probabilities_: np.ndarray | None = None
+    
+    def fit_predict(
+        self,
+        embeddings: np.ndarray,
+        **kwargs,
+    ) -> np.ndarray:
+        """
+        Fit HDBSCAN clustering.
+        
+        Parameters
+        ----------
+        embeddings : np.ndarray
+            Document embeddings (optionally reduced with UMAP first).
+            
+        Returns
+        -------
+        labels : np.ndarray
+            Cluster assignments. -1 for outliers.
+        """
+        import hdbscan
+        
+        clusterer = hdbscan.HDBSCAN(
+            min_cluster_size=self.min_cluster_size,
+            min_samples=self.min_samples,
+            metric=self.metric,
+            **kwargs,
+        )
+        
+        self.labels_ = clusterer.fit_predict(embeddings)
+        self.probabilities_ = clusterer.probabilities_
+        
+        return self.labels_

tritopic/core/embeddings.py

@@ -0,0 +1,222 @@
+"""
+Embedding Engine for TriTopic
+==============================
+
+Handles document embedding with support for multiple models:
+- Sentence-BERT models (default)
+- Instructor models (task-specific)
+- BGE models (multilingual)
+"""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+import numpy as np
+from tqdm import tqdm
+
+
+class EmbeddingEngine:
+    """
+    Generate document embeddings using transformer models.
+    
+    Supports various embedding models optimized for different use cases.
+    
+    Parameters
+    ----------
+    model_name : str
+        Name of the sentence-transformers model. Popular choices:
+        - "all-MiniLM-L6-v2": Fast, good quality (default)
+        - "all-mpnet-base-v2": Higher quality, slower
+        - "BAAI/bge-base-en-v1.5": State-of-the-art
+        - "BAAI/bge-m3": Multilingual
+        - "hkunlp/instructor-large": Task-specific (use with instruction)
+    batch_size : int
+        Batch size for encoding. Default: 32
+    device : str or None
+        Device to use ("cuda", "cpu", or None for auto).
+    show_progress : bool
+        Show progress bar. Default: True
+    """
+    
+    def __init__(
+        self,
+        model_name: str = "all-MiniLM-L6-v2",
+        batch_size: int = 32,
+        device: str | None = None,
+        show_progress: bool = True,
+    ):
+        self.model_name = model_name
+        self.batch_size = batch_size
+        self.device = device
+        self.show_progress = show_progress
+        
+        self._model = None
+        self._is_instructor = "instructor" in model_name.lower()
+    
+    def _load_model(self):
+        """Lazy load the embedding model."""
+        if self._model is not None:
+            return
+        
+        from sentence_transformers import SentenceTransformer
+        
+        self._model = SentenceTransformer(
+            self.model_name,
+            device=self.device,
+        )
+    
+    def encode(
+        self,
+        documents: list[str],
+        instruction: str | None = None,
+        normalize: bool = True,
+    ) -> np.ndarray:
+        """
+        Encode documents to embeddings.
+        
+        Parameters
+        ----------
+        documents : list[str]
+            List of document texts.
+        instruction : str, optional
+            Instruction for Instructor models (e.g., "Represent the topic of this document:").
+        normalize : bool
+            Whether to L2-normalize embeddings. Default: True
+            
+        Returns
+        -------
+        embeddings : np.ndarray
+            Document embeddings of shape (n_docs, embedding_dim).
+        """
+        self._load_model()
+        
+        # Handle instructor models
+        if self._is_instructor and instruction:
+            documents = [[instruction, doc] for doc in documents]
+        
+        # Encode in batches
+        embeddings = self._model.encode(
+            documents,
+            batch_size=self.batch_size,
+            show_progress_bar=self.show_progress,
+            normalize_embeddings=normalize,
+            convert_to_numpy=True,
+        )
+        
+        return embeddings
+    
+    def encode_with_pooling(
+        self,
+        documents: list[str],
+        pooling: Literal["mean", "max", "cls"] = "mean",
+    ) -> np.ndarray:
+        """
+        Encode with custom pooling strategy.
+        
+        Parameters
+        ----------
+        documents : list[str]
+            Document texts.
+        pooling : str
+            Pooling strategy: "mean", "max", or "cls".
+            
+        Returns
+        -------
+        embeddings : np.ndarray
+            Pooled embeddings.
+        """
+        # For now, use default pooling from model
+        # Custom pooling would require access to token-level embeddings
+        return self.encode(documents)
+    
+    @property
+    def embedding_dim(self) -> int:
+        """Get embedding dimension."""
+        self._load_model()
+        return self._model.get_sentence_embedding_dimension()
+    
+    def similarity(
+        self,
+        embeddings1: np.ndarray,
+        embeddings2: np.ndarray | None = None,
+    ) -> np.ndarray:
+        """
+        Compute cosine similarity between embeddings.
+        
+        Parameters
+        ----------
+        embeddings1 : np.ndarray
+            First set of embeddings.
+        embeddings2 : np.ndarray, optional
+            Second set. If None, compute pairwise similarity of embeddings1.
+            
+        Returns
+        -------
+        similarity : np.ndarray
+            Similarity matrix.
+        """
+        from sklearn.metrics.pairwise import cosine_similarity
+        
+        if embeddings2 is None:
+            return cosine_similarity(embeddings1)
+        return cosine_similarity(embeddings1, embeddings2)
+
+
+class MultiModelEmbedding:
+    """
+    Combine embeddings from multiple models.
+    
+    Useful for ensemble approaches where different models capture
+    different aspects of document semantics.
+    """
+    
+    def __init__(
+        self,
+        model_names: list[str],
+        weights: list[float] | None = None,
+        batch_size: int = 32,
+    ):
+        self.model_names = model_names
+        self.weights = weights or [1.0 / len(model_names)] * len(model_names)
+        self.batch_size = batch_size
+        
+        self._engines = [
+            EmbeddingEngine(name, batch_size=batch_size)
+            for name in model_names
+        ]
+    
+    def encode(
+        self,
+        documents: list[str],
+        normalize: bool = True,
+    ) -> np.ndarray:
+        """
+        Encode using all models and combine.
+        
+        Parameters
+        ----------
+        documents : list[str]
+            Document texts.
+        normalize : bool
+            Normalize final embeddings.
+            
+        Returns
+        -------
+        embeddings : np.ndarray
+            Combined embeddings (concatenated).
+        """
+        all_embeddings = []
+        
+        for engine, weight in zip(self._engines, self.weights):
+            emb = engine.encode(documents, normalize=True)
+            all_embeddings.append(emb * weight)
+        
+        # Concatenate
+        combined = np.hstack(all_embeddings)
+        
+        if normalize:
+            norms = np.linalg.norm(combined, axis=1, keepdims=True)
+            combined = combined / (norms + 1e-10)
+        
+        return combined