semaxis 0.16.0__tar.gz

semaxis-0.16.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 pillyshi
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

semaxis-0.16.0/PKG-INFO

@@ -0,0 +1,14 @@
+Metadata-Version: 2.3
+Name: semaxis
+Version: 0.16.0
+Summary: Interpretable NLI-based text features for scikit-learn
+Requires-Python: >=3.11
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: numpy (>=1.26)
+Requires-Dist: openai (>=1.30)
+Requires-Dist: scikit-learn (>=1.4)
+Requires-Dist: sentence-transformers (>=3.0)
+Requires-Dist: tiktoken (>=0.7)

semaxis-0.16.0/pyproject.toml

@@ -0,0 +1,38 @@
+[tool.poetry]
+name = "semaxis"
+version = "0.16.0"
+description = "Interpretable NLI-based text features for scikit-learn"
+authors = []
+packages = [{include = "semaxis", from = "src"}]
+
+[tool.poetry.dependencies]
+python = ">=3.11"
+openai = ">=1.30"
+tiktoken = ">=0.7"
+sentence-transformers = ">=3.0"
+scikit-learn = ">=1.4"
+numpy = ">=1.26"
+
+[tool.poetry.group.dev.dependencies]
+pytest = "*"
+ruff = "*"
+mypy = "*"
+pyalex = "^0.21"
+
+[tool.poetry.group.langchain]
+optional = true
+
+[tool.poetry.group.langchain.dependencies]
+langchain-core = ">=0.3"
+langchain-ollama = ">=0.3"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.mypy]
+python_version = "3.11"
+
+[[tool.mypy.overrides]]
+module = ["sklearn.*", "sentence_transformers.*"]
+ignore_missing_imports = true

semaxis-0.16.0/src/semaxis/__init__.py

@@ -0,0 +1,11 @@
+from .llm import LangChainLLMClient, LLMClient
+from .supervised import FeatureMeta, SupervisedTransformer
+from .unsupervised import UnsupervisedTransformer
+
+__all__ = [
+    "SupervisedTransformer",
+    "UnsupervisedTransformer",
+    "FeatureMeta",
+    "LLMClient",
+    "LangChainLLMClient",
+]

semaxis-0.16.0/src/semaxis/llm.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+import json
+import re
+from typing import Any, Protocol, runtime_checkable
+
+import tiktoken
+from openai import OpenAI
+
+
+@runtime_checkable
+class BaseLLMClient(Protocol):
+    """Protocol defining the interface required by all LLM client implementations."""
+
+    def complete(self, messages: list[dict[str, str]]) -> str: ...
+    def complete_json(self, messages: list[dict[str, str]]) -> Any: ...
+    def count_tokens(self, text: str) -> int: ...
+
+
+class LLMClient:
+    """Thin wrapper around the OpenAI chat completions API with token counting."""
+
+    def __init__(self, model: str, api_key: str | None = None) -> None:
+        self.model = model
+        self._client = OpenAI(api_key=api_key)
+        try:
+            self._encoding = tiktoken.encoding_for_model(model)
+        except KeyError:
+            self._encoding = tiktoken.get_encoding("cl100k_base")
+
+    def complete(self, messages: list[dict[str, str]]) -> str:
+        """Send a chat completion request and return the content string."""
+        response = self._client.chat.completions.create(
+            model=self.model,
+            messages=messages,  # type: ignore[arg-type]
+        )
+        return response.choices[0].message.content or ""
+
+    def complete_json(self, messages: list[dict[str, str]]) -> Any:
+        """Send a request expecting JSON output and return the parsed object."""
+        response = self._client.chat.completions.create(
+            model=self.model,
+            messages=messages,  # type: ignore[arg-type]
+            response_format={"type": "json_object"},  # type: ignore[call-overload]
+        )
+        content = response.choices[0].message.content or ""
+        return json.loads(content)
+
+    def count_tokens(self, text: str) -> int:
+        """Return the number of tokens in text using the model's tokenizer."""
+        return len(self._encoding.encode(text))
+
+
+class LangChainLLMClient:
+    """LLM client backed by any LangChain BaseChatModel.
+
+    Supports Ollama, llama.cpp, and any other LangChain-compatible provider.
+
+    Example::
+
+        from langchain_ollama import ChatOllama
+        client = LangChainLLMClient(ChatOllama(model="llama3.2", format="json"))
+    """
+
+    def __init__(self, model: Any) -> None:
+        self._model = model
+
+    def complete(self, messages: list[dict[str, str]]) -> str:
+        """Invoke the LangChain model and return the response string."""
+        from langchain_core.messages import HumanMessage, SystemMessage
+
+        lc_messages: list[SystemMessage | HumanMessage] = []
+        for m in messages:
+            role = m.get("role", "user")
+            content = m.get("content", "")
+            if role == "system":
+                lc_messages.append(SystemMessage(content=content))
+            else:
+                lc_messages.append(HumanMessage(content=content))
+
+        response = self._model.invoke(lc_messages)
+        return response.content
+
+    def complete_json(self, messages: list[dict[str, str]]) -> Any:
+        """Invoke the model and extract a JSON object from the response."""
+        content = self.complete(messages)
+        return _extract_json(content)
+
+    def count_tokens(self, text: str) -> int:
+        """Return an approximate token count for the given text."""
+        try:
+            return self._model.get_num_tokens(text)
+        except (NotImplementedError, AttributeError):
+            # Fallback: ~4 characters per token (reasonable for English)
+            return len(text) // 4
+
+
+def _extract_json(text: str) -> Any:
+    """Extract a JSON object from a string, tolerating surrounding prose.
+
+    Tries in order:
+    1. Direct json.loads (model output is clean JSON)
+    2. Extract from a markdown code block (```json ... ```)
+    3. Extract the first {...} or [...] block via regex
+    """
+    text = text.strip()
+
+    try:
+        return json.loads(text)
+    except json.JSONDecodeError:
+        pass
+
+    # Try markdown code block
+    code_block = re.search(r"```(?:json)?\s*([\s\S]+?)\s*```", text)
+    if code_block:
+        try:
+            return json.loads(code_block.group(1))
+        except json.JSONDecodeError:
+            pass
+
+    # Try first {...} or [...] block
+    brace_match = re.search(r"(\{[\s\S]*\}|\[[\s\S]*\])", text)
+    if brace_match:
+        try:
+            return json.loads(brace_match.group(1))
+        except json.JSONDecodeError:
+            pass
+
+    raise ValueError(f"Could not extract JSON from model response:\n{text}")

semaxis-0.16.0/src/semaxis/nli.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+import numpy as np
+
+
+class NLIModel:
+    """Wrapper around a sentence-transformers CrossEncoder for NLI scoring.
+
+    The model returns three logits per pair. We apply sigmoid to the entailment
+    logit and zero out pairs where the model does not predict entailment.
+
+    Args:
+        model_name: sentence-transformers CrossEncoder model name.
+        entailment_idx: Column index of the entailment class in the model output.
+            Defaults to 0, which matches ``cross-encoder/nli-deberta-v3-large``
+            (label order: [entailment, neutral, contradiction]).
+    """
+
+    def __init__(self, model_name: str, entailment_idx: int = 0) -> None:
+        from sentence_transformers import CrossEncoder
+
+        self.model_name = model_name
+        self.entailment_idx = entailment_idx
+        self._model = CrossEncoder(model_name)
+
+    def score(self, texts: list[str], hypotheses: list[str]) -> np.ndarray:
+        """Score (text, hypothesis) pairs.
+
+        Args:
+            texts: List of n texts.
+            hypotheses: List of n hypotheses, parallel to texts.
+
+        Returns:
+            np.ndarray of shape (n,) with values in [0, 1]. Pairs where the
+            model does not predict entailment are scored as 0.
+        """
+        pairs = list(zip(texts, hypotheses))
+        logits = self._model.predict(pairs)  # shape (n, 3) or (n,) depending on model
+        logits = np.array(logits)
+        if logits.ndim == 1:
+            # Binary model — return sigmoid scores directly
+            return 1.0 / (1.0 + np.exp(-logits))
+        entail_logits = logits[:, self.entailment_idx]
+        scores = 1.0 / (1.0 + np.exp(-entail_logits))
+        scores[logits.argmax(axis=1) != self.entailment_idx] = 0.0
+        return scores

semaxis-0.16.0/src/semaxis/prompts/collection_description.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+SYSTEM = """\
+You are an expert text analyst. Your task is to generate features that describe \
+properties of individual texts in a collection.
+
+Each feature must be defined by:
+- hypothesis: a declarative statement about a single text, suitable for NLI scoring \
+(e.g. "This text expresses satisfaction with the product." or \
+"This text mentions issues with build quality.")
+
+Requirements:
+- Each hypothesis must be a statement about a single text, starting with "This text"
+- Features must capture properties that are meaningfully present in some texts \
+but not all — avoid trivially universal or trivially absent properties
+- The hypothesis must be self-contained
+- Aim for diverse, non-redundant features
+
+Respond with JSON only:
+{"features": [{"hypothesis": "..."}, ...]}
+"""
+
+_USER_TEMPLATE = """\
+Here are sample texts from the collection:
+---
+{texts_block}
+
+Generate exactly {n} features that describe properties of individual texts \
+commonly found in this collection.{language_instruction}
+"""
+
+
+def build_user_message(texts: list[str], n: int, language: str | None = None) -> str:
+    texts_block = "\n---\n".join(texts) if texts else "(none)"
+    language_instruction = (
+        f"\nGenerate the hypothesis for each feature in {language}." if language else ""
+    )
+    return _USER_TEMPLATE.format(
+        texts_block=texts_block,
+        n=n,
+        language_instruction=language_instruction,
+    )

semaxis-0.16.0/src/semaxis/prompts/discriminative_features.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+SYSTEM = """\
+You are an expert text analyst. Your task is to generate features that distinguish \
+one group of texts from another.
+
+Each feature must be defined by:
+- hypothesis: a declarative statement about a single text, suitable for NLI scoring \
+(e.g. "This text expresses satisfaction with the product." or \
+"This text mentions issues with build quality.")
+
+Requirements:
+- Each hypothesis must be a statement about a single text, starting with "This text"
+- Hypotheses must capture properties that are more characteristic of the positive group \
+than the negative group — avoid properties shared equally by both groups
+- The hypothesis must be self-contained
+- Aim for diverse, non-redundant features
+
+Respond with JSON only:
+{"features": [{"hypothesis": "..."}, ...]}
+"""
+
+_USER_TEMPLATE = """\
+Positive texts (labeled "{pos_label}"):
+---
+{pos_block}
+
+Negative texts (labeled "{neg_label}"):
+---
+{neg_block}
+
+Generate exactly {n} features whose hypotheses are more likely to be true for \
+"{pos_label}" texts than "{neg_label}" texts.{language_instruction}
+"""
+
+
+def build_user_message(
+    pos_texts: list[str],
+    neg_texts: list[str],
+    pos_label: str,
+    neg_label: str,
+    n: int,
+    language: str | None = None,
+) -> str:
+    pos_block = "\n---\n".join(pos_texts) if pos_texts else "(none)"
+    neg_block = "\n---\n".join(neg_texts) if neg_texts else "(none)"
+    language_instruction = (
+        f"\nGenerate the hypothesis for each feature in {language}." if language else ""
+    )
+    return _USER_TEMPLATE.format(
+        pos_label=pos_label,
+        neg_label=neg_label,
+        pos_block=pos_block,
+        neg_block=neg_block,
+        n=n,
+        language_instruction=language_instruction,
+    )

semaxis-0.16.0/src/semaxis/sampling.py

@@ -0,0 +1,176 @@
+from __future__ import annotations
+
+import random
+from typing import Callable
+
+import numpy as np
+
+
+def sample_texts_within_budget(
+    texts: list[str],
+    token_budget: int,
+    tokenizer_fn: Callable[[str], int],
+    rng: random.Random | None = None,
+) -> list[str]:
+    """Return a random subset of texts that fits within the token budget.
+
+    Texts are shuffled, then added one by one until adding the next text
+    would exceed token_budget. The order of the returned list follows the
+    shuffled order.
+
+    Args:
+        texts: Source texts to sample from.
+        token_budget: Maximum total tokens allowed.
+        tokenizer_fn: Function that returns the token count for a single text.
+        rng: Optional Random instance for reproducibility.
+
+    Returns:
+        A list of texts whose total token count does not exceed token_budget.
+    """
+    if rng is None:
+        rng = random.Random()
+
+    indices = list(range(len(texts)))
+    rng.shuffle(indices)
+
+    selected: list[str] = []
+    total_tokens = 0
+
+    for idx in indices:
+        text = texts[idx]
+        count = tokenizer_fn(text)
+        if total_tokens + count > token_budget:
+            break
+        selected.append(text)
+        total_tokens += count
+
+    return selected
+
+
+def sample_texts_kmeans(
+    texts: list[str],
+    n: int,
+    embeddings: np.ndarray,
+    rng: random.Random | None = None,
+) -> list[str]:
+    """Return up to n texts by selecting the closest text to each K-Means centroid.
+
+    Args:
+        texts: Source texts.
+        n: Number of texts to select (capped at len(texts)).
+        embeddings: Pre-computed embedding matrix of shape (len(texts), dim).
+        rng: Optional Random instance for reproducibility.
+
+    Returns:
+        A list of up to n representative texts, one per cluster.
+    """
+    from sklearn.cluster import KMeans
+
+    n = min(n, len(texts))
+    if n == 0:
+        return []
+    if n == len(texts):
+        return list(texts)
+
+    seed = rng.randint(0, 2**31 - 1) if rng is not None else None
+    km = KMeans(n_clusters=n, random_state=seed, n_init="auto")
+    km.fit(embeddings)
+
+    selected_indices: list[int] = []
+    for center in km.cluster_centers_:
+        dists = np.linalg.norm(embeddings - center, axis=1)
+        # Exclude already-selected indices to avoid duplicates when clusters share a medoid
+        dists[selected_indices] = np.inf
+        selected_indices.append(int(np.argmin(dists)))
+
+    return [texts[i] for i in selected_indices]
+
+
+def sample_texts_votek(
+    texts: list[str],
+    n: int,
+    embeddings: np.ndarray,
+    k: int = 10,
+    rng: random.Random | None = None,
+) -> list[str]:
+    """Return up to n texts using the Vote-K algorithm (Su et al. 2022).
+
+    Vote-K balances representativeness (high vote count = many neighbours)
+    and diversity (selected texts' neighbours are suppressed from future picks).
+
+    Args:
+        texts: Source texts.
+        n: Number of texts to select (capped at len(texts)).
+        embeddings: Pre-computed embedding matrix of shape (len(texts), dim).
+        k: Number of neighbours to consider for voting and suppression.
+        rng: Optional Random instance (unused; kept for API symmetry).
+
+    Returns:
+        A list of up to n texts.
+    """
+    from sklearn.metrics.pairwise import cosine_similarity
+
+    n = min(n, len(texts))
+    if n == 0:
+        return []
+    if n == len(texts):
+        return list(texts)
+
+    k = min(k, len(texts) - 1)
+    sim = cosine_similarity(embeddings)  # (N, N)
+
+    # votes[i] = how many texts consider i among their top-k neighbours
+    votes = np.zeros(len(texts), dtype=float)
+    for i in range(len(texts)):
+        sim_row = sim[i].copy()
+        sim_row[i] = -np.inf  # exclude self
+        top_k = np.argpartition(sim_row, -k)[-k:]
+        votes[top_k] += 1.0
+
+    selected_indices: list[int] = []
+    remaining = np.ones(len(texts), dtype=bool)
+
+    while len(selected_indices) < n and remaining.any():
+        # Among remaining texts, pick the one with the most votes
+        masked_votes = np.where(remaining, votes, -np.inf)
+        best = int(np.argmax(masked_votes))
+        selected_indices.append(best)
+        remaining[best] = False
+
+        # Suppress votes of k nearest neighbours
+        sim_row = sim[best].copy()
+        sim_row[best] = -np.inf
+        top_k = np.argpartition(sim_row, -k)[-k:]
+        votes[top_k] = 0.0
+
+    return [texts[i] for i in selected_indices]
+
+
+def _estimate_n(
+    texts: list[str],
+    token_budget: int,
+    tokenizer_fn: Callable[[str], int],
+) -> int:
+    """Estimate how many texts fit in token_budget based on average token length."""
+    if not texts:
+        return 0
+    probe = texts[:min(20, len(texts))]
+    avg = sum(tokenizer_fn(t) for t in probe) / len(probe)
+    return max(1, int(token_budget / avg))
+
+
+def _trim_to_budget(
+    texts: list[str],
+    token_budget: int,
+    tokenizer_fn: Callable[[str], int],
+) -> list[str]:
+    """Trim a list of texts to fit within the token budget."""
+    result: list[str] = []
+    total = 0
+    for t in texts:
+        count = tokenizer_fn(t)
+        if total + count > token_budget:
+            break
+        result.append(t)
+        total += count
+    return result

semaxis-0.16.0/src/semaxis/supervised.py

@@ -0,0 +1,197 @@
+from __future__ import annotations
+
+import random
+from itertools import combinations
+from typing import Any, NamedTuple
+
+import numpy as np
+from sklearn.base import BaseEstimator, TransformerMixin
+from sklearn.preprocessing import LabelEncoder
+
+from .llm import BaseLLMClient, LLMClient
+from .nli import NLIModel
+from .sampling import (
+    _estimate_n,
+    _trim_to_budget,
+    sample_texts_kmeans,
+    sample_texts_votek,
+    sample_texts_within_budget,
+)
+from .prompts import discriminative_features as prompts
+
+_SAMPLE_METHODS = ("random", "kmeans", "votek")
+
+
+def _sample_group(
+    texts: list[str],
+    budget: int,
+    tokenizer_fn: Any,
+    method: str,
+    embedding_model: str,
+    rng: random.Random,
+) -> list[str]:
+    if method == "random":
+        return sample_texts_within_budget(texts, budget, tokenizer_fn, rng)
+
+    from sentence_transformers import SentenceTransformer
+    embeddings = SentenceTransformer(embedding_model).encode(
+        texts, show_progress_bar=False, convert_to_numpy=True
+    )
+    n = _estimate_n(texts, budget, tokenizer_fn)
+    if method == "kmeans":
+        sampled = sample_texts_kmeans(texts, n, embeddings, rng)
+    else:
+        sampled = sample_texts_votek(texts, n, embeddings, rng=rng)
+    return _trim_to_budget(sampled, budget, tokenizer_fn)
+
+_PROMPT_OVERHEAD = 500
+
+
+class FeatureMeta(NamedTuple):
+    positive: Any
+    negative: Any
+
+
+class SupervisedTransformer(BaseEstimator, TransformerMixin):
+    """Sklearn-compatible transformer that generates discriminative NLI features.
+
+    Fits by generating hypotheses (via LLM) that distinguish between classes,
+    then scores texts against those hypotheses using an NLI model.
+
+    Supports binary and multi-class labels (numeric or string).
+    For multi-class, use ``strategy="ovr"`` (one-vs-rest) or ``strategy="ovo"``
+    (one-vs-one). Binary classification ignores ``strategy``.
+
+    Fitted attributes:
+        classes_: Unique class labels in sorted order.
+        features_: Hypotheses as plain strings, parallel to ``feature_meta_``.
+        feature_meta_: FeatureMeta(positive, negative) for each hypothesis,
+            where negative is the original class label or ``"rest"`` for OvR.
+
+    Example::
+
+        from sklearn.pipeline import Pipeline
+        from sklearn.linear_model import LogisticRegression
+        from sklearn.model_selection import cross_val_score
+
+        pipe = Pipeline([
+            ("vect", SupervisedTransformer(llm="gpt-4o", nli_model="cross-encoder/nli-deberta-v3-large")),
+            ("clf", LogisticRegression()),
+        ])
+        cross_val_score(pipe, texts, labels, cv=5)
+    """
+
+    def __init__(
+        self,
+        llm: BaseLLMClient | str,
+        nli_model: str = "cross-encoder/nli-deberta-v3-large",
+        n_features: int = 20,
+        strategy: str = "ovr",
+        context_limit: int = 100_000,
+        language: str | None = None,
+        seed: int | None = None,
+        sample_method: str = "random",
+        embedding_model: str = "paraphrase-albert-small-v2",
+    ) -> None:
+        self.llm = llm
+        self.nli_model = nli_model
+        self.n_features = n_features
+        self.strategy = strategy
+        self.context_limit = context_limit
+        self.language = language
+        self.seed = seed
+        self.sample_method = sample_method
+        self.embedding_model = embedding_model
+
+    def fit(self, texts: list[str], y: Any) -> SupervisedTransformer:
+        """Generate discriminative hypotheses from training texts and labels.
+
+        Args:
+            texts: Training texts.
+            y: Class labels (numeric or string).
+
+        Returns:
+            self
+        """
+        if self.strategy not in ("ovr", "ovo"):
+            raise ValueError(f"strategy must be 'ovr' or 'ovo', got {self.strategy!r}")
+        if self.sample_method not in _SAMPLE_METHODS:
+            raise ValueError(f"sample_method must be one of {_SAMPLE_METHODS}, got {self.sample_method!r}")
+
+        _llm = LLMClient(self.llm) if isinstance(self.llm, str) else self.llm
+        _rng = random.Random(self.seed)
+
+        le = LabelEncoder()
+        y_enc: np.ndarray = le.fit_transform(y)
+        self.classes_ = le.classes_
+
+        n_classes = len(self.classes_)
+        budget = self.context_limit - _PROMPT_OVERHEAD
+
+        if n_classes == 2:
+            pairs: list[tuple[int, int | str]] = [(0, 1)]
+        elif self.strategy == "ovr":
+            pairs = [(i, "rest") for i in range(n_classes)]
+        else:
+            pairs = list(combinations(range(n_classes), 2))
+
+        self.features_: list[str] = []
+        self.feature_meta_: list[FeatureMeta] = []
+
+        for pos_idx, neg_idx in pairs:
+            pos_label = self.classes_[pos_idx]
+            pos_texts = [t for t, yi in zip(texts, y_enc) if yi == pos_idx]
+
+            if neg_idx == "rest":
+                neg_texts = [t for t, yi in zip(texts, y_enc) if yi != pos_idx]
+                neg_label: Any = "rest"
+            else:
+                neg_texts = [t for t, yi in zip(texts, y_enc) if yi == neg_idx]
+                neg_label = self.classes_[neg_idx]
+
+            pos_sampled = _sample_group(
+                pos_texts, budget // 2, _llm.count_tokens,
+                self.sample_method, self.embedding_model, _rng,
+            )
+            neg_sampled = _sample_group(
+                neg_texts, budget // 2, _llm.count_tokens,
+                self.sample_method, self.embedding_model, _rng,
+            )
+
+            messages = [
+                {"role": "system", "content": prompts.SYSTEM},
+                {"role": "user", "content": prompts.build_user_message(
+                    pos_texts=pos_sampled,
+                    neg_texts=neg_sampled,
+                    pos_label=str(pos_label),
+                    neg_label=str(neg_label),
+                    n=self.n_features,
+                    language=self.language,
+                )},
+            ]
+            result = _llm.complete_json(messages)
+            hypotheses = [item["hypothesis"] for item in result.get("features", [])]
+
+            self.features_.extend(hypotheses)
+            self.feature_meta_.extend(
+                FeatureMeta(positive=pos_label, negative=neg_label)
+                for _ in hypotheses
+            )
+
+        self._nli = NLIModel(self.nli_model)
+        return self
+
+    def transform(self, texts: list[str]) -> np.ndarray:
+        """Score texts against fitted hypotheses using NLI.
+
+        Args:
+            texts: Texts to score.
+
+        Returns:
+            np.ndarray of shape (n_texts, n_features) with entailment scores in [0, 1].
+        """
+        columns = [
+            self._nli.score(texts, [h] * len(texts))
+            for h in self.features_
+        ]
+        return np.column_stack(columns)

semaxis-0.16.0/src/semaxis/unsupervised.py

@@ -0,0 +1,139 @@
+from __future__ import annotations
+
+import random
+from typing import Any
+
+import numpy as np
+from sklearn.base import BaseEstimator, TransformerMixin
+
+from .llm import BaseLLMClient, LLMClient
+from .nli import NLIModel
+from .sampling import (
+    _estimate_n,
+    _trim_to_budget,
+    sample_texts_kmeans,
+    sample_texts_votek,
+    sample_texts_within_budget,
+)
+from .prompts import collection_description as prompts
+
+_SAMPLE_METHODS = ("random", "kmeans", "votek")
+
+
+def _sample_group(
+    texts: list[str],
+    budget: int,
+    tokenizer_fn: Any,
+    method: str,
+    embedding_model: str,
+    rng: random.Random,
+) -> list[str]:
+    if method == "random":
+        return sample_texts_within_budget(texts, budget, tokenizer_fn, rng)
+
+    from sentence_transformers import SentenceTransformer
+    embeddings = SentenceTransformer(embedding_model).encode(
+        texts, show_progress_bar=False, convert_to_numpy=True
+    )
+    n = _estimate_n(texts, budget, tokenizer_fn)
+    if method == "kmeans":
+        sampled = sample_texts_kmeans(texts, n, embeddings, rng)
+    else:
+        sampled = sample_texts_votek(texts, n, embeddings, rng=rng)
+    return _trim_to_budget(sampled, budget, tokenizer_fn)
+
+_PROMPT_OVERHEAD = 500
+
+
+class UnsupervisedTransformer(BaseEstimator, TransformerMixin):
+    """Sklearn-compatible transformer that generates NLI features from unlabeled texts.
+
+    Fits by generating hypotheses (via LLM) that characterize the text collection,
+    then scores texts against those hypotheses using an NLI model.
+
+    Fitted attributes:
+        features_: Hypotheses as plain strings.
+
+    Example::
+
+        from sklearn.pipeline import Pipeline
+        from sklearn.linear_model import LogisticRegression
+        from sklearn.model_selection import cross_val_score
+
+        pipe = Pipeline([
+            ("vect", UnsupervisedTransformer(llm="gpt-4o", nli_model="cross-encoder/nli-deberta-v3-large")),
+            ("clf", LogisticRegression()),
+        ])
+        cross_val_score(pipe, texts, labels, cv=5)
+    """
+
+    def __init__(
+        self,
+        llm: BaseLLMClient | str,
+        nli_model: str = "cross-encoder/nli-deberta-v3-large",
+        n_features: int = 20,
+        context_limit: int = 100_000,
+        language: str | None = None,
+        seed: int | None = None,
+        sample_method: str = "random",
+        embedding_model: str = "paraphrase-albert-small-v2",
+    ) -> None:
+        self.llm = llm
+        self.nli_model = nli_model
+        self.n_features = n_features
+        self.context_limit = context_limit
+        self.language = language
+        self.seed = seed
+        self.sample_method = sample_method
+        self.embedding_model = embedding_model
+
+    def fit(self, texts: list[str], y=None) -> UnsupervisedTransformer:
+        """Generate hypotheses from texts using LLM.
+
+        Args:
+            texts: Texts to characterize.
+            y: Ignored. Present for sklearn API compatibility.
+
+        Returns:
+            self
+        """
+        if self.sample_method not in _SAMPLE_METHODS:
+            raise ValueError(f"sample_method must be one of {_SAMPLE_METHODS}, got {self.sample_method!r}")
+
+        _llm = LLMClient(self.llm) if isinstance(self.llm, str) else self.llm
+        _rng = random.Random(self.seed)
+
+        budget = self.context_limit - _PROMPT_OVERHEAD
+        sampled = _sample_group(
+            texts, budget, _llm.count_tokens,
+            self.sample_method, self.embedding_model, _rng,
+        )
+
+        messages = [
+            {"role": "system", "content": prompts.SYSTEM},
+            {"role": "user", "content": prompts.build_user_message(
+                sampled, n=self.n_features, language=self.language
+            )},
+        ]
+        result = _llm.complete_json(messages)
+        self.features_: list[str] = [
+            item["hypothesis"] for item in result.get("features", [])
+        ]
+
+        self._nli = NLIModel(self.nli_model)
+        return self
+
+    def transform(self, texts: list[str]) -> np.ndarray:
+        """Score texts against fitted hypotheses using NLI.
+
+        Args:
+            texts: Texts to score.
+
+        Returns:
+            np.ndarray of shape (n_texts, n_features) with entailment scores in [0, 1].
+        """
+        columns = [
+            self._nli.score(texts, [h] * len(texts))
+            for h in self.features_
+        ]
+        return np.column_stack(columns)