pntx 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
pntx-0.1.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 pillyshi
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
pntx-0.1.0/PKG-INFO ADDED
@@ -0,0 +1,150 @@
1
+ Metadata-Version: 2.4
2
+ Name: pntx
3
+ Version: 0.1.0
4
+ Summary: Generate and classify text from user-defined (positive, negative) example pairs
5
+ Keywords: nlp,llm,llama.cpp,text-generation,text-classification,few-shot
6
+ Author: pillyshi
7
+ Author-email: pillyshi <pillyshi21@gmail.com>
8
+ License-Expression: MIT
9
+ License-File: LICENSE
10
+ Classifier: Development Status :: 3 - Alpha
11
+ Classifier: Intended Audience :: Developers
12
+ Classifier: Programming Language :: Python :: 3
13
+ Classifier: Programming Language :: Python :: 3.10
14
+ Classifier: Programming Language :: Python :: 3.11
15
+ Classifier: Programming Language :: Python :: 3.12
16
+ Classifier: Programming Language :: Python :: 3.13
17
+ Classifier: Topic :: Software Development :: Libraries :: Python Modules
18
+ Classifier: Topic :: Text Processing :: Linguistic
19
+ Classifier: Typing :: Typed
20
+ Requires-Dist: anthropic ; extra == 'anthropic'
21
+ Requires-Dist: sentence-transformers ; extra == 'embeddings'
22
+ Requires-Dist: llama-cpp-python ; extra == 'llama'
23
+ Requires-Python: >=3.10
24
+ Project-URL: Homepage, https://github.com/pillyshi/pntx
25
+ Project-URL: Repository, https://github.com/pillyshi/pntx
26
+ Project-URL: Issues, https://github.com/pillyshi/pntx/issues
27
+ Provides-Extra: anthropic
28
+ Provides-Extra: embeddings
29
+ Provides-Extra: llama
30
+ Description-Content-Type: text/markdown
31
+
32
+ # pntx
33
+
34
+ `pntx` is a Python library that turns a handful of user-supplied `(positive, negative)`
35
+ text pairs into:
36
+
37
+ 1. **Generation** — synthesize new text on either side of the pair.
38
+ 2. **Classification** — label arbitrary text as `positive` or `negative`.
39
+
40
+ The meaning of "positive" and "negative" is entirely up to you. It doesn't have to be
41
+ sentiment — it can be formal/casual, policy-compliant/violating, or any other contrast
42
+ you define with examples. `pntx` never interprets the pairs; it only uses them as
43
+ few-shot and scoring material.
44
+
45
+ ```python
46
+ from pntx import PNTX
47
+
48
+ model = PNTX(backend="llama", model_path="model.gguf")
49
+
50
+ pairs = [
51
+ ("The movie was fantastic", "The movie was boring"),
52
+ ("Support was quick and helpful", "Support was slow and unhelpful"),
53
+ ]
54
+ model.fit(pairs)
55
+
56
+ # Generation
57
+ texts = model.generate(
58
+ n=20,
59
+ side="positive",
60
+ temperature=1.0,
61
+ dedup=True, # filter near-duplicates (of each other and of the seed pairs)
62
+ verify=True, # self-classify and reject anything that doesn't match `side`
63
+ min_confidence=0.8, # confidence threshold used by verify
64
+ )
65
+
66
+ # Classification
67
+ result = model.classify("The staff were incredibly friendly")
68
+ result.label # "positive" | "negative"
69
+ result.confidence # float in [0.0, 1.0]
70
+ result == "positive" # True
71
+
72
+ results = model.classify_batch(texts) # batched, not a naive per-item loop
73
+ ```
74
+
75
+ ## Installation
76
+
77
+ `pntx` uses [uv](https://docs.astral.sh/uv/) for package management.
78
+
79
+ ```bash
80
+ uv add pntx # core (zero dependencies)
81
+ uv add "pntx[llama]" # + llama.cpp in-process backend
82
+ uv add "pntx[anthropic]" # + Anthropic API backend
83
+ uv add "pntx[embeddings]" # + semantic similarity for selectors
84
+ ```
85
+
86
+ The core package has no runtime dependencies. Each backend/feature lives behind its
87
+ own extra, and using one without installing it raises a clear `ImportError` with the
88
+ install command to run.
89
+
90
+ ## Backends
91
+
92
+ `pntx` runs models two ways:
93
+
94
+ - **`LlamaCppBackend`** (`pntx[llama]`) — runs a GGUF model in-process via
95
+ `llama-cpp-python`. This is the primary, most-tuned backend: classification uses
96
+ token log-probabilities directly (`score_choices`), and batched classification
97
+ reuses the shared few-shot prefix's KV cache across every item instead of
98
+ re-evaluating it per item.
99
+ - **`AnthropicBackend`** (`pntx[anthropic]`) — calls the Anthropic Messages API.
100
+ Since that API doesn't expose log-probabilities, classification asks the model to
101
+ name the label and parses it out of the response instead (confidence is then a
102
+ fixed convention value, not a calibrated probability). Batched classification runs
103
+ requests concurrently (`asyncio` + a semaphore), not in a sequential loop.
104
+
105
+ ```python
106
+ model = PNTX(backend="llama", model_path="model.gguf")
107
+ model = PNTX(backend="anthropic", model="claude-...")
108
+
109
+ # or pass a backend instance directly, e.g. for dependency injection in tests
110
+ from pntx.backends.llama import LlamaCppBackend
111
+ model = PNTX(backend=LlamaCppBackend(model_path="model.gguf"))
112
+ ```
113
+
114
+ ## Selecting exemplars
115
+
116
+ When there are more fitted pairs than comfortably fit in a prompt, a `Selector`
117
+ decides which ones to use:
118
+
119
+ - **`RandomSelector`** (default) — a uniform random subset.
120
+ - **`NearestSelector`** — picks pairs whose text is most similar to the text being
121
+ classified; dynamic, per-query selection.
122
+ - **`DiversitySelector`** — greedily picks a maximally diverse subset.
123
+
124
+ Both `NearestSelector` and `DiversitySelector` take a `similarity_fn`. It defaults to
125
+ a dependency-free character n-gram similarity (`pntx.dedup.similarity`); pass
126
+ `pntx.embeddings.cosine_similarity_fn()` (requires `pntx[embeddings]`) for semantic
127
+ similarity instead:
128
+
129
+ ```python
130
+ from pntx import PNTX
131
+ from pntx.selection import NearestSelector
132
+
133
+ model = PNTX(backend="llama", model_path="model.gguf", selector=NearestSelector())
134
+ ```
135
+
136
+ ## Development
137
+
138
+ ```bash
139
+ uv sync # install dev dependencies
140
+ uv run pytest # unit tests (integration tests are skipped by default)
141
+ uv run ruff check .
142
+ uv run mypy src tests
143
+ ```
144
+
145
+ Integration tests that hit a real model or API are opt-in:
146
+
147
+ ```bash
148
+ PNTX_LLAMA_MODEL_PATH=/path/to/model.gguf uv run pytest tests/integration
149
+ ANTHROPIC_API_KEY=... uv run pytest tests/integration/test_anthropic_backend.py
150
+ ```
pntx-0.1.0/README.md ADDED
@@ -0,0 +1,119 @@
1
+ # pntx
2
+
3
+ `pntx` is a Python library that turns a handful of user-supplied `(positive, negative)`
4
+ text pairs into:
5
+
6
+ 1. **Generation** — synthesize new text on either side of the pair.
7
+ 2. **Classification** — label arbitrary text as `positive` or `negative`.
8
+
9
+ The meaning of "positive" and "negative" is entirely up to you. It doesn't have to be
10
+ sentiment — it can be formal/casual, policy-compliant/violating, or any other contrast
11
+ you define with examples. `pntx` never interprets the pairs; it only uses them as
12
+ few-shot and scoring material.
13
+
14
+ ```python
15
+ from pntx import PNTX
16
+
17
+ model = PNTX(backend="llama", model_path="model.gguf")
18
+
19
+ pairs = [
20
+ ("The movie was fantastic", "The movie was boring"),
21
+ ("Support was quick and helpful", "Support was slow and unhelpful"),
22
+ ]
23
+ model.fit(pairs)
24
+
25
+ # Generation
26
+ texts = model.generate(
27
+ n=20,
28
+ side="positive",
29
+ temperature=1.0,
30
+ dedup=True, # filter near-duplicates (of each other and of the seed pairs)
31
+ verify=True, # self-classify and reject anything that doesn't match `side`
32
+ min_confidence=0.8, # confidence threshold used by verify
33
+ )
34
+
35
+ # Classification
36
+ result = model.classify("The staff were incredibly friendly")
37
+ result.label # "positive" | "negative"
38
+ result.confidence # float in [0.0, 1.0]
39
+ result == "positive" # True
40
+
41
+ results = model.classify_batch(texts) # batched, not a naive per-item loop
42
+ ```
43
+
44
+ ## Installation
45
+
46
+ `pntx` uses [uv](https://docs.astral.sh/uv/) for package management.
47
+
48
+ ```bash
49
+ uv add pntx # core (zero dependencies)
50
+ uv add "pntx[llama]" # + llama.cpp in-process backend
51
+ uv add "pntx[anthropic]" # + Anthropic API backend
52
+ uv add "pntx[embeddings]" # + semantic similarity for selectors
53
+ ```
54
+
55
+ The core package has no runtime dependencies. Each backend/feature lives behind its
56
+ own extra, and using one without installing it raises a clear `ImportError` with the
57
+ install command to run.
58
+
59
+ ## Backends
60
+
61
+ `pntx` runs models two ways:
62
+
63
+ - **`LlamaCppBackend`** (`pntx[llama]`) — runs a GGUF model in-process via
64
+ `llama-cpp-python`. This is the primary, most-tuned backend: classification uses
65
+ token log-probabilities directly (`score_choices`), and batched classification
66
+ reuses the shared few-shot prefix's KV cache across every item instead of
67
+ re-evaluating it per item.
68
+ - **`AnthropicBackend`** (`pntx[anthropic]`) — calls the Anthropic Messages API.
69
+ Since that API doesn't expose log-probabilities, classification asks the model to
70
+ name the label and parses it out of the response instead (confidence is then a
71
+ fixed convention value, not a calibrated probability). Batched classification runs
72
+ requests concurrently (`asyncio` + a semaphore), not in a sequential loop.
73
+
74
+ ```python
75
+ model = PNTX(backend="llama", model_path="model.gguf")
76
+ model = PNTX(backend="anthropic", model="claude-...")
77
+
78
+ # or pass a backend instance directly, e.g. for dependency injection in tests
79
+ from pntx.backends.llama import LlamaCppBackend
80
+ model = PNTX(backend=LlamaCppBackend(model_path="model.gguf"))
81
+ ```
82
+
83
+ ## Selecting exemplars
84
+
85
+ When there are more fitted pairs than comfortably fit in a prompt, a `Selector`
86
+ decides which ones to use:
87
+
88
+ - **`RandomSelector`** (default) — a uniform random subset.
89
+ - **`NearestSelector`** — picks pairs whose text is most similar to the text being
90
+ classified; dynamic, per-query selection.
91
+ - **`DiversitySelector`** — greedily picks a maximally diverse subset.
92
+
93
+ Both `NearestSelector` and `DiversitySelector` take a `similarity_fn`. It defaults to
94
+ a dependency-free character n-gram similarity (`pntx.dedup.similarity`); pass
95
+ `pntx.embeddings.cosine_similarity_fn()` (requires `pntx[embeddings]`) for semantic
96
+ similarity instead:
97
+
98
+ ```python
99
+ from pntx import PNTX
100
+ from pntx.selection import NearestSelector
101
+
102
+ model = PNTX(backend="llama", model_path="model.gguf", selector=NearestSelector())
103
+ ```
104
+
105
+ ## Development
106
+
107
+ ```bash
108
+ uv sync # install dev dependencies
109
+ uv run pytest # unit tests (integration tests are skipped by default)
110
+ uv run ruff check .
111
+ uv run mypy src tests
112
+ ```
113
+
114
+ Integration tests that hit a real model or API are opt-in:
115
+
116
+ ```bash
117
+ PNTX_LLAMA_MODEL_PATH=/path/to/model.gguf uv run pytest tests/integration
118
+ ANTHROPIC_API_KEY=... uv run pytest tests/integration/test_anthropic_backend.py
119
+ ```
@@ -0,0 +1,79 @@
1
+ [project]
2
+ name = "pntx"
3
+ version = "0.1.0"
4
+ description = "Generate and classify text from user-defined (positive, negative) example pairs"
5
+ readme = "README.md"
6
+ license = "MIT"
7
+ license-files = ["LICENSE"]
8
+ authors = [
9
+ { name = "pillyshi", email = "pillyshi21@gmail.com" }
10
+ ]
11
+ requires-python = ">=3.10"
12
+ keywords = [
13
+ "nlp",
14
+ "llm",
15
+ "llama.cpp",
16
+ "text-generation",
17
+ "text-classification",
18
+ "few-shot",
19
+ ]
20
+ classifiers = [
21
+ "Development Status :: 3 - Alpha",
22
+ "Intended Audience :: Developers",
23
+ "Programming Language :: Python :: 3",
24
+ "Programming Language :: Python :: 3.10",
25
+ "Programming Language :: Python :: 3.11",
26
+ "Programming Language :: Python :: 3.12",
27
+ "Programming Language :: Python :: 3.13",
28
+ "Topic :: Software Development :: Libraries :: Python Modules",
29
+ "Topic :: Text Processing :: Linguistic",
30
+ "Typing :: Typed",
31
+ ]
32
+ dependencies = []
33
+
34
+ [project.urls]
35
+ Homepage = "https://github.com/pillyshi/pntx"
36
+ Repository = "https://github.com/pillyshi/pntx"
37
+ Issues = "https://github.com/pillyshi/pntx/issues"
38
+
39
+ [project.optional-dependencies]
40
+ llama = [
41
+ "llama-cpp-python",
42
+ ]
43
+ anthropic = [
44
+ "anthropic",
45
+ ]
46
+ embeddings = [
47
+ "sentence-transformers",
48
+ ]
49
+
50
+ [build-system]
51
+ requires = ["uv_build>=0.11.7,<0.12.0"]
52
+ build-backend = "uv_build"
53
+
54
+ [tool.ruff]
55
+ line-length = 100
56
+ target-version = "py310"
57
+ extend-exclude = ["tmp"]
58
+
59
+ [tool.ruff.lint]
60
+ select = ["E", "F", "I", "UP", "B"]
61
+
62
+ [tool.mypy]
63
+ strict = true
64
+ python_version = "3.10"
65
+
66
+ [[tool.mypy.overrides]]
67
+ module = "sentence_transformers.*"
68
+ ignore_missing_imports = true
69
+
70
+ [tool.pytest.ini_options]
71
+ testpaths = ["tests"]
72
+
73
+ [dependency-groups]
74
+ dev = [
75
+ "mypy>=2.3.0",
76
+ "pytest>=9.1.1",
77
+ "ruff>=0.15.21",
78
+ "twine>=6.2.0",
79
+ ]
@@ -0,0 +1,6 @@
1
+ from __future__ import annotations
2
+
3
+ from .core import PNTX
4
+ from .types import ClassifyResult
5
+
6
+ __all__ = ["PNTX", "ClassifyResult"]
@@ -0,0 +1,5 @@
1
+ from __future__ import annotations
2
+
3
+ from .base import Backend, BatchBackend, BatchScoringBackend, ScoringBackend
4
+
5
+ __all__ = ["BatchBackend", "BatchScoringBackend", "Backend", "ScoringBackend"]
@@ -0,0 +1,106 @@
1
+ from __future__ import annotations
2
+
3
+ import asyncio
4
+ from typing import Any
5
+
6
+ try:
7
+ import anthropic
8
+ except ImportError as e:
9
+ raise ImportError(
10
+ "AnthropicBackend requires the 'anthropic' extra. "
11
+ "Install it with: pip install 'pntx[anthropic]'"
12
+ ) from e
13
+
14
+ _DEFAULT_CONCURRENCY = 5
15
+
16
+
17
+ class AnthropicBackend:
18
+ """Backend using the Anthropic Messages API.
19
+
20
+ Implements only ``Backend`` (plain text completion): the Messages API
21
+ doesn't expose token log-probabilities, so classification for this
22
+ backend goes through PNTX's parse-based fallback rather than
23
+ ``score_choices``. ``complete_batch`` implements ``BatchBackend`` by
24
+ running requests concurrently via asyncio, bounded by ``concurrency``.
25
+ """
26
+
27
+ def __init__(
28
+ self,
29
+ model: str,
30
+ *,
31
+ concurrency: int = _DEFAULT_CONCURRENCY,
32
+ client: anthropic.Anthropic | None = None,
33
+ async_client: anthropic.AsyncAnthropic | None = None,
34
+ **client_kwargs: Any,
35
+ ) -> None:
36
+ if concurrency < 1:
37
+ raise ValueError(f"concurrency must be >= 1, got {concurrency}")
38
+ self.model = model
39
+ self.concurrency = concurrency
40
+ self._client = client if client is not None else anthropic.Anthropic(**client_kwargs)
41
+ self._async_client = (
42
+ async_client if async_client is not None else anthropic.AsyncAnthropic(**client_kwargs)
43
+ )
44
+
45
+ def complete(
46
+ self,
47
+ prompt: str,
48
+ *,
49
+ temperature: float = 1.0,
50
+ max_tokens: int = 512,
51
+ stop: list[str] | None = None,
52
+ ) -> str:
53
+ message = self._client.messages.create(
54
+ model=self.model,
55
+ max_tokens=max_tokens,
56
+ temperature=temperature,
57
+ stop_sequences=stop or [],
58
+ messages=[{"role": "user", "content": prompt}],
59
+ )
60
+ return _extract_text(message)
61
+
62
+ def complete_batch(
63
+ self,
64
+ prompts: list[str],
65
+ *,
66
+ temperature: float = 1.0,
67
+ max_tokens: int = 512,
68
+ stop: list[str] | None = None,
69
+ ) -> list[str]:
70
+ """Complete every prompt concurrently, bounded by ``self.concurrency``."""
71
+ if not prompts:
72
+ return []
73
+ return asyncio.run(
74
+ self._complete_batch_async(
75
+ prompts, temperature=temperature, max_tokens=max_tokens, stop=stop
76
+ )
77
+ )
78
+
79
+ async def _complete_batch_async(
80
+ self,
81
+ prompts: list[str],
82
+ *,
83
+ temperature: float,
84
+ max_tokens: int,
85
+ stop: list[str] | None,
86
+ ) -> list[str]:
87
+ semaphore = asyncio.Semaphore(self.concurrency)
88
+
89
+ async def complete_one(prompt: str) -> str:
90
+ async with semaphore:
91
+ message = await self._async_client.messages.create(
92
+ model=self.model,
93
+ max_tokens=max_tokens,
94
+ temperature=temperature,
95
+ stop_sequences=stop or [],
96
+ messages=[{"role": "user", "content": prompt}],
97
+ )
98
+ return _extract_text(message)
99
+
100
+ return list(await asyncio.gather(*(complete_one(prompt) for prompt in prompts)))
101
+
102
+
103
+ def _extract_text(message: anthropic.types.Message) -> str:
104
+ return "".join(
105
+ block.text for block in message.content if isinstance(block, anthropic.types.TextBlock)
106
+ )
@@ -0,0 +1,75 @@
1
+ from __future__ import annotations
2
+
3
+ from typing import Protocol, runtime_checkable
4
+
5
+
6
+ @runtime_checkable
7
+ class Backend(Protocol):
8
+ """A text-completion backend.
9
+
10
+ Every backend (llama.cpp, Anthropic, ...) must implement at least this.
11
+ """
12
+
13
+ def complete(
14
+ self,
15
+ prompt: str,
16
+ *,
17
+ temperature: float = 1.0,
18
+ max_tokens: int = 512,
19
+ stop: list[str] | None = None,
20
+ ) -> str: ...
21
+
22
+
23
+ @runtime_checkable
24
+ class ScoringBackend(Backend, Protocol):
25
+ """A backend that can score candidate continuations directly.
26
+
27
+ This is the primary path for classification: instead of generating text
28
+ and parsing a label out of it, the backend returns a log-likelihood for
29
+ each candidate choice appended to ``prompt``.
30
+ """
31
+
32
+ def score_choices(self, prompt: str, choices: list[str]) -> list[float]:
33
+ """Return the log-likelihood of each choice in ``choices`` as a
34
+ continuation of ``prompt``."""
35
+ ...
36
+
37
+
38
+ @runtime_checkable
39
+ class BatchScoringBackend(ScoringBackend, Protocol):
40
+ """A ``ScoringBackend`` that can batch-score many queries sharing one prefix.
41
+
42
+ Used by ``classify_batch`` to warm a common prefix (e.g. the few-shot
43
+ exemplar block) once and reuse it across every query, instead of
44
+ re-evaluating it per item. Backends without this (e.g. remote API
45
+ backends) are scored one item at a time via ``score_choices``.
46
+ """
47
+
48
+ def score_choices_batch(
49
+ self, prefix: str, queries: list[str], choices: list[str]
50
+ ) -> list[list[float]]:
51
+ """For each query in ``queries``, score ``choices`` as a continuation
52
+ of ``prefix + query``. Returns one score list per query, in order."""
53
+ ...
54
+
55
+
56
+ @runtime_checkable
57
+ class BatchBackend(Backend, Protocol):
58
+ """A ``Backend`` that can complete many prompts concurrently.
59
+
60
+ Used by ``classify_batch``'s parse-based path (for backends that don't
61
+ implement ``ScoringBackend``, e.g. remote chat APIs): instead of
62
+ completing each prompt one at a time, ``complete_batch`` can run them
63
+ concurrently (e.g. via asyncio) and return results in the same order.
64
+ """
65
+
66
+ def complete_batch(
67
+ self,
68
+ prompts: list[str],
69
+ *,
70
+ temperature: float = 1.0,
71
+ max_tokens: int = 512,
72
+ stop: list[str] | None = None,
73
+ ) -> list[str]:
74
+ """Complete every prompt in ``prompts``, returning results in order."""
75
+ ...
@@ -0,0 +1,108 @@
1
+ from __future__ import annotations
2
+
3
+ from typing import Any
4
+
5
+ try:
6
+ import llama_cpp
7
+ except ImportError as e:
8
+ raise ImportError(
9
+ "LlamaCppBackend requires the 'llama' extra. "
10
+ "Install it with: pip install 'pntx[llama]'"
11
+ ) from e
12
+
13
+
14
+ class LlamaCppBackend:
15
+ """In-process backend backed by llama.cpp (via ``llama-cpp-python``).
16
+
17
+ Implements ``Backend``, ``ScoringBackend`` and ``BatchScoringBackend``.
18
+ Scoring methods reuse the KV cache of whatever prefix they share: rather
19
+ than re-evaluating a shared prefix for every downstream token span, they
20
+ eval it once and then, for each span, rewind ``n_tokens`` back to the
21
+ prefix boundary before eval'ing that span's tokens (``Llama.eval()``
22
+ trims the KV cache down to the current ``n_tokens`` before appending, so
23
+ this discards only the previous span, not the shared prefix).
24
+ """
25
+
26
+ def __init__(self, model_path: str, **kwargs: Any) -> None:
27
+ # Scoring needs per-token logits, which llama.cpp only keeps around
28
+ # when logits_all=True.
29
+ kwargs["logits_all"] = True
30
+ self._llm = llama_cpp.Llama(model_path=model_path, **kwargs)
31
+
32
+ def complete(
33
+ self,
34
+ prompt: str,
35
+ *,
36
+ temperature: float = 1.0,
37
+ max_tokens: int = 512,
38
+ stop: list[str] | None = None,
39
+ ) -> str:
40
+ result = self._llm.create_completion(
41
+ prompt,
42
+ temperature=temperature,
43
+ max_tokens=max_tokens,
44
+ stop=stop or [],
45
+ )
46
+ if not isinstance(result, dict):
47
+ raise TypeError(f"expected a non-streaming completion response, got {type(result)}")
48
+ return result["choices"][0]["text"]
49
+
50
+ def score_choices(self, prompt: str, choices: list[str]) -> list[float]:
51
+ """Return the summed log-likelihood of each choice as a continuation of ``prompt``."""
52
+ if not choices:
53
+ return []
54
+ prompt_tokens = self._tokenize(prompt, add_bos=True)
55
+ self._reset()
56
+ self._llm.eval(prompt_tokens)
57
+ return self._score_choices_at(self._llm.n_tokens, choices)
58
+
59
+ def score_choices_batch(
60
+ self, prefix: str, queries: list[str], choices: list[str]
61
+ ) -> list[list[float]]:
62
+ """For each query, score ``choices`` as a continuation of ``prefix + query``.
63
+
64
+ ``prefix`` (e.g. the few-shot exemplar block) is eval'd once and its
65
+ KV cache reused across every query, instead of re-evaluating it per
66
+ query as a naive per-item ``score_choices`` loop would.
67
+ """
68
+ if not queries:
69
+ return []
70
+ prefix_tokens = self._tokenize(prefix, add_bos=True)
71
+ self._reset()
72
+ self._llm.eval(prefix_tokens)
73
+ base = self._llm.n_tokens
74
+
75
+ results: list[list[float]] = []
76
+ for query in queries:
77
+ query_tokens = self._tokenize(query, add_bos=False)
78
+ self._llm.n_tokens = base
79
+ self._llm.eval(query_tokens)
80
+ results.append(self._score_choices_at(self._llm.n_tokens, choices))
81
+ self._llm.n_tokens = base
82
+ return results
83
+
84
+ def _tokenize(self, text: str, *, add_bos: bool) -> list[int]:
85
+ return self._llm.tokenize(text.encode("utf-8"), add_bos=add_bos)
86
+
87
+ def _reset(self) -> None:
88
+ self._llm.reset() # type: ignore[no-untyped-call] # llama_cpp.Llama.reset() has no annotations
89
+
90
+ def _score_choices_at(self, base_n_tokens: int, choices: list[str]) -> list[float]:
91
+ """Score each choice freshly on top of the KV cache at ``base_n_tokens``,
92
+ rewinding between choices so only the shared prefix up to
93
+ ``base_n_tokens`` is reused."""
94
+ scores = [
95
+ self._score_tokens(base_n_tokens, self._tokenize(choice, add_bos=False))
96
+ for choice in choices
97
+ ]
98
+ self._llm.n_tokens = base_n_tokens
99
+ return scores
100
+
101
+ def _score_tokens(self, base_n_tokens: int, tokens: list[int]) -> float:
102
+ if not tokens:
103
+ return 0.0
104
+ self._llm.n_tokens = base_n_tokens
105
+ self._llm.eval(tokens)
106
+ logits = self._llm.scores[base_n_tokens - 1 : base_n_tokens + len(tokens) - 1]
107
+ logprobs = self._llm.logits_to_logprobs(logits)
108
+ return float(sum(logprobs[j, tok] for j, tok in enumerate(tokens)))
@@ -0,0 +1,227 @@
1
+ from __future__ import annotations
2
+
3
+ import math
4
+ import warnings
5
+ from collections.abc import Iterable
6
+ from importlib import import_module
7
+ from typing import Any
8
+
9
+ from . import generate as generate_module
10
+ from . import prompts
11
+ from .backends.base import Backend, BatchBackend, BatchScoringBackend, ScoringBackend
12
+ from .selection import RandomSelector, Selector
13
+ from .types import ClassifyResult, Label, Pair
14
+
15
+ _BACKEND_REGISTRY: dict[str, tuple[str, str]] = {
16
+ "llama": ("pntx.backends.llama", "LlamaCppBackend"),
17
+ "anthropic": ("pntx.backends.anthropic", "AnthropicBackend"),
18
+ }
19
+
20
+
21
+ def _resolve_backend(backend: Backend | str, backend_kwargs: dict[str, Any]) -> Backend:
22
+ if not isinstance(backend, str):
23
+ if backend_kwargs:
24
+ raise TypeError(
25
+ "backend_kwargs are only used when backend is given as a string "
26
+ "(e.g. PNTX(backend='llama', model_path=...)); construct the "
27
+ "Backend instance directly instead"
28
+ )
29
+ return backend
30
+ try:
31
+ module_name, class_name = _BACKEND_REGISTRY[backend]
32
+ except KeyError:
33
+ raise ValueError(
34
+ f"Unknown backend {backend!r}; expected a Backend instance or one "
35
+ f"of {sorted(_BACKEND_REGISTRY)}"
36
+ ) from None
37
+ try:
38
+ module = import_module(module_name)
39
+ except ImportError as e:
40
+ raise ImportError(
41
+ f"Backend {backend!r} requires the optional '{backend}' extra. "
42
+ f"Install it with: pip install 'pntx[{backend}]'"
43
+ ) from e
44
+ return getattr(module, class_name)(**backend_kwargs) # type: ignore[no-any-return]
45
+
46
+
47
+ class PNTX:
48
+ """Generates and classifies text from user-defined (positive, negative) pairs.
49
+
50
+ The pairs' meaning is entirely up to the caller (sentiment, formality,
51
+ policy compliance, ...); this class never interprets it, it only uses the
52
+ pairs as few-shot/scoring material.
53
+ """
54
+
55
+ def __init__(
56
+ self,
57
+ backend: Backend | str,
58
+ *,
59
+ selector: Selector | None = None,
60
+ max_exemplars: int | None = None,
61
+ **backend_kwargs: Any,
62
+ ) -> None:
63
+ """Create a PNTX model.
64
+
65
+ ``backend`` is either a ready-made ``Backend`` instance, or the name
66
+ of a built-in backend (e.g. ``"llama"``) to construct lazily; any
67
+ ``backend_kwargs`` are then forwarded to that backend's constructor
68
+ (e.g. ``PNTX(backend="llama", model_path="model.gguf")``).
69
+
70
+ ``max_exemplars`` caps how many fitted pairs ``selector`` is asked to
71
+ pick for a single prompt; ``None`` means "as many as are fitted".
72
+ """
73
+ self.backend = _resolve_backend(backend, backend_kwargs)
74
+ self.selector: Selector = selector if selector is not None else RandomSelector()
75
+ self.max_exemplars = max_exemplars
76
+ self._pairs: list[Pair] = []
77
+
78
+ @property
79
+ def pairs(self) -> list[Pair]:
80
+ return list(self._pairs)
81
+
82
+ def fit(self, pairs: Iterable[Pair]) -> PNTX:
83
+ """Store the (positive, negative) pairs used as generation/classification material.
84
+
85
+ This only stores and validates ``pairs``; no training happens here.
86
+ """
87
+ pairs = list(pairs)
88
+ if not pairs:
89
+ raise ValueError("pairs must be non-empty")
90
+ self._pairs = pairs
91
+ return self
92
+
93
+ def generate(
94
+ self,
95
+ n: int,
96
+ side: Label,
97
+ *,
98
+ temperature: float = 1.0,
99
+ dedup: bool = True,
100
+ verify: bool = True,
101
+ min_confidence: float = 0.8,
102
+ max_attempts: int | None = None,
103
+ ) -> list[str]:
104
+ """Generate up to ``n`` new ``side`` texts from the fitted pairs.
105
+
106
+ If ``verify`` keeps rejecting candidates (wrong self-classified
107
+ label, or confidence below ``min_confidence``), fewer than ``n``
108
+ texts may come back; this warns rather than looping forever, capped
109
+ at ``max_attempts`` retries (default: ``max(n * 3, 3)``).
110
+ """
111
+ self._check_fitted()
112
+ if n < 0:
113
+ raise ValueError(f"n must be >= 0, got {n}")
114
+
115
+ texts = generate_module.run_generation_loop(
116
+ backend=self.backend,
117
+ pairs=self._pairs,
118
+ selector=self.selector,
119
+ exemplar_count=self._exemplar_count(),
120
+ classify=self.classify,
121
+ n=n,
122
+ side=side,
123
+ temperature=temperature,
124
+ dedup=dedup,
125
+ verify=verify,
126
+ min_confidence=min_confidence,
127
+ max_attempts=max_attempts,
128
+ )
129
+ if len(texts) < n:
130
+ warnings.warn(
131
+ f"requested n={n} {side!r} texts but only generated {len(texts)}",
132
+ UserWarning,
133
+ stacklevel=2,
134
+ )
135
+ return texts
136
+
137
+ def classify(self, text: str) -> ClassifyResult:
138
+ """Classify ``text`` as ``"positive"`` or ``"negative"``.
139
+
140
+ If the backend implements ``ScoringBackend``, this compares the
141
+ log-likelihood of each label as a continuation of the prompt
142
+ (``confidence`` is a calibrated softmax over those log-likelihoods).
143
+ Otherwise it falls back to asking the backend to write the label and
144
+ parsing it out of the response text; see
145
+ ``prompts.parse_classify_label`` for that path's confidence
146
+ convention (not a calibrated probability).
147
+ """
148
+ self._check_fitted()
149
+ exemplars = self.selector.select(self._pairs, self._exemplar_count(), query=text)
150
+ if isinstance(self.backend, ScoringBackend):
151
+ prompt = prompts.build_classify_prompt(exemplars, text)
152
+ scores = self.backend.score_choices(prompt, prompts.classify_choice_texts())
153
+ return _result_from_scores(scores)
154
+ prompt = prompts.build_classify_prompt(exemplars, text)
155
+ raw = self.backend.complete(
156
+ prompt, temperature=0.0, max_tokens=prompts.CLASSIFY_COMPLETION_MAX_TOKENS
157
+ )
158
+ return _result_from_completion(raw)
159
+
160
+ def classify_batch(self, texts: list[str]) -> list[ClassifyResult]:
161
+ """Classify each text in ``texts``; see ``classify()`` for the two
162
+ scoring strategies this dispatches between."""
163
+ self._check_fitted()
164
+ if not texts:
165
+ return []
166
+ exemplars = self.selector.select(self._pairs, self._exemplar_count(), query=None)
167
+
168
+ if isinstance(self.backend, ScoringBackend):
169
+ choices = prompts.classify_choice_texts()
170
+ prefix = prompts.build_exemplar_prefix(exemplars)
171
+ if isinstance(self.backend, BatchScoringBackend):
172
+ queries = [prompts.build_query_suffix(text) for text in texts]
173
+ all_scores = self.backend.score_choices_batch(prefix, queries, choices)
174
+ else:
175
+ # No batch-optimized path for this backend; score one prompt
176
+ # at a time. (LlamaCppBackend implements BatchScoringBackend
177
+ # and takes the branch above; a future non-batching
178
+ # ScoringBackend falls back to this.)
179
+ all_scores = [
180
+ self.backend.score_choices(prefix + prompts.build_query_suffix(text), choices)
181
+ for text in texts
182
+ ]
183
+ return [_result_from_scores(scores) for scores in all_scores]
184
+
185
+ completion_prompts = [prompts.build_classify_prompt(exemplars, text) for text in texts]
186
+ if isinstance(self.backend, BatchBackend):
187
+ raw_completions = self.backend.complete_batch(
188
+ completion_prompts,
189
+ temperature=0.0,
190
+ max_tokens=prompts.CLASSIFY_COMPLETION_MAX_TOKENS,
191
+ )
192
+ else:
193
+ # No concurrent-batch path for this backend; complete one prompt
194
+ # at a time. (AnthropicBackend implements BatchBackend and takes
195
+ # the branch above.)
196
+ raw_completions = [
197
+ self.backend.complete(
198
+ prompt, temperature=0.0, max_tokens=prompts.CLASSIFY_COMPLETION_MAX_TOKENS
199
+ )
200
+ for prompt in completion_prompts
201
+ ]
202
+ return [_result_from_completion(raw) for raw in raw_completions]
203
+
204
+ def _exemplar_count(self) -> int:
205
+ return self.max_exemplars if self.max_exemplars is not None else len(self._pairs)
206
+
207
+ def _check_fitted(self) -> None:
208
+ if not self._pairs:
209
+ raise RuntimeError("PNTX.fit(pairs) must be called before this method")
210
+
211
+
212
+ def _result_from_scores(scores: list[float]) -> ClassifyResult:
213
+ probs = _softmax(scores)
214
+ best = max(range(len(scores)), key=scores.__getitem__)
215
+ return ClassifyResult(label=prompts.CLASSIFY_LABELS[best], confidence=probs[best])
216
+
217
+
218
+ def _result_from_completion(raw: str) -> ClassifyResult:
219
+ label, confidence = prompts.parse_classify_label(raw)
220
+ return ClassifyResult(label=label, confidence=confidence)
221
+
222
+
223
+ def _softmax(scores: list[float]) -> list[float]:
224
+ top = max(scores)
225
+ exps = [math.exp(score - top) for score in scores]
226
+ total = sum(exps)
227
+ return [exp / total for exp in exps]
@@ -0,0 +1,40 @@
1
+ from __future__ import annotations
2
+
3
+ from collections.abc import Iterable
4
+
5
+ NGRAM_SIZE = 3
6
+ SIMILARITY_THRESHOLD = 0.75
7
+
8
+
9
+ def char_ngrams(text: str, n: int = NGRAM_SIZE) -> set[str]:
10
+ """Character n-grams of ``text``.
11
+
12
+ Character n-grams (rather than word n-grams) work for both Japanese,
13
+ which has no whitespace word boundaries, and English, without needing a
14
+ tokenizer or language detection.
15
+ """
16
+ if len(text) <= n:
17
+ return {text}
18
+ return {text[i : i + n] for i in range(len(text) - n + 1)}
19
+
20
+
21
+ def similarity(a: str, b: str, n: int = NGRAM_SIZE) -> float:
22
+ """Jaccard similarity between the character n-grams of ``a`` and ``b``."""
23
+ if a == b:
24
+ return 1.0
25
+ grams_a = char_ngrams(a, n)
26
+ grams_b = char_ngrams(b, n)
27
+ union = grams_a | grams_b
28
+ if not union:
29
+ return 0.0
30
+ return len(grams_a & grams_b) / len(union)
31
+
32
+
33
+ def is_near_duplicate(
34
+ text: str,
35
+ others: Iterable[str],
36
+ *,
37
+ threshold: float = SIMILARITY_THRESHOLD,
38
+ ) -> bool:
39
+ """Whether ``text`` is an approximate duplicate of any text in ``others``."""
40
+ return any(similarity(text, other) >= threshold for other in others)
@@ -0,0 +1,59 @@
1
+ from __future__ import annotations
2
+
3
+ import math
4
+ from collections.abc import Callable
5
+
6
+ try:
7
+ from sentence_transformers import SentenceTransformer
8
+ except ImportError as e:
9
+ raise ImportError(
10
+ "pntx.embeddings requires the 'embeddings' extra. "
11
+ "Install it with: pip install 'pntx[embeddings]'"
12
+ ) from e
13
+
14
+ DEFAULT_MODEL = "paraphrase-albert-small-v2"
15
+
16
+ _model_cache: dict[str, SentenceTransformer] = {}
17
+
18
+
19
+ def _load_model(model_name: str) -> SentenceTransformer:
20
+ if model_name not in _model_cache:
21
+ _model_cache[model_name] = SentenceTransformer(model_name)
22
+ return _model_cache[model_name]
23
+
24
+
25
+ def embed(texts: list[str], model_name: str = DEFAULT_MODEL) -> list[list[float]]:
26
+ """Encode ``texts`` into embedding vectors using a sentence-transformers model."""
27
+ model = _load_model(model_name)
28
+ return [vector.tolist() for vector in model.encode(texts)]
29
+
30
+
31
+ def cosine_similarity_fn(model_name: str = DEFAULT_MODEL) -> Callable[[str, str], float]:
32
+ """Build a ``(text_a, text_b) -> float`` similarity function backed by
33
+ sentence-transformers embeddings, for use as ``NearestSelector`` or
34
+ ``DiversitySelector``'s ``similarity_fn``.
35
+
36
+ Per-text embeddings are cached within the returned function, since
37
+ both selectors call it repeatedly with the same texts.
38
+ """
39
+ model = _load_model(model_name)
40
+ vector_cache: dict[str, list[float]] = {}
41
+
42
+ def embed_one(text: str) -> list[float]:
43
+ if text not in vector_cache:
44
+ vector_cache[text] = model.encode([text])[0].tolist()
45
+ return vector_cache[text]
46
+
47
+ def similarity(a: str, b: str) -> float:
48
+ return _cosine_similarity(embed_one(a), embed_one(b))
49
+
50
+ return similarity
51
+
52
+
53
+ def _cosine_similarity(a: list[float], b: list[float]) -> float:
54
+ dot = sum(x * y for x, y in zip(a, b, strict=True))
55
+ norm_a = math.sqrt(sum(x * x for x in a))
56
+ norm_b = math.sqrt(sum(y * y for y in b))
57
+ if norm_a == 0.0 or norm_b == 0.0:
58
+ return 0.0
59
+ return dot / (norm_a * norm_b)
@@ -0,0 +1,75 @@
1
+ from __future__ import annotations
2
+
3
+ from collections.abc import Callable
4
+
5
+ from . import dedup as dedup_module
6
+ from . import prompts
7
+ from .backends.base import Backend
8
+ from .selection import Selector
9
+ from .types import ClassifyResult, Label, Pair
10
+
11
+ _DEFAULT_MAX_ATTEMPTS_MULTIPLIER = 3
12
+ _MIN_DEFAULT_MAX_ATTEMPTS = 3
13
+ _TOKENS_PER_TEXT_ESTIMATE = 64
14
+
15
+
16
+ def run_generation_loop(
17
+ *,
18
+ backend: Backend,
19
+ pairs: list[Pair],
20
+ selector: Selector,
21
+ exemplar_count: int,
22
+ classify: Callable[[str], ClassifyResult],
23
+ n: int,
24
+ side: Label,
25
+ temperature: float,
26
+ dedup: bool,
27
+ verify: bool,
28
+ min_confidence: float,
29
+ max_attempts: int | None,
30
+ ) -> list[str]:
31
+ """Generate up to ``n`` new ``side`` texts, retrying failed attempts.
32
+
33
+ On each attempt, asks the backend for however many texts are still
34
+ needed, then accepts candidates that pass dedup (against both the seed
35
+ pairs and texts already accepted) and verify (self-classifies as
36
+ ``side`` with at least ``min_confidence``). Stops after ``max_attempts``
37
+ attempts even if ``n`` texts were never reached; the caller is
38
+ responsible for warning about a shortfall.
39
+ """
40
+ if n <= 0:
41
+ return []
42
+
43
+ attempts = (
44
+ max_attempts
45
+ if max_attempts is not None
46
+ else max(n * _DEFAULT_MAX_ATTEMPTS_MULTIPLIER, _MIN_DEFAULT_MAX_ATTEMPTS)
47
+ )
48
+ seed_texts = [text for pair in pairs for text in pair] if dedup else []
49
+ accepted: list[str] = []
50
+
51
+ for _attempt in range(attempts):
52
+ remaining = n - len(accepted)
53
+ if remaining <= 0:
54
+ break
55
+
56
+ exemplars = selector.select(pairs, exemplar_count, query=None)
57
+ prompt = prompts.build_generate_prompt(exemplars, side, remaining)
58
+ raw = backend.complete(
59
+ prompt,
60
+ temperature=temperature,
61
+ max_tokens=remaining * _TOKENS_PER_TEXT_ESTIMATE,
62
+ )
63
+
64
+ for candidate in prompts.parse_generated_texts(raw):
65
+ if len(accepted) >= n:
66
+ break
67
+ if dedup and dedup_module.is_near_duplicate(candidate, seed_texts + accepted):
68
+ continue
69
+ if verify:
70
+ result = classify(candidate)
71
+ if result.label != side or result.confidence < min_confidence:
72
+ continue
73
+ accepted.append(candidate)
74
+
75
+ return accepted
@@ -0,0 +1,124 @@
1
+ from __future__ import annotations
2
+
3
+ import re
4
+
5
+ from .types import NEGATIVE, POSITIVE, Label, Pair
6
+
7
+ CLASSIFY_LABELS: list[Label] = [POSITIVE, NEGATIVE]
8
+ """Fixed label order used everywhere a classify score list is produced."""
9
+
10
+ CLASSIFY_LABEL_CHOICES: dict[Label, str] = {
11
+ POSITIVE: " positive",
12
+ NEGATIVE: " negative",
13
+ }
14
+ """Text appended after the prompt to score each label (leading space matters
15
+ for most tokenizers: it keeps the label as its own token(s) rather than
16
+ merging with the preceding colon)."""
17
+
18
+ _EXEMPLAR_TEMPLATE = "Text: {text}\nLabel: {label}\n\n"
19
+ _QUERY_TEMPLATE = "Text: {text}\nLabel:"
20
+
21
+
22
+ def build_exemplar_prefix(pairs: list[Pair]) -> str:
23
+ """Render the few-shot block for ``pairs``. This is the part of the
24
+ classification prompt shared across every query, so it should come
25
+ first (see ``build_query_suffix``) to keep it a stable KV-cache prefix.
26
+ """
27
+ return "".join(
28
+ _EXEMPLAR_TEMPLATE.format(text=pos, label=POSITIVE)
29
+ + _EXEMPLAR_TEMPLATE.format(text=neg, label=NEGATIVE)
30
+ for pos, neg in pairs
31
+ )
32
+
33
+
34
+ def build_query_suffix(text: str) -> str:
35
+ """Render the query-specific tail of the classification prompt."""
36
+ return _QUERY_TEMPLATE.format(text=text)
37
+
38
+
39
+ def build_classify_prompt(pairs: list[Pair], query: str) -> str:
40
+ """Render the full classification prompt for a single ``query``."""
41
+ return build_exemplar_prefix(pairs) + build_query_suffix(query)
42
+
43
+
44
+ def classify_choice_texts() -> list[str]:
45
+ """``CLASSIFY_LABEL_CHOICES`` values in ``CLASSIFY_LABELS`` order."""
46
+ return [CLASSIFY_LABEL_CHOICES[label] for label in CLASSIFY_LABELS]
47
+
48
+
49
+ CLASSIFY_COMPLETION_MAX_TOKENS = 8
50
+ """max_tokens for the parse-based classify path: just enough for a bare label."""
51
+
52
+
53
+ def parse_classify_label(raw: str) -> tuple[Label, float]:
54
+ """Parse a freeform completion of ``build_classify_prompt`` into a label.
55
+
56
+ Used for backends without ``score_choices`` (e.g. ``AnthropicBackend``):
57
+ rather than comparing log-likelihoods, we ask the model to name the
58
+ label and look for it in the response text.
59
+
60
+ Confidence here is *not* a calibrated probability, just a fixed
61
+ convention value: ``1.0`` when exactly one label word was found, ``0.5``
62
+ when the response was ambiguous (both or neither label word present) and
63
+ we fell back to a default guess.
64
+ """
65
+ lowered = raw.lower()
66
+ positive_at = lowered.find(POSITIVE)
67
+ negative_at = lowered.find(NEGATIVE)
68
+ found_positive = positive_at != -1
69
+ found_negative = negative_at != -1
70
+
71
+ if found_positive and not found_negative:
72
+ return POSITIVE, 1.0
73
+ if found_negative and not found_positive:
74
+ return NEGATIVE, 1.0
75
+ if found_positive and found_negative:
76
+ return (POSITIVE if positive_at < negative_at else NEGATIVE), 0.5
77
+ return POSITIVE, 0.5 # neither label word found; arbitrary fallback guess
78
+
79
+
80
+ _GENERATE_PREAMBLE = (
81
+ 'Below are example pairs contrasting "positive" and "negative" texts; '
82
+ "their exact meaning is defined only by these examples.\n\n"
83
+ )
84
+ _GENERATE_EXEMPLAR_LINE = "- {label}: {text}\n"
85
+ _GENERATE_INSTRUCTION = (
86
+ "\nWrite {n} new, diverse {side} texts. Do not copy or closely paraphrase "
87
+ "the examples above. Output one text per line, numbered starting at 1, "
88
+ "with no other commentary.\n\n1. "
89
+ )
90
+
91
+ _NUMBERED_LINE = re.compile(r"^\s*\d+[.)]\s*(.+)$")
92
+
93
+
94
+ def build_generate_prompt(pairs: list[Pair], side: Label, n: int) -> str:
95
+ """Render the generation prompt asking for ``n`` new ``side`` texts.
96
+
97
+ Ends primed with ``"1. "`` so a plain text-completion backend continues
98
+ directly into the first generated item.
99
+ """
100
+ exemplars = "".join(
101
+ _GENERATE_EXEMPLAR_LINE.format(label=POSITIVE, text=pos)
102
+ + _GENERATE_EXEMPLAR_LINE.format(label=NEGATIVE, text=neg)
103
+ for pos, neg in pairs
104
+ )
105
+ return _GENERATE_PREAMBLE + exemplars + _GENERATE_INSTRUCTION.format(n=n, side=side)
106
+
107
+
108
+ def parse_generated_texts(raw: str) -> list[str]:
109
+ """Parse a numbered-list completion (continuing after a primed ``"1. "``)
110
+ into a flat list of generated texts.
111
+
112
+ This is a best-effort heuristic, not a strict parser: lines that aren't
113
+ numbered (e.g. a wrapped continuation of the previous item) are kept
114
+ as-is rather than dropped, since backends aren't guaranteed to always
115
+ produce clean single-line items.
116
+ """
117
+ lines = [line for line in raw.strip().splitlines() if line.strip()]
118
+ if not lines:
119
+ return []
120
+ texts = [lines[0].strip()]
121
+ for line in lines[1:]:
122
+ match = _NUMBERED_LINE.match(line)
123
+ texts.append(match.group(1).strip() if match else line.strip())
124
+ return [text for text in texts if text]
File without changes
@@ -0,0 +1,110 @@
1
+ from __future__ import annotations
2
+
3
+ import random
4
+ from collections.abc import Callable
5
+ from typing import Protocol, runtime_checkable
6
+
7
+ from . import dedup
8
+ from .types import Pair
9
+
10
+ SimilarityFn = Callable[[str, str], float]
11
+
12
+
13
+ @runtime_checkable
14
+ class Selector(Protocol):
15
+ """Chooses which fitted pairs to embed in a prompt.
16
+
17
+ ``query`` is the text being classified, for selectors that pick pairs
18
+ relevant to it (e.g. ``NearestSelector``); selectors that don't use it
19
+ (e.g. ``RandomSelector``) simply ignore the argument.
20
+ """
21
+
22
+ def select(self, pairs: list[Pair], k: int, query: str | None = None) -> list[Pair]: ...
23
+
24
+
25
+ class RandomSelector:
26
+ """Selects a uniform random subset of ``pairs``, ignoring ``query``."""
27
+
28
+ def __init__(self, seed: int | None = None) -> None:
29
+ self._rng = random.Random(seed)
30
+
31
+ def select(self, pairs: list[Pair], k: int, query: str | None = None) -> list[Pair]:
32
+ if k >= len(pairs):
33
+ return list(pairs)
34
+ return self._rng.sample(pairs, k)
35
+
36
+
37
+ class NearestSelector:
38
+ """Selects the ``k`` pairs whose positive or negative text is most
39
+ similar to ``query``.
40
+
41
+ ``similarity_fn`` defaults to ``dedup.similarity`` (dependency-free
42
+ character n-grams); pass e.g. ``pntx.embeddings.cosine_similarity_fn()``
43
+ for semantic similarity instead.
44
+
45
+ ``query=None`` (e.g. when called from a batch context that has no single
46
+ query to be "nearest" to) has nothing to rank by, so this falls back to
47
+ the first ``k`` pairs in ``pairs`` order; use ``RandomSelector`` or
48
+ ``DiversitySelector`` if that's not the fallback you want.
49
+ """
50
+
51
+ def __init__(self, similarity_fn: SimilarityFn = dedup.similarity) -> None:
52
+ self.similarity_fn = similarity_fn
53
+
54
+ def select(self, pairs: list[Pair], k: int, query: str | None = None) -> list[Pair]:
55
+ if k >= len(pairs):
56
+ return list(pairs)
57
+ if k <= 0:
58
+ return []
59
+ if query is None:
60
+ return list(pairs[:k])
61
+
62
+ def pair_similarity(pair: Pair) -> float:
63
+ pos, neg = pair
64
+ return max(self.similarity_fn(query, pos), self.similarity_fn(query, neg))
65
+
66
+ ranked = sorted(range(len(pairs)), key=lambda i: pair_similarity(pairs[i]), reverse=True)
67
+ return [pairs[i] for i in sorted(ranked[:k])]
68
+
69
+
70
+ class DiversitySelector:
71
+ """Greedily selects ``k`` pairs that are maximally different from each
72
+ other, ignoring ``query``.
73
+
74
+ Starts from ``pairs[0]`` and repeatedly adds whichever remaining pair
75
+ has the lowest similarity to its most-similar already-selected pair
76
+ (a farthest-point / greedy diversity heuristic), representing each pair
77
+ as its concatenated positive and negative text.
78
+
79
+ ``similarity_fn`` defaults to ``dedup.similarity`` (dependency-free
80
+ character n-grams); pass e.g. ``pntx.embeddings.cosine_similarity_fn()``
81
+ for semantic similarity instead.
82
+ """
83
+
84
+ def __init__(self, similarity_fn: SimilarityFn = dedup.similarity) -> None:
85
+ self.similarity_fn = similarity_fn
86
+
87
+ def select(self, pairs: list[Pair], k: int, query: str | None = None) -> list[Pair]:
88
+ if k >= len(pairs):
89
+ return list(pairs)
90
+ if k <= 0:
91
+ return []
92
+
93
+ representations = [f"{pos} {neg}" for pos, neg in pairs]
94
+ selected = [0]
95
+ while len(selected) < k:
96
+ best_index = -1
97
+ best_distance = -1.0
98
+ for i in range(len(pairs)):
99
+ if i in selected:
100
+ continue
101
+ similarity_to_selected = max(
102
+ self.similarity_fn(representations[i], representations[j]) for j in selected
103
+ )
104
+ distance = 1.0 - similarity_to_selected
105
+ if distance > best_distance:
106
+ best_distance = distance
107
+ best_index = i
108
+ selected.append(best_index)
109
+
110
+ return [pairs[i] for i in selected]
@@ -0,0 +1,34 @@
1
+ from __future__ import annotations
2
+
3
+ from dataclasses import dataclass
4
+ from typing import Literal
5
+
6
+ Label = Literal["positive", "negative"]
7
+
8
+ POSITIVE: Label = "positive"
9
+ NEGATIVE: Label = "negative"
10
+
11
+ Pair = tuple[str, str]
12
+ """A single (positive, negative) example pair."""
13
+
14
+
15
+ @dataclass(frozen=True, eq=False)
16
+ class ClassifyResult:
17
+ """Result of classifying a single text.
18
+
19
+ ``result == "positive"`` compares against ``label`` directly, so callers
20
+ don't need to write ``result.label == "positive"``.
21
+ """
22
+
23
+ label: Label
24
+ confidence: float
25
+
26
+ def __eq__(self, other: object) -> bool:
27
+ if isinstance(other, str):
28
+ return self.label == other
29
+ if isinstance(other, ClassifyResult):
30
+ return self.label == other.label and self.confidence == other.confidence
31
+ return NotImplemented
32
+
33
+ def __hash__(self) -> int:
34
+ return hash((self.label, self.confidence))