Focused model card; explicit generation parameters

README.md

@@ -12,8 +12,7 @@ tags:
   - non-autoregressive
   - masked-language-model
   - text-style-transfer
-  - sentiment-transfer
-  - text-rewriting
+  - data-augmentation
   - modernbert
   - emnlp2026
 language:
@@ -22,49 +21,79 @@ language:
 
 # SenseShift-base
 
-> **Generation from an encoder, with continuous control.**
-> An accepted EMNLP 2026 paper. Not a classifier, not a decoder — a masked LM
-> that *writes* text at a sentiment you dial in.
-
-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.
-
-Two things make it unusual:
-
-- **Continuous, not categorical.** Most sentiment-controlled generation offers
-  you `positive` / `negative` / `neutral`. SenseShift takes a *number* — 21
-  points on a 0.1 grid. Sweeping one sentence of a story, everything else fixed:
-
-  | target | achieved | rewritten sentence |
-  | ---: | ---: | --- |
-  | −0.9 | −0.7 | The water was bitter and made him sick. |
-  | −0.6 | −0.4 | The water was bitter and made him cough. |
-  | +0.0 | +0.0 | The water was cold and made him wet. |
-  | +0.6 | +0.7 | The water was clear and tasted very good. |
-  | +0.9 | +0.8 | His friends laughed and played games with him. |
-
-- **The generator is an encoder.** Text generation is nearly always
-  autoregressive decoding. SenseShift generates by iteratively filling masked
-  positions in a bidirectional ModernBERT, so every token it writes is
-  conditioned on the text on *both* sides. That is what lets it drop a
-  replacement sentence into the middle of a passage and have it fit.
-
-It is a masked language model fine-tuned from ModernBERT-base 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.
+Rewrite any sentence in a passage — or write a new one — at a sentiment you
+choose on a continuous scale from **−1.0 to +1.0**, in context.
 
 | Model | Base | Params |
 | --- | --- | --- |
 | **SenseShift-base** (this model) | ModernBERT-base | 150M |
-| [SenseShift-large](https://huggingface.co/shawhed/SenseShift-large) | ModernBERT-large | 396M |
+| [SenseShift-large](https://huggingface.co/shawhed/SenseShift-base) | ModernBERT-large | 396M |
+
+Accepted to EMNLP 2026.
+
+## How it works
+
+Sentiment is a **control vocabulary**: 21 special tokens `[-1.0] … [1.0]` on a
+0.1 grid, added to the ModernBERT tokenizer.
+
+Training: every sentence in a passage is scored with VADER and prefixed with its
+own sentiment token. One sentence is then masked out. The model must reconstruct
+it from the token it was given and the text on both sides — so it learns to
+write a sentence that hits a requested sentiment *and* fits its neighbours.
+
+Generation reverses this. Prefix each sentence with its current sentiment, give
+the target sentence the sentiment you *want*, mask its words, and fill the masks
+with beam search. Because the generator is a bidirectional encoder rather than a
+decoder, every token it writes is conditioned on the text before **and** after —
+which is what lets it replace a sentence mid-passage.
+
+## Examples
+
+Rewriting sentence `[3]` at **−0.8**. The model picks up "tunnel" from the
+preceding sentence:
+
+> Once upon a time, there was a little boy named Timmy. One day, Timmy went to a
+> park with his mommy. At the park, there was a big tunnel that Timmy wanted to
+> explore. **~~Timmy's mommy said he could choose whether or not to go through
+> the tunnel.~~ → But Timmy's mommy said no because it was too dangerous.**
+> Timmy was brave and decided to go through the tunnel. As Timmy crawled through
+> the tunnel, he felt mighty and strong. When he came out the other side, he was
+> so happy and proud of himself. From that day on, Timmy loved going on
+> adventures and choosing to be brave.
+
+Rewriting sentence `[2]` of a restaurant review at **+0.8**. The model carries
+the business name "Lotus 2" in from context:
+
+> I ordered delivery this evening and I must say that I was pleasantly surprised
+> at the level of customer service on the phone and at my door. I usually get my
+> Chinese take out from best wok 2 and they pale in comparison to lotus 2.
+> **~~I ordered the exact same things that I usually order from best wok and
+> everything tasted better and was of better quality and was even a little less
+> expensive.~~ → Lotus 2 are very friendly and always make sure that I get the
+> best quality of food and I get a lot of value for my money.** I am so happy
+> that I've found a new Chinese take out. And they deliver! Of course its not
+> P.F. Changs. But its well worth what you pay.
+
+## Use cases
+
+- **Writing assistant for fiction.** Adjust the emotional arc of a draft one
+  sentence at a time — darken a turning point, soften an ending — without
+  rewriting the surrounding prose.
+- **Review and copy editing.** Retune the tone of a testimonial, product
+  description or release note while keeping the concrete details intact.
+- **Data augmentation.** Generate sentiment-varied paraphrases of a corpus at
+  known target values: balance a skewed sentiment dataset, or produce minimal
+  pairs that differ in sentiment but share context.
+- **Counterfactual and robustness testing.** Probe a downstream classifier with
+  inputs where exactly one sentence's sentiment moved, holding everything else
+  fixed.
+- **Controllable-generation research.** A non-autoregressive baseline for
+  continuous attribute control, and a testbed for how far a control signal can
+  be pushed at inference time.
 
 ## Usage
 
-No install needed — the inference code ships inside this repo. Clone it and you
-have the weights and the code together:
+No install needed — the inference code ships in this repo.
 
 ```bash
 git clone https://huggingface.co/shawhed/SenseShift-base
@@ -78,138 +107,67 @@ from senseshift import SenseShift
 
 shifter = SenseShift.from_pretrained("SenseShift-base")
 
-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.")
+text = ("At the park, there was a big tunnel that Timmy wanted to explore. "
+        "Timmy's mommy said he could choose whether or not to go through the tunnel. "
+        "Timmy was brave and decided to go through the tunnel.")
 
-# 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
+# Rewrite a sentence at a target sentiment
+out = shifter.generate(text, generation_mode="rewrite", sentence_index=1, sentiment=-0.8)
+print(out.sentence)
+print(out.text)
 
-# Add a positive sentence after the last one
-out = shifter.generate(text, generation_mode="add", sentiment=0.7)
+# Write a new sentence and splice it in after sentence 2
+out = shifter.generate(text, generation_mode="add", sentence_index=2, sentiment=0.9)
 print(out.text)
 ```
 
-Or let `huggingface_hub` fetch it into the local cache instead of cloning:
+Requires `torch`, `transformers`, `huggingface-hub`, `nltk`. The VADER lexicon
+downloads itself on first use.
 
-```python
-import sys
-from huggingface_hub import snapshot_download
-
-path = snapshot_download("shawhed/SenseShift-base")
-sys.path.insert(0, path)
+### `generate` arguments
 
-from senseshift import SenseShift
-shifter = SenseShift.from_pretrained(path)
-```
+| Argument | Default | Meaning |
+| --- | --- | --- |
+| `text` | — | The passage to edit. |
+| `generation_mode` | `"rewrite"` | `"rewrite"` replaces the sentence at `sentence_index`; `"add"` inserts a new sentence after it. |
+| `sentiment` | `None` | `None` keeps the current sentiment; `"random"` draws from the grid; a number in `[-1, 1]` is snapped to the nearest 0.1. |
+| `sentence_index` | `None` | Which sentence to act on. Defaults to a random sentence (`rewrite`) or the last one (`add`). Negative indices count from the end. |
+| `num_masks` | `None` | Mask slots given to the model, i.e. roughly how long the new sentence is. Defaults to the replaced sentence's word count, or 12 for `add`. |
+| `seed` | `None` | Seeds the random index / sentiment draws. |
 
-Requirements: `torch`, `transformers`, `huggingface-hub`, `nltk` (see
-`requirements.txt`). The VADER lexicon downloads itself on first use.
+Decoding parameters — leave any at `None` to use the model's defaults:
 
-### `generate` arguments
+| Argument | Default | Meaning |
+| --- | --- | --- |
+| `beam_size` | `2` | Hypotheses kept alive. Higher is slower and usually more fluent. `num_beams` works as an alias. |
+| `top_k` | `40` | Candidate tokens considered per masked position. |
+| `temperature` | `0.8` | Below 1.0 is more conservative, above 1.0 more varied. |
+| `alpha` | `0.7` | Length-normalisation exponent; higher tolerates longer output. |
+| `gamma` | `0.05` | Diversity penalty on beams reusing the same token. |
+| `max_iters` | `30` | Cap on mask-filling steps. |
+| `min_words` | `3` | Tokens filled before the model may stop at punctuation. Raise to avoid very short rewrites. |
 
-| 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. |
+```python
+out = shifter.generate(text, sentence_index=1, sentiment=-0.8, beam_size=8, top_k=60, temperature=0.9)
+```
 
-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` (the edited passage),
+`.sentence` (what was written), `.original_sentence`, `.target_sentiment`,
+`.achieved_sentiment` (VADER of the result) and `.beams`. `str(out)` gives the
+passage.
 
-`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.
+## Training data
 
-## How it works
+TinyStories-style short children's stories, labelled per sentence with VADER.
+Despite the narrow training domain it transfers to other English prose — the
+review example above is out-of-domain.
 
-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.
-
-## Control fidelity, and how far it transfers
-
-Measured over 50 held-out passages per domain, 9 targets each, rewriting one
-randomly chosen sentence. "Achieved" is VADER on the sentence the model wrote.
-
-| target | achieved — stories *(in-domain)* | achieved — reviews *(out-of-domain)* |
-| ---: | ---: | ---: |
-| −0.9 | −0.40 | −0.44 |
-| −0.6 | −0.23 | −0.25 |
-| −0.3 | −0.14 | −0.04 |
-| +0.0 | +0.09 | +0.01 |
-| +0.3 | +0.22 | +0.14 |
-| +0.6 | +0.53 | +0.47 |
-| +0.9 | +0.75 | +0.73 |
-| **MAE** | **0.33** | **0.33** |
-
-Two things to read off this.
-
-**The response is monotone and the positive half is well calibrated** — targets
-of +0.6 and +0.9 land within about 0.15. The negative half is compressed: ask
-for −0.9 and you reliably get *more negative*, but around −0.4 rather than −0.9.
-
-**Domain transfer is essentially free.** SenseShift was trained only on
-short children's stories, yet product reviews — a domain it never saw, with
-different vocabulary, register and length — come out at an identical MAE of
-0.33. The per-target curve is the same shape. Use it outside the story domain
-with reasonable confidence.
-
-## Training data and usable range
-
-Trained on TinyStories-style children's stories, labelled per sentence with
-VADER. That corpus is lopsided, which directly shapes what the model can do:
-
-| sentiment of training sentences | share |
-| --- | ---: |
-| negative (< 0) | 15.3% |
-| neutral (= 0) | 37.8% |
-| positive (> 0) | 47.0% |
-| exactly −1.0 | **0.00%** |
-| exactly +1.0 | 0.10% |
-
-So:
-
-- **Use targets in −0.8 … +0.9.** The `[-1.0]` control token appears *zero*
-  times in training and `[1.0]` almost never; asking for ±1.0 gives an
-  undertrained token and unpredictable output. `-1.0` in particular comes back
-  near-neutral.
-- **Negative control is looser than positive**, because negative sentences are
-  a third as common in the corpus. For strongly negative text, ask for −0.9 and
-  expect roughly −0.4, or resample.
-
-The fidelity numbers above were measured on
-[SenseShift-large](https://huggingface.co/shawhed/SenseShift-large); this 150M
-variant follows the same shape but tracks the target less tightly.
-
-## Limitations
-
-- English only; trained on short children's stories, though it transfers to
-  other domains at no measured cost (see above).
-- VADER supplies the sentiment labels, so the model inherits its lexicon-based
-  view of sentiment.
-- Control is a steer, not a guarantee: MAE 0.33 over the grid, with the negative
-  half compressed and `±1.0` unusable. Single generations are noisy — check
-  `out.achieved_sentiment` and resample if you need a specific value.
-- 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`.
+Negative sentiment is under-represented in that corpus, so the negative half of
+the range is a looser steer than the positive half, and the extremes `±1.0` are
+effectively untrained. **Use targets in −0.9 … +0.9**, and read
+`out.achieved_sentiment` if you need a specific value.
 
 ## Citation
 
-Accepted to **EMNLP 2026**. A preprint and the full BibTeX entry will be linked
-here once they are public.
+Accepted to EMNLP 2026. A preprint and BibTeX entry will be linked here once
+they are public.

senseshift/pipeline.py

@@ -121,7 +121,14 @@ class SenseShift:
         sentence_index: Optional[int] = None,
         num_masks: Optional[int] = None,
         seed: Optional[int] = None,
-        **overrides,
+        beam_size: Optional[int] = None,
+        top_k: Optional[int] = None,
+        temperature: Optional[float] = None,
+        alpha: Optional[float] = None,
+        gamma: Optional[float] = None,
+        max_iters: Optional[int] = None,
+        min_words: Optional[int] = None,
+        num_beams: Optional[int] = None,
     ) -> SenseShiftOutput:
         """Rewrite one sentence, or add a new one, at a controlled sentiment.
 
@@ -140,8 +147,19 @@ class SenseShift:
                 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``.
+            beam_size: Hypotheses kept alive during decoding. Higher is slower
+                and usually more fluent. ``num_beams`` is accepted as an alias.
+            top_k: Candidate tokens considered at each masked position.
+            temperature: Softmax temperature. Below 1.0 is more conservative,
+                above 1.0 more varied.
+            alpha: Length-normalisation exponent; higher tolerates longer output.
+            gamma: Diversity penalty on beams reusing the same token.
+            max_iters: Cap on mask-filling steps.
+            min_words: Tokens that must be filled before the model may stop at
+                terminal punctuation. Raise it to prevent very short rewrites.
+
+        Any decoding argument left as ``None`` falls back to the model's
+        defaults from ``senseshift_config.json``.
 
         Returns:
             A :class:`SenseShiftOutput`; ``str(out)`` is the edited passage.
@@ -151,10 +169,23 @@ class SenseShift:
                 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))}")
+        if num_beams is not None:
+            if beam_size is not None and beam_size != num_beams:
+                raise TypeError("Pass either beam_size or num_beams, not both.")
+            beam_size = num_beams
+
+        overrides = {
+            "beam_size": beam_size, "top_k": top_k, "temperature": temperature,
+            "alpha": alpha, "gamma": gamma, "max_iters": max_iters,
+            "min_words": min_words,
+        }
+        params = {**self.generation_defaults,
+                  **{k: v for k, v in overrides.items() if v is not None}}
+
+        if params["beam_size"] < 1 or params["top_k"] < 1:
+            raise ValueError("beam_size and top_k must be >= 1.")
+        if params["temperature"] <= 0:
+            raise ValueError("temperature must be > 0.")
 
         rng = _random.Random(seed) if seed is not None else _random