Upload folder using huggingface_hub

README.md

@@ -0,0 +1,120 @@
+---
+license: apache-2.0
+base_model: answerdotai/ModernBERT-large
+library_name: transformers
+pipeline_tag: fill-mask
+tags:
+  - sentiment-control
+  - controllable-text-generation
+  - text-rewriting
+  - modernbert
+language:
+  - en
+---
+
+# SenseShift-large
+
+SenseShift rewrites a sentence — or writes a new one — at **any sentiment you ask
+for on a continuous −1.0 to +1.0 scale**, while keeping it consistent with the
+surrounding text.
+
+It is a masked language model fine-tuned from ModernBERT-large with an explicit
+**control vocabulary**: 21 special tokens `[-1.0] … [1.0]` on a 0.1 grid. At
+training time every sentence is prefixed with its own VADER sentiment token and
+one sentence is masked out, so the model learns to write a replacement that
+realises the requested sentiment in context.
+
+| Model | Base | Params |
+| --- | --- | --- |
+| [SenseShift-base](https://huggingface.co/shawhed/SenseShift-base) | ModernBERT-base | 150M |
+| **SenseShift-large** (this model) | ModernBERT-large | 396M |
+
+## Usage
+
+No install needed — the inference code ships inside this repo. Clone it and you
+have the weights and the code together:
+
+```bash
+git clone https://huggingface.co/shawhed/SenseShift-large
+```
+
+```python
+import sys
+sys.path.insert(0, "SenseShift-large")
+
+from senseshift import SenseShift
+
+shifter = SenseShift.from_pretrained("SenseShift-large")
+
+text = ("The waiter greeted us at the door. "
+        "The food arrived quickly and was still hot. "
+        "We paid the bill and walked back to the hotel.")
+
+# Rewrite sentence 1 as strongly negative
+out = shifter.generate(text, generation_mode="rewrite", sentence_index=1, sentiment=-0.9)
+print(out.sentence)   # -> "He looked sad and we felt bad."
+print(out.text)       # the full passage with that sentence swapped in
+
+# Add a positive sentence after the last one
+out = shifter.generate(text, generation_mode="add", sentiment=0.7)
+print(out.text)
+```
+
+Or let `huggingface_hub` fetch it into the local cache instead of cloning:
+
+```python
+import sys
+from huggingface_hub import snapshot_download
+
+path = snapshot_download("shawhed/SenseShift-large")
+sys.path.insert(0, path)
+
+from senseshift import SenseShift
+shifter = SenseShift.from_pretrained(path)
+```
+
+Requirements: `torch`, `transformers`, `huggingface-hub`, `nltk` (see
+`requirements.txt`). The VADER lexicon downloads itself on first use.
+
+### `generate` arguments
+
+| Argument | Meaning |
+| --- | --- |
+| `text` | The passage to edit. |
+| `generation_mode` | `"rewrite"` replaces the sentence at `sentence_index`; `"add"` inserts a new sentence right after it. |
+| `sentiment` | `None` → keep the sentiment already there. `"random"` → a random grid value, excluding the current one. A number in `[-1, 1]` → that value, snapped to the nearest 0.1. |
+| `sentence_index` | Which sentence to act on. Defaults to a random sentence (`rewrite`) or the last sentence (`add`). Negative indices count from the end. |
+| `num_masks` | How many mask slots the model gets, i.e. roughly how long the new sentence is. Defaults to the replaced sentence's word count (`rewrite`) or `12` (`add`). |
+| `seed` | Seeds the random index / sentiment draws. |
+
+Decoding can be tuned per call: `top_k`, `beam_size`, `max_iters`, `alpha`
+(length normalisation), `gamma` (beam diversity penalty), `temperature`,
+`min_words`.
+
+`generate` returns a `SenseShiftOutput` with `.text`, `.sentence`,
+`.target_sentiment`, `.achieved_sentiment` (VADER of what was actually written),
+`.beams`, and friends. `str(out)` gives the edited passage.
+
+## How it works
+
+1. Split the passage into sentences and score each with VADER.
+2. Prefix every sentence with its control token; give the target sentence the
+   *requested* token and replace its words with `[MASK]`s.
+3. Fill the masks left to right with beam search over the MLM head, with length
+   normalisation and a diversity penalty, stopping at terminal punctuation.
+4. Splice the decoded sentence back into the passage.
+
+Steps 1–4 live in the `senseshift` package, not in the weights — the checkpoint
+itself is a stock `ModernBertForMaskedLM` and can be loaded with
+`AutoModelForMaskedLM` if you want to build your own decoding loop.
+
+## Limitations
+
+- English only; trained on short narrative and review-style text.
+- VADER supplies the sentiment labels, so the model inherits its lexicon-based
+  view of sentiment. `achieved_sentiment` typically lands within ~0.2 of the
+  target rather than hitting it exactly.
+- The rewrite is length-bounded by `num_masks`, so very long sentences are
+  usually replaced by something shorter.
+- Because the whole passage is re-encoded per mask fill, generation cost grows
+  with `num_masks × beam_size`.

config.json

@@ -0,0 +1,83 @@
+{
+  "architectures": [
+    "ModernBertForMaskedLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "bos_token_id": 50281,
+  "classifier_activation": "gelu",
+  "classifier_bias": false,
+  "classifier_dropout": 0.0,
+  "classifier_pooling": "mean",
+  "cls_token_id": 50281,
+  "decoder_bias": true,
+  "deterministic_flash_attn": false,
+  "dtype": "float32",
+  "embedding_dropout": 0.0,
+  "eos_token_id": 50282,
+  "global_attn_every_n_layers": 3,
+  "gradient_checkpointing": false,
+  "hidden_activation": "gelu",
+  "hidden_size": 1024,
+  "initializer_cutoff_factor": 2.0,
+  "initializer_range": 0.02,
+  "intermediate_size": 2624,
+  "layer_norm_eps": 1e-05,
+  "layer_types": [
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention",
+    "sliding_attention",
+    "sliding_attention",
+    "full_attention"
+  ],
+  "local_attention": 128,
+  "max_position_embeddings": 8192,
+  "mlp_bias": false,
+  "mlp_dropout": 0.0,
+  "model_type": "modernbert",
+  "norm_bias": false,
+  "norm_eps": 1e-05,
+  "num_attention_heads": 16,
+  "num_hidden_layers": 28,
+  "pad_token_id": 50283,
+  "position_embedding_type": "absolute",
+  "rope_parameters": {
+    "full_attention": {
+      "rope_theta": 160000.0,
+      "rope_type": "default"
+    },
+    "sliding_attention": {
+      "rope_theta": 10000.0,
+      "rope_type": "default"
+    }
+  },
+  "sep_token_id": 50282,
+  "sparse_pred_ignore_index": -100,
+  "sparse_prediction": false,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.3.0",
+  "vocab_size": 50389
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:01cc07f8b521219991520866db120d63a888ffc3fd519c471d08d2aabd817f21
+size 1583630940

requirements.txt

@@ -0,0 +1,4 @@
+torch>=2.0
+transformers>=4.48
+huggingface-hub>=0.23
+nltk>=3.8

senseshift/__init__.py

@@ -0,0 +1,13 @@
+"""SenseShift: rewrite or extend text at any sentiment from -1.0 to 1.0."""
+
+from .pipeline import SenseShift, SenseShiftOutput
+from .text_utils import SENTIMENT_GRID, compute_vader_sentiment, split_sentences
+
+__version__ = "0.1.0"
+__all__ = [
+    "SenseShift",
+    "SenseShiftOutput",
+    "SENTIMENT_GRID",
+    "compute_vader_sentiment",
+    "split_sentences",
+]

senseshift/decoding.py

@@ -0,0 +1,111 @@
+"""Iterative left-to-right beam search over the masked positions."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
+
+import torch
+
+from .text_utils import clean_generated_text, split_sentences, strip_sentiment_marker
+
+
+@dataclass
+class Beam:
+    text: str
+    score: float
+
+
+def _terminal_token_ids(tokenizer) -> set:
+    ids = set()
+    for punct in (".", "!", "?"):
+        ids.update(tokenizer.encode(punct, add_special_tokens=False))
+    return ids
+
+
+@torch.inference_mode()
+def fill_masks_beam_search(
+    model,
+    tokenizer,
+    masked_text: str,
+    sentence_index: int,
+    top_k: int = 40,
+    beam_size: int = 2,
+    max_iters: int = 30,
+    alpha: float = 0.7,
+    gamma: float = 0.05,
+    temperature: float = 0.8,
+    min_words: int = 3,
+) -> Tuple[List[Beam], str]:
+    """Fill masks one at a time, keeping ``beam_size`` hypotheses alive.
+
+    ``alpha`` is the length-normalisation exponent, ``gamma`` penalises beams
+    that reuse a token already proposed this step, and a beam finishes early
+    once it emits terminal punctuation (after at least ``min_words`` fills).
+
+    Returns the ranked beams over the *whole* passage plus the rewritten
+    sentence extracted from the best beam.
+    """
+    device = next(model.parameters()).device
+    mask_token_id = tokenizer.mask_token_id
+    pad_token_id = tokenizer.pad_token_id
+    if pad_token_id is None:
+        pad_token_id = tokenizer.eos_token_id
+    terminal_tokens = _terminal_token_ids(tokenizer)
+
+    encoded = tokenizer(masked_text, return_tensors="pt").to(device)
+    # (ids, cumulative score, tokens filled, finished, last token)
+    beams: List[tuple] = [(encoded.input_ids[0], 0.0, 0, False, None)]
+
+    for _ in range(max_iters):
+        active = [i for i, b in enumerate(beams) if (b[0] == mask_token_id).any() and not b[3]]
+        if not active:
+            break
+
+        active_beams = [beams[i] for i in active]
+        batch_ids = torch.stack([b[0] for b in active_beams]).to(device)
+        batch_scores = torch.tensor([b[1] for b in active_beams], device=device)
+        batch_filled = torch.tensor([b[2] for b in active_beams], device=device)
+        mask_positions = (batch_ids == mask_token_id).int().argmax(dim=1)
+
+        logits = model(input_ids=batch_ids).logits
+        mask_logits = logits[torch.arange(batch_ids.shape[0]), mask_positions] / temperature
+
+        probs = torch.softmax(mask_logits, dim=-1)
+        top_probs, top_tokens = torch.topk(probs, top_k)
+        top_probs = top_probs / top_probs.sum(dim=-1, keepdim=True)
+        log_p = torch.log(top_probs + 1e-10)
+
+        candidates = [b for i, b in enumerate(beams) if i not in active]
+
+        for i in range(len(active_beams)):
+            filled = batch_filled[i].item() + 1
+            length_norm = ((5 + filled) ** alpha) / ((5 + 1) ** alpha)
+            scores = (batch_scores[i] + log_p[i]) / length_norm
+
+            for j in range(top_k):
+                token_id = top_tokens[i, j].item()
+                # discourage every beam from picking the same continuation
+                penalty = sum(1 for c in candidates if c[4] == token_id)
+                score = scores[j].item() - gamma * penalty
+
+                new_ids = batch_ids[i].clone()
+                new_ids[mask_positions[i]] = token_id
+
+                finished = token_id in terminal_tokens and filled >= min_words
+                if finished:
+                    new_ids[new_ids == mask_token_id] = pad_token_id
+
+                candidates.append((new_ids, score, filled, finished, token_id))
+
+        beams = sorted(candidates, key=lambda x: x[1], reverse=True)[:beam_size]
+
+    ranked = [Beam(tokenizer.decode(b[0], skip_special_tokens=True), float(b[1])) for b in beams]
+    best_text = ranked[0].text if ranked else ""
+
+    sentences_out = split_sentences(best_text)
+    if sentence_index < len(sentences_out):
+        best_sentence = sentences_out[sentence_index]
+    else:
+        best_sentence = sentences_out[-1] if sentences_out else best_text
+
+    return ranked, clean_generated_text(strip_sentiment_marker(best_sentence)).strip()

senseshift/masking.py

@@ -0,0 +1,112 @@
+"""Build the sentiment-annotated, partially masked prompt the model reads.
+
+The model was trained on sequences where *every* sentence is prefixed with its
+own control token, e.g.::
+
+    [0.4] The food arrived quickly. [-0.6] [MASK] [MASK] [MASK] [MASK] .
+
+Exactly one sentence carries a control token that disagrees with its text; that
+sentence is masked out and the model must write a replacement that realises the
+requested sentiment while staying consistent with its neighbours.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+from .text_utils import (
+    compute_vader_sentiment,
+    sentiment_token,
+    snap_to_grid,
+)
+
+
+@dataclass
+class MaskedPrompt:
+    masked_text: str
+    sentences: List[str]
+    sentence_index: int
+    source_sentiment: float
+    target_sentiment: float
+    sentence_sentiments: List[float] = field(default_factory=list)
+    original_sentence: Optional[str] = None
+
+
+def build_masked_prompt(
+    tokenizer,
+    text: str,
+    sentence_index: int,
+    target_sentiment: float,
+    num_masks: Optional[int] = None,
+) -> MaskedPrompt:
+    """Annotate every sentence with its control token and mask sentence ``sentence_index``.
+
+    ``num_masks=None`` uses the word count of the sentence being replaced, which
+    is the training-time convention and keeps the rewrite roughly length-matched.
+    """
+    sentences, sentiments, _ = compute_vader_sentiment(text)
+    if not sentences:
+        raise ValueError("Input text contains no sentences.")
+    if not 0 <= sentence_index < len(sentences):
+        raise IndexError(
+            f"sentence_index {sentence_index} out of range for {len(sentences)} sentences."
+        )
+
+    target_sentiment = snap_to_grid(target_sentiment)
+    source_sentiment = sentiments[sentence_index]
+
+    targets = list(sentiments)
+    targets[sentence_index] = target_sentiment
+
+    annotated: List[str] = []
+    for i, (sentence, value) in enumerate(zip(sentences, targets)):
+        token = sentiment_token(value)
+        if i == sentence_index:
+            width = len(sentence.split()) if num_masks is None else num_masks
+            width = max(1, int(width))
+            annotated.append(token + "  " + "  ".join([tokenizer.mask_token] * width))
+        else:
+            annotated.append(f"{token} {sentence}")
+
+    masked_text = " ".join(annotated).replace("  ", " ")
+
+    return MaskedPrompt(
+        masked_text=masked_text,
+        sentences=sentences,
+        sentence_index=sentence_index,
+        source_sentiment=source_sentiment,
+        target_sentiment=target_sentiment,
+        sentence_sentiments=targets,
+        original_sentence=sentences[sentence_index],
+    )
+
+
+def insert_placeholder(
+    tokenizer,
+    text: str,
+    insert_after: Optional[int],
+    num_masks: int,
+) -> tuple[str, int]:
+    """Splice an all-mask placeholder sentence into ``text``.
+
+    Returns the widened text and the index the new sentence occupies, ready to
+    be handed to :func:`build_masked_prompt`.
+    """
+    from .text_utils import split_sentences
+
+    sentences = split_sentences(text)
+    if not sentences:
+        raise ValueError("Input text contains no sentences.")
+
+    if insert_after is None:
+        insert_after = len(sentences) - 1
+    if not 0 <= insert_after < len(sentences):
+        raise IndexError(
+            f"sentence_index {insert_after} out of range for {len(sentences)} sentences."
+        )
+
+    new_index = insert_after + 1
+    placeholder = " ".join([tokenizer.mask_token] * max(1, int(num_masks))) + " ."
+    widened = list(sentences)
+    widened.insert(new_index, placeholder)
+    return " ".join(widened), new_index

senseshift/pipeline.py

@@ -0,0 +1,279 @@
+"""SenseShift — sentiment-controlled sentence rewriting and insertion."""
+from __future__ import annotations
+
+import json
+import random as _random
+from dataclasses import asdict, dataclass, field
+from typing import Any, Dict, List, Optional, Union
+
+import torch
+
+from .decoding import Beam, fill_masks_beam_search
+from .masking import build_masked_prompt, insert_placeholder
+from .text_utils import (
+    choose_random_sentiment,
+    score_sentence,
+    snap_to_grid,
+    split_sentences,
+)
+
+CONFIG_FILE = "senseshift_config.json"
+
+DEFAULT_GENERATION: Dict[str, Any] = {
+    "top_k": 40,
+    "beam_size": 2,
+    "max_iters": 30,
+    "alpha": 0.7,
+    "gamma": 0.05,
+    "temperature": 0.8,
+    "min_words": 3,
+    "add_num_masks": 12,
+}
+
+Sentiment = Union[None, str, float, int]
+
+
+@dataclass
+class SenseShiftOutput:
+    """Result of one edit."""
+
+    text: str
+    """The full passage after the edit."""
+
+    sentence: str
+    """The sentence SenseShift wrote."""
+
+    sentences: List[str]
+    sentence_index: int
+    generation_mode: str
+    target_sentiment: float
+    source_sentiment: Optional[float]
+    achieved_sentiment: float
+    original_sentence: Optional[str] = None
+    beams: List[Beam] = field(default_factory=list)
+
+    def __str__(self) -> str:  # so print(out) shows the passage
+        return self.text
+
+    def to_dict(self) -> Dict[str, Any]:
+        d = asdict(self)
+        d["beams"] = [asdict(b) for b in self.beams]
+        return d
+
+
+class SenseShift:
+    """Rewrite or extend a passage at a chosen sentiment on a [-1, 1] scale.
+
+    >>> shifter = SenseShift.from_pretrained("shawhed/SenseShift-large")
+    >>> out = shifter.generate(text, sentence_index=1, sentiment=-0.8)
+    >>> print(out.text)
+    """
+
+    def __init__(self, model, tokenizer, generation_defaults: Optional[Dict[str, Any]] = None):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.generation_defaults = {**DEFAULT_GENERATION, **(generation_defaults or {})}
+        self._validate_vocabulary()
+
+    # ------------------------------------------------------------------ load
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_id: str,
+        device: Optional[str] = None,
+        dtype: Optional["torch.dtype"] = None,
+        **kwargs,
+    ) -> "SenseShift":
+        """Load weights, tokenizer and generation defaults from the Hub or a local path."""
+        from transformers import AutoModelForMaskedLM, AutoTokenizer
+
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        tokenizer = AutoTokenizer.from_pretrained(model_id, **kwargs)
+        model = AutoModelForMaskedLM.from_pretrained(model_id, dtype=dtype, **kwargs)
+        model.to(device).eval()
+
+        return cls(model, tokenizer, generation_defaults=_load_config(model_id, **kwargs))
+
+    def _validate_vocabulary(self) -> None:
+        from .text_utils import SENTIMENT_GRID, sentiment_token
+
+        vocab = self.tokenizer.get_vocab()
+        missing = [sentiment_token(v) for v in SENTIMENT_GRID if sentiment_token(v) not in vocab]
+        if missing:
+            raise ValueError(
+                "This tokenizer is missing SenseShift control tokens "
+                f"({', '.join(missing[:5])}{'…' if len(missing) > 5 else ''}). "
+                "Load a SenseShift checkpoint, not the base ModernBERT."
+            )
+        if self.tokenizer.mask_token_id is None:
+            raise ValueError("SenseShift needs a masked-language-model tokenizer with a mask token.")
+
+    # -------------------------------------------------------------- generate
+
+    def generate(
+        self,
+        text: str,
+        generation_mode: str = "rewrite",
+        sentiment: Sentiment = None,
+        sentence_index: Optional[int] = None,
+        num_masks: Optional[int] = None,
+        seed: Optional[int] = None,
+        **overrides,
+    ) -> SenseShiftOutput:
+        """Rewrite one sentence, or add a new one, at a controlled sentiment.
+
+        Args:
+            text: The passage to edit.
+            generation_mode: ``"rewrite"`` replaces the sentence at
+                ``sentence_index``; ``"add"`` writes a new sentence and splices
+                it in directly after ``sentence_index``.
+            sentiment: ``None`` keeps the sentiment already there (for ``"add"``,
+                the sentiment of the sentence being appended to); ``"random"``
+                draws a value from the 0.1 grid, excluding the current one; a
+                number in ``[-1, 1]`` is used as-is, snapped to the nearest 0.1.
+            sentence_index: Which sentence to act on. Defaults to a random
+                sentence for ``"rewrite"`` and to the last sentence for ``"add"``.
+            num_masks: How many mask slots to give the model, i.e. roughly the
+                length of what it writes. Defaults to the replaced sentence's
+                word count (``"rewrite"``) or ``add_num_masks`` (``"add"``).
+            seed: Seed for the sentence/sentiment draws, for reproducibility.
+            **overrides: Per-call decoding overrides — ``top_k``, ``beam_size``,
+                ``max_iters``, ``alpha``, ``gamma``, ``temperature``, ``min_words``.
+
+        Returns:
+            A :class:`SenseShiftOutput`; ``str(out)`` is the edited passage.
+        """
+        if generation_mode not in ("rewrite", "add"):
+            raise ValueError(
+                f"generation_mode must be 'rewrite' or 'add', got {generation_mode!r}"
+            )
+
+        params = {**self.generation_defaults, **overrides}
+        unknown = set(overrides) - set(DEFAULT_GENERATION)
+        if unknown:
+            raise TypeError(f"Unknown generation option(s): {', '.join(sorted(unknown))}")
+
+        rng = _random.Random(seed) if seed is not None else _random
+
+        sentences = split_sentences(text)
+        if not sentences:
+            raise ValueError("Input text contains no sentences.")
+
+        anchor = self._resolve_index(sentence_index, len(sentences), generation_mode, rng)
+
+        if generation_mode == "rewrite":
+            source = score_sentence(sentences[anchor])
+            target = self._resolve_sentiment(sentiment, source, rng)
+            prompt = build_masked_prompt(
+                self.tokenizer, text, anchor, target, num_masks=num_masks
+            )
+        else:
+            source = score_sentence(sentences[anchor])
+            target = self._resolve_sentiment(sentiment, source, rng)
+            width = params["add_num_masks"] if num_masks is None else num_masks
+            widened, anchor = insert_placeholder(self.tokenizer, text, anchor, width)
+            prompt = build_masked_prompt(
+                self.tokenizer, widened, anchor, target, num_masks=width
+            )
+
+        beams, new_sentence = fill_masks_beam_search(
+            self.model,
+            self.tokenizer,
+            prompt.masked_text,
+            sentence_index=anchor,
+            top_k=params["top_k"],
+            beam_size=params["beam_size"],
+            max_iters=params["max_iters"],
+            alpha=params["alpha"],
+            gamma=params["gamma"],
+            temperature=params["temperature"],
+            min_words=params["min_words"],
+        )
+
+        if not new_sentence:
+            if generation_mode == "add":
+                raise RuntimeError("Model produced an empty sentence; nothing was inserted.")
+            new_sentence = sentences[anchor]  # fall back to the original
+
+        new_sentences = list(sentences)
+        if generation_mode == "rewrite":
+            new_sentences[anchor] = new_sentence
+            original = sentences[anchor]
+        else:
+            new_sentences.insert(anchor, new_sentence)
+            original = None
+
+        return SenseShiftOutput(
+            text=" ".join(new_sentences),
+            sentence=new_sentence,
+            sentences=new_sentences,
+            sentence_index=anchor,
+            generation_mode=generation_mode,
+            target_sentiment=target,
+            source_sentiment=source if generation_mode == "rewrite" else None,
+            achieved_sentiment=score_sentence(new_sentence),
+            original_sentence=original,
+            beams=beams,
+        )
+
+    # --------------------------------------------------------------- helpers
+
+    @staticmethod
+    def _resolve_index(
+        sentence_index: Optional[int], n: int, generation_mode: str, rng
+    ) -> int:
+        if sentence_index is None:
+            return rng.randrange(n) if generation_mode == "rewrite" else n - 1
+        if sentence_index < 0:
+            sentence_index += n
+        if not 0 <= sentence_index < n:
+            raise IndexError(f"sentence_index out of range for {n} sentences.")
+        return sentence_index
+
+    @staticmethod
+    def _resolve_sentiment(sentiment: Sentiment, source: float, rng) -> float:
+        if sentiment is None:
+            return snap_to_grid(source)
+        if isinstance(sentiment, str):
+            if sentiment.lower() != "random":
+                raise ValueError(
+                    f"sentiment string must be 'random', got {sentiment!r}"
+                )
+            return choose_random_sentiment(exclude=source, rng=rng)
+        if isinstance(sentiment, bool) or not isinstance(sentiment, (int, float)):
+            raise TypeError(
+                "sentiment must be None, 'random', or a number in [-1, 1]; "
+                f"got {type(sentiment).__name__}"
+            )
+        if not -1.0 <= float(sentiment) <= 1.0:
+            raise ValueError(f"sentiment must lie in [-1, 1], got {sentiment}")
+        return snap_to_grid(sentiment)
+
+
+def _load_config(model_id: str, **kwargs) -> Dict[str, Any]:
+    """Read ``senseshift_config.json`` from a local dir or the Hub; empty if absent."""
+    import os
+
+    local = os.path.join(model_id, CONFIG_FILE)
+    if os.path.isfile(local):
+        with open(local, encoding="utf-8") as fh:
+            return json.load(fh).get("generation", {})
+
+    try:
+        from huggingface_hub import hf_hub_download
+        from huggingface_hub.errors import EntryNotFoundError
+
+        path = hf_hub_download(
+            model_id,
+            CONFIG_FILE,
+            revision=kwargs.get("revision"),
+            token=kwargs.get("token"),
+        )
+    except Exception:
+        return {}
+
+    with open(path, encoding="utf-8") as fh:
+        return json.load(fh).get("generation", {})

senseshift/text_utils.py

@@ -0,0 +1,99 @@
+"""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)))

senseshift_config.json

@@ -0,0 +1,39 @@
+{
+  "senseshift_version": "0.1.0",
+  "base_model": "answerdotai/ModernBERT-large",
+  "control_type": "text_token",
+  "sentiment_scorer": "vader",
+  "sentiment_grid": [
+    -1.0,
+    -0.9,
+    -0.8,
+    -0.7,
+    -0.6,
+    -0.5,
+    -0.4,
+    -0.3,
+    -0.2,
+    -0.1,
+    0.0,
+    0.1,
+    0.2,
+    0.3,
+    0.4,
+    0.5,
+    0.6,
+    0.7,
+    0.8,
+    0.9,
+    1.0
+  ],
+  "generation": {
+    "top_k": 40,
+    "beam_size": 2,
+    "max_iters": 30,
+    "alpha": 0.7,
+    "gamma": 0.05,
+    "temperature": 0.8,
+    "min_words": 3,
+    "add_num_masks": 12
+  }
+}

tokenizer_config.json

@@ -0,0 +1,46 @@
+{
+  "backend": "tokenizers",
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "extra_special_tokens": [
+    "[-1.0]",
+    "[-0.9]",
+    "[-0.8]",
+    "[-0.7]",
+    "[-0.6]",
+    "[-0.5]",
+    "[-0.4]",
+    "[-0.3]",
+    "[-0.2]",
+    "[-0.1]",
+    "[0.0]",
+    "[0.1]",
+    "[0.2]",
+    "[0.3]",
+    "[0.4]",
+    "[0.5]",
+    "[0.6]",
+    "[0.7]",
+    "[0.8]",
+    "[0.9]",
+    "[1.0]"
+  ],
+  "is_local": true,
+  "mask_token": "[MASK]",
+  "max_length": 512,
+  "model_input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "model_max_length": 8192,
+  "pad_to_multiple_of": null,
+  "pad_token": "[PAD]",
+  "pad_token_type_id": 0,
+  "padding_side": "right",
+  "sep_token": "[SEP]",
+  "stride": 0,
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "truncation_side": "right",
+  "truncation_strategy": "longest_first",
+  "unk_token": "[UNK]"
+}