tritopic 0.1.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

tritopic/core/representatives.py

@@ -0,0 +1,560 @@
+"""
+Representative Documents Module
+
+Implements various methods for selecting representative documents per topic:
+- Centroid: Documents closest to topic centroid
+- Medoid: Actual document that minimizes distance to all others
+- Archetype: Documents on the convex hull (extremes/facets)
+- Diverse: Diverse set maximizing coverage
+- Hybrid: Combination of medoid + archetypes + keyword champion
+"""
+
+from typing import List, Tuple, Optional, Literal, Dict
+import numpy as np
+from dataclasses import dataclass
+
+
+@dataclass
+class RepresentativeDocument:
+    """A representative document with metadata."""
+    doc_id: int
+    text: str
+    score: float
+    method: str  # How this document was selected
+    
+
+@dataclass
+class TopicRepresentatives:
+    """Collection of representative documents for a topic."""
+    topic_id: int
+    medoid: Optional[RepresentativeDocument] = None
+    archetypes: List[RepresentativeDocument] = None
+    keyword_champion: Optional[RepresentativeDocument] = None
+    centroid_nearest: List[RepresentativeDocument] = None
+    
+    def __post_init__(self):
+        if self.archetypes is None:
+            self.archetypes = []
+        if self.centroid_nearest is None:
+            self.centroid_nearest = []
+    
+    def get_all(self) -> List[RepresentativeDocument]:
+        """Get all representative documents."""
+        result = []
+        if self.medoid:
+            result.append(self.medoid)
+        result.extend(self.archetypes)
+        if self.keyword_champion and self.keyword_champion not in result:
+            result.append(self.keyword_champion)
+        for doc in self.centroid_nearest:
+            if doc not in result:
+                result.append(doc)
+        return result
+
+
+class RepresentativeSelector:
+    """
+    Selects representative documents for topics using various methods.
+    """
+    
+    def __init__(
+        self,
+        method: Literal["centroid", "medoid", "archetype", "diverse", "hybrid"] = "hybrid",
+        n_representatives: int = 5,
+        n_archetypes: int = 4,
+        archetype_method: Literal["pcha", "convex_hull", "furthest_sum"] = "furthest_sum",
+    ):
+        """
+        Initialize the representative selector.
+        
+        Parameters
+        ----------
+        method : str
+            Selection method
+        n_representatives : int
+            Total number of representatives per topic
+        n_archetypes : int
+            Number of archetypes (for archetype/hybrid methods)
+        archetype_method : str
+            Algorithm for archetype analysis
+        """
+        self.method = method
+        self.n_representatives = n_representatives
+        self.n_archetypes = n_archetypes
+        self.archetype_method = archetype_method
+    
+    def select(
+        self,
+        embeddings: np.ndarray,
+        labels: np.ndarray,
+        documents: List[str],
+        topic_keywords: Dict[int, List[str]] = None,
+    ) -> Dict[int, TopicRepresentatives]:
+        """
+        Select representative documents for each topic.
+        
+        Parameters
+        ----------
+        embeddings : np.ndarray
+            Document embeddings
+        labels : np.ndarray
+            Topic labels
+        documents : List[str]
+            Original documents
+        topic_keywords : dict
+            Keywords per topic (for keyword champion)
+            
+        Returns
+        -------
+        Dict[int, TopicRepresentatives]
+            Representative documents per topic
+        """
+        representatives = {}
+        unique_topics = sorted([t for t in np.unique(labels) if t >= 0])
+        
+        for topic_id in unique_topics:
+            mask = labels == topic_id
+            topic_indices = np.where(mask)[0]
+            topic_embeddings = embeddings[mask]
+            topic_docs = [documents[i] for i in topic_indices]
+            
+            if len(topic_indices) < 2:
+                # Too few documents for meaningful selection
+                representatives[topic_id] = TopicRepresentatives(
+                    topic_id=topic_id,
+                    centroid_nearest=[RepresentativeDocument(
+                        doc_id=topic_indices[0],
+                        text=topic_docs[0],
+                        score=1.0,
+                        method="only_document"
+                    )] if len(topic_indices) > 0 else []
+                )
+                continue
+            
+            keywords = topic_keywords.get(topic_id, []) if topic_keywords else []
+            
+            if self.method == "centroid":
+                reps = self._select_centroid(
+                    topic_embeddings, topic_indices, topic_docs
+                )
+            elif self.method == "medoid":
+                reps = self._select_medoid(
+                    topic_embeddings, topic_indices, topic_docs
+                )
+            elif self.method == "archetype":
+                reps = self._select_archetypes(
+                    topic_embeddings, topic_indices, topic_docs
+                )
+            elif self.method == "diverse":
+                reps = self._select_diverse(
+                    topic_embeddings, topic_indices, topic_docs
+                )
+            elif self.method == "hybrid":
+                reps = self._select_hybrid(
+                    topic_embeddings, topic_indices, topic_docs, keywords
+                )
+            else:
+                raise ValueError(f"Unknown method: {self.method}")
+            
+            reps.topic_id = topic_id
+            representatives[topic_id] = reps
+        
+        return representatives
+    
+    def _select_centroid(
+        self,
+        embeddings: np.ndarray,
+        indices: np.ndarray,
+        documents: List[str],
+    ) -> TopicRepresentatives:
+        """Select documents closest to centroid."""
+        centroid = embeddings.mean(axis=0)
+        
+        # Compute distances to centroid
+        distances = np.linalg.norm(embeddings - centroid, axis=1)
+        
+        # Sort by distance
+        sorted_idx = np.argsort(distances)
+        
+        n = min(self.n_representatives, len(indices))
+        selected = []
+        
+        for i in range(n):
+            local_idx = sorted_idx[i]
+            selected.append(RepresentativeDocument(
+                doc_id=indices[local_idx],
+                text=documents[local_idx],
+                score=1.0 / (1.0 + distances[local_idx]),
+                method="centroid"
+            ))
+        
+        return TopicRepresentatives(
+            topic_id=-1,
+            centroid_nearest=selected
+        )
+    
+    def _select_medoid(
+        self,
+        embeddings: np.ndarray,
+        indices: np.ndarray,
+        documents: List[str],
+    ) -> TopicRepresentatives:
+        """Select medoid (document minimizing total distance to others)."""
+        n = len(embeddings)
+        
+        # Compute pairwise distances
+        distances = np.zeros((n, n))
+        for i in range(n):
+            for j in range(i + 1, n):
+                d = np.linalg.norm(embeddings[i] - embeddings[j])
+                distances[i, j] = d
+                distances[j, i] = d
+        
+        # Find medoid (minimum total distance)
+        total_distances = distances.sum(axis=1)
+        medoid_idx = np.argmin(total_distances)
+        
+        medoid = RepresentativeDocument(
+            doc_id=indices[medoid_idx],
+            text=documents[medoid_idx],
+            score=1.0,
+            method="medoid"
+        )
+        
+        # Fill remaining slots with centroid method
+        remaining = self.n_representatives - 1
+        if remaining > 0:
+            centroid_reps = self._select_centroid(embeddings, indices, documents)
+            # Filter out medoid
+            other_docs = [d for d in centroid_reps.centroid_nearest 
+                         if d.doc_id != medoid.doc_id][:remaining]
+        else:
+            other_docs = []
+        
+        return TopicRepresentatives(
+            topic_id=-1,
+            medoid=medoid,
+            centroid_nearest=other_docs
+        )
+    
+    def _select_archetypes(
+        self,
+        embeddings: np.ndarray,
+        indices: np.ndarray,
+        documents: List[str],
+    ) -> TopicRepresentatives:
+        """
+        Select archetypal documents using furthest-sum or convex hull.
+        
+        Archetypes represent the "extreme" examples that define the
+        boundaries of the topic - the different facets or aspects.
+        """
+        n_archetypes = min(self.n_archetypes, len(indices))
+        
+        if self.archetype_method == "furthest_sum":
+            archetype_indices = self._furthest_sum_archetypes(embeddings, n_archetypes)
+        elif self.archetype_method == "convex_hull":
+            archetype_indices = self._convex_hull_archetypes(embeddings, n_archetypes)
+        elif self.archetype_method == "pcha":
+            archetype_indices = self._pcha_archetypes(embeddings, n_archetypes)
+        else:
+            archetype_indices = self._furthest_sum_archetypes(embeddings, n_archetypes)
+        
+        archetypes = []
+        for local_idx in archetype_indices:
+            archetypes.append(RepresentativeDocument(
+                doc_id=indices[local_idx],
+                text=documents[local_idx],
+                score=1.0,
+                method="archetype"
+            ))
+        
+        return TopicRepresentatives(
+            topic_id=-1,
+            archetypes=archetypes
+        )
+    
+    def _furthest_sum_archetypes(
+        self,
+        embeddings: np.ndarray,
+        n_archetypes: int,
+    ) -> List[int]:
+        """
+        Select archetypes using the Furthest Sum algorithm.
+        
+        This iteratively selects points that maximize the sum of
+        distances to already selected points - finding extreme points.
+        """
+        n = len(embeddings)
+        
+        if n <= n_archetypes:
+            return list(range(n))
+        
+        # Start with the point furthest from centroid
+        centroid = embeddings.mean(axis=0)
+        distances_to_centroid = np.linalg.norm(embeddings - centroid, axis=1)
+        selected = [np.argmax(distances_to_centroid)]
+        
+        # Iteratively add points maximizing distance to selected set
+        for _ in range(n_archetypes - 1):
+            best_idx = -1
+            best_score = -1
+            
+            for i in range(n):
+                if i in selected:
+                    continue
+                
+                # Sum of distances to all selected points
+                total_dist = sum(
+                    np.linalg.norm(embeddings[i] - embeddings[j])
+                    for j in selected
+                )
+                
+                if total_dist > best_score:
+                    best_score = total_dist
+                    best_idx = i
+            
+            if best_idx >= 0:
+                selected.append(best_idx)
+        
+        return selected
+    
+    def _convex_hull_archetypes(
+        self,
+        embeddings: np.ndarray,
+        n_archetypes: int,
+    ) -> List[int]:
+        """
+        Select archetypes as vertices of the convex hull.
+        
+        Note: Convex hull is computed in a reduced dimensional space
+        (PCA to 2-3 dimensions) since convex hull in high dimensions
+        often includes most/all points.
+        """
+        from sklearn.decomposition import PCA
+        
+        n = len(embeddings)
+        
+        if n <= n_archetypes:
+            return list(range(n))
+        
+        # Reduce to 2D for convex hull
+        n_components = min(2, n, embeddings.shape[1])
+        pca = PCA(n_components=n_components)
+        reduced = pca.fit_transform(embeddings)
+        
+        try:
+            from scipy.spatial import ConvexHull
+            
+            if n_components == 2 and n > 3:
+                hull = ConvexHull(reduced)
+                hull_vertices = hull.vertices.tolist()
+                
+                # If we have more hull vertices than needed, select diverse subset
+                if len(hull_vertices) > n_archetypes:
+                    return self._furthest_sum_archetypes(
+                        embeddings[hull_vertices], n_archetypes
+                    )
+                else:
+                    # Fill remaining with furthest sum
+                    remaining = n_archetypes - len(hull_vertices)
+                    if remaining > 0:
+                        candidates = [i for i in range(n) if i not in hull_vertices]
+                        if candidates:
+                            extra = self._furthest_sum_archetypes(
+                                embeddings[candidates], remaining
+                            )
+                            hull_vertices.extend([candidates[i] for i in extra])
+                    return hull_vertices[:n_archetypes]
+            else:
+                # Fall back to furthest sum
+                return self._furthest_sum_archetypes(embeddings, n_archetypes)
+                
+        except Exception:
+            # Fall back to furthest sum
+            return self._furthest_sum_archetypes(embeddings, n_archetypes)
+    
+    def _pcha_archetypes(
+        self,
+        embeddings: np.ndarray,
+        n_archetypes: int,
+    ) -> List[int]:
+        """
+        Select archetypes using Principal Convex Hull Analysis (PCHA).
+        
+        PCHA finds archetypes such that all data points can be
+        expressed as convex combinations of archetypes.
+        
+        Simplified version: Uses iterative approach to find points
+        that best "span" the data.
+        """
+        # For simplicity, we use furthest sum as a proxy for PCHA
+        # Full PCHA implementation would require optimization
+        return self._furthest_sum_archetypes(embeddings, n_archetypes)
+    
+    def _select_diverse(
+        self,
+        embeddings: np.ndarray,
+        indices: np.ndarray,
+        documents: List[str],
+    ) -> TopicRepresentatives:
+        """Select diverse documents using maximal marginal relevance."""
+        n = min(self.n_representatives, len(indices))
+        
+        # Start with centroid-nearest
+        centroid = embeddings.mean(axis=0)
+        distances = np.linalg.norm(embeddings - centroid, axis=1)
+        first_idx = np.argmin(distances)
+        
+        selected = [first_idx]
+        
+        # Add documents maximizing diversity
+        lambda_param = 0.5  # Balance relevance and diversity
+        
+        for _ in range(n - 1):
+            best_idx = -1
+            best_score = float('-inf')
+            
+            for i in range(len(embeddings)):
+                if i in selected:
+                    continue
+                
+                # Relevance: closeness to centroid
+                relevance = 1.0 / (1.0 + distances[i])
+                
+                # Diversity: minimum distance to already selected
+                min_dist_to_selected = min(
+                    np.linalg.norm(embeddings[i] - embeddings[j])
+                    for j in selected
+                )
+                
+                # MMR score
+                score = lambda_param * relevance + (1 - lambda_param) * min_dist_to_selected
+                
+                if score > best_score:
+                    best_score = score
+                    best_idx = i
+            
+            if best_idx >= 0:
+                selected.append(best_idx)
+        
+        reps = []
+        for local_idx in selected:
+            reps.append(RepresentativeDocument(
+                doc_id=indices[local_idx],
+                text=documents[local_idx],
+                score=1.0 / (1.0 + distances[local_idx]),
+                method="diverse"
+            ))
+        
+        return TopicRepresentatives(
+            topic_id=-1,
+            centroid_nearest=reps
+        )
+    
+    def _select_hybrid(
+        self,
+        embeddings: np.ndarray,
+        indices: np.ndarray,
+        documents: List[str],
+        keywords: List[str],
+    ) -> TopicRepresentatives:
+        """
+        Hybrid selection combining:
+        1. Medoid (most typical document)
+        2. Archetypes (facets/extremes)
+        3. Keyword Champion (highest keyword density)
+        """
+        # 1. Get medoid
+        n = len(embeddings)
+        distances = np.zeros((n, n))
+        for i in range(n):
+            for j in range(i + 1, n):
+                d = np.linalg.norm(embeddings[i] - embeddings[j])
+                distances[i, j] = d
+                distances[j, i] = d
+        
+        total_distances = distances.sum(axis=1)
+        medoid_idx = np.argmin(total_distances)
+        
+        medoid = RepresentativeDocument(
+            doc_id=indices[medoid_idx],
+            text=documents[medoid_idx],
+            score=1.0,
+            method="medoid"
+        )
+        
+        # 2. Get archetypes (excluding medoid)
+        n_arch = min(self.n_archetypes, n - 1)
+        if n_arch > 0:
+            # Filter out medoid
+            other_indices = [i for i in range(n) if i != medoid_idx]
+            other_embeddings = embeddings[other_indices]
+            
+            arch_local = self._furthest_sum_archetypes(other_embeddings, n_arch)
+            arch_indices = [other_indices[i] for i in arch_local]
+            
+            archetypes = [
+                RepresentativeDocument(
+                    doc_id=indices[i],
+                    text=documents[i],
+                    score=1.0,
+                    method="archetype"
+                )
+                for i in arch_indices
+            ]
+        else:
+            archetypes = []
+        
+        # 3. Get keyword champion
+        keyword_champion = None
+        if keywords:
+            keyword_set = set(w.lower() for w in keywords[:10])
+            best_score = 0
+            best_idx = -1
+            
+            for i, doc in enumerate(documents):
+                doc_words = set(doc.lower().split())
+                score = len(keyword_set & doc_words) / len(keyword_set)
+                if score > best_score:
+                    best_score = score
+                    best_idx = i
+            
+            if best_idx >= 0 and best_idx != medoid_idx and best_idx not in arch_indices:
+                keyword_champion = RepresentativeDocument(
+                    doc_id=indices[best_idx],
+                    text=documents[best_idx],
+                    score=best_score,
+                    method="keyword_champion"
+                )
+        
+        return TopicRepresentatives(
+            topic_id=-1,
+            medoid=medoid,
+            archetypes=archetypes,
+            keyword_champion=keyword_champion
+        )
+
+
+def get_representative_documents(
+    topic_representatives: Dict[int, TopicRepresentatives],
+    n_per_topic: int = 5,
+) -> Dict[int, List[Tuple[int, str, str]]]:
+    """
+    Get flattened list of representative documents per topic.
+    
+    Returns
+    -------
+    Dict[int, List[Tuple[int, str, str]]]
+        topic_id -> [(doc_id, text, method), ...]
+    """
+    result = {}
+    
+    for topic_id, reps in topic_representatives.items():
+        all_reps = reps.get_all()[:n_per_topic]
+        result[topic_id] = [
+            (r.doc_id, r.text, r.method)
+            for r in all_reps
+        ]
+    
+    return result