social-jev 0.1.0a1__tar.gz

social_jev-0.1.0a1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Daniel Göttlich
+
+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.

social_jev-0.1.0a1/PKG-INFO

@@ -0,0 +1,69 @@
+Metadata-Version: 2.4
+Name: social-jev
+Version: 0.1.0a1
+Summary: Jev-like classification foundation models for the social sciences
+License-Expression: MIT
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.13
+Requires-Dist: transformers<6,>=5.17
+Requires-Dist: pydantic>=2
+Dynamic: license-file
+
+# SocialJev
+
+SocialJev is a Jev-like classification foundation model for the social sciences,
+inspired by TypeSafe Jev. It is trained to combine the flexibility of modern
+language models pretrained on trillions of tokens of text with the speed of
+traditional lightweight NLP methods. Define your task through a simple typed
+interface and SocialJev returns typed outputs.
+
+**Note**
+SocialJev is under active development. The API and model interfaces may change,
+including breaking changes between releases.
+
+## Getting started
+
+```bash
+uv pip install social-jev
+```
+
+Load a model from Hugging Face and classify a batch of documents:
+
+```python
+from social_jev import Query, Request, SocialJev
+
+model = SocialJev.from_pretrained("social-jev/social-jev-0.6B-v1")
+
+query = Query(
+    text="Which policy area does this document concern?",
+    options=["Transport", "Education", "Healthcare"],
+)
+
+results = model.query(
+    [
+        Request(state="The council expanded the bus network.", queries=[query]),
+        Request(state="The government recruited more teachers.", queries=[query]),
+    ],
+    batch_size=16,
+)
+
+print(results[0][0].model_dump())
+```
+
+Example result:
+
+```python
+{
+    "choice": "Transport",
+    "probabilities": {
+        "Transport": 0.91,
+        "Education": 0.03,
+        "Healthcare": 0.06,
+    },
+}
+```
+
+Results follow document order, with one result per query. The probabilities above
+are illustrative.

social_jev-0.1.0a1/README.md

@@ -0,0 +1,56 @@
+# SocialJev
+
+SocialJev is a Jev-like classification foundation model for the social sciences,
+inspired by TypeSafe Jev. It is trained to combine the flexibility of modern
+language models pretrained on trillions of tokens of text with the speed of
+traditional lightweight NLP methods. Define your task through a simple typed
+interface and SocialJev returns typed outputs.
+
+**Note**
+SocialJev is under active development. The API and model interfaces may change,
+including breaking changes between releases.
+
+## Getting started
+
+```bash
+uv pip install social-jev
+```
+
+Load a model from Hugging Face and classify a batch of documents:
+
+```python
+from social_jev import Query, Request, SocialJev
+
+model = SocialJev.from_pretrained("social-jev/social-jev-0.6B-v1")
+
+query = Query(
+    text="Which policy area does this document concern?",
+    options=["Transport", "Education", "Healthcare"],
+)
+
+results = model.query(
+    [
+        Request(state="The council expanded the bus network.", queries=[query]),
+        Request(state="The government recruited more teachers.", queries=[query]),
+    ],
+    batch_size=16,
+)
+
+print(results[0][0].model_dump())
+```
+
+Example result:
+
+```python
+{
+    "choice": "Transport",
+    "probabilities": {
+        "Transport": 0.91,
+        "Education": 0.03,
+        "Healthcare": 0.06,
+    },
+}
+```
+
+Results follow document order, with one result per query. The probabilities above
+are illustrative.

social_jev-0.1.0a1/pyproject.toml

@@ -0,0 +1,114 @@
+[project]
+name = "social-jev"
+version = "0.1.0a1"
+description = "Jev-like classification foundation models for the social sciences"
+license = "MIT"
+license-files = ["LICENSE"]
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = ["torch>=2.13", "transformers>=5.17,<6", "pydantic>=2"]
+
+[dependency-groups]
+dev = [
+    "actionlint-py>=1.7.12.24",
+    "deptry>=0.25.1",
+    "mutmut>=3.7.0",
+    "prek>=0.5.2",
+    "pytest>=9.1.1",
+    "pytest-cov>=7.1.0",
+    "pytest-randomly>=5.0.0",
+    "ruff>=0.16.7",
+    "shellcheck-py>=0.11.0.1",
+    "ty>=0.0.80",
+    "vulture>=2.16",
+    "zizmor>=1.30.1",
+]
+
+[project.scripts]
+
+[build-system]
+requires = ["setuptools>=84.0.0"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.package-dir]
+"" = "src"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = [
+    "--strict-config",
+    "--strict-markers",
+    "--cov=src",
+    "--cov=ci",
+    "--cov-branch",
+    "--cov-report=term-missing",
+    "--cov-fail-under=80",
+]
+xfail_strict = true
+filterwarnings = ["error"]
+
+
+[tool.ruff]
+line-length = 100
+
+[tool.ruff.lint]
+extend-select = [
+    "A",      # Prevent shadowing Python built-ins
+    "ARG",    # Unused function arguments
+    "B",      # Bug-prone Python patterns
+    "C4",     # Better comprehensions
+    "C90",    # McCabe complexity
+    "E",      # pycodestyle errors
+    "ERA",    # Commented-out code
+    "F",      # Pyflakes, including stricter import checks
+    "N",      # PEP 8 naming conventions
+    "PERF",   # Performance anti-patterns
+    "PGH003", # Require specific type: ignore codes
+    "PGH004", # Require specific noqa codes
+    "PL",     # Pylint-derived code quality checks
+    "PT",     # Pytest-specific checks
+    "PTH",    # Prefer pathlib
+    "RET",    # Cleaner return control flow
+    "RUF",    # Ruff-specific correctness/style checks
+    "SIM",    # Simplify unnecessarily complex constructs
+    "TC",     # Type-checking import hygiene
+    "TRY",    # Exception-handling structure
+    "UP013",  # Modern TypedDict syntax
+    "UP015",  # Remove redundant open() modes
+    "UP042",  # Prefer StrEnum
+    "W",      # pycodestyle warnings
+]
+
+# These conflict with, or are redundant with, `ruff format`.
+ignore = [
+    "E111",
+    "E114",
+    "E117",
+    "W191",
+]
+
+[tool.ruff.lint.mccabe]
+max-complexity = 6
+
+[tool.ruff.lint.per-file-ignores]
+"tests/**/*.py" = [
+    "PLR2004", # Literal values are normal and useful in assertions
+]
+
+[tool.check-tests]
+fail = ["TST001", "TST002", "TST003", "TST004", "TST005"]
+warn = []
+
+[tool.mutmut]
+source_paths = ["src", "ci"]
+# mutation subsets must not fail whole-project coverage or vary their execution order
+pytest_add_cli_args = ["--no-cov", "-p", "no:randomly"]
+# protect explicit local incremental runs when dependencies change
+on_dependency_change = "rerun"
+
+[tool.setuptools.package-data]
+social_jev = ["py.typed", "templates/*.jinja"]

social_jev-0.1.0a1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

social_jev-0.1.0a1/src/social_jev/__init__.py

@@ -0,0 +1,5 @@
+from .data import Processor
+from .model import SocialJev
+from .types import Example, Query, Request, Result
+
+__all__ = ["Example", "Processor", "Query", "Request", "Result", "SocialJev"]

social_jev-0.1.0a1/src/social_jev/attention.py

@@ -0,0 +1,137 @@
+"""Packed visibility masks and logical token positions."""
+
+from typing import TYPE_CHECKING
+
+import torch
+from torch.nn.attention.flex_attention import create_block_mask
+
+if TYPE_CHECKING:
+    from .data import Batch
+
+
+def block_mask(ranges):
+    """Build a FlexAttention block mask from allowed token ranges."""
+    batch, length = ranges.shape[:2]
+
+    def allowed(b, _h, q, k):
+        """Return whether a query token may attend to a key token."""
+        valid = (q < length) & (k < length)
+        bounds = ranges[b, q.clamp(max=length - 1)]
+        return valid & (
+            ((k >= bounds[0, 0]) & (k < bounds[0, 1]))
+            | ((k >= bounds[1, 0]) & (k < bounds[1, 1]))
+            | ((k >= bounds[2, 0]) & (k < bounds[2, 1]))
+        )
+
+    return create_block_mask(allowed, batch, None, length, length, device=ranges.device)
+
+
+def message_spans(boundaries, roles, answer_length):
+    """Expose the appended answer prefix without changing stored message roles."""
+    start = 0
+    for index, (boundary, role) in enumerate(zip(boundaries, roles, strict=True)):
+        end = boundary + 1
+        final_option = role == "option" and (
+            index + 1 == len(roles) or roles[index + 1] != "option"
+        )
+        if answer_length and final_option:
+            yield start, end - answer_length, role
+            yield end - answer_length, end, "answer"
+        else:
+            yield start, end, role
+        start = end
+
+
+class Positions:
+    """Compute logical token positions for isolated queries and options.
+
+    Args:
+        batch: Padded tokens and their message boundaries.
+        answer_length: Number of tokens in each appended answer prefix.
+
+    Examples:
+        >>> positions = Positions(batch, answer_length=2)  # doctest: +SKIP
+        >>> positions.value.shape == batch.input_ids.shape  # doctest: +SKIP
+        True
+    """
+
+    def __init__(  # noqa: C901 -- One branch per message role.
+        self, batch: "Batch", *, answer_length: int = 0
+    ):
+        """Build logical positions for all tokens in the batch."""
+        positions = []
+        for boundaries, roles in zip(batch.message_boundaries, batch.roles, strict=True):
+            sequence = []
+            state_length = query_length = longest_example = longest_option = 0
+            for start, end, role in message_spans(boundaries, roles, answer_length):
+                length = end - start
+                if role == "state":
+                    state_length = length
+                    query_length = longest_example = longest_option = 0
+                    offset = 0
+                elif role == "query":
+                    query_length = length
+                    longest_example = longest_option = 0
+                    offset = state_length
+                elif role == "example":
+                    offset = state_length + query_length
+                    longest_example = max(longest_example, length)
+                elif role == "answer":
+                    offset = state_length + query_length + longest_example + longest_option
+                else:  # option; Batch validates supported roles.
+                    longest_option = max(longest_option, length)
+                    offset = state_length + query_length + longest_example
+                sequence.extend(range(offset, offset + length))
+            padding = batch.input_ids.shape[1] - len(sequence)
+            positions.append(sequence + [0] * padding)
+        self.value = torch.tensor(positions, dtype=torch.long, device=batch.input_ids.device)
+
+
+class AttentionMask:
+    """Compute causal visibility within documents, queries, and candidate answers.
+
+    Args:
+        batch: Padded tokens and their message boundaries.
+        answer_length: Number of tokens in each appended answer prefix.
+
+    Examples:
+        >>> mask = AttentionMask(batch, answer_length=2)  # doctest: +SKIP
+        >>> block_mask = mask.value  # doctest: +SKIP
+    """
+
+    def __init__(  # noqa: C901 -- One branch per message role.
+        self, batch: "Batch", *, answer_length: int = 0
+    ):
+        """Build the visibility mask for all tokens in the batch."""
+        batch_size, length = batch.input_ids.shape
+        tokens = torch.arange(length)
+        ranges = torch.zeros((batch_size, length, 3, 2), dtype=torch.int32)
+        ranges[:, :, 0, 0] = tokens
+        ranges[:, :, 0, 1] = tokens + 1
+        for index, (boundaries, roles) in enumerate(
+            zip(batch.message_boundaries, batch.roles, strict=True)
+        ):
+            state = query = (0, 0)
+            context_end = 0
+            for start, end, role in message_spans(boundaries, roles, answer_length):
+                ranges[index, start:end, 0, 0] = start
+                ranges[index, start:end, 0, 1] = torch.arange(start + 1, end + 1)
+                if role == "state":
+                    state = (start, end)
+                    query = (0, 0)
+                    context_end = 0
+                elif role == "query":
+                    query = (start, end)
+                    context_end = end
+                    ranges[index, start:end, 1] = torch.tensor(state)
+                elif role == "example":
+                    ranges[index, start:end, 1] = torch.tensor(state)
+                    ranges[index, start:end, 2] = torch.tensor(query)
+                    context_end = end
+                elif role == "answer":
+                    ranges[index, start:end, 1] = torch.tensor(state)
+                    ranges[index, start:end, 2] = torch.tensor([query[0], start])
+                else:
+                    ranges[index, start:end, 1] = torch.tensor(state)
+                    ranges[index, start:end, 2] = torch.tensor([query[0], context_end])
+        self.value = block_mask(ranges.to(batch.input_ids.device))

social_jev-0.1.0a1/src/social_jev/config.py

@@ -0,0 +1,60 @@
+"""Hugging Face configuration for SocialJev LM-head inference."""
+
+from string import ascii_uppercase
+from typing import ClassVar
+
+from transformers import AutoConfig, PretrainedConfig
+
+
+class SocialJevConfig(PretrainedConfig):
+    """Hugging Face configuration for an LM-head SocialJev checkpoint.
+
+    Args:
+        backbone_config: Backbone configuration object or serialized configuration.
+        decision_head: Saved head configuration; only lm_head is supported.
+        option_codes: Single-token answer codes assigned to candidate options.
+        role_start: Special token starting a message.
+        role_end: Special token ending a message.
+        **kwargs: Additional Hugging Face configuration settings.
+
+    Examples:
+        >>> config = SocialJevConfig(option_codes="ABC")
+        >>> config.decision_head["type"]
+        'lm_head'
+    """
+
+    model_type = "social_jev"
+    """Hugging Face identifier for the SocialJev architecture."""
+    sub_configs: ClassVar = {"backbone_config": AutoConfig}
+    """Configuration types used by the nested backbone."""
+
+    def __init__(
+        self,
+        backbone_config=None,
+        decision_head=None,
+        option_codes=ascii_uppercase,
+        role_start="<|im_start|>",
+        role_end="<|im_end|>",
+        **kwargs,
+    ):
+        """Initialize the backbone, message delimiters, and answer codes."""
+        kwargs.pop("tie_word_embeddings", None)
+        super().__init__(**kwargs)
+        self.tie_word_embeddings = False
+        if backbone_config is None:
+            backbone_config = {"model_type": "qwen3"}
+        if isinstance(backbone_config, dict):
+            values = dict(backbone_config)
+            backbone_config = AutoConfig.for_model(values.pop("model_type"), **values)
+        self.backbone_config = backbone_config
+        self.backbone_config.tie_word_embeddings = False
+        if decision_head is not None and decision_head.get("type") != "lm_head":
+            message = "Only lm_head checkpoints are supported"
+            raise ValueError(message)
+        self.decision_head = {"type": "lm_head", "dims": [], "activation": "gelu"}
+        self.option_codes = option_codes
+        self.role_start = role_start
+        self.role_end = role_end
+
+
+SocialJevConfig.register_for_auto_class()

social_jev-0.1.0a1/src/social_jev/data.py

@@ -0,0 +1,362 @@
+"""Pull-based inference preparation, exposed through Processor."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, replace
+from importlib.resources import files
+from itertools import islice
+from typing import TYPE_CHECKING, Any, cast
+
+import torch
+
+from .types import Request
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedTokenizerBase
+
+    from .config import SocialJevConfig
+
+
+@dataclass
+class Sample:
+    """Formatted request messages before tokenization."""
+
+    messages: list[dict[str, Any]]
+    """Role-tagged messages in request order."""
+
+    @classmethod
+    def from_request(cls, request: Request) -> Sample:
+        """Format a typed request into role-tagged messages."""
+        messages: list[dict[str, Any]] = [{"role": "state", "content": request.state}]
+        for query in request.queries:
+            messages.append({"role": "query", "content": query.text})
+            messages.extend(
+                {"role": "example", "content": example.model_dump()} for example in query.examples
+            )
+            messages.extend({"role": "option", "content": option} for option in query.options)
+        return cls(messages)
+
+
+@dataclass
+class TokenizedSample:
+    """Tokenized request with message boundaries and roles."""
+
+    ids: list[int]
+    """Token IDs in message order."""
+    message_boundaries: list[int]
+    """Inclusive final token index of each message."""
+    roles: list[str]
+    """Message roles aligned with message boundaries."""
+
+    @classmethod
+    def from_sample(cls, sample: Sample, processor: Processor) -> TokenizedSample:
+        """Tokenize formatted messages and record their boundaries and roles."""
+        config = processor.config
+        ids = cast(
+            "list[int]",
+            processor.tokenizer.apply_chat_template(
+                sample.messages,
+                role_start=config.role_start,
+                role_end=config.role_end,
+                option_codes=config.option_codes,
+                tokenize=True,
+                return_dict=False,
+                add_generation_prompt=False,
+            ),
+        )
+        roles = [message["role"] for message in sample.messages]
+        closings = iter(i for i, token in enumerate(ids) if token == processor.end_id)
+        boundaries = []
+        for index, role in enumerate(roles):
+            if role == "example":
+                next(closings)
+                next(closings)
+            boundary = next(closings)
+            if role == "option" and (index + 1 == len(roles) or roles[index + 1] != "option"):
+                end = boundary + 1 + len(processor.answer_ids)
+                if ids[boundary + 1 : end] != processor.answer_ids:
+                    message = "Tokenizer template must append the LM-head answer prefix"
+                    raise ValueError(message)
+                boundary = end - 1
+            boundaries.append(boundary)
+        if len(ids) > config.backbone_config.max_position_embeddings:
+            message = "Request exceeds model context capacity"
+            raise ValueError(message)
+        return cls(ids, boundaries, roles)
+
+    def pad(self, length: int, pad_id: int) -> TokenizedSample:
+        """Return a copy padded to the requested token length."""
+        return replace(self, ids=self.ids + [pad_id] * (length - len(self.ids)))
+
+
+@dataclass
+class Batch:
+    """Padded CPU token tensors and aligned message metadata."""
+
+    input_ids: torch.Tensor
+    """Padded token IDs with shape (documents, tokens)."""
+    message_boundaries: list[list[int]]
+    """Inclusive message end indices for each document."""
+    roles: list[list[str]]
+    """Message roles for each document, aligned with its boundaries."""
+
+    @classmethod
+    def from_samples(cls, samples: list[TokenizedSample]) -> Batch:
+        """Collate padded samples into a CPU batch."""
+        return cls(
+            torch.tensor([sample.ids for sample in samples], dtype=torch.long),
+            [sample.message_boundaries for sample in samples],
+            [sample.roles for sample in samples],
+        )
+
+    def to(self, device: torch.device | str) -> Batch:
+        """Return a copy with token tensors transferred to the requested device."""
+        return replace(self, input_ids=self.input_ids.to(device))
+
+    @property
+    def option_positions(self) -> torch.Tensor:
+        """Return the batch row and final token index of every option message."""
+        return torch.tensor(
+            [
+                (row, end)
+                for row, (boundaries, roles) in enumerate(
+                    zip(self.message_boundaries, self.roles, strict=True)
+                )
+                for end, role in zip(boundaries, roles, strict=True)
+                if role == "option"
+            ],
+            dtype=torch.long,
+            device=self.input_ids.device,
+        )
+
+    @property
+    def option_indices(self) -> torch.Tensor:
+        """Group option indices by query, using -1 for unavailable candidates."""
+        groups: list[list[int]] = []
+        index = 0
+        for roles in self.roles:
+            for role in roles:
+                if role == "query":
+                    groups.append([])
+                elif role == "option":
+                    groups[-1].append(index)
+                    index += 1
+        width = max(map(len, groups))
+        return torch.tensor(
+            [group + [-1] * (width - len(group)) for group in groups],
+            dtype=torch.long,
+            device=self.input_ids.device,
+        )
+
+
+class Step[T](ABC):
+    """Define the pull-based interface shared by inference pipeline stages.
+
+    Examples:
+        >>> source = RequestStep([])
+        >>> isinstance(source, Step)
+        True
+        >>> source.next() is None
+        True
+    """
+
+    @abstractmethod
+    def next(self) -> T | None:
+        """Pull the next item, or return None when the upstream source is exhausted."""
+        ...
+
+
+class RequestStep(Step[Request]):
+    """Read typed requests in their original order.
+
+    Args:
+        requests: Documents and queries to pass through the pipeline.
+
+    Examples:
+        >>> source = RequestStep([])
+        >>> source.next() is None
+        True
+    """
+
+    def __init__(self, requests: list[Request]):
+        """Initialize the ordered request source."""
+        self.requests = iter(requests)
+
+    def next(self) -> Request | None:
+        """Pull the next item, or return None when the upstream source is exhausted."""
+        return next(self.requests, None)
+
+
+class FormatStep(Step[Sample]):
+    """Orchestrate conversion of requests into formatted samples.
+
+    Args:
+        parent: Source of typed requests.
+
+    Examples:
+        >>> formatted = FormatStep(RequestStep([]))
+        >>> formatted.next() is None
+        True
+    """
+
+    def __init__(self, parent: Step[Request]):
+        """Attach the request source to the formatting stage."""
+        self.parent = parent
+
+    def next(self) -> Sample | None:
+        """Pull the next item, or return None when the upstream source is exhausted."""
+        request = self.parent.next()
+        if request is None:
+            return None
+        return Sample.from_request(request)
+
+
+class TokenizeStep(Step[TokenizedSample]):
+    """Orchestrate tokenization of formatted samples.
+
+    Args:
+        parent: Source of formatted samples.
+        processor: Tokenizer and checkpoint settings used for tokenization.
+
+    Examples:
+        >>> formatted = FormatStep(RequestStep([]))
+        >>> tokenized = TokenizeStep(formatted, processor)  # doctest: +SKIP
+        >>> tokenized.next() is None  # doctest: +SKIP
+        True
+    """
+
+    def __init__(self, parent: Step[Sample], processor: Processor):
+        """Attach the sample source and tokenizer settings."""
+        self.parent = parent
+        self.processor = processor
+
+    def next(self) -> TokenizedSample | None:
+        """Pull the next item, or return None when the upstream source is exhausted."""
+        sample = self.parent.next()
+        if sample is None:
+            return None
+        return TokenizedSample.from_sample(sample, self.processor)
+
+
+class PadStep(Step[list[TokenizedSample]]):
+    """Collect requests and orchestrate padding to the group's longest sequence.
+
+    Args:
+        parent: Source of tokenized samples.
+        batch_size: Maximum number of documents per group.
+        pad_id: Token ID used for padding.
+
+    Examples:
+        >>> padded = PadStep(tokenized, batch_size=16, pad_id=0)  # doctest: +SKIP
+        >>> padded.next() is None  # doctest: +SKIP
+        True
+    """
+
+    def __init__(self, parent: Step[TokenizedSample], batch_size: int, pad_id: int):
+        """Configure the upstream source, batch size, and padding token."""
+        if batch_size < 1:
+            message = "batch_size must be positive"
+            raise ValueError(message)
+        self.parent = parent
+        self.batch_size = batch_size
+        self.pad_id = pad_id
+
+    def next(self) -> list[TokenizedSample] | None:
+        """Pull the next item, or return None when the upstream source is exhausted."""
+        samples = list(islice(iter(self.parent.next, None), self.batch_size))
+        if not samples:
+            return None
+        length = max(len(sample.ids) for sample in samples)
+        return [sample.pad(length, self.pad_id) for sample in samples]
+
+
+class CollateStep(Step[Batch]):
+    """Orchestrate construction of CPU batches from padded sample groups.
+
+    Args:
+        parent: Source of padded sample groups.
+
+    Examples:
+        >>> batches = CollateStep(padded)  # doctest: +SKIP
+        >>> batch = batches.next()  # doctest: +SKIP
+    """
+
+    def __init__(self, parent: Step[list[TokenizedSample]]):
+        """Attach the padded sample source to the collator."""
+        self.parent = parent
+
+    def next(self) -> Batch | None:
+        """Pull the next item, or return None when the upstream source is exhausted."""
+        samples = self.parent.next()
+        if samples is None:
+            return None
+        return Batch.from_samples(samples)
+
+
+class Processor:
+    """Expose the pull-based request preparation chain.
+
+    Args:
+        tokenizer: Tokenizer containing the checkpoint's chat template.
+        config: Model configuration defining message roles and option codes.
+
+    Examples:
+        >>> from social_jev import Query, Request
+        >>> processor = Processor(tokenizer, config)  # doctest: +SKIP
+        >>> request = Request(state="The bus network expanded.", queries=[
+        ...     Query(text="Policy area?", options=["Transport", "Education"])
+        ... ])
+        >>> batch = processor.encode(request)  # doctest: +SKIP
+        >>> batch.input_ids.shape[0]  # doctest: +SKIP
+        1
+    """
+
+    def __init__(self, tokenizer: PreTrainedTokenizerBase, config: SocialJevConfig):
+        """Initialize the tokenizer settings and validate single-token answer codes."""
+        self.tokenizer = tokenizer
+        self.config = config
+        self.end_id = cast("int", tokenizer.convert_tokens_to_ids(config.role_end))
+        self.answer_ids = tokenizer.encode(config.role_start + "answer\n", add_special_tokens=False)
+        self.code_token_ids: list[int] = []
+        for code in config.option_codes:
+            encoded = tokenizer.encode(
+                config.role_start + "answer\n" + code, add_special_tokens=False
+            )
+            if encoded[:-1] != self.answer_ids or encoded[-1] == tokenizer.unk_token_id:
+                message = "Each option code must be one known answer token"
+                raise ValueError(message)
+            self.code_token_ids.append(encoded[-1])
+        if len(set(self.code_token_ids)) != len(self.code_token_ids):
+            message = "Option codes must have distinct token IDs"
+            raise ValueError(message)
+
+    @staticmethod
+    def template() -> str:
+        """Read the bundled LM-head chat template from package resources."""
+        return files("social_jev").joinpath("templates/messages.jinja").read_text()
+
+    def batches(self, requests: list[Request], *, batch_size: int = 16) -> CollateStep:
+        """Assemble a fresh pull-based chain that produces padded CPU batches."""
+        if any(
+            len(query.options) > len(self.code_token_ids)
+            for request in requests
+            for query in request.queries
+        ):
+            message = "Query has more options than available option codes"
+            raise ValueError(message)
+        source = RequestStep(requests)
+        formatted = FormatStep(source)
+        tokenized = TokenizeStep(formatted, self)
+        padded = PadStep(tokenized, batch_size, self.end_id)
+        return CollateStep(padded)
+
+    def encode(self, requests: Request | list[Request]) -> Batch:
+        """Encode all supplied requests into one padded CPU batch."""
+        items = [requests] if isinstance(requests, Request) else requests
+        if not items:
+            message = "encode requires at least one request"
+            raise ValueError(message)
+        batch = self.batches(items, batch_size=len(items)).next()
+        assert batch is not None
+        return batch

social_jev-0.1.0a1/src/social_jev/model.py

@@ -0,0 +1,145 @@
+"""Hugging Face LM-head loading and offline batched inference."""
+
+from typing import overload
+
+import torch
+from torch.nn.attention.flex_attention import create_mask
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    PreTrainedModel,
+    PreTrainedTokenizerBase,
+)
+
+from .attention import AttentionMask, Positions
+from .config import SocialJevConfig
+from .data import Batch, Processor
+from .types import Request, Result
+
+
+class SocialJev(PreTrainedModel):
+    """Classify documents with a Hugging Face LM-head checkpoint.
+
+    Args:
+        config: Backbone and LM-head inference configuration.
+
+    Examples:
+        >>> from social_jev import Query, Request
+        >>> model = SocialJev.from_pretrained("social-jev/social-jev-0.6B-v1")  # doctest: +SKIP
+        >>> request = Request(state="The bus network expanded.", queries=[
+        ...     Query(text="Policy area?", options=["Transport", "Education"])
+        ... ])
+        >>> results = model.query([request], batch_size=16)  # doctest: +SKIP
+        >>> results[0][0].choice  # doctest: +SKIP
+        'Transport'
+    """
+
+    config_class = SocialJevConfig
+    """Hugging Face configuration class used to load this model."""
+    base_model_prefix = "backbone"
+    """Checkpoint prefix for backbone parameters."""
+    _supports_sdpa = True
+    """Declare support for PyTorch scaled dot-product attention."""
+    _supports_flex_attn = True
+    """Declare support for PyTorch FlexAttention."""
+
+    def __init__(self, config: SocialJevConfig):
+        """Initialize the backbone and LM head from the checkpoint configuration."""
+        super().__init__(config)
+        causal = AutoModelForCausalLM.from_config(
+            config.backbone_config,
+            attn_implementation=config._attn_implementation or "sdpa",
+        )
+        self.backbone = causal.base_model
+        self.decision_head = causal.get_output_embeddings()
+        self.post_init()
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        """Load model and tokenizer assets using the same Hugging Face loading options."""
+        loaded = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+        model = loaded[0] if isinstance(loaded, tuple) else loaded
+        assert isinstance(model, cls)
+        tokenizer_options = {
+            key: kwargs[key]
+            for key in (
+                "revision",
+                "cache_dir",
+                "local_files_only",
+                "token",
+                "force_download",
+                "subfolder",
+                "trust_remote_code",
+            )
+            if key in kwargs
+        }
+        tokenizer = AutoTokenizer.from_pretrained(
+            pretrained_model_name_or_path,
+            config=model.config.backbone_config,
+            **tokenizer_options,
+        )
+        assert isinstance(tokenizer, PreTrainedTokenizerBase)
+        model.processor = Processor(tokenizer, model.config)
+        return loaded
+
+    def forward(self, batch: Batch) -> torch.Tensor:
+        """Score candidate answer codes for every query in the batch."""
+        answer_length = len(self.processor.answer_ids)
+        mask = AttentionMask(batch, answer_length=answer_length).value
+        if self.backbone.config._attn_implementation != "flex_attention":
+            size, length = batch.input_ids.shape
+            mask = create_mask(
+                mask.mask_mod, size, 1, length, length, device=str(batch.input_ids.device)
+            )
+        hidden = self.backbone(
+            input_ids=batch.input_ids,
+            position_ids=Positions(batch, answer_length=answer_length).value,
+            attention_mask={"full_attention": mask},
+            use_cache=False,
+        ).last_hidden_state
+        indices = batch.option_indices
+        positions = batch.option_positions[indices.max(dim=-1).values]
+        answers = hidden[positions[:, 0], positions[:, 1]]
+        codes = torch.tensor(self.processor.code_token_ids, device=hidden.device)
+        scores = self.decision_head(answers)[:, codes[: indices.shape[1]]]
+        return scores.masked_fill(indices < 0, -torch.inf)
+
+    @overload
+    def query(self, requests: Request, *, batch_size: int = 16) -> list[Result]:
+        """Classify requests in batches and return results in document and query order."""
+
+    @overload
+    def query(self, requests: list[Request], *, batch_size: int = 16) -> list[list[Result]]:
+        """Classify requests in batches and return results in document and query order."""
+
+    @torch.inference_mode()
+    def query(
+        self, requests: Request | list[Request], *, batch_size: int = 16
+    ) -> list[Result] | list[list[Result]]:
+        """Classify requests in batches and return results in document and query order."""
+        self.eval()
+        single = not isinstance(requests, list)
+        items = [requests] if single else requests
+        chain = self.processor.batches(items, batch_size=batch_size)
+        pending = iter(items)
+        results: list[list[Result]] = []
+        while (batch := chain.next()) is not None:
+            rows = iter(self(batch.to(self.device)))
+            for _ in batch.roles:
+                request = next(pending)
+                answers = []
+                for query in request.queries:
+                    probabilities = next(rows)[: len(query.options)].float().softmax(0).cpu()
+                    answers.append(
+                        Result(
+                            choice=query.options[int(probabilities.argmax())],
+                            probabilities=dict(
+                                zip(query.options, probabilities.tolist(), strict=True)
+                            ),
+                        )
+                    )
+                results.append(answers)
+        return results[0] if single else results
+
+
+SocialJev.register_for_auto_class("AutoModel")

social_jev-0.1.0a1/src/social_jev/templates/messages.jinja

@@ -0,0 +1,24 @@
+{%- set option_codes = option_codes|default('ABCDEFGHIJKLMNOPQRSTUVWXYZ') %}
+{%- set role_start = role_start|default('<|im_start|>') %}
+{%- set role_end = role_end|default('<|im_end|>') %}
+{%- set counter = namespace(option=0) %}
+{%- for message in messages %}
+{%- if message.role == 'query' %}{%- set counter.option = 0 %}{%- endif %}
+{{- role_start + message.role + '\n' }}
+{%- if message.role == 'example' %}
+{{- role_start + 'state\n' + message.content.state + role_end }}
+{{- role_start + 'option\n' + message.content.choice + role_end }}
+{%- elif message.role == 'option' %}
+{%- if counter.option >= option_codes|length %}
+{{- raise_exception('More options than available letter codes') }}
+{%- endif %}
+{{- option_codes[counter.option] + ': ' + message.content }}
+{%- set counter.option = counter.option + 1 %}
+{%- else %}
+{{- message.content }}
+{%- endif %}
+{{- role_end }}
+{%- if message.role == 'option' and (loop.last or messages[loop.index].role != 'option') %}
+{{- role_start + 'answer\n' }}
+{%- endif %}
+{%- endfor %}

social_jev-0.1.0a1/src/social_jev/types.py

@@ -0,0 +1,90 @@
+"""Typed document coding requests and results."""
+
+from pydantic import BaseModel, Field, model_validator
+
+
+class Example(BaseModel):
+    """A labelled document illustrating a candidate answer.
+
+    Args:
+        state: Example document text.
+        choice: Answer matching one of the query's options.
+
+    Examples:
+        >>> example = Example(state="The bus network expanded.", choice="Transport")
+        >>> example.choice
+        'Transport'
+    """
+
+    state: str
+    """Text of the labelled example document."""
+    choice: str
+    """Candidate answer assigned to the example."""
+
+
+class Query(BaseModel):
+    """A classification question with candidate answers and optional examples.
+
+    Args:
+        text: Question describing the classification task.
+        options: Nonempty list of distinct candidate answers.
+        examples: Labelled documents whose choices appear in options.
+
+    Examples:
+        >>> query = Query(text="Policy area?", options=["Transport", "Education"])
+        >>> query.options[0]
+        'Transport'
+    """
+
+    text: str
+    """Question defining the classification task."""
+    options: list[str] = Field(min_length=1)
+    """Distinct candidate answers in their supplied order."""
+    examples: list[Example] = Field(default_factory=list)
+    """Optional labelled examples for this query."""
+
+    @model_validator(mode="after")
+    def validate_options(self):
+        """Require distinct options and example choices drawn from those options."""
+        assert len(self.options) == len(set(self.options))
+        assert all(example.choice in self.options for example in self.examples)
+        return self
+
+
+class Request(BaseModel):
+    """A document and the classification questions to ask about it.
+
+    Args:
+        state: Document text to classify.
+        queries: Nonempty list of questions about the document.
+
+    Examples:
+        >>> query = Query(text="Policy area?", options=["Transport", "Education"])
+        >>> request = Request(state="The bus network expanded.", queries=[query])
+        >>> len(request.queries)
+        1
+    """
+
+    state: str
+    """Document text to classify."""
+    queries: list[Query] = Field(min_length=1)
+    """Questions to answer about the document, in output order."""
+
+
+class Result(BaseModel):
+    """The selected answer and probabilities over a query's candidate answers.
+
+    Args:
+        choice: Candidate answer with the highest probability.
+        probabilities: Candidate answers mapped to their normalized probabilities.
+
+    Examples:
+        >>> result = Result(choice="Transport", probabilities={"Transport": 0.9, "Education": 0.1})
+        >>> result.choice
+        'Transport'
+    """
+
+    choice: str
+    """Candidate answer with the highest probability."""
+    probabilities: dict[str, float]
+    """Probabilities normalized over the supplied candidate answers."""

social_jev-0.1.0a1/src/social_jev.egg-info/PKG-INFO

@@ -0,0 +1,69 @@
+Metadata-Version: 2.4
+Name: social-jev
+Version: 0.1.0a1
+Summary: Jev-like classification foundation models for the social sciences
+License-Expression: MIT
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.13
+Requires-Dist: transformers<6,>=5.17
+Requires-Dist: pydantic>=2
+Dynamic: license-file
+
+# SocialJev
+
+SocialJev is a Jev-like classification foundation model for the social sciences,
+inspired by TypeSafe Jev. It is trained to combine the flexibility of modern
+language models pretrained on trillions of tokens of text with the speed of
+traditional lightweight NLP methods. Define your task through a simple typed
+interface and SocialJev returns typed outputs.
+
+**Note**
+SocialJev is under active development. The API and model interfaces may change,
+including breaking changes between releases.
+
+## Getting started
+
+```bash
+uv pip install social-jev
+```
+
+Load a model from Hugging Face and classify a batch of documents:
+
+```python
+from social_jev import Query, Request, SocialJev
+
+model = SocialJev.from_pretrained("social-jev/social-jev-0.6B-v1")
+
+query = Query(
+    text="Which policy area does this document concern?",
+    options=["Transport", "Education", "Healthcare"],
+)
+
+results = model.query(
+    [
+        Request(state="The council expanded the bus network.", queries=[query]),
+        Request(state="The government recruited more teachers.", queries=[query]),
+    ],
+    batch_size=16,
+)
+
+print(results[0][0].model_dump())
+```
+
+Example result:
+
+```python
+{
+    "choice": "Transport",
+    "probabilities": {
+        "Transport": 0.91,
+        "Education": 0.03,
+        "Healthcare": 0.06,
+    },
+}
+```
+
+Results follow document order, with one result per query. The probabilities above
+are illustrative.

social_jev-0.1.0a1/src/social_jev.egg-info/SOURCES.txt

@@ -0,0 +1,16 @@
+LICENSE
+README.md
+pyproject.toml
+src/social_jev/__init__.py
+src/social_jev/attention.py
+src/social_jev/config.py
+src/social_jev/data.py
+src/social_jev/model.py
+src/social_jev/py.typed
+src/social_jev/types.py
+src/social_jev.egg-info/PKG-INFO
+src/social_jev.egg-info/SOURCES.txt
+src/social_jev.egg-info/dependency_links.txt
+src/social_jev.egg-info/requires.txt
+src/social_jev.egg-info/top_level.txt
+src/social_jev/templates/messages.jinja

social_jev-0.1.0a1/src/social_jev.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

social_jev-0.1.0a1/src/social_jev.egg-info/requires.txt

@@ -0,0 +1,3 @@
+torch>=2.13
+transformers<6,>=5.17
+pydantic>=2

social_jev-0.1.0a1/src/social_jev.egg-info/top_level.txt

@@ -0,0 +1 @@
+social_jev