muscriptor 0.1.0__py3-none-any.whl

muscriptor/utils/audio.py

@@ -0,0 +1,112 @@
+"""Audio loading and resampling utilities. WAV is handled by the stdlib;
+other formats fall back to `soundfile`."""
+
+import wave
+from pathlib import Path
+from typing import IO
+
+import numpy as np
+import torch
+
+from muscriptor.utils.resample import resample_frac
+
+
+def _read_wav_file(source) -> tuple[torch.Tensor, int]:
+    """Load a PCM WAV file using the stdlib `wave` module.
+
+    `source` may be a filesystem path or a binary file-like object.
+
+    Returns:
+        (wav, sr) where wav has shape [C, T] and is float32 in [-1, 1].
+    """
+    if hasattr(source, "read"):
+        opened = wave.open(source, "rb")
+    else:
+        opened = wave.open(str(source), "rb")
+    with opened as wf:
+        n_channels = wf.getnchannels()
+        sr = wf.getframerate()
+        sampwidth = wf.getsampwidth()
+        n_frames = wf.getnframes()
+        raw = wf.readframes(n_frames)
+
+    if sampwidth == 1:
+        data = np.frombuffer(raw, dtype=np.uint8).astype(np.float32)
+        data = (data - 128.0) / 128.0
+    elif sampwidth == 2:
+        data = np.frombuffer(raw, dtype=np.int16).astype(np.float32) / 32768.0
+    elif sampwidth == 3:
+        bytes_ = np.frombuffer(raw, dtype=np.uint8).reshape(-1, 3)
+        as_int32 = (
+            bytes_[:, 0].astype(np.int32)
+            | (bytes_[:, 1].astype(np.int32) << 8)
+            | (bytes_[:, 2].astype(np.int32) << 16)
+        )
+        as_int32 = np.where(as_int32 >= (1 << 23), as_int32 - (1 << 24), as_int32)
+        data = as_int32.astype(np.float32) / float(1 << 23)
+    elif sampwidth == 4:
+        data = np.frombuffer(raw, dtype=np.int32).astype(np.float32) / float(1 << 31)
+    else:
+        raise ValueError(f"Unsupported WAV sample width: {sampwidth} bytes")
+
+    data = data.reshape(-1, n_channels)
+    return torch.from_numpy(np.ascontiguousarray(data.T)), sr
+
+
+def _read_non_wav_file(source: str | Path | IO[bytes]) -> tuple[torch.Tensor, int]:
+    """Load a non-WAV audio file using `soundfile`.
+
+    `source` may be a filesystem path or a binary file-like object (e.g. an
+    ``io.BytesIO`` of an uploaded file), since libsndfile reads either.
+
+    Returns:
+        (wav, sr) where wav has shape [C, T] and is float32 in [-1, 1].
+    """
+    try:
+        import soundfile as sf
+    except ImportError as e:
+        raise ImportError(
+            "soundfile is required to read non-WAV audio files. "
+            "Install with: `pip install soundfile` or `uvx --with soundfile`"
+        ) from e
+
+    target = str(source) if isinstance(source, (str, Path)) else source
+    data, sample_rate = sf.read(target, dtype="float32")
+    if data.ndim == 1:
+        data = data[:, None]
+    wav = torch.from_numpy(np.ascontiguousarray(data.T))
+    return wav, sample_rate
+
+
+def resample(
+    waveform: torch.Tensor,
+    orig_freq: int,
+    new_freq: int,
+) -> torch.Tensor:
+    """Sinc resampler via julius `resample_frac`. Operates along the last dim."""
+    if orig_freq == new_freq:
+        return waveform
+    return resample_frac(waveform, int(orig_freq), int(new_freq))
+
+
+def load_audio(path: str | Path, target_sr: int = 16000) -> torch.Tensor:
+    """Load an audio file and return a mono float32 tensor at target_sr.
+
+    PCM WAV files are read with the stdlib `wave` module. Other formats (mp3,
+    flac, ogg, m4a, …) are decoded via `soundfile`. Dispatch is by content, not
+    file extension, so misnamed files (e.g. an MP3 upload saved as .wav) still
+    load.
+
+    Returns:
+        Tensor of shape [1, T] at target_sr.
+    """
+    filepath = Path(path)
+    try:
+        wav, sr = _read_wav_file(str(filepath))
+    except (wave.Error, EOFError):
+        wav, sr = _read_non_wav_file(str(filepath))
+    if wav.shape[0] > 1:
+        wav = wav.mean(dim=0, keepdim=True)
+    if sr != target_sr:
+        wav = resample(wav, sr, target_sr)
+    return wav

muscriptor/utils/auralization.py

@@ -0,0 +1,145 @@
+"""FluidSynth-based MIDI auralization.
+
+Synthesizes a MIDI file with FluidSynth and blends the result with the
+original audio into a stereo mix (L = original, R = synthesis).
+
+Requires:
+  - fluidsynth on the system PATH
+  - soundfile Python package (already a muscriptor dependency)
+"""
+
+import os
+import subprocess
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+
+from muscriptor.utils.audio import load_audio
+
+_DEFAULT_SOUNDFONT = Path(__file__).parent.parent.parent / "MuseScore_General.sf2"
+_SAMPLE_RATE = 44100
+
+
+def _load_mono_44k(path: Path) -> np.ndarray:
+    """Return a mono float32 numpy array at 44100 Hz for any audio file."""
+    wav = load_audio(str(path), target_sr=_SAMPLE_RATE)  # [1, T]
+    return wav[0].numpy()
+
+
+def _resolve_soundfont(soundfont_path: str | Path | None) -> Path:
+    soundfont_path = Path(soundfont_path) if soundfont_path else _DEFAULT_SOUNDFONT
+    if not soundfont_path.exists():
+        raise FileNotFoundError(
+            f"SoundFont not found: {soundfont_path}\n"
+            "Pass --soundfont /path/to/file.sf2 or ensure the bundled "
+            "MuseScore_General.sf2 is present in the muscriptor project root."
+        )
+    return soundfont_path
+
+
+def _synthesize_midi(midi_path: Path, soundfont_path: Path) -> np.ndarray:
+    """Render a MIDI file with FluidSynth → mono float32 array at 44100 Hz.
+
+    Raises:
+        RuntimeError: If fluidsynth is not available or fails.
+    """
+    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
+        synth_tmp = tmp.name
+    try:
+        # Options must precede the positional soundfont/MIDI arguments:
+        # fluidsynth >= 2.5 silently ignores trailing options (exit 0, no
+        # output file written).
+        result = subprocess.run(
+            [
+                "fluidsynth", "-ni",
+                "-F", synth_tmp,
+                "-r", str(_SAMPLE_RATE),
+                str(soundfont_path), str(midi_path),
+            ],
+            capture_output=True,
+        )
+        if result.returncode != 0:
+            raise RuntimeError(
+                f"fluidsynth failed (exit {result.returncode}).\n"
+                "Ensure fluidsynth is installed and the SoundFont path is correct.\n"
+                f"stderr: {result.stderr.decode(errors='replace')}"
+            )
+        synth_audio, _ = sf.read(synth_tmp, dtype="float32")
+        if synth_audio.ndim > 1:
+            synth_audio = synth_audio.mean(axis=1)
+        return synth_audio
+    finally:
+        if os.path.exists(synth_tmp):
+            os.remove(synth_tmp)
+
+
+def synthesize(
+    midi_path: str | Path,
+    output_path: str | Path,
+    soundfont_path: str | Path | None = None,
+) -> None:
+    """Render just the transcription: MIDI → mono WAV via FluidSynth.
+
+    Args:
+        midi_path: Path to the MIDI file to synthesize.
+        output_path: Destination WAV file path.
+        soundfont_path: Path to a ``.sf2`` SoundFont file.
+            Defaults to the bundled MuseScore_General.sf2.
+
+    Raises:
+        RuntimeError: If fluidsynth is not available or fails.
+        FileNotFoundError: If the SoundFont file is not found.
+    """
+    soundfont = _resolve_soundfont(soundfont_path)
+    synth_audio = _synthesize_midi(Path(midi_path), soundfont)
+    sf.write(str(output_path), synth_audio, _SAMPLE_RATE)
+
+
+def auralize(
+    midi_path: str | Path,
+    original_audio_path: str | Path,
+    output_path: str | Path,
+    soundfont_path: str | Path | None = None,
+) -> None:
+    """Create a stereo auralization of a transcription.
+
+    Left channel:  original audio
+    Right channel: FluidSynth MIDI synthesis (RMS-matched to original)
+
+    Args:
+        midi_path: Path to the MIDI file to synthesize.
+        original_audio_path: Path to the source audio file (any format soundfile supports).
+        output_path: Destination WAV file path.
+        soundfont_path: Path to a ``.sf2`` SoundFont file.
+            Defaults to the bundled MuseScore_General.sf2.
+
+    Raises:
+        RuntimeError: If fluidsynth is not available or fails.
+        FileNotFoundError: If the SoundFont file is not found.
+    """
+    original_audio_path = Path(original_audio_path)
+    output_path = Path(output_path)
+    soundfont = _resolve_soundfont(soundfont_path)
+
+    # 1. Synthesize MIDI via FluidSynth
+    synth_audio = _synthesize_midi(Path(midi_path), soundfont)
+
+    # 2. Load original audio at 44100 Hz mono
+    original_audio = _load_mono_44k(original_audio_path)
+
+    # 3. Pad both to the same length
+    length = max(len(original_audio), len(synth_audio))
+    original_audio = np.pad(original_audio, (0, length - len(original_audio)))
+    synth_audio = np.pad(synth_audio, (0, length - len(synth_audio)))
+
+    # 4. RMS-normalize synthesis to match the original's loudness
+    rms_orig = np.sqrt(np.mean(original_audio ** 2))
+    rms_synth = np.sqrt(np.mean(synth_audio ** 2))
+    if rms_synth > 1e-8:
+        synth_audio = synth_audio * (rms_orig / rms_synth)
+
+    # 5. Assemble stereo array [T, 2] and write WAV
+    stereo = np.stack([original_audio, synth_audio], axis=1)
+    sf.write(str(output_path), stereo, _SAMPLE_RATE)

muscriptor/utils/download.py

@@ -0,0 +1,68 @@
+"""Weight download utility with caching in ~/.cache/muscriptor/."""
+
+import hashlib
+import urllib.request
+from pathlib import Path
+from huggingface_hub import hf_hub_download
+from huggingface_hub.errors import HfHubHTTPError
+from huggingface_hub.utils import EntryNotFoundError
+
+
+_CACHE_DIR = Path.home() / ".cache" / "muscriptor"
+
+
+def download_if_necessary(url: str | Path) -> Path:
+    """Resolve a weights location to a local file, downloading if necessary.
+
+    Args:
+        url: Where to find the weights:
+            - ``hf://<repo_id>/<path/in/repo>`` — downloaded via huggingface_hub.
+            - ``http(s)://…`` — fetched with a plain HTTP GET and cached under
+              the cache dir (filename prefixed with a hash of the URL).
+            - anything else (a local path, as ``str`` or ``Path``) — used as-is;
+              nothing is downloaded, but the file must already exist.
+
+    Returns:
+        Path to the local file.
+    """
+    if isinstance(url, str) and url.startswith("hf://"):
+        org, name, hf_filename = url[len("hf://") :].split("/", 2)
+        cached = hf_hub_download(repo_id=f"{org}/{name}", filename=hf_filename)
+        return Path(cached)
+
+    if isinstance(url, str) and url.startswith(("http://", "https://")):
+        # Prefix the cache filename with a hash of the URL so two different URLs
+        # that share a filename don't map to the same file.
+        _CACHE_DIR.mkdir(parents=True, exist_ok=True)
+        filename = url.split("/")[-1].split("?")[0]
+        url_hash = hashlib.sha256(url.encode()).hexdigest()[:8]
+        dest = _CACHE_DIR / f"{url_hash}_{filename}"
+        if dest.exists():
+            return dest
+        print(f"Downloading {filename} …")
+        urllib.request.urlretrieve(url, dest)
+        return dest
+
+    # Local file — nothing to download, just check it's there.
+    path = Path(url)
+    if not path.exists():
+        raise FileNotFoundError(f"weights file not found: {path}")
+    return path
+
+
+def download_companion(url: str | Path, filename: str) -> Path | None:
+    """Best-effort fetch of a sibling file from the same ``hf://`` repo.
+
+    Used to grab a model's ``config.json`` next to its weights. Returns the
+    local path, or ``None`` if ``url`` isn't an ``hf://`` URL or the file can't
+    be fetched — repo/file missing, gated, or offline (so callers can fall back
+    to other detection schemes rather than failing the whole load).
+    """
+    if not (isinstance(url, str) and url.startswith("hf://")):
+        return None
+    org, name, _ = url[len("hf://") :].split("/", 2)
+    try:
+        cached = hf_hub_download(repo_id=f"{org}/{name}", filename=filename)
+    except (EntryNotFoundError, HfHubHTTPError):
+        return None
+    return Path(cached)

muscriptor/utils/midi.py

@@ -0,0 +1,25 @@
+"""MIDI output utilities."""
+
+from pathlib import Path
+
+from muscriptor.tokenizer.notes import Note, note2note_event, note_event2midi
+
+
+def notes_to_midi(notes: list[Note], velocity: int = 100, tempo_bpm: int = 120):
+    """Convert a list of Note objects to a mido MidiFile."""
+    note_events = note2note_event(notes)
+    tempo_us = int(60_000_000 / tempo_bpm)
+    return note_event2midi(
+        note_events, output_file=None, velocity=velocity, tempo=tempo_us
+    )
+
+
+def save_midi(
+    notes: list[Note],
+    path: str | Path,
+    velocity: int = 100,
+    tempo_bpm: int = 120,
+) -> None:
+    """Save a list of Note objects as a MIDI file."""
+    midi = notes_to_midi(notes, velocity=velocity, tempo_bpm=tempo_bpm)
+    midi.save(str(path))

muscriptor/utils/resample.py

@@ -0,0 +1,191 @@
+# File under the MIT license, see https://github.com/adefossez/julius/LICENSE for details.
+# Author: adefossez, 2020
+"""
+Differentiable, Pytorch based resampling.
+Implementation of Julius O. Smith algorithm for resampling.
+See https://ccrma.stanford.edu/~jos/resample/ for details.
+This implementation is specially optimized for when new_sr / old_sr is a fraction
+with a small numerator and denominator when removing the gcd (e.g. new_sr = 700, old_sr = 500).
+
+Very similar to [bmcfee/resampy](https://github.com/bmcfee/resampy) except this implementation
+is optimized for the case mentioned before, while resampy is slower but more general.
+
+"""
+
+import math
+from typing import Optional
+
+import torch
+from torch.nn import functional as F
+
+
+def sinc(x: torch.Tensor) -> torch.Tensor:
+    """sin(x) / x, with the limit value 1 at x == 0.
+
+    __Warning__: the input is not multiplied by `pi`!
+    """
+    return torch.where(
+        x == 0,
+        torch.tensor(1.0, device=x.device, dtype=x.dtype),
+        torch.sin(x) / x,
+    )
+
+
+class ResampleFrac(torch.nn.Module):
+    """
+    Resampling from the sample rate `old_sr` to `new_sr`.
+    """
+
+    def __init__(
+        self, old_sr: int, new_sr: int, zeros: int = 24, rolloff: float = 0.945
+    ):
+        """
+        Args:
+            old_sr (int): sample rate of the input signal x.
+            new_sr (int): sample rate of the output.
+            zeros (int): number of zero crossing to keep in the sinc filter.
+            rolloff (float): use a lowpass filter that is `rolloff * new_sr / 2`,
+                to ensure sufficient margin due to the imperfection of the FIR filter used.
+                Lowering this value will reduce anti-aliasing, but will reduce some of the
+                highest frequencies.
+
+        Shape:
+
+            - Input: `[*, T]`
+            - Output: `[*, T']` with `T' = int(new_sr * T / old_sr)
+
+
+        .. caution::
+            After dividing `old_sr` and `new_sr` by their GCD, both should be small
+            for this implementation to be fast.
+
+        >>> import torch
+        >>> resample = ResampleFrac(4, 5)
+        >>> x = torch.randn(1000)
+        >>> print(len(resample(x)))
+        1250
+        """
+        super().__init__()
+        if not isinstance(old_sr, int) or not isinstance(new_sr, int):
+            raise ValueError("old_sr and new_sr should be integers")
+        gcd = math.gcd(old_sr, new_sr)
+        self.old_sr = old_sr // gcd
+        self.new_sr = new_sr // gcd
+        self.zeros = zeros
+        self.rolloff = rolloff
+
+        self._init_kernels()
+
+    def _init_kernels(self):
+        if self.old_sr == self.new_sr:
+            return
+
+        kernels = []
+        sr = min(self.new_sr, self.old_sr)
+        # rolloff will perform antialiasing filtering by removing the highest frequencies.
+        # At first I thought I only needed this when downsampling, but when upsampling
+        # you will get edge artifacts without this, the edge is equivalent to zero padding,
+        # which will add high freq artifacts.
+        sr *= self.rolloff
+
+        # The key idea of the algorithm is that x(t) can be exactly reconstructed from x[i] (tensor)
+        # using the sinc interpolation formula:
+        #   x(t) = sum_i x[i] sinc(pi * old_sr * (i / old_sr - t))
+        # We can then sample the function x(t) with a different sample rate:
+        #    y[j] = x(j / new_sr)
+        # or,
+        #    y[j] = sum_i x[i] sinc(pi * old_sr * (i / old_sr - j / new_sr))
+
+        # We see here that y[j] is the convolution of x[i] with a specific filter, for which
+        # we take an FIR approximation, stopping when we see at least `zeros` zeros crossing.
+        # But y[j+1] is going to have a different set of weights and so on, until y[j + new_sr].
+        # Indeed:
+        # y[j + new_sr] = sum_i x[i] sinc(pi * old_sr * ((i / old_sr - (j + new_sr) / new_sr))
+        #               = sum_i x[i] sinc(pi * old_sr * ((i - old_sr) / old_sr - j / new_sr))
+        #               = sum_i x[i + old_sr] sinc(pi * old_sr * (i / old_sr - j / new_sr))
+        # so y[j+new_sr] uses the same filter as y[j], but on a shifted version of x by `old_sr`.
+        # This will explain the F.conv1d after, with a stride of old_sr.
+        self._width = math.ceil(self.zeros * self.old_sr / sr)
+        # If old_sr is still big after GCD reduction, most filters will be very unbalanced, i.e.,
+        # they will have a lot of almost zero values to the left or to the right...
+        # There is probably a way to evaluate those filters more efficiently, but this is kept for
+        # future work.
+        idx = torch.arange(-self._width, self._width + self.old_sr).float()
+        for i in range(self.new_sr):
+            t = (-i / self.new_sr + idx / self.old_sr) * sr
+            t = t.clamp_(-self.zeros, self.zeros)
+            t *= math.pi
+            window = torch.cos(t / self.zeros / 2) ** 2
+            kernel = sinc(t) * window
+            # Renormalize kernel to ensure a constant signal is preserved.
+            kernel.div_(kernel.sum())
+            kernels.append(kernel)
+
+        self.register_buffer("kernel", torch.stack(kernels).view(self.new_sr, 1, -1))
+
+    def forward(
+        self, x: torch.Tensor, output_length: Optional[int] = None, full: bool = False
+    ):
+        """
+        Resample x.
+        Args:
+            x (Tensor): signal to resample, time should be the last dimension
+            output_length (None or int): This can be set to the desired output length
+                (last dimension). Allowed values are between 0 and
+                ceil(length * new_sr / old_sr). When None (default) is specified, the
+                floored output length will be used. In order to select the largest possible
+                size, use the `full` argument.
+            full (bool): return the longest possible output from the input. This can be useful
+                if you chain resampling operations, and want to give the `output_length` only
+                for the last one, while passing `full=True` to all the other ones.
+        """
+        if self.old_sr == self.new_sr:
+            return x
+        shape = x.shape
+        length = x.shape[-1]
+        x = x.reshape(-1, length)
+        x = F.pad(
+            x[:, None], (self._width, self._width + self.old_sr), mode="replicate"
+        )
+        ys = F.conv1d(x, self.kernel, stride=self.old_sr)  # type: ignore
+        y = ys.transpose(1, 2).reshape(list(shape[:-1]) + [-1])
+
+        float_output_length = torch.as_tensor(self.new_sr * length / self.old_sr)
+        max_output_length = torch.ceil(float_output_length).long()
+        default_output_length = torch.floor(float_output_length).long()
+
+        if output_length is None:
+            applied_output_length = max_output_length if full else default_output_length
+        elif output_length < 0 or output_length > max_output_length:
+            raise ValueError(f"output_length must be between 0 and {max_output_length}")
+        else:
+            applied_output_length = torch.tensor(output_length)
+            if full:
+                raise ValueError("You cannot pass both full=True and output_length")
+        return y[..., :applied_output_length]  # type: ignore
+
+    def __repr__(self):
+        return (
+            f"ResampleFrac(old_sr={self.old_sr}, new_sr={self.new_sr}, "
+            f"zeros={self.zeros}, rolloff={self.rolloff})"
+        )
+
+
+def resample_frac(
+    x: torch.Tensor,
+    old_sr: int,
+    new_sr: int,
+    zeros: int = 24,
+    rolloff: float = 0.945,
+    output_length: Optional[int] = None,
+    full: bool = False,
+):
+    """
+    Functional version of `ResampleFrac`, refer to its documentation for more information.
+
+    ..warning::
+        If you call repeatidly this functions with the same sample rates, then the
+        resampling kernel will be recomputed everytime. For best performance, you should use
+        and cache an instance of `ResampleFrac`.
+    """
+    return ResampleFrac(old_sr, new_sr, zeros, rolloff).to(x)(x, output_length, full)

muscriptor/utils/sampling.py

@@ -0,0 +1,89 @@
+"""Sampling utilities for autoregressive generation."""
+
+import torch
+
+
+def length_to_mask(lengths: torch.Tensor, max_len: int | None = None) -> torch.Tensor:
+    """Convert a tensor of sequence lengths to a boolean mask."""
+    assert len(lengths.shape) == 1
+    final_length = int(lengths.max().item()) if not max_len else max_len
+    final_length = max(final_length, 1)
+    return torch.arange(final_length, device=lengths.device)[None, :] < lengths[:, None]
+
+
+def multinomial(
+    input: torch.Tensor, num_samples: int, replacement: bool = False, *, generator=None
+) -> torch.Tensor:
+    """torch.multinomial with arbitrary number of dimensions."""
+    input_ = input.reshape(-1, input.shape[-1])
+    output_ = torch.multinomial(
+        input_, num_samples=num_samples, replacement=replacement, generator=generator
+    )
+    output = output_.reshape(*list(input.shape[:-1]), -1)
+    return output
+
+
+def sample_top_k(probs: torch.Tensor, k: int, num_samples: int = 1) -> torch.Tensor:
+    """Sample from top-k probabilities."""
+    top_k_value, _ = torch.topk(probs, k, dim=-1)
+    min_value_top_k = top_k_value[..., [-1]]
+    probs = probs * (probs >= min_value_top_k).float()
+    probs = probs / probs.sum(dim=-1, keepdim=True)
+    return multinomial(probs, num_samples=num_samples)
+
+
+def sample_top_p(probs: torch.Tensor, p: float, num_samples: int = 1) -> torch.Tensor:
+    """Sample from nucleus (top-p) distribution."""
+    probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
+    probs_sum = torch.cumsum(probs_sort, dim=-1)
+    mask = probs_sum - probs_sort > p
+    probs_sort = probs_sort * (~mask).float()
+    probs_sort = probs_sort / probs_sort.sum(dim=-1, keepdim=True)
+    next_token = multinomial(probs_sort, num_samples=num_samples)
+    return torch.gather(probs_idx, -1, next_token)
+
+
+def sample_from_probs(
+    probs: torch.Tensor, top_p: float = 0.0, top_k: int = 0
+) -> torch.Tensor:
+    """Sample one token from probs, optionally filtered by top-p or top-k."""
+    if top_p > 0.0:
+        return sample_top_p(probs, top_p)
+    if top_k > 0:
+        return sample_top_k(probs, top_k)
+    return multinomial(probs, num_samples=1)
+
+
+def sample_stratified(
+    probs: torch.Tensor,
+    special_token: int,
+    first_temp: float,
+    second_temp: float = 1.0,
+    top_p: float = 0.0,
+    top_k: int = 0,
+) -> torch.Tensor:
+    """Stratified sampling: first decide special vs. non-special, then sample among non-special."""
+    eps = 1e-12
+    probs_special = probs[..., special_token : special_token + 1].clamp(
+        min=eps, max=1 - eps
+    )
+    logits_two = torch.cat(
+        [torch.log(probs_special), torch.log(1 - probs_special)], dim=-1
+    )
+    logits_two = logits_two / max(first_temp, eps)
+    probs_two = torch.softmax(logits_two, dim=-1)
+    probs_special_temp = probs_two[..., 0:1]
+    next_token_is_special = torch.rand_like(probs_special_temp).lt(probs_special_temp)
+
+    denom = (1 - probs_special).clamp(min=eps)
+    new_probs = probs.clone() / denom
+    new_probs[..., special_token] = 0.0
+    if second_temp > 0:
+        log_new = torch.log(new_probs.clamp(min=eps)) / second_temp
+        new_probs = torch.softmax(log_new, dim=-1)
+
+    next_token = sample_from_probs(new_probs, top_p=top_p, top_k=top_k)
+
+    return torch.where(
+        next_token_is_special, torch.full_like(next_token, special_token), next_token
+    )