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

tritopic/labeling.py

@@ -0,0 +1,313 @@
+"""
+LLM-powered topic labeling for TriTopic.
+
+Supports multiple LLM providers for generating human-readable topic labels.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional
+
+
+class BaseLLMLabeler(ABC):
+    """Abstract base class for LLM-based topic labelers."""
+    
+    def __init__(
+        self,
+        language: str = "english",
+        max_label_words: int = 5,
+        include_context: bool = True
+    ):
+        self.language = language
+        self.max_label_words = max_label_words
+        self.include_context = include_context
+    
+    @abstractmethod
+    def _call_llm(self, prompt: str) -> str:
+        """Call the LLM with the given prompt."""
+        pass
+    
+    def generate_label(
+        self,
+        keywords: List[str],
+        representative_docs: List[str],
+        context: Optional[str] = None
+    ) -> str:
+        """
+        Generate a human-readable label for a topic.
+        
+        Parameters
+        ----------
+        keywords : List[str]
+            Top keywords for the topic.
+        representative_docs : List[str]
+            Representative document excerpts.
+        context : str, optional
+            Additional context about the corpus.
+            
+        Returns
+        -------
+        str
+            Generated topic label.
+        """
+        prompt = self._build_prompt(keywords, representative_docs, context)
+        response = self._call_llm(prompt)
+        return self._parse_response(response)
+    
+    def _build_prompt(
+        self,
+        keywords: List[str],
+        representative_docs: List[str],
+        context: Optional[str] = None
+    ) -> str:
+        """Build the prompt for label generation."""
+        keyword_str = ", ".join(keywords[:10])
+        
+        docs_str = ""
+        for i, doc in enumerate(representative_docs[:3], 1):
+            # Truncate long documents
+            doc_preview = doc[:300] + "..." if len(doc) > 300 else doc
+            docs_str += f"\n{i}. \"{doc_preview}\""
+        
+        language_instruction = ""
+        if self.language.lower() != "english":
+            language_instruction = f"\nIMPORTANT: Generate the label in {self.language}."
+        
+        context_str = ""
+        if context and self.include_context:
+            context_str = f"\nCorpus context: {context}"
+        
+        prompt = f"""Generate a concise, descriptive label for this topic.
+
+Keywords: {keyword_str}
+
+Representative documents:{docs_str}
+{context_str}
+{language_instruction}
+Requirements:
+- Maximum {self.max_label_words} words
+- Be specific and descriptive
+- Capture the main theme
+- Use title case
+
+Respond with ONLY the label, nothing else."""
+        
+        return prompt
+    
+    def _parse_response(self, response: str) -> str:
+        """Parse and clean the LLM response."""
+        # Clean up the response
+        label = response.strip()
+        
+        # Remove quotes if present
+        label = label.strip('"\'')
+        
+        # Remove any prefix like "Label:" or "Topic:"
+        for prefix in ["Label:", "Topic:", "Title:"]:
+            if label.lower().startswith(prefix.lower()):
+                label = label[len(prefix):].strip()
+        
+        # Truncate if too long
+        words = label.split()
+        if len(words) > self.max_label_words:
+            label = " ".join(words[:self.max_label_words])
+        
+        return label
+
+
+class AnthropicLabeler(BaseLLMLabeler):
+    """Topic labeler using Anthropic's Claude models."""
+    
+    def __init__(
+        self,
+        api_key: str,
+        model: str = "claude-sonnet-4-20250514",
+        language: str = "english",
+        max_label_words: int = 5,
+        **kwargs
+    ):
+        super().__init__(language, max_label_words, **kwargs)
+        self.api_key = api_key
+        self.model = model
+        self._client = None
+    
+    @property
+    def client(self):
+        if self._client is None:
+            try:
+                from anthropic import Anthropic
+                self._client = Anthropic(api_key=self.api_key)
+            except ImportError:
+                raise ImportError(
+                    "Anthropic package required. Install with: pip install anthropic"
+                )
+        return self._client
+    
+    def _call_llm(self, prompt: str) -> str:
+        message = self.client.messages.create(
+            model=self.model,
+            max_tokens=100,
+            messages=[
+                {"role": "user", "content": prompt}
+            ]
+        )
+        return message.content[0].text
+
+
+class OpenAILabeler(BaseLLMLabeler):
+    """Topic labeler using OpenAI's GPT models."""
+    
+    def __init__(
+        self,
+        api_key: str,
+        model: str = "gpt-4o",
+        language: str = "english",
+        max_label_words: int = 5,
+        **kwargs
+    ):
+        super().__init__(language, max_label_words, **kwargs)
+        self.api_key = api_key
+        self.model = model
+        self._client = None
+    
+    @property
+    def client(self):
+        if self._client is None:
+            try:
+                from openai import OpenAI
+                self._client = OpenAI(api_key=self.api_key)
+            except ImportError:
+                raise ImportError(
+                    "OpenAI package required. Install with: pip install openai"
+                )
+        return self._client
+    
+    def _call_llm(self, prompt: str) -> str:
+        response = self.client.chat.completions.create(
+            model=self.model,
+            max_tokens=100,
+            messages=[
+                {"role": "user", "content": prompt}
+            ]
+        )
+        return response.choices[0].message.content
+
+
+class LLMLabeler:
+    """
+    Unified interface for LLM-based topic labeling.
+    
+    Supports multiple providers: Anthropic (Claude), OpenAI (GPT).
+    
+    Parameters
+    ----------
+    provider : str
+        LLM provider: "anthropic" or "openai".
+    api_key : str
+        API key for the provider.
+    model : str, optional
+        Specific model to use. Defaults vary by provider.
+    language : str
+        Output language for labels.
+    max_label_words : int
+        Maximum words in generated labels.
+        
+    Examples
+    --------
+    >>> labeler = LLMLabeler(
+    ...     provider="anthropic",
+    ...     api_key="your-api-key",
+    ...     language="german"
+    ... )
+    >>> model.generate_labels(labeler)
+    """
+    
+    def __init__(
+        self,
+        provider: str = "anthropic",
+        api_key: Optional[str] = None,
+        model: Optional[str] = None,
+        language: str = "english",
+        max_label_words: int = 5,
+        **kwargs
+    ):
+        provider = provider.lower()
+        
+        if provider == "anthropic":
+            self._labeler = AnthropicLabeler(
+                api_key=api_key,
+                model=model or "claude-sonnet-4-20250514",
+                language=language,
+                max_label_words=max_label_words,
+                **kwargs
+            )
+        elif provider == "openai":
+            self._labeler = OpenAILabeler(
+                api_key=api_key,
+                model=model or "gpt-4o",
+                language=language,
+                max_label_words=max_label_words,
+                **kwargs
+            )
+        else:
+            raise ValueError(f"Unknown provider: {provider}. Use 'anthropic' or 'openai'.")
+    
+    def generate_label(
+        self,
+        keywords: List[str],
+        representative_docs: List[str],
+        context: Optional[str] = None
+    ) -> str:
+        """
+        Generate a label for a topic.
+        
+        Parameters
+        ----------
+        keywords : List[str]
+            Topic keywords.
+        representative_docs : List[str]
+            Representative documents.
+        context : str, optional
+            Corpus context.
+            
+        Returns
+        -------
+        str
+            Generated label.
+        """
+        return self._labeler.generate_label(keywords, representative_docs, context)
+
+
+class KeywordLabeler:
+    """
+    Simple labeler that combines top keywords into a label.
+    
+    No LLM required - uses keyword-based heuristics.
+    """
+    
+    def __init__(
+        self,
+        n_keywords: int = 3,
+        separator: str = " & ",
+        capitalize: bool = True
+    ):
+        self.n_keywords = n_keywords
+        self.separator = separator
+        self.capitalize = capitalize
+    
+    def generate_label(
+        self,
+        keywords: List[str],
+        representative_docs: Optional[List[str]] = None,
+        context: Optional[str] = None
+    ) -> str:
+        """Generate label from keywords."""
+        top_keywords = keywords[:self.n_keywords]
+        
+        if self.capitalize:
+            top_keywords = [kw.title() for kw in top_keywords]
+        
+        return self.separator.join(top_keywords)