senseshift/text_utils.py

"""Sentence segmentation, VADER scoring and output cleanup.

Self-contained copies of the helpers the research repo keeps in
``generate_utils.py``, so the released package does not depend on it.
"""
from __future__ import annotations

import re
from functools import lru_cache
from typing import List, Sequence, Tuple

# The control vocabulary the model was trained with: 21 tokens on a 0.1 grid.
SENTIMENT_GRID: Tuple[float, ...] = tuple(round(i / 10, 1) + 0.0 for i in range(-10, 11))


def sentiment_token(value: float) -> str:
    """Map a sentiment value to the special token the model expects."""
    return f"[{snap_to_grid(value)}]"


def snap_to_grid(value: float) -> float:
    """Clamp to [-1, 1] and round to the nearest 0.1 (never returns -0.0)."""
    value = float(value)
    if value != value:  # NaN
        raise ValueError("sentiment must be a real number, got NaN")
    value = max(-1.0, min(1.0, value))
    return round(value, 1) + 0.0


@lru_cache(maxsize=1)
def _analyzer():
    import nltk
    from nltk.sentiment import SentimentIntensityAnalyzer

    try:
        nltk.data.find("sentiment/vader_lexicon.zip")
    except LookupError:
        nltk.download("vader_lexicon", quiet=True)
    return SentimentIntensityAnalyzer()


def split_sentences(text: str) -> List[str]:
    parts = re.split(r"(?<=[.!?])\s+", text.strip())
    return [p.strip() for p in parts if p.strip()]


def score_sentence(sentence: str) -> float:
    return snap_to_grid(_analyzer().polarity_scores(sentence)["compound"])


def compute_vader_sentiment(text: str) -> Tuple[List[str], List[float], float]:
    """Return (sentences, per-sentence sentiment on the 0.1 grid, overall)."""
    sentences = split_sentences(text)
    sentiments = [score_sentence(s) for s in sentences]
    overall = snap_to_grid(_analyzer().polarity_scores(text)["compound"])
    return sentences, sentiments, overall


def strip_sentiment_marker(text: str) -> str:
    """Drop a leading ``[0.3]`` style control token."""
    return re.sub(r"^\s*\[[+-]?\d+(\.\d+)?\]\s*", "", text)


def clean_generated_text(text: str) -> str:
    """Detokenisation cleanup for text decoded out of the MLM."""
    cleaned = re.sub(r"\b(\w+)\s+##(\w+)", r"\1\2", text)
    cleaned = re.sub(r"<[^>]*>", "", cleaned)
    cleaned = re.sub(r"\[[+-]?\d+(\.\d+)?\]", " ", cleaned)
    cleaned = re.sub(r"\s+", " ", cleaned).strip()

    parts = [p.strip() for p in re.split(r"\s{2,}", cleaned) if p.strip()]
    if not parts:
        return cleaned

    result_parts = []
    for p in parts:
        if p and p[-1] not in ".!?":
            p += "."
        result_parts.append(p)

    out = " ".join(result_parts)
    out = re.sub(r"\s+", " ", out).strip()
    out = re.sub(r"\s+([,.;:!?])", r"\1", out)
    out = re.sub(r"\s*'\s*", "'", out)
    out = re.sub(r"\.\.+", ".", out)
    out = re.sub(r'"', "", out)
    return out


def choose_random_sentiment(exclude: float | None = None, rng=None) -> float:
    """Pick a grid value, optionally excluding the current one."""
    import random as _random

    rng = rng or _random
    options: Sequence[float] = SENTIMENT_GRID
    if exclude is not None:
        exclude = snap_to_grid(exclude)
        options = [v for v in SENTIMENT_GRID if v != exclude]
    return float(rng.choice(list(options)))