pntx 0.1.0__py3-none-any.whl

pntx/__init__.py

@@ -0,0 +1,6 @@
+from __future__ import annotations
+
+from .core import PNTX
+from .types import ClassifyResult
+
+__all__ = ["PNTX", "ClassifyResult"]

pntx/backends/__init__.py

@@ -0,0 +1,5 @@
+from __future__ import annotations
+
+from .base import Backend, BatchBackend, BatchScoringBackend, ScoringBackend
+
+__all__ = ["BatchBackend", "BatchScoringBackend", "Backend", "ScoringBackend"]

pntx/backends/anthropic.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+import asyncio
+from typing import Any
+
+try:
+    import anthropic
+except ImportError as e:
+    raise ImportError(
+        "AnthropicBackend requires the 'anthropic' extra. "
+        "Install it with: pip install 'pntx[anthropic]'"
+    ) from e
+
+_DEFAULT_CONCURRENCY = 5
+
+
+class AnthropicBackend:
+    """Backend using the Anthropic Messages API.
+
+    Implements only ``Backend`` (plain text completion): the Messages API
+    doesn't expose token log-probabilities, so classification for this
+    backend goes through PNTX's parse-based fallback rather than
+    ``score_choices``. ``complete_batch`` implements ``BatchBackend`` by
+    running requests concurrently via asyncio, bounded by ``concurrency``.
+    """
+
+    def __init__(
+        self,
+        model: str,
+        *,
+        concurrency: int = _DEFAULT_CONCURRENCY,
+        client: anthropic.Anthropic | None = None,
+        async_client: anthropic.AsyncAnthropic | None = None,
+        **client_kwargs: Any,
+    ) -> None:
+        if concurrency < 1:
+            raise ValueError(f"concurrency must be >= 1, got {concurrency}")
+        self.model = model
+        self.concurrency = concurrency
+        self._client = client if client is not None else anthropic.Anthropic(**client_kwargs)
+        self._async_client = (
+            async_client if async_client is not None else anthropic.AsyncAnthropic(**client_kwargs)
+        )
+
+    def complete(
+        self,
+        prompt: str,
+        *,
+        temperature: float = 1.0,
+        max_tokens: int = 512,
+        stop: list[str] | None = None,
+    ) -> str:
+        message = self._client.messages.create(
+            model=self.model,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            stop_sequences=stop or [],
+            messages=[{"role": "user", "content": prompt}],
+        )
+        return _extract_text(message)
+
+    def complete_batch(
+        self,
+        prompts: list[str],
+        *,
+        temperature: float = 1.0,
+        max_tokens: int = 512,
+        stop: list[str] | None = None,
+    ) -> list[str]:
+        """Complete every prompt concurrently, bounded by ``self.concurrency``."""
+        if not prompts:
+            return []
+        return asyncio.run(
+            self._complete_batch_async(
+                prompts, temperature=temperature, max_tokens=max_tokens, stop=stop
+            )
+        )
+
+    async def _complete_batch_async(
+        self,
+        prompts: list[str],
+        *,
+        temperature: float,
+        max_tokens: int,
+        stop: list[str] | None,
+    ) -> list[str]:
+        semaphore = asyncio.Semaphore(self.concurrency)
+
+        async def complete_one(prompt: str) -> str:
+            async with semaphore:
+                message = await self._async_client.messages.create(
+                    model=self.model,
+                    max_tokens=max_tokens,
+                    temperature=temperature,
+                    stop_sequences=stop or [],
+                    messages=[{"role": "user", "content": prompt}],
+                )
+                return _extract_text(message)
+
+        return list(await asyncio.gather(*(complete_one(prompt) for prompt in prompts)))
+
+
+def _extract_text(message: anthropic.types.Message) -> str:
+    return "".join(
+        block.text for block in message.content if isinstance(block, anthropic.types.TextBlock)
+    )

pntx/backends/base.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+
+@runtime_checkable
+class Backend(Protocol):
+    """A text-completion backend.
+
+    Every backend (llama.cpp, Anthropic, ...) must implement at least this.
+    """
+
+    def complete(
+        self,
+        prompt: str,
+        *,
+        temperature: float = 1.0,
+        max_tokens: int = 512,
+        stop: list[str] | None = None,
+    ) -> str: ...
+
+
+@runtime_checkable
+class ScoringBackend(Backend, Protocol):
+    """A backend that can score candidate continuations directly.
+
+    This is the primary path for classification: instead of generating text
+    and parsing a label out of it, the backend returns a log-likelihood for
+    each candidate choice appended to ``prompt``.
+    """
+
+    def score_choices(self, prompt: str, choices: list[str]) -> list[float]:
+        """Return the log-likelihood of each choice in ``choices`` as a
+        continuation of ``prompt``."""
+        ...
+
+
+@runtime_checkable
+class BatchScoringBackend(ScoringBackend, Protocol):
+    """A ``ScoringBackend`` that can batch-score many queries sharing one prefix.
+
+    Used by ``classify_batch`` to warm a common prefix (e.g. the few-shot
+    exemplar block) once and reuse it across every query, instead of
+    re-evaluating it per item. Backends without this (e.g. remote API
+    backends) are scored one item at a time via ``score_choices``.
+    """
+
+    def score_choices_batch(
+        self, prefix: str, queries: list[str], choices: list[str]
+    ) -> list[list[float]]:
+        """For each query in ``queries``, score ``choices`` as a continuation
+        of ``prefix + query``. Returns one score list per query, in order."""
+        ...
+
+
+@runtime_checkable
+class BatchBackend(Backend, Protocol):
+    """A ``Backend`` that can complete many prompts concurrently.
+
+    Used by ``classify_batch``'s parse-based path (for backends that don't
+    implement ``ScoringBackend``, e.g. remote chat APIs): instead of
+    completing each prompt one at a time, ``complete_batch`` can run them
+    concurrently (e.g. via asyncio) and return results in the same order.
+    """
+
+    def complete_batch(
+        self,
+        prompts: list[str],
+        *,
+        temperature: float = 1.0,
+        max_tokens: int = 512,
+        stop: list[str] | None = None,
+    ) -> list[str]:
+        """Complete every prompt in ``prompts``, returning results in order."""
+        ...

pntx/backends/llama.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+from typing import Any
+
+try:
+    import llama_cpp
+except ImportError as e:
+    raise ImportError(
+        "LlamaCppBackend requires the 'llama' extra. "
+        "Install it with: pip install 'pntx[llama]'"
+    ) from e
+
+
+class LlamaCppBackend:
+    """In-process backend backed by llama.cpp (via ``llama-cpp-python``).
+
+    Implements ``Backend``, ``ScoringBackend`` and ``BatchScoringBackend``.
+    Scoring methods reuse the KV cache of whatever prefix they share: rather
+    than re-evaluating a shared prefix for every downstream token span, they
+    eval it once and then, for each span, rewind ``n_tokens`` back to the
+    prefix boundary before eval'ing that span's tokens (``Llama.eval()``
+    trims the KV cache down to the current ``n_tokens`` before appending, so
+    this discards only the previous span, not the shared prefix).
+    """
+
+    def __init__(self, model_path: str, **kwargs: Any) -> None:
+        # Scoring needs per-token logits, which llama.cpp only keeps around
+        # when logits_all=True.
+        kwargs["logits_all"] = True
+        self._llm = llama_cpp.Llama(model_path=model_path, **kwargs)
+
+    def complete(
+        self,
+        prompt: str,
+        *,
+        temperature: float = 1.0,
+        max_tokens: int = 512,
+        stop: list[str] | None = None,
+    ) -> str:
+        result = self._llm.create_completion(
+            prompt,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            stop=stop or [],
+        )
+        if not isinstance(result, dict):
+            raise TypeError(f"expected a non-streaming completion response, got {type(result)}")
+        return result["choices"][0]["text"]
+
+    def score_choices(self, prompt: str, choices: list[str]) -> list[float]:
+        """Return the summed log-likelihood of each choice as a continuation of ``prompt``."""
+        if not choices:
+            return []
+        prompt_tokens = self._tokenize(prompt, add_bos=True)
+        self._reset()
+        self._llm.eval(prompt_tokens)
+        return self._score_choices_at(self._llm.n_tokens, choices)
+
+    def score_choices_batch(
+        self, prefix: str, queries: list[str], choices: list[str]
+    ) -> list[list[float]]:
+        """For each query, score ``choices`` as a continuation of ``prefix + query``.
+
+        ``prefix`` (e.g. the few-shot exemplar block) is eval'd once and its
+        KV cache reused across every query, instead of re-evaluating it per
+        query as a naive per-item ``score_choices`` loop would.
+        """
+        if not queries:
+            return []
+        prefix_tokens = self._tokenize(prefix, add_bos=True)
+        self._reset()
+        self._llm.eval(prefix_tokens)
+        base = self._llm.n_tokens
+
+        results: list[list[float]] = []
+        for query in queries:
+            query_tokens = self._tokenize(query, add_bos=False)
+            self._llm.n_tokens = base
+            self._llm.eval(query_tokens)
+            results.append(self._score_choices_at(self._llm.n_tokens, choices))
+        self._llm.n_tokens = base
+        return results
+
+    def _tokenize(self, text: str, *, add_bos: bool) -> list[int]:
+        return self._llm.tokenize(text.encode("utf-8"), add_bos=add_bos)
+
+    def _reset(self) -> None:
+        self._llm.reset()  # type: ignore[no-untyped-call]  # llama_cpp.Llama.reset() has no annotations
+
+    def _score_choices_at(self, base_n_tokens: int, choices: list[str]) -> list[float]:
+        """Score each choice freshly on top of the KV cache at ``base_n_tokens``,
+        rewinding between choices so only the shared prefix up to
+        ``base_n_tokens`` is reused."""
+        scores = [
+            self._score_tokens(base_n_tokens, self._tokenize(choice, add_bos=False))
+            for choice in choices
+        ]
+        self._llm.n_tokens = base_n_tokens
+        return scores
+
+    def _score_tokens(self, base_n_tokens: int, tokens: list[int]) -> float:
+        if not tokens:
+            return 0.0
+        self._llm.n_tokens = base_n_tokens
+        self._llm.eval(tokens)
+        logits = self._llm.scores[base_n_tokens - 1 : base_n_tokens + len(tokens) - 1]
+        logprobs = self._llm.logits_to_logprobs(logits)
+        return float(sum(logprobs[j, tok] for j, tok in enumerate(tokens)))

pntx/core.py

@@ -0,0 +1,227 @@
+from __future__ import annotations
+
+import math
+import warnings
+from collections.abc import Iterable
+from importlib import import_module
+from typing import Any
+
+from . import generate as generate_module
+from . import prompts
+from .backends.base import Backend, BatchBackend, BatchScoringBackend, ScoringBackend
+from .selection import RandomSelector, Selector
+from .types import ClassifyResult, Label, Pair
+
+_BACKEND_REGISTRY: dict[str, tuple[str, str]] = {
+    "llama": ("pntx.backends.llama", "LlamaCppBackend"),
+    "anthropic": ("pntx.backends.anthropic", "AnthropicBackend"),
+}
+
+
+def _resolve_backend(backend: Backend | str, backend_kwargs: dict[str, Any]) -> Backend:
+    if not isinstance(backend, str):
+        if backend_kwargs:
+            raise TypeError(
+                "backend_kwargs are only used when backend is given as a string "
+                "(e.g. PNTX(backend='llama', model_path=...)); construct the "
+                "Backend instance directly instead"
+            )
+        return backend
+    try:
+        module_name, class_name = _BACKEND_REGISTRY[backend]
+    except KeyError:
+        raise ValueError(
+            f"Unknown backend {backend!r}; expected a Backend instance or one "
+            f"of {sorted(_BACKEND_REGISTRY)}"
+        ) from None
+    try:
+        module = import_module(module_name)
+    except ImportError as e:
+        raise ImportError(
+            f"Backend {backend!r} requires the optional '{backend}' extra. "
+            f"Install it with: pip install 'pntx[{backend}]'"
+        ) from e
+    return getattr(module, class_name)(**backend_kwargs)  # type: ignore[no-any-return]
+
+
+class PNTX:
+    """Generates and classifies text from user-defined (positive, negative) pairs.
+
+    The pairs' meaning is entirely up to the caller (sentiment, formality,
+    policy compliance, ...); this class never interprets it, it only uses the
+    pairs as few-shot/scoring material.
+    """
+
+    def __init__(
+        self,
+        backend: Backend | str,
+        *,
+        selector: Selector | None = None,
+        max_exemplars: int | None = None,
+        **backend_kwargs: Any,
+    ) -> None:
+        """Create a PNTX model.
+
+        ``backend`` is either a ready-made ``Backend`` instance, or the name
+        of a built-in backend (e.g. ``"llama"``) to construct lazily; any
+        ``backend_kwargs`` are then forwarded to that backend's constructor
+        (e.g. ``PNTX(backend="llama", model_path="model.gguf")``).
+
+        ``max_exemplars`` caps how many fitted pairs ``selector`` is asked to
+        pick for a single prompt; ``None`` means "as many as are fitted".
+        """
+        self.backend = _resolve_backend(backend, backend_kwargs)
+        self.selector: Selector = selector if selector is not None else RandomSelector()
+        self.max_exemplars = max_exemplars
+        self._pairs: list[Pair] = []
+
+    @property
+    def pairs(self) -> list[Pair]:
+        return list(self._pairs)
+
+    def fit(self, pairs: Iterable[Pair]) -> PNTX:
+        """Store the (positive, negative) pairs used as generation/classification material.
+
+        This only stores and validates ``pairs``; no training happens here.
+        """
+        pairs = list(pairs)
+        if not pairs:
+            raise ValueError("pairs must be non-empty")
+        self._pairs = pairs
+        return self
+
+    def generate(
+        self,
+        n: int,
+        side: Label,
+        *,
+        temperature: float = 1.0,
+        dedup: bool = True,
+        verify: bool = True,
+        min_confidence: float = 0.8,
+        max_attempts: int | None = None,
+    ) -> list[str]:
+        """Generate up to ``n`` new ``side`` texts from the fitted pairs.
+
+        If ``verify`` keeps rejecting candidates (wrong self-classified
+        label, or confidence below ``min_confidence``), fewer than ``n``
+        texts may come back; this warns rather than looping forever, capped
+        at ``max_attempts`` retries (default: ``max(n * 3, 3)``).
+        """
+        self._check_fitted()
+        if n < 0:
+            raise ValueError(f"n must be >= 0, got {n}")
+
+        texts = generate_module.run_generation_loop(
+            backend=self.backend,
+            pairs=self._pairs,
+            selector=self.selector,
+            exemplar_count=self._exemplar_count(),
+            classify=self.classify,
+            n=n,
+            side=side,
+            temperature=temperature,
+            dedup=dedup,
+            verify=verify,
+            min_confidence=min_confidence,
+            max_attempts=max_attempts,
+        )
+        if len(texts) < n:
+            warnings.warn(
+                f"requested n={n} {side!r} texts but only generated {len(texts)}",
+                UserWarning,
+                stacklevel=2,
+            )
+        return texts
+
+    def classify(self, text: str) -> ClassifyResult:
+        """Classify ``text`` as ``"positive"`` or ``"negative"``.
+
+        If the backend implements ``ScoringBackend``, this compares the
+        log-likelihood of each label as a continuation of the prompt
+        (``confidence`` is a calibrated softmax over those log-likelihoods).
+        Otherwise it falls back to asking the backend to write the label and
+        parsing it out of the response text; see
+        ``prompts.parse_classify_label`` for that path's confidence
+        convention (not a calibrated probability).
+        """
+        self._check_fitted()
+        exemplars = self.selector.select(self._pairs, self._exemplar_count(), query=text)
+        if isinstance(self.backend, ScoringBackend):
+            prompt = prompts.build_classify_prompt(exemplars, text)
+            scores = self.backend.score_choices(prompt, prompts.classify_choice_texts())
+            return _result_from_scores(scores)
+        prompt = prompts.build_classify_prompt(exemplars, text)
+        raw = self.backend.complete(
+            prompt, temperature=0.0, max_tokens=prompts.CLASSIFY_COMPLETION_MAX_TOKENS
+        )
+        return _result_from_completion(raw)
+
+    def classify_batch(self, texts: list[str]) -> list[ClassifyResult]:
+        """Classify each text in ``texts``; see ``classify()`` for the two
+        scoring strategies this dispatches between."""
+        self._check_fitted()
+        if not texts:
+            return []
+        exemplars = self.selector.select(self._pairs, self._exemplar_count(), query=None)
+
+        if isinstance(self.backend, ScoringBackend):
+            choices = prompts.classify_choice_texts()
+            prefix = prompts.build_exemplar_prefix(exemplars)
+            if isinstance(self.backend, BatchScoringBackend):
+                queries = [prompts.build_query_suffix(text) for text in texts]
+                all_scores = self.backend.score_choices_batch(prefix, queries, choices)
+            else:
+                # No batch-optimized path for this backend; score one prompt
+                # at a time. (LlamaCppBackend implements BatchScoringBackend
+                # and takes the branch above; a future non-batching
+                # ScoringBackend falls back to this.)
+                all_scores = [
+                    self.backend.score_choices(prefix + prompts.build_query_suffix(text), choices)
+                    for text in texts
+                ]
+            return [_result_from_scores(scores) for scores in all_scores]
+
+        completion_prompts = [prompts.build_classify_prompt(exemplars, text) for text in texts]
+        if isinstance(self.backend, BatchBackend):
+            raw_completions = self.backend.complete_batch(
+                completion_prompts,
+                temperature=0.0,
+                max_tokens=prompts.CLASSIFY_COMPLETION_MAX_TOKENS,
+            )
+        else:
+            # No concurrent-batch path for this backend; complete one prompt
+            # at a time. (AnthropicBackend implements BatchBackend and takes
+            # the branch above.)
+            raw_completions = [
+                self.backend.complete(
+                    prompt, temperature=0.0, max_tokens=prompts.CLASSIFY_COMPLETION_MAX_TOKENS
+                )
+                for prompt in completion_prompts
+            ]
+        return [_result_from_completion(raw) for raw in raw_completions]
+
+    def _exemplar_count(self) -> int:
+        return self.max_exemplars if self.max_exemplars is not None else len(self._pairs)
+
+    def _check_fitted(self) -> None:
+        if not self._pairs:
+            raise RuntimeError("PNTX.fit(pairs) must be called before this method")
+
+
+def _result_from_scores(scores: list[float]) -> ClassifyResult:
+    probs = _softmax(scores)
+    best = max(range(len(scores)), key=scores.__getitem__)
+    return ClassifyResult(label=prompts.CLASSIFY_LABELS[best], confidence=probs[best])
+
+
+def _result_from_completion(raw: str) -> ClassifyResult:
+    label, confidence = prompts.parse_classify_label(raw)
+    return ClassifyResult(label=label, confidence=confidence)
+
+
+def _softmax(scores: list[float]) -> list[float]:
+    top = max(scores)
+    exps = [math.exp(score - top) for score in scores]
+    total = sum(exps)
+    return [exp / total for exp in exps]

pntx/dedup.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+NGRAM_SIZE = 3
+SIMILARITY_THRESHOLD = 0.75
+
+
+def char_ngrams(text: str, n: int = NGRAM_SIZE) -> set[str]:
+    """Character n-grams of ``text``.
+
+    Character n-grams (rather than word n-grams) work for both Japanese,
+    which has no whitespace word boundaries, and English, without needing a
+    tokenizer or language detection.
+    """
+    if len(text) <= n:
+        return {text}
+    return {text[i : i + n] for i in range(len(text) - n + 1)}
+
+
+def similarity(a: str, b: str, n: int = NGRAM_SIZE) -> float:
+    """Jaccard similarity between the character n-grams of ``a`` and ``b``."""
+    if a == b:
+        return 1.0
+    grams_a = char_ngrams(a, n)
+    grams_b = char_ngrams(b, n)
+    union = grams_a | grams_b
+    if not union:
+        return 0.0
+    return len(grams_a & grams_b) / len(union)
+
+
+def is_near_duplicate(
+    text: str,
+    others: Iterable[str],
+    *,
+    threshold: float = SIMILARITY_THRESHOLD,
+) -> bool:
+    """Whether ``text`` is an approximate duplicate of any text in ``others``."""
+    return any(similarity(text, other) >= threshold for other in others)

pntx/embeddings.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import math
+from collections.abc import Callable
+
+try:
+    from sentence_transformers import SentenceTransformer
+except ImportError as e:
+    raise ImportError(
+        "pntx.embeddings requires the 'embeddings' extra. "
+        "Install it with: pip install 'pntx[embeddings]'"
+    ) from e
+
+DEFAULT_MODEL = "paraphrase-albert-small-v2"
+
+_model_cache: dict[str, SentenceTransformer] = {}
+
+
+def _load_model(model_name: str) -> SentenceTransformer:
+    if model_name not in _model_cache:
+        _model_cache[model_name] = SentenceTransformer(model_name)
+    return _model_cache[model_name]
+
+
+def embed(texts: list[str], model_name: str = DEFAULT_MODEL) -> list[list[float]]:
+    """Encode ``texts`` into embedding vectors using a sentence-transformers model."""
+    model = _load_model(model_name)
+    return [vector.tolist() for vector in model.encode(texts)]
+
+
+def cosine_similarity_fn(model_name: str = DEFAULT_MODEL) -> Callable[[str, str], float]:
+    """Build a ``(text_a, text_b) -> float`` similarity function backed by
+    sentence-transformers embeddings, for use as ``NearestSelector`` or
+    ``DiversitySelector``'s ``similarity_fn``.
+
+    Per-text embeddings are cached within the returned function, since
+    both selectors call it repeatedly with the same texts.
+    """
+    model = _load_model(model_name)
+    vector_cache: dict[str, list[float]] = {}
+
+    def embed_one(text: str) -> list[float]:
+        if text not in vector_cache:
+            vector_cache[text] = model.encode([text])[0].tolist()
+        return vector_cache[text]
+
+    def similarity(a: str, b: str) -> float:
+        return _cosine_similarity(embed_one(a), embed_one(b))
+
+    return similarity
+
+
+def _cosine_similarity(a: list[float], b: list[float]) -> float:
+    dot = sum(x * y for x, y in zip(a, b, strict=True))
+    norm_a = math.sqrt(sum(x * x for x in a))
+    norm_b = math.sqrt(sum(y * y for y in b))
+    if norm_a == 0.0 or norm_b == 0.0:
+        return 0.0
+    return dot / (norm_a * norm_b)

pntx/generate.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+
+from . import dedup as dedup_module
+from . import prompts
+from .backends.base import Backend
+from .selection import Selector
+from .types import ClassifyResult, Label, Pair
+
+_DEFAULT_MAX_ATTEMPTS_MULTIPLIER = 3
+_MIN_DEFAULT_MAX_ATTEMPTS = 3
+_TOKENS_PER_TEXT_ESTIMATE = 64
+
+
+def run_generation_loop(
+    *,
+    backend: Backend,
+    pairs: list[Pair],
+    selector: Selector,
+    exemplar_count: int,
+    classify: Callable[[str], ClassifyResult],
+    n: int,
+    side: Label,
+    temperature: float,
+    dedup: bool,
+    verify: bool,
+    min_confidence: float,
+    max_attempts: int | None,
+) -> list[str]:
+    """Generate up to ``n`` new ``side`` texts, retrying failed attempts.
+
+    On each attempt, asks the backend for however many texts are still
+    needed, then accepts candidates that pass dedup (against both the seed
+    pairs and texts already accepted) and verify (self-classifies as
+    ``side`` with at least ``min_confidence``). Stops after ``max_attempts``
+    attempts even if ``n`` texts were never reached; the caller is
+    responsible for warning about a shortfall.
+    """
+    if n <= 0:
+        return []
+
+    attempts = (
+        max_attempts
+        if max_attempts is not None
+        else max(n * _DEFAULT_MAX_ATTEMPTS_MULTIPLIER, _MIN_DEFAULT_MAX_ATTEMPTS)
+    )
+    seed_texts = [text for pair in pairs for text in pair] if dedup else []
+    accepted: list[str] = []
+
+    for _attempt in range(attempts):
+        remaining = n - len(accepted)
+        if remaining <= 0:
+            break
+
+        exemplars = selector.select(pairs, exemplar_count, query=None)
+        prompt = prompts.build_generate_prompt(exemplars, side, remaining)
+        raw = backend.complete(
+            prompt,
+            temperature=temperature,
+            max_tokens=remaining * _TOKENS_PER_TEXT_ESTIMATE,
+        )
+
+        for candidate in prompts.parse_generated_texts(raw):
+            if len(accepted) >= n:
+                break
+            if dedup and dedup_module.is_near_duplicate(candidate, seed_texts + accepted):
+                continue
+            if verify:
+                result = classify(candidate)
+                if result.label != side or result.confidence < min_confidence:
+                    continue
+            accepted.append(candidate)
+
+    return accepted

pntx/prompts.py

@@ -0,0 +1,124 @@
+from __future__ import annotations
+
+import re
+
+from .types import NEGATIVE, POSITIVE, Label, Pair
+
+CLASSIFY_LABELS: list[Label] = [POSITIVE, NEGATIVE]
+"""Fixed label order used everywhere a classify score list is produced."""
+
+CLASSIFY_LABEL_CHOICES: dict[Label, str] = {
+    POSITIVE: " positive",
+    NEGATIVE: " negative",
+}
+"""Text appended after the prompt to score each label (leading space matters
+for most tokenizers: it keeps the label as its own token(s) rather than
+merging with the preceding colon)."""
+
+_EXEMPLAR_TEMPLATE = "Text: {text}\nLabel: {label}\n\n"
+_QUERY_TEMPLATE = "Text: {text}\nLabel:"
+
+
+def build_exemplar_prefix(pairs: list[Pair]) -> str:
+    """Render the few-shot block for ``pairs``. This is the part of the
+    classification prompt shared across every query, so it should come
+    first (see ``build_query_suffix``) to keep it a stable KV-cache prefix.
+    """
+    return "".join(
+        _EXEMPLAR_TEMPLATE.format(text=pos, label=POSITIVE)
+        + _EXEMPLAR_TEMPLATE.format(text=neg, label=NEGATIVE)
+        for pos, neg in pairs
+    )
+
+
+def build_query_suffix(text: str) -> str:
+    """Render the query-specific tail of the classification prompt."""
+    return _QUERY_TEMPLATE.format(text=text)
+
+
+def build_classify_prompt(pairs: list[Pair], query: str) -> str:
+    """Render the full classification prompt for a single ``query``."""
+    return build_exemplar_prefix(pairs) + build_query_suffix(query)
+
+
+def classify_choice_texts() -> list[str]:
+    """``CLASSIFY_LABEL_CHOICES`` values in ``CLASSIFY_LABELS`` order."""
+    return [CLASSIFY_LABEL_CHOICES[label] for label in CLASSIFY_LABELS]
+
+
+CLASSIFY_COMPLETION_MAX_TOKENS = 8
+"""max_tokens for the parse-based classify path: just enough for a bare label."""
+
+
+def parse_classify_label(raw: str) -> tuple[Label, float]:
+    """Parse a freeform completion of ``build_classify_prompt`` into a label.
+
+    Used for backends without ``score_choices`` (e.g. ``AnthropicBackend``):
+    rather than comparing log-likelihoods, we ask the model to name the
+    label and look for it in the response text.
+
+    Confidence here is *not* a calibrated probability, just a fixed
+    convention value: ``1.0`` when exactly one label word was found, ``0.5``
+    when the response was ambiguous (both or neither label word present) and
+    we fell back to a default guess.
+    """
+    lowered = raw.lower()
+    positive_at = lowered.find(POSITIVE)
+    negative_at = lowered.find(NEGATIVE)
+    found_positive = positive_at != -1
+    found_negative = negative_at != -1
+
+    if found_positive and not found_negative:
+        return POSITIVE, 1.0
+    if found_negative and not found_positive:
+        return NEGATIVE, 1.0
+    if found_positive and found_negative:
+        return (POSITIVE if positive_at < negative_at else NEGATIVE), 0.5
+    return POSITIVE, 0.5  # neither label word found; arbitrary fallback guess
+
+
+_GENERATE_PREAMBLE = (
+    'Below are example pairs contrasting "positive" and "negative" texts; '
+    "their exact meaning is defined only by these examples.\n\n"
+)
+_GENERATE_EXEMPLAR_LINE = "- {label}: {text}\n"
+_GENERATE_INSTRUCTION = (
+    "\nWrite {n} new, diverse {side} texts. Do not copy or closely paraphrase "
+    "the examples above. Output one text per line, numbered starting at 1, "
+    "with no other commentary.\n\n1. "
+)
+
+_NUMBERED_LINE = re.compile(r"^\s*\d+[.)]\s*(.+)$")
+
+
+def build_generate_prompt(pairs: list[Pair], side: Label, n: int) -> str:
+    """Render the generation prompt asking for ``n`` new ``side`` texts.
+
+    Ends primed with ``"1. "`` so a plain text-completion backend continues
+    directly into the first generated item.
+    """
+    exemplars = "".join(
+        _GENERATE_EXEMPLAR_LINE.format(label=POSITIVE, text=pos)
+        + _GENERATE_EXEMPLAR_LINE.format(label=NEGATIVE, text=neg)
+        for pos, neg in pairs
+    )
+    return _GENERATE_PREAMBLE + exemplars + _GENERATE_INSTRUCTION.format(n=n, side=side)
+
+
+def parse_generated_texts(raw: str) -> list[str]:
+    """Parse a numbered-list completion (continuing after a primed ``"1. "``)
+    into a flat list of generated texts.
+
+    This is a best-effort heuristic, not a strict parser: lines that aren't
+    numbered (e.g. a wrapped continuation of the previous item) are kept
+    as-is rather than dropped, since backends aren't guaranteed to always
+    produce clean single-line items.
+    """
+    lines = [line for line in raw.strip().splitlines() if line.strip()]
+    if not lines:
+        return []
+    texts = [lines[0].strip()]
+    for line in lines[1:]:
+        match = _NUMBERED_LINE.match(line)
+        texts.append(match.group(1).strip() if match else line.strip())
+    return [text for text in texts if text]

pntx/selection.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+import random
+from collections.abc import Callable
+from typing import Protocol, runtime_checkable
+
+from . import dedup
+from .types import Pair
+
+SimilarityFn = Callable[[str, str], float]
+
+
+@runtime_checkable
+class Selector(Protocol):
+    """Chooses which fitted pairs to embed in a prompt.
+
+    ``query`` is the text being classified, for selectors that pick pairs
+    relevant to it (e.g. ``NearestSelector``); selectors that don't use it
+    (e.g. ``RandomSelector``) simply ignore the argument.
+    """
+
+    def select(self, pairs: list[Pair], k: int, query: str | None = None) -> list[Pair]: ...
+
+
+class RandomSelector:
+    """Selects a uniform random subset of ``pairs``, ignoring ``query``."""
+
+    def __init__(self, seed: int | None = None) -> None:
+        self._rng = random.Random(seed)
+
+    def select(self, pairs: list[Pair], k: int, query: str | None = None) -> list[Pair]:
+        if k >= len(pairs):
+            return list(pairs)
+        return self._rng.sample(pairs, k)
+
+
+class NearestSelector:
+    """Selects the ``k`` pairs whose positive or negative text is most
+    similar to ``query``.
+
+    ``similarity_fn`` defaults to ``dedup.similarity`` (dependency-free
+    character n-grams); pass e.g. ``pntx.embeddings.cosine_similarity_fn()``
+    for semantic similarity instead.
+
+    ``query=None`` (e.g. when called from a batch context that has no single
+    query to be "nearest" to) has nothing to rank by, so this falls back to
+    the first ``k`` pairs in ``pairs`` order; use ``RandomSelector`` or
+    ``DiversitySelector`` if that's not the fallback you want.
+    """
+
+    def __init__(self, similarity_fn: SimilarityFn = dedup.similarity) -> None:
+        self.similarity_fn = similarity_fn
+
+    def select(self, pairs: list[Pair], k: int, query: str | None = None) -> list[Pair]:
+        if k >= len(pairs):
+            return list(pairs)
+        if k <= 0:
+            return []
+        if query is None:
+            return list(pairs[:k])
+
+        def pair_similarity(pair: Pair) -> float:
+            pos, neg = pair
+            return max(self.similarity_fn(query, pos), self.similarity_fn(query, neg))
+
+        ranked = sorted(range(len(pairs)), key=lambda i: pair_similarity(pairs[i]), reverse=True)
+        return [pairs[i] for i in sorted(ranked[:k])]
+
+
+class DiversitySelector:
+    """Greedily selects ``k`` pairs that are maximally different from each
+    other, ignoring ``query``.
+
+    Starts from ``pairs[0]`` and repeatedly adds whichever remaining pair
+    has the lowest similarity to its most-similar already-selected pair
+    (a farthest-point / greedy diversity heuristic), representing each pair
+    as its concatenated positive and negative text.
+
+    ``similarity_fn`` defaults to ``dedup.similarity`` (dependency-free
+    character n-grams); pass e.g. ``pntx.embeddings.cosine_similarity_fn()``
+    for semantic similarity instead.
+    """
+
+    def __init__(self, similarity_fn: SimilarityFn = dedup.similarity) -> None:
+        self.similarity_fn = similarity_fn
+
+    def select(self, pairs: list[Pair], k: int, query: str | None = None) -> list[Pair]:
+        if k >= len(pairs):
+            return list(pairs)
+        if k <= 0:
+            return []
+
+        representations = [f"{pos} {neg}" for pos, neg in pairs]
+        selected = [0]
+        while len(selected) < k:
+            best_index = -1
+            best_distance = -1.0
+            for i in range(len(pairs)):
+                if i in selected:
+                    continue
+                similarity_to_selected = max(
+                    self.similarity_fn(representations[i], representations[j]) for j in selected
+                )
+                distance = 1.0 - similarity_to_selected
+                if distance > best_distance:
+                    best_distance = distance
+                    best_index = i
+            selected.append(best_index)
+
+        return [pairs[i] for i in selected]

pntx/types.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+Label = Literal["positive", "negative"]
+
+POSITIVE: Label = "positive"
+NEGATIVE: Label = "negative"
+
+Pair = tuple[str, str]
+"""A single (positive, negative) example pair."""
+
+
+@dataclass(frozen=True, eq=False)
+class ClassifyResult:
+    """Result of classifying a single text.
+
+    ``result == "positive"`` compares against ``label`` directly, so callers
+    don't need to write ``result.label == "positive"``.
+    """
+
+    label: Label
+    confidence: float
+
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, str):
+            return self.label == other
+        if isinstance(other, ClassifyResult):
+            return self.label == other.label and self.confidence == other.confidence
+        return NotImplemented
+
+    def __hash__(self) -> int:
+        return hash((self.label, self.confidence))

pntx-0.1.0.dist-info/METADATA

@@ -0,0 +1,150 @@
+Metadata-Version: 2.4
+Name: pntx
+Version: 0.1.0
+Summary: Generate and classify text from user-defined (positive, negative) example pairs
+Keywords: nlp,llm,llama.cpp,text-generation,text-classification,few-shot
+Author: pillyshi
+Author-email: pillyshi <pillyshi21@gmail.com>
+License-Expression: MIT
+License-File: LICENSE
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Topic :: Text Processing :: Linguistic
+Classifier: Typing :: Typed
+Requires-Dist: anthropic ; extra == 'anthropic'
+Requires-Dist: sentence-transformers ; extra == 'embeddings'
+Requires-Dist: llama-cpp-python ; extra == 'llama'
+Requires-Python: >=3.10
+Project-URL: Homepage, https://github.com/pillyshi/pntx
+Project-URL: Repository, https://github.com/pillyshi/pntx
+Project-URL: Issues, https://github.com/pillyshi/pntx/issues
+Provides-Extra: anthropic
+Provides-Extra: embeddings
+Provides-Extra: llama
+Description-Content-Type: text/markdown
+
+# pntx
+
+`pntx` is a Python library that turns a handful of user-supplied `(positive, negative)`
+text pairs into:
+
+1. **Generation** — synthesize new text on either side of the pair.
+2. **Classification** — label arbitrary text as `positive` or `negative`.
+
+The meaning of "positive" and "negative" is entirely up to you. It doesn't have to be
+sentiment — it can be formal/casual, policy-compliant/violating, or any other contrast
+you define with examples. `pntx` never interprets the pairs; it only uses them as
+few-shot and scoring material.
+
+```python
+from pntx import PNTX
+
+model = PNTX(backend="llama", model_path="model.gguf")
+
+pairs = [
+    ("The movie was fantastic", "The movie was boring"),
+    ("Support was quick and helpful", "Support was slow and unhelpful"),
+]
+model.fit(pairs)
+
+# Generation
+texts = model.generate(
+    n=20,
+    side="positive",
+    temperature=1.0,
+    dedup=True,          # filter near-duplicates (of each other and of the seed pairs)
+    verify=True,         # self-classify and reject anything that doesn't match `side`
+    min_confidence=0.8,  # confidence threshold used by verify
+)
+
+# Classification
+result = model.classify("The staff were incredibly friendly")
+result.label        # "positive" | "negative"
+result.confidence   # float in [0.0, 1.0]
+result == "positive"  # True
+
+results = model.classify_batch(texts)  # batched, not a naive per-item loop
+```
+
+## Installation
+
+`pntx` uses [uv](https://docs.astral.sh/uv/) for package management.
+
+```bash
+uv add pntx              # core (zero dependencies)
+uv add "pntx[llama]"     # + llama.cpp in-process backend
+uv add "pntx[anthropic]" # + Anthropic API backend
+uv add "pntx[embeddings]" # + semantic similarity for selectors
+```
+
+The core package has no runtime dependencies. Each backend/feature lives behind its
+own extra, and using one without installing it raises a clear `ImportError` with the
+install command to run.
+
+## Backends
+
+`pntx` runs models two ways:
+
+- **`LlamaCppBackend`** (`pntx[llama]`) — runs a GGUF model in-process via
+  `llama-cpp-python`. This is the primary, most-tuned backend: classification uses
+  token log-probabilities directly (`score_choices`), and batched classification
+  reuses the shared few-shot prefix's KV cache across every item instead of
+  re-evaluating it per item.
+- **`AnthropicBackend`** (`pntx[anthropic]`) — calls the Anthropic Messages API.
+  Since that API doesn't expose log-probabilities, classification asks the model to
+  name the label and parses it out of the response instead (confidence is then a
+  fixed convention value, not a calibrated probability). Batched classification runs
+  requests concurrently (`asyncio` + a semaphore), not in a sequential loop.
+
+```python
+model = PNTX(backend="llama", model_path="model.gguf")
+model = PNTX(backend="anthropic", model="claude-...")
+
+# or pass a backend instance directly, e.g. for dependency injection in tests
+from pntx.backends.llama import LlamaCppBackend
+model = PNTX(backend=LlamaCppBackend(model_path="model.gguf"))
+```
+
+## Selecting exemplars
+
+When there are more fitted pairs than comfortably fit in a prompt, a `Selector`
+decides which ones to use:
+
+- **`RandomSelector`** (default) — a uniform random subset.
+- **`NearestSelector`** — picks pairs whose text is most similar to the text being
+  classified; dynamic, per-query selection.
+- **`DiversitySelector`** — greedily picks a maximally diverse subset.
+
+Both `NearestSelector` and `DiversitySelector` take a `similarity_fn`. It defaults to
+a dependency-free character n-gram similarity (`pntx.dedup.similarity`); pass
+`pntx.embeddings.cosine_similarity_fn()` (requires `pntx[embeddings]`) for semantic
+similarity instead:
+
+```python
+from pntx import PNTX
+from pntx.selection import NearestSelector
+
+model = PNTX(backend="llama", model_path="model.gguf", selector=NearestSelector())
+```
+
+## Development
+
+```bash
+uv sync                          # install dev dependencies
+uv run pytest                    # unit tests (integration tests are skipped by default)
+uv run ruff check .
+uv run mypy src tests
+```
+
+Integration tests that hit a real model or API are opt-in:
+
+```bash
+PNTX_LLAMA_MODEL_PATH=/path/to/model.gguf uv run pytest tests/integration
+ANTHROPIC_API_KEY=... uv run pytest tests/integration/test_anthropic_backend.py
+```

pntx-0.1.0.dist-info/RECORD

@@ -0,0 +1,17 @@
+pntx/__init__.py,sha256=l-XtqdBxBj3v6XzAU_tz-OFK7QJz4nIUlYdhn-2NLng,131
+pntx/backends/__init__.py,sha256=ST7pfXWirR1Xx9w32cVlQrrTv2bG85WV-T-h7jSp3B8,193
+pntx/backends/anthropic.py,sha256=cKfjhfNH3zvGkw-Hj7Z7KwFQdtfTUfSAa54_RFF1nfk,3396
+pntx/backends/base.py,sha256=5soJhmtmKFocb7Nt5VtcGhnz5tWmPLYaCknK6bci544,2433
+pntx/backends/llama.py,sha256=78zJ-0EdV2Fps27N32H1LLarELvkU3c7iVj-774MnQw,4286
+pntx/core.py,sha256=W158hpICAqESnVWqp9qBBRX74fX7OOcfi-VMWUkZ6o4,9076
+pntx/dedup.py,sha256=5yBaDzbNqmHl2qCVc6u55-BADOzmclZlQg6oeC6mPVw,1161
+pntx/embeddings.py,sha256=y1dIzfD8CMXwooICz2n9bbilABoXrc0qX5a2KsU7tkg,2026
+pntx/generate.py,sha256=Ee_jbr1bK5s1xbrcu-0fPsUJBljaI5XZL4SbZXirh9s,2371
+pntx/prompts.py,sha256=Tze3SGoOw6leneNDSSEAyeBa_FbF5ZVopmlsOkV5Nps,4704
+pntx/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+pntx/selection.py,sha256=L55L7WHtKef-zbB1tJAWqjgYHOqVqBkCazvzw2BnshI,3955
+pntx/types.py,sha256=C110SSPFa9oR5JLIVHLhudd5wAFdGA-AARMyzXW49xs,919
+pntx-0.1.0.dist-info/licenses/LICENSE,sha256=CsKg5iZ5CavN1hLssUVZLR3KwXDysPkEzHHAE3-uwa0,1065
+pntx-0.1.0.dist-info/WHEEL,sha256=WvwXFgRajeoYkfRVmDhkP4Qlqo31Mk687zIO2QQoFmw,80
+pntx-0.1.0.dist-info/METADATA,sha256=7mSnCG7O-Joe5y3vaUWcvmBgzL6es5EkTm8MWeALtlA,5602
+pntx-0.1.0.dist-info/RECORD,,

pntx-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: uv 0.11.7
+Root-Is-Purelib: true
+Tag: py3-none-any

pntx-0.1.0.dist-info/licenses/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.