File size: 3,304 Bytes
df10fc7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | """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)))
|