transcript.py

"""The one place that defines the chat transcript format.

Training data, the Streamlit app and the evaluator all go through this module,
so what the model sees at inference is exactly what it saw in training.

Format (each message is its own line, continuation lines have no label):

    --- Wed 14 Jan 2026, 11:41 PM ---
    Yash: kya kar rhi
    Divyanshi: Abhi toh bola 8?
    Divyanshi: Phir se kuch sunnan h?🤧

The prompt ends with "Divyanshi:" and the model writes the rest of the burst
until it starts a "Yash:" line (the stop sequence).
"""

from __future__ import annotations

import re
from datetime import datetime
from typing import Any
from zoneinfo import ZoneInfo

DAYS = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]
MONTHS = ["Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"]

MEDIA_PLACEHOLDERS = {
    "reel", "post", "story", "profile", "link", "photo", "video", "voice note", "file",
}
PLACEHOLDER_RE = re.compile(r"^\[(?:%s)\]$" % "|".join(re.escape(p) for p in MEDIA_PLACEHOLDERS))

HEADER_LABEL = "---"


def is_placeholder(text: str) -> bool:
    return bool(PLACEHOLDER_RE.match(text.strip()))


def format_ts(ts_ms: int, tz: str) -> str:
    dt = datetime.fromtimestamp(ts_ms / 1000, tz=ZoneInfo(tz))
    return f"{DAYS[dt.weekday()]} {dt.day} {MONTHS[dt.month - 1]} {dt.year}, {dt.hour % 12 or 12}:{dt.minute:02d} {'AM' if dt.hour < 12 else 'PM'}"


def messages_to_lines(messages: list[dict[str, Any]], gap_minutes: int) -> list[dict[str, Any]]:
    """Messages ({role: me|them, text, ts}) -> transcript lines, inserting a header per session."""
    lines: list[dict[str, Any]] = []
    prev_ts: int | None = None
    for m in messages:
        if prev_ts is None or m["ts"] - prev_ts >= gap_minutes * 60_000:
            lines.append({"kind": "header", "text": "", "ts": m["ts"]})
        lines.append({"kind": m["role"], "text": m["text"], "ts": m["ts"]})
        prev_ts = m["ts"]
    return lines


def _parts(line: dict[str, Any], cfg: dict[str, Any]) -> tuple[str, str]:
    """(label, body): the label is the speaker marker, the body is the rest of the line incl. newline."""
    if line["kind"] == "header":
        return HEADER_LABEL, f" {format_ts(line['ts'], cfg['timezone'])} ---\n"
    alias = cfg["target_alias"] if line["kind"] == "them" else cfg["my_alias"]
    return f"{alias}:", f" {line['text']}\n"


def render_lines(lines: list[dict[str, Any]], cfg: dict[str, Any]) -> str:
    return "".join("".join(_parts(l, cfg)) for l in lines)


def build_segments(
    lines: list[dict[str, Any]],
    start: int,
    end: int,
    train_from: int,
    cfg: dict[str, Any],
) -> list[tuple[str, bool]]:
    """Render lines[start:end] as (text, trainable) segments.

    Loss goes only on the target's message text (plus its newline) and on the speaker
    label that follows one of the target's lines, which teaches when a burst ends.
    Lines before `train_from` are context only. Placeholder-only target lines are
    never trained so the model doesn't learn to send media.
    """
    segs: list[tuple[str, bool]] = []
    if lines[start]["kind"] != "header":
        label, body = _parts({"kind": "header", "text": "", "ts": lines[start]["ts"]}, cfg)
        segs.append((label + body, False))
    for i in range(start, end):
        line = lines[i]
        label, body = _parts(line, cfg)
        live = i >= train_from
        prev_them = i > 0 and lines[i - 1]["kind"] == "them"
        segs.append((label, live and prev_them))
        if line["kind"] == "them":
            segs.append((body, live and not is_placeholder(line["text"])))
        else:
            segs.append((body, False))
    merged: list[tuple[str, bool]] = []
    for text, flag in segs:
        if merged and merged[-1][1] == flag:
            merged[-1] = (merged[-1][0] + text, flag)
        else:
            merged.append((text, flag))
    return merged


def build_prompt(history: list[dict[str, Any]], cfg: dict[str, Any]) -> str:
    """Prompt for the next target burst. history = [{role, text, ts}], newest last.

    The caller includes a final message (or a bare timestamp) so the header for the
    current session exists; the text is trimmed to `context_chars` from the oldest side.
    """
    lines = messages_to_lines(history, cfg["session_gap_minutes"])
    budget = cfg.get("context_chars", 5500)
    start, used = len(lines), 0
    while start > 0:
        size = sum(len(p) for p in _parts(lines[start - 1], cfg))
        if used + size > budget:
            break
        used += size
        start -= 1
    start = min(start, len(lines) - 1) if lines else 0
    kept = lines[start:]
    if kept and kept[0]["kind"] != "header":
        kept.insert(0, {"kind": "header", "text": "", "ts": kept[0]["ts"]})
    return render_lines(kept, cfg) + f"{cfg['target_alias']}:"


def parse_reply(raw: str, cfg: dict[str, Any], max_burst: int = 6) -> list[str]:
    """Split raw model output into separate messages, stopping when anyone else speaks."""
    me_label = f"{cfg['my_alias']}:"
    them_label = f"{cfg['target_alias']}:"
    messages: list[list[str]] = [[]]
    first = True
    for line in raw.split("\n"):
        if first:
            messages[-1].append(line[1:] if line.startswith(" ") else line)
            first = False
        elif line.startswith(me_label) or line.startswith(HEADER_LABEL):
            break
        elif line.startswith(them_label):
            rest = line[len(them_label):]
            messages.append([rest[1:] if rest.startswith(" ") else rest])
        else:
            messages[-1].append(line)
    out = []
    for parts in messages:
        text = "\n".join(parts).strip()
        if text and not is_placeholder(text):
            out.append(text)
    return out[:max_burst]


def tokenize_segments(
    tokenizer: Any,
    segments: list[tuple[str, bool]],
    max_len: int,
) -> dict[str, list[int]] | None:
    """Tokenize and mask: labels are -100 except on trainable segments."""
    text = "".join(s for s, _ in segments)
    spans, pos = [], 0
    for s, flag in segments:
        if flag:
            spans.append((pos, pos + len(s)))
        pos += len(s)
    enc = tokenizer(text, add_special_tokens=False, return_offsets_mapping=True)
    ids, offsets = list(enc["input_ids"]), list(enc["offset_mapping"])
    bos = getattr(tokenizer, "bos_token_id", None)
    if bos is not None:
        ids.insert(0, bos)
        offsets.insert(0, (-1, -1))
    labels = []
    for tok, (a, _b) in zip(ids, offsets):
        trained = a >= 0 and any(s <= a < e for s, e in spans)
        labels.append(tok if trained else -100)
    ids, labels = ids[:max_len], labels[:max_len]
    if all(l == -100 for l in labels):
        return None
    return {"input_ids": ids, "labels": labels, "attention_mask": [1] * len(ids)}