contextattr 0.1.0__py3-none-any.whl

contextattr/prompt/io.py

@@ -0,0 +1,222 @@
+"""Reading and writing prompts.jsonl, one prompt per line (docs/interface.md §2).
+
+Reading checks every field and names the line and the place of any error. Writing gives each group its units
+explicitly, so a file read back gives equal prompts whatever splitter made the groups.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from collections.abc import Iterable
+
+from .split import KINDS, split
+from .units import Group, Message, Prompt, Response, Unit
+
+_PROMPT_FIELDS = {"id", "messages", "text", "reference", "responses"}
+_MESSAGE_FIELDS = {"role", "content"}
+_UNIT_FIELDS = {"unit", "text", "sep", "tag", "replace", "split"}
+_GROUP_FIELDS = {"group", "units", "lead", "sep", "tag", "replace"}
+_RESPONSE_FIELDS = {"id", "text", "label", "span"}
+
+
+class _FormatError(ValueError):
+    """An error already prefixed with its place in the prompt."""
+
+
+def read_prompts(path) -> list[Prompt]:
+    """Every prompt of a JSON Lines file; blank lines are skipped and prompt ids must be unique."""
+    prompts, lines = [], {}
+    with open(path, encoding="utf-8-sig") as f:     # -sig: a byte-order mark some editors add is skipped
+        for number, line in enumerate(f, 1):
+            if not line.strip():
+                continue
+            where = f"{os.fspath(path)} line {number}"
+            try:
+                data = json.loads(line)
+            except json.JSONDecodeError as e:
+                raise ValueError(f"{where}: not valid JSON ({e.msg}, column {e.colno})") from None
+            prompt = prompt_from_dict(data, where)
+            if prompt.id in lines:
+                raise ValueError(f"{where}: prompt id {prompt.id!r} is already used on line {lines[prompt.id]}")
+            lines[prompt.id] = number
+            prompts.append(prompt)
+    return prompts
+
+
+def write_prompts(prompts: Iterable[Prompt], path) -> None:
+    """Write prompts as JSON Lines, replacing the file only once every line is written."""
+    rows, seen = [], set()
+    for p in prompts:
+        if not isinstance(p, Prompt):
+            raise TypeError(f"write_prompts takes Prompt objects, not {type(p).__name__}")
+        if p.id in seen:
+            raise ValueError(f"prompt id {p.id!r} is used twice")
+        seen.add(p.id)
+        rows.append(json.dumps(prompt_to_dict(p), ensure_ascii=False))
+    tmp = f"{os.fspath(path)}.tmp"
+    with open(tmp, "w", encoding="utf-8", newline="\n") as f:
+        f.writelines(row + "\n" for row in rows)
+    os.replace(tmp, path)
+
+
+def prompt_from_dict(data, where: str = "prompt") -> Prompt:
+    """One prompt from its JSON object, every field checked."""
+    try:
+        _object(data, "the prompt")
+        _fields(data, _PROMPT_FIELDS, {"id"}, "the prompt")
+        messages = text = None
+        if "messages" in data:
+            if not isinstance(data["messages"], list):
+                raise _FormatError("messages must be a list")
+            messages = [_at(f"messages[{k}]", _message, m, f"messages[{k}]") for k, m in enumerate(data["messages"])]
+        if "text" in data:
+            text = _content(data["text"], "text")
+        responses = data.get("responses", [])
+        if not isinstance(responses, list):
+            raise _FormatError("responses must be a list")
+        return Prompt(data["id"], messages=messages, text=text, reference=data.get("reference"),
+                      responses=tuple(_at(f"responses[{k}]", _response, r, f"responses[{k}]")
+                                      for k, r in enumerate(responses)))
+    except (TypeError, ValueError) as e:
+        raise ValueError(f"{where}: {e}") from None
+
+
+def prompt_to_dict(prompt: Prompt) -> dict:
+    """The JSON object of a prompt; prompt_from_dict gives it back equal."""
+    d: dict = {"id": prompt.id}
+    if prompt.messages is not None:
+        d["messages"] = [{"role": m.role, "content": _content_out(m.content)} for m in prompt.messages]
+    else:
+        d["text"] = _content_out(prompt.text)
+    if prompt.reference is not None:
+        d["reference"] = prompt.reference
+    if prompt.responses:
+        d["responses"] = [_response_out(r) for r in prompt.responses]
+    return d
+
+
+# -- reading ----------------------------------------------------------------------------------------------------------
+
+def _at(where: str, fn, *args):
+    try:
+        return fn(*args)
+    except _FormatError:
+        raise
+    except (TypeError, ValueError) as e:
+        raise _FormatError(f"{where}: {e}") from None
+
+
+def _object(value, where: str) -> None:
+    if not isinstance(value, dict):
+        raise _FormatError(f"{where} must be a JSON object, not {type(value).__name__}")
+
+
+def _fields(obj: dict, allowed: set, required: set, where: str) -> None:
+    unknown = sorted(set(obj) - allowed)
+    if unknown:
+        raise _FormatError(f"{where}: unknown field {', '.join(map(repr, unknown))} "
+                           f"(allowed: {', '.join(sorted(allowed))})")
+    missing = sorted(required - set(obj))
+    if missing:
+        raise _FormatError(f"{where}: missing field {', '.join(map(repr, missing))}")
+
+
+def _message(value, where: str) -> Message:
+    _object(value, where)
+    _fields(value, _MESSAGE_FIELDS, _MESSAGE_FIELDS, where)
+    return Message(value["role"], _content(value["content"], f"{where}.content"))
+
+
+def _content(value, where: str):
+    if isinstance(value, str):
+        return value
+    if not isinstance(value, list):
+        raise _FormatError(f"{where} must be a string or a list, not {type(value).__name__}")
+    return [_at(f"{where}[{k}]", _piece, p, f"{where}[{k}]") for k, p in enumerate(value)]
+
+
+def _piece(value, where: str):
+    if isinstance(value, str):
+        return value
+    _object(value, where)
+    if "unit" in value and "group" in value:
+        raise _FormatError(f"{where}: a piece is a unit or a group, not both")
+    if "unit" in value:
+        _fields(value, _UNIT_FIELDS, {"unit", "text"}, where)
+        if "split" in value:
+            if value["split"] not in KINDS:
+                raise _FormatError(f"{where}: split must be one of {', '.join(KINDS)}, not {value['split']!r}")
+            return split(value["text"], value["split"], id=value["unit"], sep=value.get("sep", ""),
+                         tag=value.get("tag"), replace=value.get("replace"))
+        return _unit(value)
+    if "group" in value:
+        _fields(value, _GROUP_FIELDS, {"group", "units"}, where)
+        if not isinstance(value["units"], list):
+            raise _FormatError(f"{where}: units must be a list")
+        units = []
+        for k, u in enumerate(value["units"]):
+            _object(u, f"{where}.units[{k}]")
+            _fields(u, _UNIT_FIELDS - {"split"}, {"unit", "text"}, f"{where}.units[{k}]")
+            units.append(_at(f"{where}.units[{k}]", _unit, u))
+        return Group(value["group"], tuple(units), lead=value.get("lead", ""), sep=value.get("sep", ""),
+                     tag=value.get("tag"), replace=value.get("replace"))
+    raise _FormatError(f"{where}: an object piece needs a 'unit' or a 'group' field")
+
+
+def _unit(value: dict) -> Unit:
+    return Unit(value["unit"], value["text"], sep=value.get("sep", ""), tag=value.get("tag"),
+                replace=value.get("replace"))
+
+
+def _response(value, where: str) -> Response:
+    _object(value, where)
+    _fields(value, _RESPONSE_FIELDS, {"id", "text"}, where)
+    span = value.get("span")
+    if span is not None and not isinstance(span, list):
+        raise _FormatError(f"{where}: span must be a list [start, end]")
+    return Response(value["id"], value["text"], label=value.get("label"),
+                    span=tuple(span) if span is not None else None)
+
+
+# -- writing ----------------------------------------------------------------------------------------------------------
+
+def _content_out(pieces):
+    if len(pieces) == 1 and isinstance(pieces[0], str):
+        return pieces[0]
+    return [_piece_out(p) for p in pieces]
+
+
+def _piece_out(p):
+    if isinstance(p, str):
+        return p
+    if isinstance(p, Unit):
+        return _unit_out(p, None, None)
+    d = {"group": p.id, "units": [_unit_out(u, p.tag, p.replace) for u in p.units]}
+    for name in ("lead", "sep"):
+        if getattr(p, name):
+            d[name] = getattr(p, name)
+    for name in ("tag", "replace"):
+        if getattr(p, name) is not None:
+            d[name] = getattr(p, name)
+    return d
+
+
+def _unit_out(u: Unit, group_tag, group_replace) -> dict:
+    d = {"unit": u.id, "text": u.text}
+    if u.sep:
+        d["sep"] = u.sep
+    if u.tag is not None and u.tag != group_tag:
+        d["tag"] = u.tag
+    if u.replace is not None and u.replace != group_replace:
+        d["replace"] = u.replace
+    return d
+
+
+def _response_out(r: Response) -> dict:
+    d = {"id": r.id, "text": r.text}
+    if r.label is not None:
+        d["label"] = r.label
+    if r.span is not None:
+        d["span"] = list(r.span)
+    return d

contextattr/prompt/render.py

@@ -0,0 +1,70 @@
+"""A mask → the prompt as the model sees it: chat messages or text, removed units deleted or replaced (§2, §6)."""
+
+from __future__ import annotations
+
+from .units import Group, Prompt, Unit
+
+
+def as_mask(mask, n: int) -> tuple[bool, ...]:
+    """Check a mask: n entries, each 1 or True (keep) or 0 or False (remove). None keeps every unit."""
+    if mask is None:
+        return (True,) * n
+    if isinstance(mask, (str, bytes)):
+        raise TypeError("a mask is a sequence of 0/1 or True/False, not a string")
+    try:
+        values = list(mask)
+    except TypeError:
+        raise TypeError(f"a mask is a sequence of 0/1 or True/False, not {type(mask).__name__}") from None
+    if len(values) != n:
+        raise ValueError(f"the mask has {len(values)} entries, the prompt {n} units")
+    out = []
+    for k, v in enumerate(values):
+        if v is True or v is False:
+            out.append(v)
+            continue
+        try:
+            ok = not isinstance(v, (str, bytes)) and bool(v == 0 or v == 1)
+        except (TypeError, ValueError):  # e.g. a row of a 2-D array
+            ok = False
+        if not ok:
+            raise ValueError(f"mask entry {k} is {v!r}: use 0/1 or True/False")
+        out.append(bool(v == 1))
+    return tuple(out)
+
+
+def render(prompt: Prompt, mask=None, replace: str | None = None):
+    """The prompt with only the units `mask` keeps. A removed unit is deleted together with its `sep`, or, when it
+    has its own `replace` text or `replace` is given, becomes that text followed by its `sep`. A group's `lead` and
+    `sep` stay while any of its units shows. Chat prompts give [{"role", "content"}] dicts, text prompts a string."""
+    if replace is not None and not isinstance(replace, str):
+        raise TypeError(f"replace must be a string or None, not {type(replace).__name__}")
+    keep = iter(as_mask(mask, prompt.n))
+    if prompt.messages is not None:
+        return [{"role": m.role, "content": _join(m.content, keep, replace)} for m in prompt.messages]
+    return _join(prompt.text, keep, replace)
+
+
+def _join(pieces, keep, replace: str | None) -> str:
+    parts: list[str] = []
+    for p in pieces:
+        if isinstance(p, str):
+            parts.append(p)
+        elif isinstance(p, Unit):
+            parts.append(_unit(p, next(keep), replace)[0])
+        elif isinstance(p, Group):
+            shown = [_unit(u, next(keep), replace) for u in p.units]
+            if any(visible for _, visible in shown):
+                parts.append(p.lead)
+                parts.extend(text for text, _ in shown)
+                parts.append(p.sep)
+    return "".join(parts)
+
+
+def _unit(u: Unit, kept: bool, replace: str | None) -> tuple[str, bool]:
+    """The unit's text in the rendered prompt, and whether it shows (kept, or replaced)."""
+    if kept:
+        return u.text + u.sep, True
+    r = u.replace if u.replace is not None else replace
+    if r is None:
+        return "", False
+    return r + u.sep, True

contextattr/prompt/split.py

@@ -0,0 +1,153 @@
+"""Cutting a text into units that join back to it exactly (docs/interface.md §2).
+
+Every splitter only chooses where units start and end: the text before the first unit becomes the group's `lead`,
+and the text after each unit, up to the next one, becomes that unit's `sep`, so nothing is ever lost or added.
+"""
+
+from __future__ import annotations
+
+import operator
+import re
+from collections.abc import Callable, Iterable
+
+from .units import Group, Unit
+
+KINDS = ("word", "sentence", "line", "paragraph")
+
+_LINE_BREAK = re.compile(r"\r\n|\r|\n")
+_WORD = re.compile(r"\S+")
+_CLOSERS = "\"'”’»)]"
+_OPENERS = "\"'“‘«([{"
+_LATIN_END = re.compile(r"[.!?…]+[\"'”’»)\]]*(?=\s)")  # terminal punctuation, closing quotes or brackets, a space
+_CJK_END = re.compile(r"[。！？]+[」』”’）)]*")  # full-width terminals need no space after them
+_NUMBER_NEXT = re.compile(r"\(?\d")
+_DOTTED = re.compile(r"(?:[^\W\d_]\.)+[^\W\d_]")  # U.S, e.g, i.e, a.m: what precedes the final period
+# A period after these never ends a sentence: titles and short forms that precede a name or a word.
+_ABBREV = frozenset("mr mrs ms mx dr prof st mt gen gov sen rep rev lt col sgt capt cmdr adm maj vs cf al".split())
+# A period after these does not end a sentence when a number follows: "Fig. 3", "No. 5", "Eq. (2)", "Jan. 5".
+_ABBREV_NUMBER = frozenset("no nos fig figs eq eqs sec secs vol vols pp p ch art ref tab approx "
+                           "jan feb mar apr jun jul aug sep sept oct nov dec".split())
+
+
+def split(text: str, by: str | Callable[[str], Iterable[tuple[int, int]]] = "sentence", *, id: str,
+          sep: str = "", tag: str | None = None, replace: str | None = None) -> Group:
+    """Cut `text` into units `<id>.0`, `<id>.1`, … held in a Group with `sep` after it. The units with their
+    separators give back `text` exactly. `by`: "word", "sentence", "line", "paragraph", or a function returning the
+    units' (start, end) character spans, sorted and not overlapping."""
+    if not isinstance(text, str):
+        raise TypeError(f"split: text must be a string, not {type(text).__name__}")
+    if callable(by):
+        spans = _custom_spans(text, by)
+    elif by in _SPLITTERS:
+        spans = _SPLITTERS[by](text)
+    else:
+        raise ValueError(f"split: by must be one of {', '.join(KINDS)} or a function, not {by!r}")
+    if not spans:
+        raise ValueError(f"split: no units in {text[:40]!r} (by {by if isinstance(by, str) else 'the function'})")
+    units = []
+    for k, (start, end) in enumerate(spans):
+        nxt = spans[k + 1][0] if k + 1 < len(spans) else len(text)
+        units.append(Unit(f"{id}.{k}", text[start:end], sep=text[end:nxt]))
+    return Group(id, tuple(units), lead=text[: spans[0][0]], sep=sep, tag=tag, replace=replace)
+
+
+def _word_spans(text: str) -> list[tuple[int, int]]:
+    """Runs of non-space characters."""
+    return [(m.start(), m.end()) for m in _WORD.finditer(text)]
+
+
+def _line_spans(text: str) -> list[tuple[int, int]]:
+    """Lines with text, each from the line's start (its indentation included) to its last non-space character."""
+    spans = []
+    start = 0
+    for m in [*_LINE_BREAK.finditer(text), None]:
+        end = m.start() if m else len(text)
+        content = text[start:end].rstrip()
+        if content.strip():
+            spans.append((start, start + len(content)))
+        if m:
+            start = m.end()
+    return spans
+
+
+def _paragraph_spans(text: str) -> list[tuple[int, int]]:
+    """Runs of lines with text that no blank line separates, from the first line's start to the last's last
+    non-space character."""
+    spans: list[tuple[int, int]] = []
+    for a, b in _line_spans(text):
+        if spans and len(_LINE_BREAK.findall(text, spans[-1][1], a)) == 1:
+            spans[-1] = (spans[-1][0], b)
+        else:
+            spans.append((a, b))
+    return spans
+
+
+def _sentence_spans(text: str) -> list[tuple[int, int]]:
+    """Sentences within each line (a line break always ends one), without surrounding spaces."""
+    spans = []
+    for a, b in _line_spans(text):
+        a = _skip_space(text, a, b)
+        start = a
+        for end in _sentence_ends(text, a, b):
+            spans.append((start, end))
+            start = _skip_space(text, end, b)
+        spans.append((start, b))
+    return spans
+
+
+def _sentence_ends(text: str, a: int, b: int) -> list[int]:
+    """Positions inside text[a:b] where a sentence ends and another begins on the same line."""
+    ends = set()
+    for m in _LATIN_END.finditer(text, a, b):
+        nxt = _skip_space(text, m.end(), b)
+        if nxt < b and not text[nxt].islower() and not _abbreviation(text, a, m.start(), m.group(), nxt, b):
+            ends.add(m.end())
+    for m in _CJK_END.finditer(text, a, b):
+        if _skip_space(text, m.end(), b) < b:
+            ends.add(m.end())
+    return sorted(ends)
+
+
+def _abbreviation(text: str, line_start: int, p: int, run: str, nxt: int, b: int) -> bool:
+    """Whether the single period at p closes an abbreviation rather than a sentence."""
+    if run.rstrip(_CLOSERS) != ".":
+        return False
+    q = p
+    while q > line_start and not text[q - 1].isspace():
+        q -= 1
+    word = text[q:p].lstrip(_OPENERS)
+    if not word:
+        return False
+    return ((len(word) == 1 and word.isupper())                                    # an initial: "J. K. Rowling"
+            or _DOTTED.fullmatch(word) is not None                                 # "U.S.", "e.g.", "a.m."
+            or word.lower() in _ABBREV                                             # "Dr. Smith", "et al. (2020)"
+            or (word.lower() in _ABBREV_NUMBER and _NUMBER_NEXT.match(text, nxt, b) is not None)   # "Fig. 3"
+            or ((word.isdigit() or (len(word) == 1 and word.isalpha())) and q == line_start))   # "1. Mix", "a. Add"
+
+
+def _skip_space(text: str, i: int, b: int) -> int:
+    while i < b and text[i].isspace():
+        i += 1
+    return i
+
+
+def _custom_spans(text: str, fn: Callable[[str], Iterable[tuple[int, int]]]) -> list[tuple[int, int]]:
+    spans = []
+    prev = 0
+    for k, span in enumerate(fn(text)):
+        try:
+            start, end = span
+            if isinstance(start, bool) or isinstance(end, bool):
+                raise TypeError
+            start, end = operator.index(start), operator.index(end)
+        except (TypeError, ValueError):
+            raise TypeError(f"split: the function gave {span!r} as span {k}; a span is (start, end) integers") from None
+        if not prev <= start < end <= len(text):
+            raise ValueError(f"split: span {k} is ({start}, {end}); spans must be sorted, not overlapping and not "
+                             f"empty, within 0..{len(text)}")
+        spans.append((start, end))
+        prev = end
+    return spans
+
+
+_SPLITTERS = {"word": _word_spans, "sentence": _sentence_spans, "line": _line_spans, "paragraph": _paragraph_spans}