SenseShift-base / README.md
shawhed's picture
Fix cross-link to SenseShift-large
b70b4fb verified
|
Raw
History Blame Contribute Delete
7.5 kB
---
license: apache-2.0
base_model: answerdotai/ModernBERT-base
library_name: transformers
pipeline_tag: text-generation
inference: false
tags:
- sentiment-control
- continuous-control
- controllable-text-generation
- encoder-generation
- non-autoregressive
- masked-language-model
- text-style-transfer
- data-augmentation
- modernbert
- emnlp2026
language:
- en
---
# SenseShift-base
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 |
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 in this repo.
```bash
git clone https://huggingface.co/shawhed/SenseShift-base
```
```python
import sys
sys.path.insert(0, "SenseShift-base")
from senseshift import SenseShift
shifter = SenseShift.from_pretrained("SenseShift-base")
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 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)
# 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)
```
Requires `torch`, `transformers`, `huggingface-hub`, `nltk`. The VADER lexicon
downloads itself on first use.
### `generate` arguments
| 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. |
Decoding parameters β€” leave any at `None` to use the model's defaults:
| 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. |
```python
out = shifter.generate(text, sentence_index=1, sentiment=-0.8, beam_size=8, top_k=60, temperature=0.9)
```
`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.
## Training data
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.
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 BibTeX entry will be linked here once
they are public.