Publish BioVoice-TTS sparse energy checkpoint and model card

README.md

@@ -0,0 +1,136 @@
+---
+language:
+- en
+license: other
+library_name: pytorch
+pipeline_tag: text-to-speech
+tags:
+- text-to-speech
+- voice-cloning
+- sparse-attention
+- low-rank
+- cpu-first
+- ljspeech
+- biovoice-tts
+datasets:
+- keithito/lj_speech
+---
+
+# BioVoice-TTS Sparse Energy Voice Model
+
+BioVoice-TTS is a CPU-first text-to-speech and voice-cloning research model built around the same sparse-energy design principles as Bio-LLM/SSET:
+
+- low-rank Q/K/V projections
+- causal sparse candidate attention
+- local, memory, landmark, and content candidate routing
+- laminar excitatory/inhibitory refinement
+- explicit speaker conditioning
+- explicit duration, pitch, and energy prediction
+- sparse acoustic decoding with an energy/gated mel head
+
+This upload contains a compact LJSpeech-trained text-to-mel checkpoint. The optimizer state and intermediate checkpoints were intentionally omitted to keep the repository small.
+
+## Files
+
+- `model.safetensors`: model-only weights converted from local checkpoint `step_8000.pt`
+- `config.json`: BioVoice-TTS architecture and training config
+- `tokenizer.json`: phoneme/character tokenizer used for the LJSpeech run
+- `training_summary.json`: checkpoint metrics summary
+- `training_metrics_step_8000.json`: raw metrics saved with the selected checkpoint
+- `bio_voice_tts/`: model, audio feature, dataset, training, inference, streaming, and vocoder code
+- `bio_llm/`: shared sparse-energy language-model utilities used by the project
+
+## Checkpoint Metrics
+
+Selected checkpoint: `step_8000`
+
+| Metric | Value |
+|---|---:|
+| loss | 1.7827 |
+| mel_loss / mel_mae | 1.6115 |
+| duration_loss | 0.0077 |
+| pitch_loss | 0.3701 |
+| energy_loss | 1.3269 |
+| speaker_cosine proxy | 1.0000 |
+
+These are internal training metrics from the local run, not standardized MOS, WER, speaker-verification EER, or cross-model benchmark scores.
+
+## Architecture Path
+
+The real forward path is:
+
+1. Reference mel -> `SpeakerEncoder`
+2. Text tokens -> `SemanticEncoder`
+3. Semantic states + speaker latent -> FiLM conditioning
+4. Duration predictor -> length regulation
+5. Pitch and energy predictors -> frame-level controls
+6. Frame states + speaker + pitch + energy -> `SparseAcousticDecoder`
+7. Acoustic energy/gating head -> mel spectrogram
+8. Optional `SparseNeuralVocoder` -> waveform
+
+The model is not a wrapper around Tacotron, FastSpeech, VITS, XTTS, StyleTTS, or F5-TTS. It is a custom sparse-energy TTS architecture. Some config fields are reserved or scaffolded and are not fully wired yet; see the limitations section.
+
+## Minimal Loading Example
+
+```python
+import json
+import torch
+from safetensors.torch import load_file
+
+from bio_voice_tts import BioVoiceConfig, BioVoiceTTS
+
+def merge_dataclass(instance, payload):
+    for key, value in payload.items():
+        current = getattr(instance, key)
+        if hasattr(current, "__dataclass_fields__") and isinstance(value, dict):
+            merge_dataclass(current, value)
+        else:
+            setattr(instance, key, value)
+    return instance
+
+config = merge_dataclass(BioVoiceConfig(), json.load(open("config.json")))
+model = BioVoiceTTS(config)
+model.load_state_dict(load_file("model.safetensors"), strict=False)
+model.eval()
+
+token_ids = torch.randint(0, config.semantic.vocab_size, (1, 16))
+reference_mel = torch.randn(1, 128, config.audio.n_mels)
+
+with torch.no_grad():
+    outputs = model(token_ids, reference_mel)
+
+print(outputs["mel"].shape)
+```
+
+For waveform synthesis, pass `outputs["mel"]` through `bio_voice_tts.vocoder.sparse_vocoder.SparseNeuralVocoder`. A separately trained vocoder checkpoint is recommended for production-quality audio.
+
+## Real Comparison Snapshot
+
+This table compares design and deployment tradeoffs, not universal audio quality. BioVoice-TTS has not yet been evaluated with a public MOS/WER/EER benchmark suite against these systems.
+
+| Model | Publicly known strength | Where BioVoice-TTS is different | Where BioVoice-TTS is currently weaker |
+|---|---|---|---|
+| Coqui XTTS-v2 | Mature multilingual voice cloning; model card states 17 languages and cloning from a short reference clip. Source: https://huggingface.co/coqui/XTTS-v2 | BioVoice-TTS is smaller here and built around sparse low-rank CPU-first modules rather than a large ready-to-use multilingual stack. | XTTS-v2 is more production-ready, multilingual, and widely tested. BioVoice-TTS currently has a single-speaker LJSpeech checkpoint and needs more evaluation. |
+| StyleTTS 2 | Paper reports human-level/super-human judged naturalness on LJSpeech/VCTK settings using style diffusion and adversarial training with speech language models. Source: https://arxiv.org/abs/2306.07691 | BioVoice-TTS avoids diffusion/style sampling and emphasizes interpretable sparse memory, laminar refinement, and CPU-oriented low-rank compute. | StyleTTS 2 has stronger published quality claims. BioVoice-TTS does not yet have MOS evidence at that level. |
+| F5-TTS | Flow-matching DiT system; paper emphasizes a simpler non-autoregressive design without explicit duration model/text encoder/phoneme alignment. Source: https://arxiv.org/abs/2410.06885 | BioVoice-TTS intentionally keeps explicit duration, pitch, energy, speaker, memory, and acoustic components for control and interpretability. | F5-TTS is a stronger modern zero-shot baseline for naturalness/voice cloning. BioVoice-TTS needs more training scale and public listening tests. |
+| OpenVoice | Instant voice cloning from a short reference; paper focuses on flexible tone-color cloning and multilingual generation. Source: https://arxiv.org/abs/2312.01479 | BioVoice-TTS is an end-to-end sparse text-to-mel architecture with an optional neural vocoder, not mainly a tone-color conversion stack. | OpenVoice has a clearer instant-cloning product path. BioVoice-TTS still needs stronger speaker verification and cloning evaluations. |
+| Older Tacotron/FastSpeech-style systems | Stable, well-known TTS baselines. | BioVoice-TTS has richer sparse attention, memory routing, laminar refinement, and explicit energy/gated acoustic decoding. | Older systems may have simpler tooling and better documented recipes; BioVoice-TTS is research-stage. |
+
+## Limitations
+
+- This upload is primarily a research checkpoint and code release.
+- It is trained on LJSpeech-style single-speaker data, not a broad multilingual/multi-speaker corpus.
+- No standardized MOS, WER, speaker EER, latency benchmark, or human preference study is included yet.
+- Some architecture config fields are currently scaffolded or hard-coded rather than fully parameterized.
+- The published checkpoint is text-to-mel. For final audio quality, use or train a matching vocoder checkpoint.
+- Do not use for impersonation, fraud, or cloning a voice without consent.
+
+## Suggested Evaluation Before Production
+
+- MOS or MUSHRA-style listening test against XTTS-v2, StyleTTS 2, F5-TTS, and OpenVoice
+- Speaker similarity with a real speaker verification model
+- WER using an ASR model to verify intelligibility
+- CPU latency at multiple text lengths
+- Long-form stability and pronunciation tests
+- Ablations for sparse memory, laminar refinement, and acoustic energy gates
+

bio_llm/__init__.py

@@ -0,0 +1 @@
+"""Bio-LLM package."""

bio_llm/model/__init__.py

@@ -0,0 +1 @@
+"""Model components for the Structured Sparse Energy Transformer."""

bio_llm/model/candidate_retrieval.py

@@ -0,0 +1,61 @@
+import math
+
+import torch
+from torch import nn
+
+
+class CandidateRetriever(nn.Module):
+    """Two-stage candidate filtering to avoid a full-vocabulary energy pass."""
+
+    def __init__(self, vocab_size: int, d_model: int, stage1_dim: int, stage1_k: int, stage2_k: int):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.stage1_k = stage1_k
+        self.stage2_k = stage2_k
+        self.query_low = nn.Parameter(torch.empty(d_model, stage1_dim))
+        self.vocab_low = nn.Parameter(torch.empty(vocab_size, stage1_dim))
+        nn.init.normal_(self.query_low, mean=0.0, std=1.0 / math.sqrt(d_model))
+        nn.init.normal_(self.vocab_low, mean=0.0, std=1.0 / math.sqrt(stage1_dim))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        embedding_weight: torch.Tensor,
+        target_ids: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        batch_size, seq_len, _ = hidden_states.shape
+        low_queries = torch.matmul(hidden_states, self.query_low)
+        low_scores = torch.matmul(low_queries, self.vocab_low.transpose(0, 1))
+        stage1_k = min(self.stage1_k, self.vocab_size)
+        coarse_candidates = torch.topk(low_scores, k=stage1_k, dim=-1).indices
+        final_candidates = torch.zeros(
+            batch_size, seq_len, self.stage2_k, dtype=torch.long, device=hidden_states.device
+        )
+
+        for batch_index in range(batch_size):
+            for position in range(seq_len):
+                candidate_ids = coarse_candidates[batch_index, position]
+                candidate_embeddings = embedding_weight[candidate_ids]
+                fine_scores = torch.matmul(candidate_embeddings, hidden_states[batch_index, position])
+
+                if target_ids is not None:
+                    target_id = int(target_ids[batch_index, position].item())
+                    if target_id not in candidate_ids.tolist():
+                        candidate_ids = torch.cat(
+                            [candidate_ids[:-1], torch.tensor([target_id], device=hidden_states.device)]
+                        )
+                        candidate_embeddings = embedding_weight[candidate_ids]
+                        fine_scores = torch.matmul(candidate_embeddings, hidden_states[batch_index, position])
+
+                stage2_k = min(self.stage2_k, candidate_ids.numel())
+                top_indices = torch.topk(fine_scores, k=stage2_k).indices
+                chosen = candidate_ids[top_indices]
+
+                if target_ids is not None:
+                    target_id = int(target_ids[batch_index, position].item())
+                    if target_id not in chosen.tolist():
+                        chosen[-1] = target_id
+
+                final_candidates[batch_index, position, :stage2_k] = chosen[:stage2_k]
+
+        return final_candidates

bio_llm/model/embedding.py

@@ -0,0 +1,18 @@
+import torch
+from torch import nn
+
+
+class TokenEmbedding(nn.Module):
+    """Maps token ids to continuous vectors x_t = W_embed[token]."""
+
+    def __init__(self, vocab_size: int, d_model: int):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, d_model)
+        nn.init.normal_(self.embedding.weight, mean=0.0, std=0.02)
+
+    @property
+    def weight(self) -> torch.Tensor:
+        return self.embedding.weight
+
+    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
+        return self.embedding(token_ids)

bio_llm/model/energy_head.py

@@ -0,0 +1,96 @@
+import math
+
+import torch
+from torch import nn
+
+
+class FactorizedTransitionBias(nn.Module):
+    """Approximates transition_bias(prev_token, y) with a low-rank factorization."""
+
+    def __init__(self, vocab_size: int, rank: int):
+        super().__init__()
+        self.prev_factor = nn.Embedding(vocab_size, rank)
+        self.next_factor = nn.Embedding(vocab_size, rank)
+        nn.init.normal_(self.prev_factor.weight, mean=0.0, std=1.0 / math.sqrt(rank))
+        nn.init.normal_(self.next_factor.weight, mean=0.0, std=1.0 / math.sqrt(rank))
+
+    def forward(self, prev_tokens: torch.Tensor, candidate_ids: torch.Tensor) -> torch.Tensor:
+        prev_repr = self.prev_factor(prev_tokens).unsqueeze(-2)
+        next_repr = self.next_factor(candidate_ids)
+        return torch.sum(prev_repr * next_repr, dim=-1)
+
+
+class EnergyHead(nn.Module):
+    """Computes E(y) and P(y) on a sparse candidate set."""
+
+    def __init__(self, vocab_size: int, d_model: int, transition_rank: int):
+        super().__init__()
+        self.transition_bias = FactorizedTransitionBias(vocab_size, transition_rank)
+        self.d_model = d_model
+        self.log_temperature = nn.Parameter(torch.zeros(1))
+        self.temperature_floor = 0.5
+
+    def _context_embedding_summary(
+        self,
+        input_ids: torch.Tensor,
+        attention_indices: torch.Tensor,
+        attention_weights: torch.Tensor,
+        embedding_weight: torch.Tensor,
+    ) -> torch.Tensor:
+        batch_size, seq_len, _ = attention_indices.shape
+        summary = torch.zeros(batch_size, seq_len, self.d_model, device=input_ids.device, dtype=embedding_weight.dtype)
+        for batch_index in range(batch_size):
+            for position in range(seq_len):
+                token_positions = attention_indices[batch_index, position]
+                token_ids = input_ids[batch_index, token_positions]
+                token_embeddings = embedding_weight[token_ids]
+                weights = attention_weights[batch_index, position].unsqueeze(-1)
+                summary[batch_index, position] = torch.sum(token_embeddings * weights, dim=0)
+        return summary
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        input_ids: torch.Tensor,
+        prev_tokens: torch.Tensor,
+        candidate_ids: torch.Tensor,
+        attention_indices: torch.Tensor,
+        attention_weights: torch.Tensor,
+        embedding_weight: torch.Tensor,
+    ) -> dict[str, torch.Tensor]:
+        candidate_embeddings = embedding_weight[candidate_ids]
+        context_summary = self._context_embedding_summary(
+            input_ids=input_ids,
+            attention_indices=attention_indices,
+            attention_weights=attention_weights,
+            embedding_weight=embedding_weight,
+        )
+
+        e_sim = -torch.sum(hidden_states.unsqueeze(-2) * candidate_embeddings, dim=-1)
+        e_ctx = -torch.sum(context_summary.unsqueeze(-2) * candidate_embeddings, dim=-1)
+        e_mem = -self.transition_bias(prev_tokens, candidate_ids)
+
+        confidence = attention_weights.max(dim=-1).values
+        g1 = confidence
+        g2 = 1.0 - torch.abs(0.5 - confidence)
+        g3 = 1.0 - confidence
+        energies = g1.unsqueeze(-1) * e_sim + g2.unsqueeze(-1) * e_ctx + g3.unsqueeze(-1) * e_mem
+
+        temperature = self.temperature_floor + torch.nn.functional.softplus(self.log_temperature)
+        scaled_energies = energies / temperature
+
+        log_probs = -scaled_energies - torch.logsumexp(-scaled_energies, dim=-1, keepdim=True)
+        probabilities = torch.exp(log_probs)
+        return {
+            "energies": energies,
+            "scaled_energies": scaled_energies,
+            "log_probs": log_probs,
+            "probabilities": probabilities,
+            "temperature": temperature,
+            "e_sim": e_sim,
+            "e_ctx": e_ctx,
+            "e_mem": e_mem,
+            "g1": g1,
+            "g2": g2,
+            "g3": g3,
+        }

bio_llm/model/laminar_layer.py

@@ -0,0 +1,40 @@
+import torch
+from torch import nn
+
+
+class LaminarRefinement(nn.Module):
+    """Applies h_i <- h_i + eta (E_i - I_i) for a few sparse iterations."""
+
+    def __init__(self, steps: int = 2, eta: float = 0.1):
+        super().__init__()
+        self.steps = steps
+        self.eta = eta
+
+    def _weighted_sum(
+        self,
+        states: torch.Tensor,
+        attention_indices: torch.Tensor,
+        attention_weights: torch.Tensor,
+    ) -> torch.Tensor:
+        batch_size, seq_len, _ = states.shape
+        mixed = torch.zeros_like(states)
+        for batch_index in range(batch_size):
+            for position in range(seq_len):
+                indices = attention_indices[batch_index, position]
+                weights = attention_weights[batch_index, position].unsqueeze(-1)
+                mixed[batch_index, position] = torch.sum(states[batch_index, indices] * weights, dim=0)
+        return mixed
+
+    def forward(
+        self,
+        states: torch.Tensor,
+        attention_indices: torch.Tensor,
+        attention_weights: torch.Tensor,
+    ) -> torch.Tensor:
+        refined = states
+        for _ in range(self.steps):
+            excitatory = self._weighted_sum(refined, attention_indices, attention_weights)
+            prefix_totals = refined.cumsum(dim=1)
+            inhibitory = prefix_totals - excitatory
+            refined = refined + self.eta * (excitatory - inhibitory)
+        return refined

bio_llm/model/low_rank_qkv.py

@@ -0,0 +1,35 @@
+import math
+
+import torch
+from torch import nn
+
+
+class LowRankLinear(nn.Module):
+    """Computes xW with W = UV^T without materializing the full matrix."""
+
+    def __init__(self, in_dim: int, out_dim: int, rank: int):
+        super().__init__()
+        self.left = nn.Parameter(torch.empty(in_dim, rank))
+        self.right = nn.Parameter(torch.empty(out_dim, rank))
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        nn.init.normal_(self.left, mean=0.0, std=1.0 / math.sqrt(self.left.size(0)))
+        nn.init.normal_(self.right, mean=0.0, std=1.0 / math.sqrt(self.right.size(1)))
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        hidden = torch.matmul(inputs, self.left)
+        return torch.matmul(hidden, self.right.transpose(0, 1))
+
+
+class LowRankQKV(nn.Module):
+    """Builds efficient Q, K, V projections with low-rank factors."""
+
+    def __init__(self, d_model: int, rank: int):
+        super().__init__()
+        self.q_proj = LowRankLinear(d_model, d_model, rank)
+        self.k_proj = LowRankLinear(d_model, d_model, rank)
+        self.v_proj = LowRankLinear(d_model, d_model, rank)
+
+    def forward(self, inputs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return self.q_proj(inputs), self.k_proj(inputs), self.v_proj(inputs)

bio_llm/model/model.py

@@ -0,0 +1,98 @@
+import torch
+from torch import nn
+
+from bio_llm.model.candidate_retrieval import CandidateRetriever
+from bio_llm.model.embedding import TokenEmbedding
+from bio_llm.model.energy_head import EnergyHead
+from bio_llm.model.laminar_layer import LaminarRefinement
+from bio_llm.model.sparse_attention import SparseAttention
+from bio_llm.utils.config import SSETConfig
+
+
+class StructuredSparseEnergyTransformer(nn.Module):
+    """CPU-oriented language model with sparse interaction and energy decoding."""
+
+    def __init__(self, config: SSETConfig):
+        super().__init__()
+        self.config = config
+        self.embedding = TokenEmbedding(config.vocab_size, config.d_model)
+        self.sparse_attention = SparseAttention(
+            d_model=config.d_model,
+            rank=config.low_rank,
+            max_seq_len=config.max_seq_len,
+            top_k=config.attention_top_k,
+            local_window=config.local_window,
+            memory_candidates=config.memory_candidates,
+            landmark_count=config.landmark_count,
+            content_memory_candidates=config.content_memory_candidates,
+        )
+        self.laminar = LaminarRefinement(steps=config.laminar_steps, eta=config.laminar_eta)
+        self.retriever = CandidateRetriever(
+            vocab_size=config.vocab_size,
+            d_model=config.d_model,
+            stage1_dim=config.stage1_dim,
+            stage1_k=config.retrieval_stage1_k,
+            stage2_k=config.retrieval_stage2_k,
+        )
+        self.energy_head = EnergyHead(
+            vocab_size=config.vocab_size,
+            d_model=config.d_model,
+            transition_rank=config.transition_rank,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        target_ids: torch.Tensor | None = None,
+        attention_mode: str | None = None,
+    ) -> dict[str, torch.Tensor]:
+        embedded = self.embedding(input_ids)
+        attention_state = self.sparse_attention(embedded, mode=attention_mode or self.config.attention_mode)
+        refined_states = self.laminar(
+            states=attention_state.context,
+            attention_indices=attention_state.attention_indices,
+            attention_weights=attention_state.attention_weights,
+        )
+        candidate_ids = self.retriever(
+            hidden_states=refined_states,
+            embedding_weight=self.embedding.weight,
+            target_ids=target_ids,
+        )
+        energy_outputs = self.energy_head(
+            hidden_states=refined_states,
+            input_ids=input_ids,
+            prev_tokens=input_ids,
+            candidate_ids=candidate_ids,
+            attention_indices=attention_state.attention_indices,
+            attention_weights=attention_state.attention_weights,
+            embedding_weight=self.embedding.weight,
+        )
+        return {
+            "hidden_states": refined_states,
+            "candidate_ids": candidate_ids,
+            "attention_indices": attention_state.attention_indices,
+            "attention_weights": attention_state.attention_weights,
+            "confidence": attention_state.confidence,
+            **energy_outputs,
+        }
+
+    @torch.no_grad()
+    def generate(
+        self,
+        prompt_ids: list[int],
+        eos_id: int,
+        max_new_tokens: int = 12,
+        attention_mode: str | None = None,
+    ) -> list[int]:
+        generated = list(prompt_ids)
+        for _ in range(max_new_tokens):
+            window = generated[-self.config.max_seq_len :]
+            input_ids = torch.tensor([window], dtype=torch.long)
+            outputs = self.forward(input_ids, attention_mode=attention_mode)
+            last_probabilities = outputs["probabilities"][0, -1]
+            last_candidates = outputs["candidate_ids"][0, -1]
+            next_token = int(last_candidates[last_probabilities.argmax()].item())
+            generated.append(next_token)
+            if next_token == eos_id:
+                break
+        return generated

bio_llm/model/sparse_attention.py

@@ -0,0 +1,143 @@
+from dataclasses import dataclass
+import math
+
+import torch
+from torch import nn
+
+from .low_rank_qkv import LowRankQKV
+
+
+@dataclass
+class SparseAttentionState:
+    context: torch.Tensor
+    attention_weights: torch.Tensor
+    attention_indices: torch.Tensor
+    confidence: torch.Tensor
+
+
+def sparsemax(scores: torch.Tensor) -> torch.Tensor:
+    """Sparsemax is a simple entmax-like normalization with exact zeros."""
+
+    if scores.numel() == 1:
+        return torch.ones_like(scores)
+
+    sorted_scores, _ = torch.sort(scores, descending=True)
+    cumulative = torch.cumsum(sorted_scores, dim=0) - 1.0
+    steps = torch.arange(1, scores.numel() + 1, device=scores.device, dtype=scores.dtype)
+    support = sorted_scores - cumulative / steps > 0
+    support_size = int(max(1, support.sum().item()))
+    tau = cumulative[support_size - 1] / steps[support_size - 1]
+    output = torch.clamp(scores - tau, min=0.0)
+    total = output.sum()
+    if total <= 0:
+        return torch.softmax(scores, dim=0)
+    return output / total
+
+
+class SparseAttention(nn.Module):
+    """Causal sparse interaction with top-k or sparsemax normalization."""
+
+    def __init__(
+        self,
+        d_model: int,
+        rank: int,
+        max_seq_len: int,
+        top_k: int,
+        local_window: int,
+        memory_candidates: int,
+        landmark_count: int,
+        content_memory_candidates: int,
+    ):
+        super().__init__()
+        self.qkv = LowRankQKV(d_model, rank)
+        self.top_k = top_k
+        self.local_window = local_window
+        self.memory_candidates = memory_candidates
+        self.landmark_count = landmark_count
+        self.content_memory_candidates = content_memory_candidates
+        self.max_candidates = local_window + memory_candidates + landmark_count + content_memory_candidates
+        self.channel_logits = nn.Parameter(torch.tensor([1.0, 0.4, 0.4]))
+        self.memory_bias = nn.Parameter(torch.zeros(max_seq_len))
+        self.landmark_logits = nn.Parameter(torch.zeros(max_seq_len))
+        self.scale = 1.0 / math.sqrt(d_model)
+
+    def _candidate_positions(
+        self,
+        position: int,
+        query: torch.Tensor,
+        prior_keys: torch.Tensor,
+    ) -> list[int]:
+        local_start = max(0, position - self.local_window + 1)
+        local_positions = list(range(local_start, position + 1))
+        memory_budget = min(self.memory_candidates, position + 1)
+        top_distances = torch.topk(self.memory_bias[: position + 1], k=memory_budget).indices.tolist()
+        memory_positions = [position - int(distance) for distance in top_distances]
+        landmark_budget = min(self.landmark_count, position + 1)
+        top_landmarks = torch.topk(self.landmark_logits[: position + 1], k=landmark_budget).indices.tolist()
+
+        content_positions: list[int] = []
+        if position > 0 and self.content_memory_candidates > 0:
+            similarity = torch.matmul(prior_keys, query) * self.scale
+            content_budget = min(self.content_memory_candidates, similarity.numel())
+            content_positions = torch.topk(similarity, k=content_budget).indices.tolist()
+
+        merged = sorted(set(local_positions + memory_positions + top_landmarks + content_positions))
+        return merged[: self.max_candidates]
+
+    def _normalize(self, scores: torch.Tensor, mode: str) -> torch.Tensor:
+        if mode == "topk":
+            keep = min(self.top_k, scores.numel())
+            values, indices = torch.topk(scores, k=keep)
+            weights = torch.zeros_like(scores)
+            weights[indices] = torch.softmax(values, dim=0)
+            return weights
+        if mode == "sparsemax":
+            return sparsemax(scores)
+        raise ValueError(f"Unsupported sparse attention mode: {mode}")
+
+    def forward(self, inputs: torch.Tensor, mode: str = "topk") -> SparseAttentionState:
+        queries, keys, values = self.qkv(inputs)
+        batch_size, seq_len, d_model = queries.shape
+
+        attention_indices = torch.zeros(
+            batch_size, seq_len, self.max_candidates, dtype=torch.long, device=inputs.device
+        )
+        attention_weights = torch.zeros(
+            batch_size, seq_len, self.max_candidates, dtype=inputs.dtype, device=inputs.device
+        )
+        context = torch.zeros(batch_size, seq_len, d_model, dtype=inputs.dtype, device=inputs.device)
+        channel_weights = torch.softmax(self.channel_logits, dim=0)
+
+        for batch_index in range(batch_size):
+            for position in range(seq_len):
+                prior_keys = keys[batch_index, : position + 1]
+                candidates = self._candidate_positions(position, queries[batch_index, position], prior_keys)
+                candidate_tensor = torch.tensor(candidates, dtype=torch.long, device=inputs.device)
+                candidate_count = candidate_tensor.numel()
+                attention_indices[batch_index, position, :candidate_count] = candidate_tensor
+
+                q_i = queries[batch_index, position]
+                k_subset = keys[batch_index, candidate_tensor]
+                v_subset = values[batch_index, candidate_tensor]
+
+                content_term = torch.matmul(k_subset, q_i) * self.scale
+                distances = (position - candidate_tensor).to(inputs.dtype)
+                positional_term = -distances
+                memory_term = self.memory_bias[position - candidate_tensor]
+                scores = (
+                    channel_weights[0] * content_term
+                    + channel_weights[1] * positional_term
+                    + channel_weights[2] * memory_term
+                )
+
+                normalized = self._normalize(scores, mode=mode)
+                attention_weights[batch_index, position, :candidate_count] = normalized
+                context[batch_index, position] = torch.sum(normalized.unsqueeze(-1) * v_subset, dim=0)
+
+        confidence = attention_weights.max(dim=-1).values
+        return SparseAttentionState(
+            context=context,
+            attention_weights=attention_weights,
+            attention_indices=attention_indices,
+            confidence=confidence,
+        )

bio_llm/training/__init__.py

@@ -0,0 +1 @@
+"""Training utilities for the Structured Sparse Energy Transformer."""

bio_llm/training/chat_dataset.py

@@ -0,0 +1,183 @@
+from __future__ import annotations
+
+import re
+from collections.abc import Iterable, Sequence
+
+
+USER_ROLES = {"user", "human", "prompter", "customer", "client"}
+ASSISTANT_ROLES = {"assistant", "gpt", "bot", "model", "response"}
+SYSTEM_ROLES = {"system", "moderator"}
+RAW_TURN_PATTERN = re.compile(
+    r"<start_of_turn>\s*(?P<role>[a-zA-Z_]+)\s*\n(?P<content>.*?)<end_of_turn>",
+    flags=re.DOTALL,
+)
+
+
+def normalize_role(role: str) -> str:
+    role_key = role.strip().lower()
+    if role_key in USER_ROLES:
+        return "User"
+    if role_key in ASSISTANT_ROLES:
+        return "Assistant"
+    if role_key in SYSTEM_ROLES:
+        return "System"
+    if not role.strip():
+        return "User"
+    return role.strip().title()
+
+
+def _extract_text(message: object) -> tuple[str, str]:
+    if isinstance(message, dict):
+        role = str(message.get("role", message.get("from", message.get("speaker", ""))))
+        content = message.get("content", message.get("value", message.get("text", "")))
+        return role, str(content)
+
+    if isinstance(message, (list, tuple)) and len(message) >= 2:
+        return str(message[0]), str(message[1])
+
+    return "", str(message)
+
+
+def extract_messages(example: dict[str, object]) -> list[tuple[str, str]]:
+    if "messages" in example and isinstance(example["messages"], Sequence) and not isinstance(example["messages"], (str, bytes)):
+        return [_extract_text(message) for message in example["messages"]]
+
+    if "conversations" in example and isinstance(example["conversations"], Sequence) and not isinstance(
+        example["conversations"], (str, bytes)
+    ):
+        return [_extract_text(message) for message in example["conversations"]]
+
+    if "dialog" in example and isinstance(example["dialog"], Sequence) and not isinstance(example["dialog"], (str, bytes)):
+        return [_extract_text(message) for message in example["dialog"]]
+
+    if "prompt" in example and "response" in example:
+        return [("user", str(example["prompt"])), ("assistant", str(example["response"]))]
+
+    if "instruction" in example and "output" in example:
+        return [("user", str(example["instruction"])), ("assistant", str(example["output"]))]
+
+    if "input" in example and "output" in example:
+        return [("user", str(example["input"])), ("assistant", str(example["output"]))]
+
+    if "raw_text_content" in example:
+        text = str(example["raw_text_content"])
+        messages = [(match.group("role"), match.group("content")) for match in RAW_TURN_PATTERN.finditer(text)]
+        if messages:
+            return messages
+
+    return []
+
+
+def _clean_text(text: str) -> str:
+    return " ".join(str(text).split()).strip()
+
+
+def conversation_to_text(
+    example: dict[str, object],
+    max_turns: int | None = None,
+    max_message_chars: int | None = None,
+) -> str:
+    messages = extract_messages(example)
+    if max_turns is not None:
+        messages = messages[: max(1, max_turns)]
+
+    lines: list[str] = []
+    for role, content in messages:
+        clean_content = _clean_text(content)
+        if max_message_chars is not None:
+            clean_content = clean_content[: max(1, max_message_chars)].strip()
+        if not clean_content:
+            continue
+        lines.append(f"{normalize_role(role)}: {clean_content}")
+    return "\n".join(lines).strip()
+
+
+def conversation_to_pair(
+    example: dict[str, object],
+    max_turns: int | None = None,
+    max_message_chars: int | None = None,
+) -> dict[str, str] | None:
+    messages = extract_messages(example)
+    if max_turns is not None:
+        messages = messages[: max(1, max_turns)]
+
+    if not messages:
+        return None
+
+    first_user_index = None
+    first_assistant_index = None
+    for index, (role, content) in enumerate(messages):
+        role_key = role.strip().lower()
+        if first_user_index is None and role_key in USER_ROLES:
+            first_user_index = index
+        elif first_user_index is not None and role_key in ASSISTANT_ROLES:
+            first_assistant_index = index
+            break
+
+    if first_user_index is None or first_assistant_index is None:
+        return None
+
+    prompt_parts: list[str] = []
+    for role, content in messages[:first_assistant_index]:
+        clean_content = _clean_text(content)
+        if max_message_chars is not None:
+            clean_content = clean_content[: max(1, max_message_chars)].strip()
+        if not clean_content:
+            continue
+        prompt_parts.append(f"{normalize_role(role)}: {clean_content}")
+
+    assistant_role, assistant_content = messages[first_assistant_index]
+    assistant_text = _clean_text(assistant_content)
+    if max_message_chars is not None:
+        assistant_text = assistant_text[: max(1, max_message_chars)].strip()
+    if not prompt_parts or not assistant_text:
+        return None
+
+    prompt_text = "\n".join(prompt_parts + [f"{normalize_role(assistant_role)}:"]).strip()
+    sentence_text = f"{prompt_text} {assistant_text}".strip()
+    return {
+        "input": prompt_text,
+        "next_word": assistant_text,
+        "sentence": sentence_text,
+    }
+
+
+def build_chat_corpus(
+    examples: Iterable[dict[str, object]],
+    max_turns: int | None = None,
+    max_message_chars: int | None = None,
+) -> str:
+    conversations = [
+        conversation_to_text(
+            example,
+            max_turns=max_turns,
+            max_message_chars=max_message_chars,
+        )
+        for example in examples
+    ]
+    conversations = [conversation for conversation in conversations if conversation]
+    if not conversations:
+        raise ValueError("No conversational text could be built from the dataset examples.")
+    return "\n\n".join(conversations)
+
+
+def build_chat_pairs(
+    examples: Iterable[dict[str, object]],
+    max_turns: int | None = None,
+    max_message_chars: int | None = None,
+) -> list[dict[str, str]]:
+    pairs = [
+        pair
+        for example in examples
+        if (
+            pair := conversation_to_pair(
+                example,
+                max_turns=max_turns,
+                max_message_chars=max_message_chars,
+            )
+        )
+        is not None
+    ]
+    if not pairs:
+        raise ValueError("No conversational prompt/response pairs could be built from the dataset examples.")
+    return pairs

bio_llm/training/interview_dataset.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+import json
+import random
+from pathlib import Path
+
+
+def load_interview_examples(path: str | Path) -> list[dict[str, str]]:
+    payload = json.loads(Path(path).read_text(encoding="utf-8"))
+    if not isinstance(payload, list):
+        raise ValueError(f"Expected a JSON list in {path}, got {type(payload).__name__}.")
+
+    examples: list[dict[str, str]] = []
+    for index, item in enumerate(payload, start=1):
+        if not isinstance(item, dict):
+            raise ValueError(f"Expected object items in {path} at index {index}, got {type(item).__name__}.")
+        for field in ("question", "answer"):
+            if field not in item:
+                raise ValueError(f"Missing field {field!r} in {path} at index {index}.")
+        examples.append(
+            {
+                "domain": str(item.get("domain", "")).strip(),
+                "question": str(item["question"]).strip(),
+                "answer": str(item["answer"]).strip(),
+            }
+        )
+
+    filtered = [example for example in examples if example["question"] and example["answer"]]
+    if not filtered:
+        raise ValueError(f"No usable interview examples found in {path}.")
+    return filtered
+
+
+def format_interview_prompt(example: dict[str, str], include_domain: bool = True) -> str:
+    lines: list[str] = []
+    domain = example.get("domain", "").strip()
+    if include_domain and domain:
+        lines.append(f"Domain: {domain}")
+    lines.append(f"Question: {example['question'].strip()}")
+    lines.append("Answer:")
+    return "\n".join(lines).strip()
+
+
+def build_interview_corpus(
+    examples: list[dict[str, str]],
+    include_domain: bool = True,
+) -> str:
+    documents = [
+        f"{format_interview_prompt(example, include_domain=include_domain)} {example['answer'].strip()}".strip()
+        for example in examples
+    ]
+    documents = [document for document in documents if document]
+    if not documents:
+        raise ValueError("No interview training documents could be built from the dataset.")
+    return "\n\n".join(documents)
+
+
+def build_interview_pairs(
+    examples: list[dict[str, str]],
+    include_domain: bool = True,
+) -> list[dict[str, str]]:
+    pairs = [
+        {
+            "input": format_interview_prompt(example, include_domain=include_domain),
+            "next_word": example["answer"].strip(),
+            "sentence": (
+                f"{format_interview_prompt(example, include_domain=include_domain)} {example['answer'].strip()}".strip()
+            ),
+        }
+        for example in examples
+        if example["question"].strip() and example["answer"].strip()
+    ]
+    if not pairs:
+        raise ValueError("No interview prompt/answer pairs could be built from the dataset.")
+    return pairs
+
+
+def split_interview_examples(
+    examples: list[dict[str, str]],
+    eval_fraction: float,
+    seed: int,
+) -> tuple[list[dict[str, str]], list[dict[str, str]]]:
+    shuffled = list(examples)
+    random.Random(seed).shuffle(shuffled)
+
+    if len(shuffled) < 2 or eval_fraction <= 0:
+        return shuffled, []
+
+    eval_count = max(1, int(round(len(shuffled) * eval_fraction)))
+    eval_count = min(eval_count, len(shuffled) - 1)
+    train_examples = shuffled[:-eval_count]
+    eval_examples = shuffled[-eval_count:]
+    return train_examples, eval_examples
+
+
+def dump_pairs_jsonl(pairs: list[dict[str, str]], path: str | Path) -> None:
+    output_path = Path(path)
+    output_path.write_text(
+        "\n".join(json.dumps(pair, ensure_ascii=True) for pair in pairs),
+        encoding="utf-8",
+    )

bio_llm/training/loss.py

@@ -0,0 +1,65 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class StructuredEnergyLoss(nn.Module):
+    def __init__(
+        self,
+        margin: float,
+        margin_lambda: float,
+        ignore_index: int | None = None,
+        label_smoothing: float = 0.0,
+        hard_negative_count: int = 1,
+    ):
+        super().__init__()
+        self.margin = margin
+        self.margin_lambda = margin_lambda
+        self.ignore_index = ignore_index
+        self.label_smoothing = max(0.0, label_smoothing)
+        self.hard_negative_count = max(1, hard_negative_count)
+
+    def forward(
+        self,
+        energies: torch.Tensor,
+        log_probs: torch.Tensor,
+        candidate_ids: torch.Tensor,
+        target_ids: torch.Tensor,
+    ) -> tuple[torch.Tensor, dict[str, float]]:
+        target_mask = candidate_ids.eq(target_ids.unsqueeze(-1))
+        positive_log_probs = log_probs.masked_fill(~target_mask, float("-inf")).max(dim=-1).values
+        positive_energies = energies.masked_fill(~target_mask, float("inf")).min(dim=-1).values
+        negative_energies = energies.masked_fill(target_mask, float("inf"))
+        negative_scores = torch.where(
+            torch.isinf(negative_energies),
+            torch.full_like(negative_energies, float("-inf")),
+            -negative_energies,
+        )
+        hard_negative_count = min(self.hard_negative_count, negative_scores.size(-1))
+        hard_negative_energies = -torch.topk(negative_scores, k=hard_negative_count, dim=-1).values
+        hard_negative_energies = torch.where(
+            torch.isinf(hard_negative_energies),
+            positive_energies.detach().unsqueeze(-1),
+            hard_negative_energies,
+        )
+
+        if self.label_smoothing > 0:
+            uniform_log_probs = log_probs.mean(dim=-1)
+            nll = -(1.0 - self.label_smoothing) * positive_log_probs - self.label_smoothing * uniform_log_probs
+        else:
+            nll = -positive_log_probs
+
+        margin_term = F.relu(self.margin + positive_energies.unsqueeze(-1) - hard_negative_energies).mean(dim=-1)
+
+        if self.ignore_index is not None:
+            valid = target_ids.ne(self.ignore_index)
+            nll = nll[valid]
+            margin_term = margin_term[valid]
+
+        loss = nll.mean() + self.margin_lambda * margin_term.mean()
+        stats = {
+            "loss": float(loss.item()),
+            "nll": float(nll.mean().item()),
+            "margin": float(margin_term.mean().item()),
+        }
+        return loss, stats

bio_llm/training/metrics.py

@@ -0,0 +1,291 @@
+from __future__ import annotations
+
+import math
+import re
+import time
+
+import torch
+from torch.utils.data import DataLoader
+
+from bio_llm.model.model import StructuredSparseEnergyTransformer
+from bio_llm.utils.tokenizer import Tokenizer
+
+
+def _valid_mask(target_ids: torch.Tensor, ignore_index: int | None) -> torch.Tensor:
+    if ignore_index is None:
+        return torch.ones_like(target_ids, dtype=torch.bool)
+    return target_ids.ne(ignore_index)
+
+
+def _gather_target_log_probs(
+    log_probs: torch.Tensor,
+    candidate_ids: torch.Tensor,
+    target_ids: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    target_mask = candidate_ids.eq(target_ids.unsqueeze(-1))
+    positive_log_probs = log_probs.masked_fill(~target_mask, float("-inf")).max(dim=-1).values
+    found_target = target_mask.any(dim=-1)
+    return positive_log_probs, found_target
+
+
+def compute_prediction_metrics(
+    log_probs: torch.Tensor,
+    candidate_ids: torch.Tensor,
+    target_ids: torch.Tensor,
+    ignore_index: int | None = None,
+    top_k: int = 3,
+) -> dict[str, float]:
+    valid = _valid_mask(target_ids, ignore_index)
+    positive_log_probs, found_target = _gather_target_log_probs(log_probs, candidate_ids, target_ids)
+    effective = valid & found_target
+
+    if effective.sum().item() == 0:
+        return {
+            "candidate_nll": 0.0,
+            "candidate_perplexity": 1.0,
+            "top1_accuracy": 0.0,
+            "topk_accuracy": 0.0,
+            "coverage": 0.0,
+        }
+
+    top1_predictions = candidate_ids.gather(
+        -1,
+        log_probs.argmax(dim=-1, keepdim=True),
+    ).squeeze(-1)
+    top1_correct = top1_predictions.eq(target_ids) & effective
+
+    capped_top_k = min(top_k, candidate_ids.size(-1))
+    topk_candidates = torch.topk(log_probs, k=capped_top_k, dim=-1).indices
+    topk_predictions = candidate_ids.gather(-1, topk_candidates)
+    topk_correct = topk_predictions.eq(target_ids.unsqueeze(-1)).any(dim=-1) & effective
+
+    nll = -positive_log_probs[effective]
+    mean_nll = nll.mean().item()
+    return {
+        "candidate_nll": mean_nll,
+        "candidate_perplexity": math.exp(mean_nll),
+        "top1_accuracy": top1_correct.float().sum().item() / effective.sum().item(),
+        "topk_accuracy": topk_correct.float().sum().item() / effective.sum().item(),
+        "coverage": found_target[valid].float().mean().item(),
+    }
+
+
+def compute_retrieval_recall(
+    candidate_ids: torch.Tensor,
+    target_ids: torch.Tensor,
+    ignore_index: int | None = None,
+) -> float:
+    valid = _valid_mask(target_ids, ignore_index)
+    if valid.sum().item() == 0:
+        return 0.0
+    recalled = candidate_ids.eq(target_ids.unsqueeze(-1)).any(dim=-1)
+    return recalled[valid].float().mean().item()
+
+
+@torch.no_grad()
+def evaluate_model(
+    model: StructuredSparseEnergyTransformer,
+    loader: DataLoader,
+    ignore_index: int | None = None,
+    attention_mode: str | None = None,
+    top_k: int = 3,
+) -> dict[str, float]:
+    was_training = model.training
+    model.eval()
+
+    total_tokens = 0
+    total_nll = 0.0
+    total_top1 = 0.0
+    total_topk = 0.0
+    total_coverage = 0.0
+    total_retrieval = 0.0
+
+    for input_ids, target_ids in loader:
+        scored_outputs = model(input_ids, target_ids=target_ids, attention_mode=attention_mode)
+        retrieval_outputs = model(input_ids, target_ids=None, attention_mode=attention_mode)
+
+        valid = _valid_mask(target_ids, ignore_index)
+        token_count = int(valid.sum().item())
+        if token_count == 0:
+            continue
+
+        prediction_metrics = compute_prediction_metrics(
+            log_probs=scored_outputs["log_probs"],
+            candidate_ids=scored_outputs["candidate_ids"],
+            target_ids=target_ids,
+            ignore_index=ignore_index,
+            top_k=top_k,
+        )
+        retrieval_recall = compute_retrieval_recall(
+            candidate_ids=retrieval_outputs["candidate_ids"],
+            target_ids=target_ids,
+            ignore_index=ignore_index,
+        )
+
+        total_tokens += token_count
+        total_nll += prediction_metrics["candidate_nll"] * token_count
+        total_top1 += prediction_metrics["top1_accuracy"] * token_count
+        total_topk += prediction_metrics["topk_accuracy"] * token_count
+        total_coverage += prediction_metrics["coverage"] * token_count
+        total_retrieval += retrieval_recall * token_count
+
+    if was_training:
+        model.train()
+
+    if total_tokens == 0:
+        return {
+            "candidate_nll": 0.0,
+            "candidate_perplexity": 1.0,
+            "top1_accuracy": 0.0,
+            "topk_accuracy": 0.0,
+            "candidate_coverage": 0.0,
+            "retrieval_recall": 0.0,
+        }
+
+    mean_nll = total_nll / total_tokens
+    return {
+        "candidate_nll": mean_nll,
+        "candidate_perplexity": math.exp(mean_nll),
+        "top1_accuracy": total_top1 / total_tokens,
+        "topk_accuracy": total_topk / total_tokens,
+        "candidate_coverage": total_coverage / total_tokens,
+        "retrieval_recall": total_retrieval / total_tokens,
+    }
+
+
+@torch.no_grad()
+def benchmark_attention_mode(
+    model: StructuredSparseEnergyTransformer,
+    input_ids: torch.Tensor,
+    attention_mode: str,
+    warmup: int,
+    runs: int,
+) -> dict[str, float]:
+    latencies_ms: list[float] = []
+    attention_density: list[float] = []
+    active_edges: list[float] = []
+    tokens_per_batch = float(input_ids.numel())
+
+    for _ in range(warmup):
+        model(input_ids, attention_mode=attention_mode)
+
+    for _ in range(runs):
+        start = time.perf_counter()
+        outputs = model(input_ids, attention_mode=attention_mode)
+        elapsed_ms = (time.perf_counter() - start) * 1000.0
+        latencies_ms.append(elapsed_ms)
+
+        nonzero = outputs["attention_weights"].gt(0).float()
+        attention_density.append(nonzero.mean().item())
+        active_edges.append(nonzero.sum(dim=-1).mean().item())
+
+    mean_latency_ms = sum(latencies_ms) / len(latencies_ms)
+    variance = sum((value - mean_latency_ms) ** 2 for value in latencies_ms) / len(latencies_ms)
+    return {
+        "latency_ms_mean": mean_latency_ms,
+        "latency_ms_std": math.sqrt(variance),
+        "tokens_per_second": tokens_per_batch / (mean_latency_ms / 1000.0),
+        "attention_density": sum(attention_density) / len(attention_density),
+        "attention_sparsity": 1.0 - (sum(attention_density) / len(attention_density)),
+        "active_edges_per_token": sum(active_edges) / len(active_edges),
+    }
+
+
+WORD_PATTERN = re.compile(r"\w+|[^\w\s]", re.UNICODE)
+
+
+def _surface_from_tokens(tokens: list[str]) -> str:
+    output: list[str] = []
+    for token in tokens:
+        if not output:
+            output.append(token)
+            continue
+        if re.match(r"^[^\w\s]+$", token):
+            output.append(token)
+            continue
+        if re.match(r"^[^\w\s]+$", output[-1]):
+            output.append(" ")
+            output.append(token)
+            continue
+        output.append(" ")
+        output.append(token)
+    return "".join(output).strip()
+
+
+def sentence_to_prompt_completion(sentence: str) -> tuple[str, str]:
+    tokens = WORD_PATTERN.findall(sentence.strip())
+    if len(tokens) < 3:
+        raise ValueError(f"Sentence too short for completion evaluation: {sentence!r}")
+
+    split_index = len(tokens) - 1
+    if re.match(r"^[^\w\s]+$", tokens[-1]) and len(tokens) >= 3:
+        split_index = len(tokens) - 2
+
+    prompt_tokens = tokens[:split_index]
+    completion_tokens = tokens[split_index:]
+    return _surface_from_tokens(prompt_tokens), _surface_from_tokens(completion_tokens)
+
+
+@torch.no_grad()
+def evaluate_sentence_completions(
+    model: StructuredSparseEnergyTransformer,
+    tokenizer: Tokenizer,
+    sentences: list[str],
+    attention_mode: str | None = None,
+) -> dict[str, object]:
+    was_training = model.training
+    model.eval()
+
+    total_examples = 0
+    exact_matches = 0
+    token_matches = 0
+    total_completion_tokens = 0
+    examples: list[dict[str, str | bool]] = []
+
+    for sentence in sentences:
+        sentence = sentence.strip()
+        if not sentence:
+            continue
+
+        prompt_text, completion_text = sentence_to_prompt_completion(sentence)
+        prompt_ids = tokenizer.encode(prompt_text, add_bos=True)
+        target_ids = tokenizer.encode(completion_text, add_eos=True)
+        generated_ids = model.generate(
+            prompt_ids=prompt_ids,
+            eos_id=tokenizer.eos_id,
+            max_new_tokens=len(target_ids),
+            attention_mode=attention_mode,
+        )
+        predicted_completion_ids = generated_ids[len(prompt_ids) : len(prompt_ids) + len(target_ids)]
+        predicted_completion_text = tokenizer.decode(predicted_completion_ids)
+
+        exact_match = predicted_completion_ids == target_ids
+        exact_matches += int(exact_match)
+        total_examples += 1
+
+        common = min(len(predicted_completion_ids), len(target_ids))
+        token_matches += sum(
+            int(predicted_completion_ids[index] == target_ids[index])
+            for index in range(common)
+        )
+        total_completion_tokens += len(target_ids)
+
+        if len(examples) < 5:
+            examples.append(
+                {
+                    "prompt": prompt_text,
+                    "target_completion": completion_text,
+                    "predicted_completion": predicted_completion_text,
+                    "exact_match": exact_match,
+                }
+            )
+
+    if was_training:
+        model.train()
+
+    return {
+        "sentence_count": total_examples,
+        "exact_match_accuracy": exact_matches / max(1, total_examples),
+        "completion_token_accuracy": token_matches / max(1, total_completion_tokens),
+        "examples": examples,
+    }

bio_llm/training/recipes.py

@@ -0,0 +1,133 @@
+from __future__ import annotations
+
+import re
+from collections.abc import Iterable
+
+
+WORD_PATTERN = re.compile(r"\w+|[^\w\s]", re.UNICODE)
+
+
+CURRICULUM_PREFIXES: list[str] = [
+    "",
+    "in practice",
+    "for cpu tests",
+    "with sparse attention",
+    "during small runs",
+    "for quick checks",
+    "in this setup",
+    "on a tiny corpus",
+    "for compact models",
+    "under limited data",
+    "during fine tuning",
+    "for sentence completion",
+    "in the benchmark",
+    "for next word tests",
+    "with curriculum learning",
+    "on cpu",
+    "for retrieval tests",
+    "with energy scoring",
+    "for short prompts",
+    "during validation",
+    "for local windows",
+    "under sparse decoding",
+    "for repeated examples",
+    "with simple sentences",
+]
+
+
+def _complexity_key(example: dict[str, str]) -> tuple[int, int, int, str]:
+    prompt_tokens = WORD_PATTERN.findall(example["input"].strip())
+    sentence_tokens = WORD_PATTERN.findall(example["sentence"].strip())
+    target_tokens = WORD_PATTERN.findall(example["next_word"].strip())
+    return (len(sentence_tokens), len(prompt_tokens), len(target_tokens), example["input"].strip().lower())
+
+
+def augment_pair_examples(
+    examples: Iterable[dict[str, str]],
+    prefix_limit: int | None = None,
+) -> list[dict[str, str]]:
+    prefixes = CURRICULUM_PREFIXES if prefix_limit is None else CURRICULUM_PREFIXES[: max(1, prefix_limit)]
+    augmented: list[dict[str, str]] = []
+    seen: set[tuple[str, str, str]] = set()
+
+    for example in examples:
+        prompt = example["input"].strip()
+        next_word = example["next_word"].strip()
+        sentence = example["sentence"].strip()
+        for prefix in prefixes:
+            prefix_text = prefix.strip()
+            if prefix_text:
+                prompt_text = f"{prefix_text} {prompt}"
+                sentence_text = f"{prefix_text} {sentence}"
+            else:
+                prompt_text = prompt
+                sentence_text = sentence
+
+            key = (prompt_text.lower(), next_word.lower(), sentence_text.lower())
+            if key in seen:
+                continue
+            seen.add(key)
+            augmented.append(
+                {
+                    "input": prompt_text,
+                    "next_word": next_word,
+                    "sentence": sentence_text,
+                }
+            )
+
+    return augmented
+
+
+def order_examples_for_curriculum(examples: Iterable[dict[str, str]]) -> list[dict[str, str]]:
+    return sorted((dict(example) for example in examples), key=_complexity_key)
+
+
+def build_training_text(
+    examples: Iterable[dict[str, str]],
+    repeat_factor: int = 1,
+    curriculum: bool = True,
+) -> str:
+    ordered_examples = order_examples_for_curriculum(examples) if curriculum else [dict(example) for example in examples]
+    sentences = [example["sentence"].strip() for example in ordered_examples if example["sentence"].strip()]
+    if not sentences:
+        raise ValueError("No sentences available to build the training corpus.")
+    repeated = sentences * max(1, repeat_factor)
+    return " ".join(repeated)
+
+
+def split_phase_epochs(total_epochs: int) -> tuple[int, int, int]:
+    if total_epochs <= 1:
+        return max(1, total_epochs), 0, 0
+    if total_epochs == 2:
+        return 1, 1, 0
+
+    pretrain = max(1, round(total_epochs * 0.5))
+    freeze = max(1, round(total_epochs * 0.25))
+    finetune = total_epochs - pretrain - freeze
+
+    if finetune < 1:
+        finetune = 1
+    while pretrain + freeze + finetune > total_epochs and pretrain > 1:
+        pretrain -= 1
+    while pretrain + freeze + finetune > total_epochs and freeze > 1:
+        freeze -= 1
+    while pretrain + freeze + finetune > total_epochs and finetune > 1:
+        finetune -= 1
+    return pretrain, freeze, finetune
+
+
+def _set_requires_grad(module, requires_grad: bool) -> None:
+    for parameter in module.parameters():
+        parameter.requires_grad = requires_grad
+
+
+def freeze_backbone(model) -> None:
+    _set_requires_grad(model.embedding, False)
+    _set_requires_grad(model.sparse_attention, False)
+    _set_requires_grad(model.laminar, False)
+    _set_requires_grad(model.retriever, False)
+    _set_requires_grad(model.energy_head, True)
+
+
+def unfreeze_all(model) -> None:
+    _set_requires_grad(model, True)

bio_llm/training/trainer.py

@@ -0,0 +1,308 @@
+import json
+import time
+from dataclasses import replace
+from pathlib import Path
+
+import torch
+from torch.utils.data import DataLoader, TensorDataset
+
+from bio_llm.model.model import StructuredSparseEnergyTransformer
+from bio_llm.training.loss import StructuredEnergyLoss
+from bio_llm.training.metrics import evaluate_model
+from bio_llm.utils.config import SSETConfig
+from bio_llm.utils.tokenizer import BPETokenizer, Tokenizer, load_tokenizer
+
+
+def set_seed(seed: int) -> None:
+    torch.manual_seed(seed)
+
+
+def read_corpus(path: str | None) -> str:
+    if path is None:
+        base_corpus = (
+            "my name is sam . "
+            "i am ready . "
+            "we are happy . "
+            "the cat is small . "
+            "the sky is blue ."
+        )
+        return " ".join([base_corpus] * 4)
+    return Path(path).read_text(encoding="utf-8")
+
+
+def make_language_model_dataset(token_ids: list[int], seq_len: int, pad_id: int) -> TensorDataset:
+    windows: list[list[int]] = []
+    stride = max(1, seq_len // 2)
+    if len(token_ids) <= seq_len:
+        padded = token_ids + [pad_id] * (seq_len + 1 - len(token_ids))
+        windows.append(padded[: seq_len + 1])
+    else:
+        for start in range(0, len(token_ids) - 1, stride):
+            chunk = token_ids[start : start + seq_len + 1]
+            if len(chunk) < 2:
+                continue
+            if len(chunk) < seq_len + 1:
+                chunk = chunk + [pad_id] * (seq_len + 1 - len(chunk))
+            windows.append(chunk)
+            if start + seq_len + 1 >= len(token_ids):
+                break
+
+    if not windows:
+        empty = torch.empty((0, seq_len + 1), dtype=torch.long)
+        return TensorDataset(empty[:, :-1], empty[:, 1:])
+
+    data = torch.tensor(windows, dtype=torch.long)
+    return TensorDataset(data[:, :-1], data[:, 1:])
+
+
+def build_tokenizer_from_text(text: str, config: SSETConfig) -> Tokenizer:
+    return BPETokenizer.build(
+        [text],
+        vocab_size=config.tokenizer_vocab_size,
+        min_frequency=config.tokenizer_min_frequency,
+    )
+
+
+def split_token_ids_for_holdout(
+    token_ids: list[int],
+    holdout_fraction: float,
+) -> tuple[list[int], list[int]]:
+    if len(token_ids) < 2 or holdout_fraction <= 0:
+        return token_ids, []
+
+    holdout_size = max(1, int(round(len(token_ids) * holdout_fraction)))
+    holdout_size = min(holdout_size, len(token_ids) - 1)
+    split_index = len(token_ids) - holdout_size
+    train_ids = token_ids[:split_index]
+    holdout_ids = token_ids[split_index:]
+    if not train_ids:
+        return token_ids, token_ids
+    return train_ids, holdout_ids
+
+
+def fit_model(
+    model: StructuredSparseEnergyTransformer,
+    tokenizer: Tokenizer,
+    runtime_config: SSETConfig,
+    text: str,
+    evaluate_each_epoch: bool = True,
+    shuffle_train: bool = True,
+    verbose: bool = False,
+    log_interval: int = 1,
+) -> list[dict[str, float | int]]:
+    set_seed(runtime_config.seed)
+    model.config = runtime_config
+    token_ids = tokenizer.encode(text, add_bos=True, add_eos=True)
+    train_token_ids, holdout_token_ids = split_token_ids_for_holdout(token_ids, runtime_config.holdout_fraction)
+    train_dataset = make_language_model_dataset(train_token_ids, runtime_config.seq_len, tokenizer.pad_id)
+    holdout_dataset = make_language_model_dataset(holdout_token_ids, runtime_config.seq_len, tokenizer.pad_id)
+    if len(train_dataset) == 0:
+        raise ValueError("The training corpus is too short to build language-model windows.")
+
+    train_loader = DataLoader(train_dataset, batch_size=runtime_config.batch_size, shuffle=shuffle_train)
+    holdout_loader = DataLoader(holdout_dataset, batch_size=runtime_config.batch_size, shuffle=False)
+
+    trainable_parameters = [parameter for parameter in model.parameters() if parameter.requires_grad]
+    if not trainable_parameters:
+        raise ValueError("No trainable parameters are available for optimization.")
+
+    optimizer = torch.optim.Adam(trainable_parameters, lr=runtime_config.learning_rate)
+    loss_fn = StructuredEnergyLoss(
+        margin=runtime_config.margin,
+        margin_lambda=runtime_config.margin_lambda,
+        ignore_index=tokenizer.pad_id,
+        label_smoothing=runtime_config.label_smoothing,
+        hard_negative_count=runtime_config.hard_negative_count,
+    )
+
+    history: list[dict[str, float | int]] = []
+    model.train()
+    if verbose:
+        print("Training setup:")
+        print(f"  total_tokens={len(token_ids)}")
+        print(f"  train_tokens={len(train_token_ids)}")
+        print(f"  holdout_tokens={len(holdout_token_ids)}")
+        print(f"  train_windows={len(train_dataset)}")
+        print(f"  holdout_windows={len(holdout_dataset)}")
+        print(f"  batch_size={runtime_config.batch_size}")
+        print(f"  epochs={runtime_config.epochs}")
+        print(f"  learning_rate={runtime_config.learning_rate}")
+        print(f"  attention_mode={runtime_config.attention_mode}")
+    for epoch in range(runtime_config.epochs):
+        epoch_wall_start = time.perf_counter()
+        epoch_cpu_start = time.process_time()
+        epoch_loss = 0.0
+        batches = 0
+        train_tokens = 0
+        train_examples = 0
+        batch_count = len(train_loader)
+        if verbose:
+            print(f"\n[train] epoch {epoch + 1}/{runtime_config.epochs} started")
+        for input_ids, target_ids in train_loader:
+            optimizer.zero_grad(set_to_none=True)
+            outputs = model(input_ids, target_ids=target_ids, attention_mode=runtime_config.attention_mode)
+            loss, stats = loss_fn(
+                energies=outputs["energies"],
+                log_probs=outputs["log_probs"],
+                candidate_ids=outputs["candidate_ids"],
+                target_ids=target_ids,
+            )
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), runtime_config.grad_clip)
+            optimizer.step()
+            epoch_loss += stats["loss"]
+            batches += 1
+            train_tokens += int(input_ids.numel())
+            train_examples += int(input_ids.size(0))
+            if verbose and (batches % max(1, log_interval) == 0 or batches == batch_count):
+                average_loss = epoch_loss / max(1, batches)
+                print(
+                    f"[train] epoch {epoch + 1}/{runtime_config.epochs} "
+                    f"batch {batches}/{batch_count} "
+                    f"batch_loss={stats['loss']:.4f} avg_loss={average_loss:.4f} "
+                    f"examples_seen={train_examples}"
+                )
+
+        train_seconds = time.perf_counter() - epoch_wall_start
+        train_cpu_seconds = time.process_time() - epoch_cpu_start
+
+        epoch_record: dict[str, float | int] = {
+            "epoch": epoch + 1,
+            "loss": epoch_loss / max(1, batches),
+            "train_seconds": train_seconds,
+            "train_cpu_seconds": train_cpu_seconds,
+            "train_tokens": train_tokens,
+            "train_examples": train_examples,
+            "train_tokens_per_second": train_tokens / max(train_seconds, 1e-12),
+            "train_examples_per_second": train_examples / max(train_seconds, 1e-12),
+            "train_cpu_to_wall_ratio": train_cpu_seconds / max(train_seconds, 1e-12),
+            "holdout_tokens": len(holdout_token_ids),
+            "holdout_examples": len(holdout_dataset),
+        }
+        if evaluate_each_epoch:
+            if verbose:
+                print(f"[eval] epoch {epoch + 1}/{runtime_config.epochs} holdout evaluation started")
+            eval_wall_start = time.perf_counter()
+            epoch_metrics = evaluate_model(
+                model=model,
+                loader=holdout_loader,
+                ignore_index=tokenizer.pad_id,
+                attention_mode=runtime_config.attention_mode,
+                top_k=min(3, runtime_config.retrieval_stage2_k),
+            )
+            epoch_record.update({f"validation_{key}": value for key, value in epoch_metrics.items()})
+            epoch_record["evaluation_seconds"] = time.perf_counter() - eval_wall_start
+            if verbose:
+                print(
+                    f"[eval] epoch {epoch + 1}/{runtime_config.epochs} "
+                    f"candidate_perplexity={epoch_metrics['candidate_perplexity']:.4f} "
+                    f"top1_accuracy={epoch_metrics['top1_accuracy']:.4f} "
+                    f"topk_accuracy={epoch_metrics['topk_accuracy']:.4f} "
+                    f"retrieval_recall={epoch_metrics['retrieval_recall']:.4f}"
+                )
+        if verbose:
+            print(
+                f"[train] epoch {epoch + 1}/{runtime_config.epochs} completed "
+                f"loss={epoch_record['loss']:.4f} "
+                f"train_seconds={train_seconds:.2f}"
+            )
+        history.append(epoch_record)
+
+    model.eval()
+    return history
+
+
+def save_checkpoint(
+    model: StructuredSparseEnergyTransformer,
+    tokenizer: Tokenizer,
+    output_dir: str | Path,
+    history: list[dict[str, float | int]] | None = None,
+) -> Path:
+    output_path = Path(output_dir)
+    output_path.mkdir(parents=True, exist_ok=True)
+    checkpoint_path = output_path / "sset.pt"
+    torch.save({"config": model.config.to_dict(), "state_dict": model.state_dict()}, checkpoint_path)
+    tokenizer.save(output_path / "tokenizer.json")
+    if history is not None:
+        (output_path / "history.json").write_text(json.dumps(history, indent=2), encoding="utf-8")
+    return checkpoint_path
+
+
+def _load_state_dict_with_resize(
+    model: StructuredSparseEnergyTransformer,
+    state_dict: dict[str, torch.Tensor],
+) -> None:
+    current_state = model.state_dict()
+    merged_state: dict[str, torch.Tensor] = dict(current_state)
+
+    for name, source in state_dict.items():
+        if name not in current_state:
+            continue
+
+        target = current_state[name]
+        if source.shape == target.shape:
+            merged_state[name] = source
+            continue
+
+        if source.ndim == target.ndim and source.ndim > 0:
+            if source.shape[:-1] == target.shape[:-1] and source.shape[-1] <= target.shape[-1]:
+                resized = target.clone()
+                resized[..., : source.shape[-1]] = source
+                merged_state[name] = resized
+                continue
+
+            if source.shape[1:] == target.shape[1:] and source.shape[0] <= target.shape[0]:
+                resized = target.clone()
+                resized[: source.shape[0]] = source
+                merged_state[name] = resized
+                continue
+
+    model.load_state_dict(merged_state)
+
+
+def load_checkpoint(
+    checkpoint_path: str | Path,
+    tokenizer_path: str | Path | None = None,
+    config_overrides: dict[str, object] | None = None,
+) -> tuple[StructuredSparseEnergyTransformer, Tokenizer]:
+    checkpoint = torch.load(checkpoint_path, map_location="cpu")
+    config = SSETConfig(**checkpoint["config"])
+    if config_overrides:
+        overrides = dict(config_overrides)
+        if "max_seq_len" in overrides and "seq_len" not in overrides:
+            overrides["seq_len"] = overrides["max_seq_len"]
+        config = replace(config, **overrides)
+    model = StructuredSparseEnergyTransformer(config)
+    _load_state_dict_with_resize(model, checkpoint["state_dict"])
+    model.eval()
+    resolved_tokenizer_path = (
+        Path(tokenizer_path) if tokenizer_path is not None else Path(checkpoint_path).with_name("tokenizer.json")
+    )
+    tokenizer = load_tokenizer(resolved_tokenizer_path)
+    return model, tokenizer
+
+
+def train_model(
+    config: SSETConfig,
+    text: str,
+    tokenizer: Tokenizer | None = None,
+    evaluate_each_epoch: bool = True,
+    shuffle_train: bool = True,
+    verbose: bool = False,
+    log_interval: int = 1,
+) -> tuple[StructuredSparseEnergyTransformer, Tokenizer, list[dict[str, float | int]]]:
+    set_seed(config.seed)
+    tokenizer = tokenizer or build_tokenizer_from_text(text, config)
+    runtime_config = replace(config, vocab_size=tokenizer.vocab_size)
+    model = StructuredSparseEnergyTransformer(runtime_config)
+    history = fit_model(
+        model=model,
+        tokenizer=tokenizer,
+        runtime_config=runtime_config,
+        text=text,
+        evaluate_each_epoch=evaluate_each_epoch,
+        shuffle_train=shuffle_train,
+        verbose=verbose,
+        log_interval=log_interval,
+    )
+    return model, tokenizer, history

bio_llm/utils/__init__.py

@@ -0,0 +1 @@
+"""Utility helpers for the Structured Sparse Energy Transformer."""

bio_llm/utils/config.py

@@ -0,0 +1,37 @@
+from dataclasses import asdict, dataclass
+
+
+@dataclass
+class SSETConfig:
+    vocab_size: int = 0
+    d_model: int = 128
+    low_rank: int = 16
+    stage1_dim: int = 32
+    max_seq_len: int = 32
+    attention_top_k: int = 6
+    local_window: int = 12
+    memory_candidates: int = 6
+    landmark_count: int = 4
+    content_memory_candidates: int = 4
+    retrieval_stage1_k: int = 48
+    retrieval_stage2_k: int = 12
+    laminar_steps: int = 2
+    laminar_eta: float = 0.1
+    transition_rank: int = 32
+    attention_mode: str = "sparsemax"
+    learning_rate: float = 3e-3
+    batch_size: int = 8
+    epochs: int = 8
+    seq_len: int = 32
+    holdout_fraction: float = 0.1
+    margin: float = 0.25
+    margin_lambda: float = 0.2
+    label_smoothing: float = 0.05
+    hard_negative_count: int = 2
+    grad_clip: float = 1.0
+    seed: int = 7
+    tokenizer_vocab_size: int = 256
+    tokenizer_min_frequency: int = 2
+
+    def to_dict(self) -> dict:
+        return asdict(self)

bio_llm/utils/export.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+
+from bio_llm.training.trainer import load_checkpoint
+
+
+def export_checkpoint_to_npz(
+    checkpoint_path: str | Path,
+    output_path: str | Path,
+    tokenizer_path: str | Path | None = None,
+    max_seq_len: int | None = None,
+) -> Path:
+    config_overrides = {"max_seq_len": max_seq_len} if max_seq_len is not None else None
+    model, tokenizer = load_checkpoint(
+        checkpoint_path=checkpoint_path,
+        tokenizer_path=tokenizer_path,
+        config_overrides=config_overrides,
+    )
+
+    arrays: dict[str, np.ndarray] = {}
+    parameter_names: list[str] = []
+    for name, tensor in model.state_dict().items():
+        safe_name = name.replace(".", "__")
+        arrays[safe_name] = tensor.detach().cpu().numpy()
+        parameter_names.append(name)
+
+    arrays["__parameter_names__"] = np.array(parameter_names, dtype=object)
+    arrays["__config_json__"] = np.array(json.dumps(model.config.to_dict(), ensure_ascii=True))
+    arrays["__tokenizer_json__"] = np.array(
+        json.dumps(
+            {
+                "type": "bpe" if hasattr(tokenizer, "merges") else "simple",
+                "vocab": tokenizer.id_to_token,
+                "merges": [list(pair) for pair in getattr(tokenizer, "merges", [])],
+            },
+            ensure_ascii=True,
+        )
+    )
+
+    resolved_output = Path(output_path)
+    resolved_output.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(resolved_output, **arrays)
+    return resolved_output

bio_llm/utils/tokenizer.py

@@ -0,0 +1,277 @@
+import json
+import re
+from collections import Counter
+from pathlib import Path
+from typing import Iterable, List, Sequence
+
+
+class SimpleTokenizer:
+    """A small word-and-punctuation tokenizer for CPU-only experiments."""
+
+    PAD = "<pad>"
+    BOS = "<bos>"
+    EOS = "<eos>"
+    UNK = "<unk>"
+    TOKEN_PATTERN = re.compile(r"\w+|[^\w\s]", re.UNICODE)
+
+    def __init__(self, vocab: List[str]):
+        self.id_to_token = vocab
+        self.token_to_id = {token: index for index, token in enumerate(vocab)}
+
+    @classmethod
+    def build(cls, texts: Iterable[str], min_freq: int = 1) -> "SimpleTokenizer":
+        counter: Counter[str] = Counter()
+        for text in texts:
+            counter.update(cls.tokenize(text))
+
+        vocab = [cls.PAD, cls.BOS, cls.EOS, cls.UNK]
+        for token, freq in counter.most_common():
+            if freq >= min_freq and token not in vocab:
+                vocab.append(token)
+        return cls(vocab)
+
+    @staticmethod
+    def tokenize(text: str) -> List[str]:
+        return SimpleTokenizer.TOKEN_PATTERN.findall(text)
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.id_to_token)
+
+    @property
+    def pad_id(self) -> int:
+        return self.token_to_id[self.PAD]
+
+    @property
+    def bos_id(self) -> int:
+        return self.token_to_id[self.BOS]
+
+    @property
+    def eos_id(self) -> int:
+        return self.token_to_id[self.EOS]
+
+    @property
+    def unk_id(self) -> int:
+        return self.token_to_id[self.UNK]
+
+    def encode(self, text: str, add_bos: bool = False, add_eos: bool = False) -> List[int]:
+        tokens = self.tokenize(text)
+        ids = [self.token_to_id.get(token, self.unk_id) for token in tokens]
+        if add_bos:
+            ids.insert(0, self.bos_id)
+        if add_eos:
+            ids.append(self.eos_id)
+        return ids
+
+    def decode(self, token_ids: Iterable[int], skip_special_tokens: bool = True) -> str:
+        tokens: List[str] = []
+        specials = {self.PAD, self.BOS, self.EOS, self.UNK}
+        for token_id in token_ids:
+            token = self.id_to_token[int(token_id)]
+            if skip_special_tokens and token in specials:
+                continue
+            tokens.append(token)
+
+        output = []
+        for token in tokens:
+            if output and re.match(r"\w", token) and re.match(r"\w", output[-1][-1]):
+                output.append(" ")
+            elif output and token not in {".", ",", "!", "?", ":", ";", "'", '"', ")"} and output[-1] not in {"(", '"'}:
+                output.append(" ")
+            output.append(token)
+        return "".join(output).strip()
+
+    def save(self, path: str | Path) -> None:
+        payload = {"vocab": self.id_to_token}
+        Path(path).write_text(json.dumps(payload, indent=2), encoding="utf-8")
+
+    @classmethod
+    def load(cls, path: str | Path) -> "SimpleTokenizer":
+        payload = json.loads(Path(path).read_text(encoding="utf-8"))
+        return cls(payload["vocab"])
+
+
+class BPETokenizer:
+    """A compact BPE tokenizer with greedy longest-match encoding."""
+
+    PAD = "<pad>"
+    BOS = "<bos>"
+    EOS = "<eos>"
+    UNK = "<unk>"
+    END_OF_WORD = "</w>"
+    TOKEN_PATTERN = re.compile(r"\w+|[^\w\s]", re.UNICODE)
+
+    def __init__(self, vocab: Sequence[str], merges: Sequence[list[str] | tuple[str, str]]):
+        self.id_to_token = list(vocab)
+        self.token_to_id = {token: index for index, token in enumerate(self.id_to_token)}
+        self.merges = [tuple(pair) for pair in merges]
+        self.merge_ranks = {pair: index for index, pair in enumerate(self.merges)}
+
+    @classmethod
+    def build(
+        cls,
+        texts: Iterable[str],
+        vocab_size: int = 256,
+        min_frequency: int = 2,
+    ) -> "BPETokenizer":
+        words = Counter()
+        for text in texts:
+            words.update(cls.TOKEN_PATTERN.findall(text))
+
+        word_pieces = {
+            word: tuple(list(word) + [cls.END_OF_WORD])
+            for word, frequency in words.items()
+            if frequency >= 1
+        }
+        merges: list[tuple[str, str]] = []
+        special_tokens = [cls.PAD, cls.BOS, cls.EOS, cls.UNK]
+        symbol_vocab = {symbol for pieces in word_pieces.values() for symbol in pieces}
+
+        while len(symbol_vocab) + len(special_tokens) < vocab_size:
+            pair_counts: Counter[tuple[str, str]] = Counter()
+            for word, pieces in word_pieces.items():
+                frequency = words[word]
+                for index in range(len(pieces) - 1):
+                    pair_counts[(pieces[index], pieces[index + 1])] += frequency
+
+            if not pair_counts:
+                break
+
+            best_pair, best_frequency = pair_counts.most_common(1)[0]
+            if best_frequency < min_frequency:
+                break
+
+            merged_symbol = "".join(best_pair)
+            merges.append(best_pair)
+            updated: dict[str, tuple[str, ...]] = {}
+            for word, pieces in word_pieces.items():
+                new_pieces: list[str] = []
+                index = 0
+                while index < len(pieces):
+                    if index < len(pieces) - 1 and (pieces[index], pieces[index + 1]) == best_pair:
+                        new_pieces.append(merged_symbol)
+                        index += 2
+                    else:
+                        new_pieces.append(pieces[index])
+                        index += 1
+                updated[word] = tuple(new_pieces)
+            word_pieces = updated
+            symbol_vocab = {symbol for pieces in word_pieces.values() for symbol in pieces}
+
+        vocab = special_tokens + sorted(symbol_vocab)
+        return cls(vocab=vocab, merges=merges)
+
+    @staticmethod
+    def tokenize(text: str) -> List[str]:
+        return BPETokenizer.TOKEN_PATTERN.findall(text)
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.id_to_token)
+
+    @property
+    def pad_id(self) -> int:
+        return self.token_to_id[self.PAD]
+
+    @property
+    def bos_id(self) -> int:
+        return self.token_to_id[self.BOS]
+
+    @property
+    def eos_id(self) -> int:
+        return self.token_to_id[self.EOS]
+
+    @property
+    def unk_id(self) -> int:
+        return self.token_to_id[self.UNK]
+
+    def _apply_merges(self, word: str) -> list[str]:
+        pieces = list(word) + [self.END_OF_WORD]
+        if len(pieces) == 1:
+            return pieces
+
+        while True:
+            candidates = []
+            for index in range(len(pieces) - 1):
+                pair = (pieces[index], pieces[index + 1])
+                if pair in self.merge_ranks:
+                    candidates.append((self.merge_ranks[pair], index, pair))
+            if not candidates:
+                break
+
+            _, merge_index, pair = min(candidates)
+            pieces = pieces[:merge_index] + ["".join(pair)] + pieces[merge_index + 2 :]
+        return pieces
+
+    def encode(self, text: str, add_bos: bool = False, add_eos: bool = False) -> List[int]:
+        ids: list[int] = []
+        if add_bos:
+            ids.append(self.bos_id)
+        for token in self.tokenize(text):
+            if re.match(r"\w+", token):
+                pieces = self._apply_merges(token)
+            else:
+                pieces = [token + self.END_OF_WORD]
+                if pieces[0] not in self.token_to_id:
+                    pieces = [token, self.END_OF_WORD]
+            for piece in pieces:
+                ids.append(self.token_to_id.get(piece, self.unk_id))
+        if add_eos:
+            ids.append(self.eos_id)
+        return ids
+
+    def decode(self, token_ids: Iterable[int], skip_special_tokens: bool = True) -> str:
+        specials = {self.PAD, self.BOS, self.EOS, self.UNK}
+        words: list[str] = []
+        current = ""
+        for token_id in token_ids:
+            token = self.id_to_token[int(token_id)]
+            if skip_special_tokens and token in specials:
+                continue
+            if token == self.END_OF_WORD:
+                if current:
+                    words.append(current)
+                    current = ""
+                continue
+            if token.endswith(self.END_OF_WORD):
+                current += token[: -len(self.END_OF_WORD)]
+                words.append(current)
+                current = ""
+            else:
+                current += token
+
+        if current:
+            words.append(current)
+
+        output: list[str] = []
+        for word in words:
+            if not output:
+                output.append(word)
+            elif re.match(r"^[^\w\s]+$", word):
+                output.append(word)
+            elif re.match(r"^[^\w\s]+$", output[-1]):
+                output.append(" ")
+                output.append(word)
+            else:
+                output.append(" ")
+                output.append(word)
+        return "".join(output).replace("  ", " ").strip()
+
+    def save(self, path: str | Path) -> None:
+        payload = {"type": "bpe", "vocab": self.id_to_token, "merges": [list(pair) for pair in self.merges]}
+        Path(path).write_text(json.dumps(payload, indent=2), encoding="utf-8")
+
+    @classmethod
+    def load(cls, path: str | Path) -> "BPETokenizer":
+        payload = json.loads(Path(path).read_text(encoding="utf-8"))
+        return cls(payload["vocab"], payload.get("merges", []))
+
+
+Tokenizer = SimpleTokenizer | BPETokenizer
+
+
+def load_tokenizer(path: str | Path) -> Tokenizer:
+    payload = json.loads(Path(path).read_text(encoding="utf-8"))
+    if payload.get("type") == "bpe":
+        return BPETokenizer(payload["vocab"], payload.get("merges", []))
+    return SimpleTokenizer(payload["vocab"])

bio_voice_tts/TRAINING_INFRASTRUCTURE.md

@@ -0,0 +1,964 @@
+# BioVoice-TTS Training Infrastructure
+## Sparse Energy-Based Voice Cloning Foundation Model
+
+This document describes the complete training infrastructure, data pipeline, code blueprint, and staged systems plan for `BioVoice-TTS`, a CPU-first sparse voice cloning foundation model built around:
+
+- low-rank QKV projections
+- sparse temporal attention
+- laminar refinement
+- energy-based decoding
+- multi-scale memory
+- FiLM speaker conditioning
+- sparse acoustic decoding
+- lightweight sparse neural vocoding
+- streaming-first inference and training
+
+The code accompanying this document lives under `bio_voice_tts/` and is organized as a modular PyTorch training stack with production-oriented boundaries.
+
+---
+
+# 1. Complete Repository Structure
+
+```text
+bio_voice_tts/
+├── __init__.py
+├── TRAINING_INFRASTRUCTURE.md
+├── audio/
+│   ├── __init__.py
+│   ├── stft.py
+│   ├── mel.py
+│   ├── phoneme.py
+│   └── features.py
+├── benchmarks/
+│   ├── __init__.py
+│   └── benchmark_cpu.py
+├── configs/
+│   ├── base.yaml
+│   ├── speaker_encoder.yaml
+│   ├── acoustic_decoder.yaml
+│   ├── vocoder.yaml
+│   └── datasets.yaml
+├── datasets/
+│   ├── __init__.py
+│   ├── manifest.py
+│   ├── speaker_dataset.py
+│   └── tts_dataset.py
+├── evaluation/
+│   ├── __init__.py
+│   ├── metrics.py
+│   └── evaluate.py
+├── inference/
+│   ├── __init__.py
+│   ├── synthesize.py
+│   ├── clone_voice.py
+│   └── realtime_stream.py
+├── model/
+│   ├── __init__.py
+│   ├── low_rank_qkv.py
+│   ├── sparse_attention.py
+│   ├── laminar.py
+│   ├── memory.py
+│   ├── speaker_encoder.py
+│   ├── semantic_encoder.py
+│   ├── prosody.py
+│   ├── acoustic_decoder.py
+│   ├── mel_generator.py
+│   └── biovoice_tts.py
+├── preprocessing/
+│   ├── __init__.py
+│   ├── preprocessing.py
+│   ├── stft.py
+│   ├── mel.py
+│   └── phoneme.py
+├── scripts/
+│   ├── prepare_manifest.py
+│   ├── train_speaker.py
+│   ├── train_tts.py
+│   └── train_vocoder.py
+├── streaming/
+│   ├── __init__.py
+│   ├── cache.py
+│   └── realtime_stream.py
+├── training/
+│   ├── __init__.py
+│   ├── losses.py
+│   ├── checkpoint.py
+│   ├── distributed.py
+│   ├── trainer.py
+│   ├── speaker_trainer.py
+│   ├── acoustic_trainer.py
+│   └── vocoder_trainer.py
+├── utils/
+│   ├── __init__.py
+│   ├── config.py
+│   ├── logging.py
+│   ├── seed.py
+│   └── device.py
+└── vocoder/
+    ├── __init__.py
+    ├── sparse_vocoder.py
+    └── discriminator.py
+```
+
+## Module Roles
+
+`audio/`
+
+- Implements mathematically explicit STFT, mel projection, text normalization, fallback phonemization, and CPU-friendly feature extraction.
+- `features.py` is the main acoustic preprocessing engine used by training, evaluation, and inference.
+
+`datasets/`
+
+- Standardizes JSONL manifests into `ManifestEntry` records.
+- Splits data loading into speaker-centric triplet sampling and TTS-centric token/mel loading.
+- Keeps batching memory-safe by padding only to batch-local maxima.
+
+`preprocessing/`
+
+- Houses the manifest preprocessing pipeline and file-level wrappers around signal processing modules.
+- Intended for offline feature extraction so CPU training spends less time in repeated IO and FFT work.
+
+`model/`
+
+- Contains the sparse semantic core, speaker encoder, prosody heads, acoustic decoder, and full assembly model.
+- Preserves the original SSET principles by keeping low-rank, sparse, and laminar modules isolated and reusable.
+
+`training/`
+
+- Contains reusable losses, checkpoint management, optimizer scheduling, distributed setup, and specialized trainers.
+- Supports staged training: speaker encoder first, then text-to-mel, then vocoder, then joint refinement.
+
+`vocoder/`
+
+- Implements a lightweight causal sparse vocoder plus a compact waveform discriminator for adversarial refinement.
+
+`streaming/`
+
+- Encapsulates cache management and chunk-wise synthesis, so streaming logic does not pollute core model code.
+
+`inference/`
+
+- Exposes offline cloning and realtime-style chunked synthesis entrypoints.
+
+`evaluation/`
+
+- Measures mel error, pseudo-MOS, and speaker consistency proxies.
+- Intended to be expanded with ASR-backed WER, speaker verification EER, and MOSNet-like learned estimators.
+
+`benchmarks/`
+
+- Tracks mean CPU latency, peak memory, and mel throughput.
+- Used to compare sparse vs dense variants under fixed input sizes.
+
+`configs/`
+
+- YAML-driven configuration surface for datasets, speaker encoder, text-to-mel, and vocoder recipes.
+
+`scripts/`
+
+- Command-line entrypoints for preprocessing and staged training.
+
+---
+
+# 2. Data Sources
+
+The datasets below are the best starting points for a CPU-first voice cloning foundation model. Sizes and speaker counts are approximate and should be verified against the exact release you download.
+
+## Recommended Dataset Roles
+
+- Best for speaker encoder:
+  - VoxCeleb
+  - VCTK
+  - MLS
+  - Common Voice
+- Best for single-speaker clean TTS bootstrapping:
+  - LJSpeech
+- Best for multi-speaker TTS:
+  - LibriTTS
+  - VCTK
+  - MLS
+- Best for multilingual synthesis:
+  - MLS
+  - Common Voice
+  - AISHELL for Mandarin
+  - Hindi subsets from Common Voice and Indic speech resources
+- Best for streaming speech and future duplex STT/TTS:
+  - GigaSpeech
+  - Fisher English
+  - Common Voice
+
+## Dataset Table
+
+| Dataset | Primary Use | Approx. Size | Speakers | Typical SR | Transcript Format | License / Access | Official Link |
+|---|---|---:|---:|---:|---|---|---|
+| LibriTTS | Clean multi-speaker TTS | ~585 hours | ~2.4k | 24 kHz | aligned text / normalized text | CC BY 4.0 | https://openslr.org/60/ |
+| VCTK | Multi-speaker cloning / accents | ~44 hours | 100+ | 48 kHz | per-utterance text | Edinburgh DataShare terms | https://datashare.ed.ac.uk/handle/10283/3443 |
+| LJSpeech | Single-speaker TTS bootstrap | ~24 hours | 1 | 22.05 kHz | metadata CSV | Public domain derivatives / project terms | https://keithito.com/LJ-Speech-Dataset/ |
+| Common Voice | Multilingual speech | release-dependent, large | very large | 48 kHz source | TSV/CSV + clips | CC0 | https://commonvoice.mozilla.org/ |
+| VoxCeleb | Speaker encoder pretraining | >2k hours | 7k+ | 16 kHz common recipes | speaker ids + metadata | research access, verify terms | https://www.robots.ox.ac.uk/~vgg/data/voxceleb/ |
+| AISHELL-1 | Mandarin ASR / multilingual speech | ~178 hours | 400 | 16 kHz | Kaldi-style transcripts | Apache 2.0 | https://openslr.org/33/ |
+| MLS | Multilingual LibriSpeech-style corpus | tens of thousands of hours across languages | many thousands | 16 kHz | transcript files | CC BY 4.0 | https://openslr.org/94/ |
+| GigaSpeech | large-scale speech + transcripts | 10k+ hours by subset | many | 16 kHz | JSON / segment metadata | see repo license and corpus terms | https://github.com/SpeechColab/GigaSpeech |
+| Fisher English | conversational speech / streaming STT | ~2k hours | many thousands | 8 kHz | LDC transcripts | LDC licensed | https://catalog.ldc.upenn.edu/LDC2004T19 |
+| Hindi: Common Voice Hindi | multilingual Hindi speech | release-dependent | many | 48 kHz source | TSV/CSV | CC0 | https://commonvoice.mozilla.org/ |
+| Hindi: IndicVoices / AI4Bharat resources | Indian language speech | release-dependent | large | varies | manifest / JSON / TSV | verify individual dataset license | https://ai4bharat.iitm.ac.in/ |
+
+## Dataset Preparation Strategy
+
+1. Convert all audio to a canonical sample rate:
+   - `24000` Hz for TTS
+   - `16000` Hz auxiliary branch for speaker verification if needed
+2. Normalize transcripts:
+   - lowercase
+   - Unicode normalization
+   - punctuation pruning
+   - numeric expansion where appropriate
+3. Build a canonical JSONL manifest.
+4. Precompute mel, log-energy, and pitch features.
+5. Store speaker-balanced splits.
+
+## Download and Preparation Scripts
+
+The codebase expects a JSONL manifest. Example workflow:
+
+```bash
+python -m bio_voice_tts.scripts.prepare_manifest \
+  --manifest data/train_manifest.jsonl \
+  --feature-dir data/features \
+  --sample-rate 24000 \
+  --n-mels 80
+```
+
+For large sources:
+
+- keep original archives in `data/raw/`
+- export normalized manifests into `data/manifests/`
+- store precomputed features in `data/features/<speaker_id>/`
+
+---
+
+# 3. Dataset Format
+
+## Canonical Layout
+
+```text
+datasets/
+├── speaker_001/
+│   ├── audio/
+│   │   ├── utt_0001.wav
+│   │   └── utt_0002.wav
+│   ├── transcript.txt
+│   └── metadata.json
+└── manifests/
+    ├── train_manifest.jsonl
+    └── eval_manifest.jsonl
+```
+
+## `metadata.json`
+
+```json
+{
+  "speaker_id": "001",
+  "gender": "female",
+  "language": "en",
+  "sample_rate": 22050
+}
+```
+
+## JSONL Manifest Schema
+
+```json
+{"audio_path":"datasets/speaker_001/audio/utt_0001.wav","text":"welcome to the interview","speaker_id":"001","language":"en","sample_rate":24000}
+{"audio_path":"datasets/speaker_001/audio/utt_0002.wav","text":"thank you for joining us","speaker_id":"001","language":"en","sample_rate":24000}
+```
+
+## Normalization Rules
+
+- transcript normalization:
+  - lowercase
+  - whitespace collapse
+  - remove unsupported symbols
+- punctuation cleaning:
+  - keep sentence boundaries and pauses when useful
+- phoneme conversion:
+  - language-specific G2P if available
+  - fallback grapheme-to-token path for fast experiments
+- silence trimming:
+  - trim leading/trailing silence with energy threshold or VAD
+- audio normalization:
+  - peak or RMS normalization
+- mel extraction:
+  - fixed FFT/hop/window parameters per experiment family
+
+---
+
+# 4. Audio Preprocessing
+
+The code files are:
+
+- `audio/stft.py`
+- `audio/mel.py`
+- `audio/phoneme.py`
+- `audio/features.py`
+- `preprocessing/preprocessing.py`
+
+## Mathematics
+
+STFT:
+
+\[
+\text{STFT}(m,\omega)=\sum_{n=0}^{N-1}x[n]w[n-mH]e^{-j\omega n}
+\]
+
+Power spectrogram:
+
+\[
+P(m,\omega)=|\text{STFT}(m,\omega)|^2
+\]
+
+Mel projection:
+
+\[
+M_{m,f}=\log\left(\epsilon + \sum_{\omega}H_f(\omega)P(m,\omega)\right)
+\]
+
+Mel scale:
+
+\[
+\text{mel}(f)=2595\log_{10}\left(1+\frac{f}{700}\right)
+\]
+
+Peak normalization:
+
+\[
+\widetilde{x}[n]=\frac{x[n]}{\max(\epsilon,\max_k |x[k]|)}
+\]
+
+Energy:
+
+\[
+e_m=\log\left(\epsilon + \frac{1}{F}\sum_{\omega}P(m,\omega)\right)
+\]
+
+Pitch:
+
+\[
+p_m = \log(1 + F_0(m))
+\]
+
+## Implementation Notes
+
+- `AudioFeatureExtractor` uses `torchaudio` for loading, resampling, and pitch detection.
+- `build_mel_filterbank` explicitly constructs the mel basis rather than hiding it behind a monolithic transform.
+- `normalize_text` and `naive_phonemize` provide a CPU-safe fallback path.
+- preprocessing is offline by default to reduce repeated FFT work during training.
+
+## Silence and VAD
+
+Current code includes energy-based trimming. Production upgrades should add:
+
+- WebRTC VAD
+- framewise voiced/unvoiced masks
+- language-specific pause retention logic
+
+## Alignment
+
+The current blueprint uses duration heuristics as a bootstrap. For full training:
+
+- use Montreal Forced Aligner
+- or an internal CTC/attention aligner
+- export token or phoneme durations into the manifest or feature cache
+
+---
+
+# 5. Speaker Encoder Training
+
+Files:
+
+- `model/speaker_encoder.py`
+- `datasets/speaker_dataset.py`
+- `training/speaker_trainer.py`
+
+## Architecture
+
+1. Convolutional frontend over mel frames.
+2. Sparse temporal encoder.
+3. Laminar refinement.
+4. Attentive statistics pooling.
+5. L2-normalized speaker projection.
+
+Mathematically:
+
+\[
+z_s = f_{\text{speaker}}(M)
+\]
+
+where:
+
+- \(M \in \mathbb{R}^{B \times T \times F_{\text{mel}}}\)
+- \(z_s \in \mathbb{R}^{B \times d_s}\)
+
+## Losses
+
+Triplet loss:
+
+\[
+L = \max(0, d(a,p) - d(a,n) + m)
+\]
+
+where \(d(\cdot,\cdot)\) is cosine or angular distance.
+
+Contrastive / supervised contrastive auxiliary loss:
+
+\[
+\mathcal{L}_{\text{con}} = -\log \frac{\sum_{p \in P(i)} \exp(\text{sim}(z_i,z_p)/\tau)}
+{\sum_{j \ne i}\exp(\text{sim}(z_i,z_j)/\tau)}
+\]
+
+## Batching Strategy
+
+- sample anchor, positive, negative triplets by speaker id
+- keep variable-length batch padding local
+- periodically mine hard negatives from an embedding memory bank
+
+## Hard Negative Mining
+
+Practical CPU-first approach:
+
+1. embed a speaker minibatch
+2. compute cosine matrix
+3. choose negatives with highest non-matching cosine
+
+## Augmentation Pipeline
+
+- additive noise
+- room impulse responses
+- mild codec artifacts
+- small gain shifts
+
+Use conservative augmentation to avoid collapsing speaker identity.
+
+---
+
+# 6. Semantic Sparse Encoder Training
+
+Files:
+
+- `datasets/tts_dataset.py`
+- `model/low_rank_qkv.py`
+- `model/sparse_attention.py`
+- `model/laminar.py`
+- `model/memory.py`
+- `model/semantic_encoder.py`
+
+## Tokenization
+
+The current code includes a lightweight phoneme-like tokenizer built from normalized text. Production upgrades should provide:
+
+- phoneme tokenizer
+- multilingual symbol tables
+- byte fallback for out-of-vocabulary text
+
+## Tensor Shapes
+
+- tokens: \(u \in \mathbb{N}^{B \times T_{\text{txt}}}\)
+- embeddings: \(X \in \mathbb{R}^{B \times T_{\text{txt}} \times d}\)
+- hidden states: \(H \in \mathbb{R}^{B \times T_{\text{txt}} \times d}\)
+
+## Low-Rank QKV
+
+\[
+W_Q = U_QV_Q^\top,\quad W_K = U_KV_K^\top,\quad W_V = U_VV_V^\top
+\]
+
+with \(U \in \mathbb{R}^{d \times r}\), \(V \in \mathbb{R}^{d \times r}\), \(r \ll d\).
+
+## Sparse Attention
+
+For query \(i\), candidate set:
+
+\[
+\mathcal{C}_i = \mathcal{C}^{\text{local}}_i \cup \mathcal{C}^{\text{memory}}_i \cup \mathcal{C}^{\text{landmark}}_i \cup \mathcal{C}^{\text{content}}_i
+\]
+
+Score:
+
+\[
+S_{ij}=w_1(Q_i^\top K_j)+w_2P_{ij}+w_3M_{ij}+w_4R_{ij}
+\]
+
+Normalization:
+
+- top-k sparse softmax
+- sparsemax
+
+## Laminar Refinement
+
+\[
+h_i \leftarrow h_i + \eta(E_i-I_i)
+\]
+
+This stage improves stability while retaining shallow compute depth.
+
+## CPU Optimization Notes
+
+- local window candidates are contiguous and cache-friendly
+- content retrieval is bounded to a small top-k
+- memory summaries compress long-range context
+- low-rank projections reduce parameter bandwidth
+
+---
+
+# 7. Prosody Modeling
+
+Files:
+
+- `model/prosody.py`
+
+## Heads
+
+- duration predictor
+- pitch predictor
+- energy predictor
+- FiLM speaker conditioning
+- length regulator
+
+## Equations
+
+Duration:
+
+\[
+\widehat{d}_t = \text{softplus}(w_d^\top h_t + b_d)
+\]
+
+Pitch:
+
+\[
+\widehat{p}_t = w_p^\top h_t + b_p
+\]
+
+Energy:
+
+\[
+\widehat{e}_t = w_e^\top h_t + b_e
+\]
+
+Speaker FiLM:
+
+\[
+h'_t = \gamma(z_s) \odot h_t + \beta(z_s)
+\]
+
+Length regulation:
+
+\[
+\widetilde{H} = \text{Expand}(H, \widehat{d})
+\]
+
+## Rhythm Modeling
+
+Prosody is modeled explicitly so the acoustic decoder does not need to discover:
+
+- alignment
+- stress
+- speaking rate
+- energy contour
+
+from scratch with dense frame attention.
+
+---
+
+# 8. Acoustic Decoder Training
+
+Files:
+
+- `model/acoustic_decoder.py`
+- `model/mel_generator.py`
+- `training/acoustic_trainer.py`
+
+## Architecture
+
+```text
+semantic latent
++ speaker latent
++ pitch
++ energy
+-> sparse acoustic decoder
+-> energy head
+-> mel spectrogram
+```
+
+The decoder uses sparse temporal attention again, but at the frame-expanded acoustic level.
+
+## Teacher Forcing and Scheduled Sampling
+
+Current blueprint uses target durations, target pitch, and target energy when provided.
+
+Recommended progression:
+
+1. full teacher forcing with ground-truth durations
+2. predicted duration with teacher-forced prosody
+3. scheduled sampling on duration and prosody
+4. chunk-wise streaming teacher forcing
+
+## Losses
+
+- mel L1
+- duration loss
+- pitch loss
+- energy loss
+- optional STFT or alignment losses
+
+Alignment loss can be added as:
+
+\[
+\mathcal{L}_{\text{align}} = \|A - A^\*\|_1
+\]
+
+where \(A^\*\) is an external aligner target.
+
+---
+
+# 9. Sparse Vocoder Training
+
+Files:
+
+- `vocoder/sparse_vocoder.py`
+- `vocoder/discriminator.py`
+- `training/vocoder_trainer.py`
+
+## Architecture
+
+- causal transposed-convolution upsamplers
+- sparse residual blocks
+- compact adversarial discriminator
+
+Autoregressive probability:
+
+\[
+P(x)=\prod_t P(x_t \mid x_{<t}, M, z_s)
+\]
+
+The provided generator is a practical hybrid: conditioning is mel-driven and causal, but the structure is simplified for CPU deployment.
+
+## Losses
+
+- multi-resolution STFT loss
+- waveform adversarial loss
+- optional feature matching loss
+
+Generator loss:
+
+\[
+\mathcal{L}_G = \lambda_{\text{stft}}\mathcal{L}_{\text{MR-STFT}} + \lambda_{\text{gan}}\mathcal{L}_{\text{GAN}}^G
+\]
+
+Discriminator loss:
+
+\[
+\mathcal{L}_D = \mathbb{E}[\max(0,1-D(x))] + \mathbb{E}[\max(0,1+D(\widehat{x}))]
+\]
+
+---
+
+# 10. Training Orchestration
+
+Files:
+
+- `training/trainer.py`
+- `training/checkpoint.py`
+- `training/distributed.py`
+
+## Optimizer and Scheduling
+
+- `AdamW`
+- cosine decay
+- warmup
+- gradient clipping
+- gradient accumulation
+
+## Checkpointing
+
+`CheckpointManager` saves:
+
+- model weights
+- optimizer state
+- global step
+- metrics snapshot
+
+## Distributed Training
+
+The current blueprint includes `gloo`-based initialization for CPU-compatible distributed experiments.
+
+Recommended CPU DDP usage:
+
+```bash
+torchrun --nproc_per_node=4 -m bio_voice_tts.scripts.train_tts --config bio_voice_tts/configs/base.yaml
+```
+
+## Curriculum and Freezing
+
+Recommended staging:
+
+1. pretrain speaker encoder
+2. freeze speaker encoder and train semantic + prosody + acoustic stack
+3. train vocoder on ground-truth mels
+4. unfreeze speaker encoder partially for joint refinement
+5. enable streaming chunk curriculum
+
+---
+
+# 11. Streaming Training
+
+Files:
+
+- `streaming/cache.py`
+- `streaming/realtime_stream.py`
+
+## Chunked Training
+
+Simulate streaming by splitting sequences into chunks of size \(C\):
+
+\[
+Y = [Y^{(1)}, Y^{(2)}, \dots, Y^{(N)}]
+\]
+
+where each chunk is decoded with a bounded cache and bounded history.
+
+## KV-Cache Simulation
+
+For each chunk:
+
+- preserve previous mel summaries
+- preserve acoustic hidden summaries
+- bound cached frames to `streaming_cache_frames`
+
+## Latency Optimization
+
+- fixed chunk sizes
+- fixed sparse candidate counts
+- reuse speaker embedding
+- keep cache tensors small and contiguous
+
+---
+
+# 12. Configuration System
+
+Files:
+
+- `configs/base.yaml`
+- `configs/speaker_encoder.yaml`
+- `configs/acoustic_decoder.yaml`
+- `configs/vocoder.yaml`
+- `configs/datasets.yaml`
+- `utils/config.py`
+
+The config loader merges YAML payloads into nested dataclasses so training recipes remain type-safe and script-friendly.
+
+Example:
+
+```yaml
+training:
+  epochs: 100
+  batch_size: 4
+  learning_rate: 0.0002
+semantic:
+  d_model: 256
+  low_rank: 32
+```
+
+---
+
+# 13. Benchmarking
+
+Files:
+
+- `benchmarks/benchmark_cpu.py`
+
+## Metrics
+
+- mean CPU latency
+- peak memory usage
+- mel frames per second
+- chunk throughput
+- streaming cache overhead
+
+## Dense vs Sparse Comparison
+
+Dense attention cost:
+
+\[
+O(n^2d)
+\]
+
+Sparse attention cost:
+
+\[
+O(nkd)
+\]
+
+Benchmark protocol:
+
+1. fix token length
+2. fix mel frame length
+3. run sparse and dense variants
+4. compare latency and peak memory
+
+---
+
+# 14. Evaluation
+
+Files:
+
+- `evaluation/metrics.py`
+- `evaluation/evaluate.py`
+
+## Current Metrics
+
+- mel MAE
+- pseudo-MOS proxy
+- speaker cosine similarity proxy
+
+## Recommended Production Metrics
+
+- MOSNet or neural MOS proxy
+- ASR-backed intelligibility and WER
+- speaker verification cosine / EER
+- multi-resolution STFT metrics
+- streaming first-chunk latency
+- real-time factor
+
+WER compatibility is especially important for voice cloning because a model can sound speaker-correct while still degrading lexical accuracy.
+
+---
+
+# 15. Inference System
+
+Files:
+
+- `inference/synthesize.py`
+- `inference/clone_voice.py`
+- `inference/realtime_stream.py`
+
+Example:
+
+```bash
+python -m bio_voice_tts.inference.clone_voice \
+  --config bio_voice_tts/configs/base.yaml \
+  --speaker sample.wav \
+  --text "welcome to the interview" \
+  --out out.wav
+```
+
+Streaming:
+
+```bash
+python -m bio_voice_tts.streaming.realtime_stream \
+  --config bio_voice_tts/configs/base.yaml \
+  --speaker sample.wav \
+  --text "welcome to the interview. this is a streaming test." \
+  --out stream.wav
+```
+
+---
+
+# 16. Training Roadmap
+
+## Stage 1: Speaker Encoder
+
+- data: VoxCeleb + VCTK + MLS speaker-balanced subsets
+- objective: triplet + contrastive
+- hardware: 16-32 CPU cores or modest single GPU
+- CPU estimate: batch 4-16 depending on mel length
+- duration: 1-5 days on CPU-only, much less on GPU
+
+## Stage 2: Semantic Encoder
+
+- data: LibriTTS + MLS text/transcript-aligned corpora
+- objective: stable sparse semantic representations
+- batch: 4-8 CPU-friendly
+- focus: tokenizer quality, sparse attention stability
+
+## Stage 3: Prosody Modules
+
+- data: forced-aligned durations + extracted F0/energy
+- objective: duration/pitch/energy prediction
+- strategy: freeze speaker encoder initially
+
+## Stage 4: Acoustic Decoder
+
+- objective: text+speaker+prosody to mel
+- training: teacher forcing then scheduled sampling
+- CPU estimate: batch 2-6 depending on frame length
+
+## Stage 5: Vocoder
+
+- train on ground-truth mels first
+- use multi-resolution STFT heavily
+- add GAN gradually
+- CPU estimate: slowest stage; consider GPU assist if available
+
+## Stage 6: Joint Finetuning
+
+- partially unfreeze speaker encoder
+- align text-to-mel with vocoder feedback
+- use lower LR and aggressive checkpointing
+
+## Stage 7: Streaming Optimization
+
+- chunk curriculum
+- cache simulation
+- latency-oriented benchmarking
+- prune or quantize dense remnants
+
+## Practical Hardware Expectations
+
+- minimal CPU research box:
+  - 16 cores
+  - 64 GB RAM
+- comfortable CPU box:
+  - 32 to 64 cores
+  - 128 GB RAM
+- optional GPU acceleration:
+  - 16 GB VRAM helps with vocoder and adversarial training
+
+CPU-first does not mean CPU-only forever. It means the architecture remains deployable and trainable in low-memory environments and does not assume giant dense accelerators.
+
+---
+
+# 17. Future Research
+
+- multilingual cloning with language-conditioned phoneme spaces
+- emotional speech latents disentangled from speaker identity
+- sparse diffusion or consistency vocoders
+- unified speech foundation models with shared acoustic latents
+- TTS/STT shared sparse front-ends
+- duplex speech agents with interruption handling
+- persistent speaker memory that updates over repeated interactions
+
+---
+
+## Implementation Status
+
+The code shipped with this blueprint is intentionally structured as a strong research foundation rather than a fully benchmarked production release. It already provides:
+
+- real PyTorch modules
+- low-rank sparse attention
+- laminar refinement
+- speaker encoder training scaffolding
+- text-to-mel training scaffolding
+- sparse vocoder scaffolding
+- chunked streaming cache support
+- YAML configs
+- benchmarking and evaluation hooks
+
+The main gaps for a full-scale deployment are:
+
+- forced alignment integration
+- multilingual phonemization backends
+- large-scale distributed dataset recipes
+- richer vocoder feature matching
+- real ASR-backed intelligibility evaluation
+- quantization/export tooling
+
+Those are natural next steps, not architectural blockers.

bio_voice_tts/__init__.py

@@ -0,0 +1,6 @@
+"""BioVoice-TTS package."""
+
+from .model.biovoice_tts import BioVoiceTTS
+from .utils.config import BioVoiceConfig, load_config
+
+__all__ = ["BioVoiceTTS", "BioVoiceConfig", "load_config"]

bio_voice_tts/audio/__init__.py

@@ -0,0 +1 @@
+"""Audio feature extraction primitives."""

bio_voice_tts/audio/features.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torchaudio
+from scipy.io import wavfile
+
+from .mel import build_mel_filterbank
+from .stft import STFT, STFTConfig
+
+
+@dataclass
+class FeatureBatch:
+    waveform: torch.Tensor
+    mel: torch.Tensor
+    log_energy: torch.Tensor
+    pitch: torch.Tensor
+    voiced_mask: torch.Tensor
+
+
+class AudioFeatureExtractor:
+    """CPU-first feature extraction with minimal temporary allocations."""
+
+    def __init__(
+        self,
+        sample_rate: int = 24000,
+        n_fft: int = 1024,
+        win_length: int = 1024,
+        hop_length: int = 256,
+        n_mels: int = 80,
+        f_min: float = 0.0,
+        f_max: float = 12000.0,
+        trim_db: float = 32.0,
+    ):
+        self.sample_rate = sample_rate
+        self.trim_db = trim_db
+        self.stft = STFT(STFTConfig(sample_rate, n_fft, win_length, hop_length))
+        self.mel_basis = build_mel_filterbank(sample_rate, n_fft, n_mels, f_min, f_max)
+
+    def load_audio(self, path: str) -> torch.Tensor:
+        try:
+            waveform, sample_rate = torchaudio.load(path)
+            waveform = waveform.mean(dim=0, keepdim=True)
+        except (ImportError, OSError):
+            sample_rate, waveform_np = wavfile.read(path)
+            waveform_np = np.asarray(waveform_np)
+            if waveform_np.ndim == 2:
+                waveform_np = waveform_np.mean(axis=1)
+            if np.issubdtype(waveform_np.dtype, np.integer):
+                waveform_np = waveform_np.astype(np.float32) / np.iinfo(waveform_np.dtype).max
+            else:
+                waveform_np = waveform_np.astype(np.float32)
+            waveform = torch.from_numpy(waveform_np).unsqueeze(0)
+        if sample_rate != self.sample_rate:
+            waveform = torchaudio.functional.resample(waveform, sample_rate, self.sample_rate)
+        return waveform.squeeze(0)
+
+    def trim_silence(self, waveform: torch.Tensor) -> torch.Tensor:
+        energy = waveform.abs()
+        threshold = energy.max() * 10 ** (-self.trim_db / 20.0)
+        mask = energy > threshold
+        if not torch.any(mask):
+            return waveform
+        start = int(mask.float().argmax().item())
+        end = int(mask.numel() - mask.flip(0).float().argmax().item())
+        return waveform[start:end]
+
+    def compute_pitch(self, waveform: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        frame_time = self.stft.config.hop_length / self.sample_rate
+        pitch = torchaudio.functional.detect_pitch_frequency(
+            waveform.unsqueeze(0),
+            self.sample_rate,
+            frame_time=frame_time,
+        ).squeeze(0)
+        voiced_mask = pitch > 1.0
+        return pitch.log1p(), voiced_mask
+
+    def compute_features(self, waveform: torch.Tensor) -> FeatureBatch:
+        waveform = self.trim_silence(waveform)
+        waveform = waveform / waveform.abs().max().clamp_min(1e-6)
+        spec = self.stft.forward(waveform)
+        if spec.dim() == 3:
+            spec = spec.squeeze(0)
+        power_spec = spec.abs().pow(2.0)
+        mel = torch.matmul(self.mel_basis.to(power_spec.device), power_spec)
+        mel = torch.log(mel.clamp_min(1e-5)).transpose(0, 1)
+        log_energy = torch.log(power_spec.mean(dim=0).clamp_min(1e-5))
+        pitch, voiced_mask = self.compute_pitch(waveform)
+        frame_count = mel.size(0)
+        pitch = F.pad(pitch, (0, max(0, frame_count - pitch.numel())))[:frame_count]
+        voiced_mask = F.pad(voiced_mask.float(), (0, max(0, frame_count - voiced_mask.numel())))[:frame_count].bool()
+        return FeatureBatch(
+            waveform=waveform,
+            mel=mel,
+            log_energy=log_energy[:frame_count],
+            pitch=pitch,
+            voiced_mask=voiced_mask,
+        )

bio_voice_tts/audio/mel.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+import numpy as np
+import torch
+
+
+def hz_to_mel(freq_hz: np.ndarray) -> np.ndarray:
+    return 2595.0 * np.log10(1.0 + freq_hz / 700.0)
+
+
+def mel_to_hz(mel: np.ndarray) -> np.ndarray:
+    return 700.0 * (10 ** (mel / 2595.0) - 1.0)
+
+
+def build_mel_filterbank(
+    sample_rate: int,
+    n_fft: int,
+    n_mels: int,
+    f_min: float,
+    f_max: float,
+) -> torch.Tensor:
+    mel_min = hz_to_mel(np.array([f_min], dtype=np.float64))[0]
+    mel_max = hz_to_mel(np.array([f_max], dtype=np.float64))[0]
+    mel_points = np.linspace(mel_min, mel_max, n_mels + 2)
+    hz_points = mel_to_hz(mel_points)
+    bins = np.floor((n_fft + 1) * hz_points / sample_rate).astype(np.int64)
+    fb = np.zeros((n_mels, n_fft // 2 + 1), dtype=np.float32)
+    for idx in range(1, n_mels + 1):
+        left, center, right = bins[idx - 1], bins[idx], bins[idx + 1]
+        center = max(center, left + 1)
+        right = max(right, center + 1)
+        for fft_bin in range(left, center):
+            fb[idx - 1, fft_bin] = (fft_bin - left) / max(1, center - left)
+        for fft_bin in range(center, right):
+            fb[idx - 1, fft_bin] = (right - fft_bin) / max(1, right - center)
+    return torch.from_numpy(fb)

bio_voice_tts/audio/phoneme.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+import re
+
+
+def normalize_text(text: str) -> str:
+    text = text.lower().strip()
+    text = re.sub(r"\s+", " ", text)
+    text = re.sub(r"[^a-z0-9,;:!?.' -]", "", text)
+    return text
+
+
+def naive_phonemize(text: str) -> list[str]:
+    """Fallback grapheme-to-phoneme approximation for CPU-only pipelines."""
+
+    normalized = normalize_text(text)
+    phonemes: list[str] = []
+    for token in normalized.split():
+        phonemes.extend(list(token))
+        phonemes.append("|")
+    return phonemes[:-1] if phonemes else []

bio_voice_tts/audio/stft.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass
+class STFTConfig:
+    sample_rate: int = 24000
+    n_fft: int = 1024
+    win_length: int = 1024
+    hop_length: int = 256
+
+
+class STFT(torch.nn.Module):
+    r"""Short-time Fourier transform.
+
+    X(m, \omega) = \sum_{n=0}^{N-1} x[n] w[n - mH] e^{-j \omega n}
+    """
+
+    def __init__(self, config: STFTConfig):
+        super().__init__()
+        self.config = config
+        self.register_buffer("window", torch.hann_window(config.win_length), persistent=False)
+
+    def forward(self, waveform: torch.Tensor) -> torch.Tensor:
+        if waveform.dim() == 1:
+            waveform = waveform.unsqueeze(0)
+        return torch.stft(
+            waveform,
+            n_fft=self.config.n_fft,
+            hop_length=self.config.hop_length,
+            win_length=self.config.win_length,
+            window=self.window,
+            center=True,
+            return_complex=True,
+        )
+
+    def magnitude(self, waveform: torch.Tensor) -> torch.Tensor:
+        return self.forward(waveform).abs()

bio_voice_tts/benchmarks/__init__.py

@@ -0,0 +1 @@
+"""Benchmark helpers."""

bio_voice_tts/benchmarks/benchmark_cpu.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+import json
+import time
+import tracemalloc
+from pathlib import Path
+
+import torch
+
+from bio_voice_tts.model.biovoice_tts import BioVoiceTTS
+
+
+def benchmark_model(model: BioVoiceTTS, token_ids: torch.Tensor, reference_mel: torch.Tensor, iterations: int = 5) -> dict[str, float]:
+    latencies = []
+    tracemalloc.start()
+    for _ in range(iterations):
+        start = time.perf_counter()
+        with torch.no_grad():
+            _ = model(token_ids, reference_mel)
+        latencies.append(time.perf_counter() - start)
+    _, peak = tracemalloc.get_traced_memory()
+    tracemalloc.stop()
+    mel_frames = float(reference_mel.size(1))
+    return {
+        "mean_latency_seconds": sum(latencies) / len(latencies),
+        "peak_memory_mb": peak / (1024 * 1024),
+        "mel_frames_per_second": mel_frames / max(1e-6, sum(latencies) / len(latencies)),
+    }
+
+
+def save_benchmark(payload: dict[str, float], path: str) -> None:
+    Path(path).write_text(json.dumps(payload, indent=2), encoding="utf-8")

bio_voice_tts/configs/acoustic_decoder.yaml

@@ -0,0 +1,14 @@
+semantic:
+  d_model: 256
+  low_rank: 32
+  top_k: 12
+acoustic:
+  d_model: 256
+  low_rank: 32
+  top_k: 24
+  local_window: 48
+  chunk_size: 24
+training:
+  epochs: 100
+  batch_size: 4
+  output_dir: artifacts/biovoice_acoustic

bio_voice_tts/configs/base.yaml

@@ -0,0 +1,43 @@
+audio:
+  sample_rate: 24000
+  n_fft: 1024
+  win_length: 1024
+  hop_length: 256
+  n_mels: 80
+dataset:
+  train_manifest: data/train_manifest.jsonl
+  eval_manifest: data/eval_manifest.jsonl
+  feature_dir: data/features
+semantic:
+  vocab_size: 4096
+  d_model: 256
+  num_heads: 4
+  low_rank: 32
+  top_k: 12
+  local_window: 32
+  max_positions: 512
+speaker:
+  input_dim: 80
+  conv_channels: 128
+  embedding_dim: 192
+prosody:
+  d_model: 256
+  hidden_dim: 128
+acoustic:
+  d_model: 256
+  speaker_dim: 192
+  n_mels: 80
+  low_rank: 32
+  top_k: 24
+  local_window: 48
+  chunk_size: 24
+vocoder:
+  n_mels: 80
+  channels: 128
+  residual_layers: 6
+training:
+  epochs: 10
+  batch_size: 2
+  learning_rate: 0.0002
+  grad_accum_steps: 2
+  output_dir: artifacts/biovoice_base

bio_voice_tts/configs/datasets.yaml

@@ -0,0 +1,6 @@
+datasets:
+  train_manifest: data/train_manifest.jsonl
+  eval_manifest: data/eval_manifest.jsonl
+  feature_dir: data/features
+  sample_rate: 24000
+  transcript_format: jsonl

bio_voice_tts/configs/speaker_encoder.yaml

@@ -0,0 +1,8 @@
+training:
+  epochs: 20
+  batch_size: 8
+  learning_rate: 0.001
+  output_dir: artifacts/biovoice_speaker
+speaker:
+  conv_channels: 128
+  embedding_dim: 192

bio_voice_tts/configs/vocoder.yaml

@@ -0,0 +1,9 @@
+vocoder:
+  channels: 128
+  residual_layers: 6
+  upsample_scales: [8, 5, 3, 2]
+training:
+  epochs: 120
+  batch_size: 8
+  learning_rate: 0.0002
+  output_dir: artifacts/biovoice_vocoder

bio_voice_tts/datasets/__init__.py

@@ -0,0 +1 @@
+"""Manifest-driven datasets for TTS and speaker training."""

bio_voice_tts/datasets/manifest.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+
+
+@dataclass
+class ManifestEntry:
+    audio_path: str
+    text: str
+    speaker_id: str
+    language: str = "en"
+    sample_rate: int = 24000
+    duration_seconds: float | None = None
+    feature_path: str | None = None
+
+
+def load_manifest(path: str | Path) -> list[ManifestEntry]:
+    entries: list[ManifestEntry] = []
+    with Path(path).open("r", encoding="utf-8") as handle:
+        for line in handle:
+            if not line.strip():
+                continue
+            entries.append(ManifestEntry(**json.loads(line)))
+    return entries

bio_voice_tts/datasets/speaker_dataset.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import random
+from collections import defaultdict
+
+import torch
+from torch.utils.data import Dataset
+
+from bio_voice_tts.audio.features import AudioFeatureExtractor
+from bio_voice_tts.datasets.manifest import ManifestEntry
+
+
+class SpeakerDataset(Dataset):
+    """Triplet-friendly dataset that samples anchor/positive/negative mels."""
+
+    def __init__(self, entries: list[ManifestEntry], extractor: AudioFeatureExtractor):
+        self.entries = entries
+        self.extractor = extractor
+        self.by_speaker: dict[str, list[ManifestEntry]] = defaultdict(list)
+        for entry in entries:
+            self.by_speaker[entry.speaker_id].append(entry)
+        self.speakers = sorted(self.by_speaker)
+
+    def __len__(self) -> int:
+        return len(self.entries)
+
+    def _mel(self, entry: ManifestEntry) -> torch.Tensor:
+        return self.extractor.compute_features(self.extractor.load_audio(entry.audio_path)).mel
+
+    def __getitem__(self, index: int) -> dict[str, torch.Tensor | str]:
+        anchor = self.entries[index]
+        positive_pool = [entry for entry in self.by_speaker[anchor.speaker_id] if entry.audio_path != anchor.audio_path]
+        positive = random.choice(positive_pool or self.by_speaker[anchor.speaker_id])
+        negative_speaker = random.choice([speaker for speaker in self.speakers if speaker != anchor.speaker_id])
+        negative = random.choice(self.by_speaker[negative_speaker])
+        return {
+            "anchor_mel": self._mel(anchor),
+            "positive_mel": self._mel(positive),
+            "negative_mel": self._mel(negative),
+            "speaker_id": anchor.speaker_id,
+        }
+
+
+def pad_mels(batch: list[dict[str, torch.Tensor | str]]) -> dict[str, torch.Tensor]:
+    def _pad(key: str) -> torch.Tensor:
+        tensors = [item[key] for item in batch]
+        assert all(isinstance(tensor, torch.Tensor) for tensor in tensors)
+        max_len = max(tensor.size(0) for tensor in tensors)  # type: ignore[arg-type]
+        padded = []
+        for tensor in tensors:  # type: ignore[assignment]
+            pad = max_len - tensor.size(0)
+            padded.append(torch.nn.functional.pad(tensor, (0, 0, 0, pad)))
+        return torch.stack(padded)
+
+    return {
+        "anchor_mel": _pad("anchor_mel"),
+        "positive_mel": _pad("positive_mel"),
+        "negative_mel": _pad("negative_mel"),
+    }

bio_voice_tts/datasets/tts_dataset.py

@@ -0,0 +1,133 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+from torch.utils.data import Dataset
+
+from bio_voice_tts.audio.features import AudioFeatureExtractor
+from bio_voice_tts.audio.phoneme import naive_phonemize, normalize_text
+from bio_voice_tts.datasets.manifest import ManifestEntry
+
+
+@dataclass
+class TokenizerState:
+    stoi: dict[str, int]
+    itos: list[str]
+    pad_id: int = 0
+    bos_id: int = 1
+    eos_id: int = 2
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.itos)
+
+    def encode(self, text: str) -> list[int]:
+        phonemes = naive_phonemize(normalize_text(text))
+        ids = [self.bos_id]
+        ids.extend(self.stoi.get(token, self.stoi["<unk>"]) for token in phonemes)
+        ids.append(self.eos_id)
+        return ids
+
+
+def build_tokenizer(entries: list[ManifestEntry]) -> TokenizerState:
+    vocab = ["<pad>", "<bos>", "<eos>", "<unk>"]
+    seen = set(vocab)
+    for entry in entries:
+        for token in naive_phonemize(normalize_text(entry.text)):
+            if token not in seen:
+                seen.add(token)
+                vocab.append(token)
+    return TokenizerState(stoi={token: idx for idx, token in enumerate(vocab)}, itos=vocab)
+
+
+class TTSDataset(Dataset):
+    def __init__(self, entries: list[ManifestEntry], extractor: AudioFeatureExtractor, tokenizer: TokenizerState):
+        self.entries = entries
+        self.extractor = extractor
+        self.tokenizer = tokenizer
+
+    def _load_features(self, entry: ManifestEntry) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if entry.feature_path:
+            feature_path = Path(entry.feature_path)
+            if feature_path.exists():
+                payload = torch.load(feature_path, map_location="cpu")
+                mel = payload["mel"].float()
+                pitch = payload["pitch"].float()
+                energy = payload["energy"].float()
+                frame_count = min(mel.size(0), pitch.size(0), energy.size(0))
+                return mel[:frame_count], pitch[:frame_count], energy[:frame_count]
+
+        features = self.extractor.compute_features(self.extractor.load_audio(entry.audio_path))
+        frame_count = min(features.mel.size(0), features.pitch.size(0), features.log_energy.size(0))
+        return features.mel[:frame_count], features.pitch[:frame_count], features.log_energy[:frame_count]
+
+    def __len__(self) -> int:
+        return len(self.entries)
+
+    def __getitem__(self, index: int) -> dict[str, torch.Tensor | str]:
+        entry = self.entries[index]
+        mel, pitch, energy = self._load_features(entry)
+        token_ids = torch.tensor(self.tokenizer.encode(entry.text), dtype=torch.long)
+        duration = max(1, mel.size(0) // max(1, token_ids.numel() - 2))
+        durations = torch.full((token_ids.numel(),), duration, dtype=torch.long)
+        return {
+            "token_ids": token_ids,
+            "mel": mel,
+            "pitch": pitch,
+            "energy": energy,
+            "durations": durations,
+            "speaker_id": entry.speaker_id,
+            "text": entry.text,
+        }
+
+
+def tts_collate(batch: list[dict[str, torch.Tensor | str]]) -> dict[str, torch.Tensor]:
+    token_tensors = [item["token_ids"] for item in batch]
+    mel_tensors = [item["mel"] for item in batch]
+    pitch_tensors = [item["pitch"] for item in batch]
+    energy_tensors = [item["energy"] for item in batch]
+    duration_tensors = [item["durations"] for item in batch]
+    assert all(isinstance(tensor, torch.Tensor) for tensor in token_tensors + mel_tensors + pitch_tensors + energy_tensors + duration_tensors)
+
+    max_tokens = max(tensor.size(0) for tensor in token_tensors)  # type: ignore[arg-type]
+    max_frames = max(tensor.size(0) for tensor in mel_tensors)  # type: ignore[arg-type]
+    n_mels = mel_tensors[0].size(1)  # type: ignore[index]
+
+    token_batch = torch.zeros(len(batch), max_tokens, dtype=torch.long)
+    duration_batch = torch.zeros(len(batch), max_tokens, dtype=torch.long)
+    mel_batch = torch.zeros(len(batch), max_frames, n_mels)
+    pitch_batch = torch.zeros(len(batch), max_frames)
+    energy_batch = torch.zeros(len(batch), max_frames)
+    token_lengths = torch.zeros(len(batch), dtype=torch.long)
+    frame_lengths = torch.zeros(len(batch), dtype=torch.long)
+
+    for idx, item in enumerate(batch):
+        tokens = item["token_ids"]
+        mel = item["mel"]
+        pitch = item["pitch"]
+        energy = item["energy"]
+        durations = item["durations"]
+        assert isinstance(tokens, torch.Tensor)
+        assert isinstance(mel, torch.Tensor)
+        assert isinstance(pitch, torch.Tensor)
+        assert isinstance(energy, torch.Tensor)
+        assert isinstance(durations, torch.Tensor)
+        token_batch[idx, : tokens.size(0)] = tokens
+        duration_batch[idx, : durations.size(0)] = durations
+        mel_batch[idx, : mel.size(0)] = mel
+        pitch_batch[idx, : pitch.size(0)] = pitch
+        energy_batch[idx, : energy.size(0)] = energy
+        token_lengths[idx] = tokens.size(0)
+        frame_lengths[idx] = mel.size(0)
+
+    return {
+        "token_ids": token_batch,
+        "durations": duration_batch,
+        "mel": mel_batch,
+        "pitch": pitch_batch,
+        "energy": energy_batch,
+        "token_lengths": token_lengths,
+        "frame_lengths": frame_lengths,
+    }

bio_voice_tts/evaluation/__init__.py

@@ -0,0 +1 @@
+"""Evaluation metrics and CLI helpers."""

bio_voice_tts/evaluation/evaluate.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+
+from bio_voice_tts.datasets.tts_dataset import TTSDataset, build_tokenizer, tts_collate
+from bio_voice_tts.evaluation.metrics import mel_mae, pseudo_mos
+from bio_voice_tts.model.biovoice_tts import BioVoiceTTS
+from bio_voice_tts.utils.config import load_config
+
+
+@torch.no_grad()
+def evaluate_model(config_path: str, checkpoint_path: str, entries, extractor) -> dict[str, float]:
+    config = load_config(config_path)
+    tokenizer = build_tokenizer(entries)
+    config.semantic.vocab_size = tokenizer.vocab_size
+    dataset = TTSDataset(entries, extractor, tokenizer)
+    batch = tts_collate([dataset[idx] for idx in range(min(2, len(dataset)))])
+    model = BioVoiceTTS(config)
+    state = torch.load(checkpoint_path, map_location="cpu")
+    model.load_state_dict(state["model"], strict=False)
+    model.eval()
+    outputs = model(batch["token_ids"], batch["mel"], batch["durations"], batch["pitch"], batch["energy"])
+    frame_limit = min(outputs["mel"].size(1), batch["mel"].size(1))
+    mae = mel_mae(outputs["mel"][:, :frame_limit], batch["mel"][:, :frame_limit])
+    mos = pseudo_mos(mae, torch.tensor(0.9))
+    return {"mel_mae": float(mae.item()), "pseudo_mos": float(mos.item())}
+
+
+def save_evaluation(payload: dict[str, float], output: str) -> None:
+    Path(output).write_text(json.dumps(payload, indent=2), encoding="utf-8")

bio_voice_tts/evaluation/metrics.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+
+def mel_mae(predicted: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    return (predicted - target).abs().mean()
+
+
+def speaker_cosine_similarity(predicted: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+    return F.cosine_similarity(predicted, target, dim=-1).mean()
+
+
+def pseudo_mos(mel_error: torch.Tensor, speaker_similarity: torch.Tensor) -> torch.Tensor:
+    score = 5.0 - 2.0 * mel_error.clamp(max=1.5) + 0.5 * speaker_similarity.clamp(min=0.0, max=1.0)
+    return score.clamp(1.0, 5.0)

bio_voice_tts/inference/__init__.py

@@ -0,0 +1 @@
+"""Inference entrypoints for cloning and synthesis."""

bio_voice_tts/inference/clone_voice.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+import argparse
+
+import torch
+
+from bio_voice_tts.inference.synthesize import build_dummy_tokenizer, synthesize
+from bio_voice_tts.model.biovoice_tts import BioVoiceTTS
+from bio_voice_tts.utils.config import load_config
+from bio_voice_tts.vocoder.sparse_vocoder import SparseNeuralVocoder
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Clone a voice with BioVoice-TTS.")
+    parser.add_argument("--config", required=True)
+    parser.add_argument("--checkpoint", required=False)
+    parser.add_argument("--vocoder-checkpoint", required=False)
+    parser.add_argument("--speaker", required=True)
+    parser.add_argument("--text", required=True)
+    parser.add_argument("--out", default="out.wav")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    config = load_config(args.config)
+    model = BioVoiceTTS(config)
+    vocoder = SparseNeuralVocoder(
+        n_mels=config.vocoder.n_mels,
+        channels=config.vocoder.channels,
+        residual_layers=config.vocoder.residual_layers,
+        upsample_scales=config.vocoder.upsample_scales,
+    )
+    if args.checkpoint:
+        state = torch.load(args.checkpoint, map_location="cpu")
+        model.load_state_dict(state["model"], strict=False)
+    if args.vocoder_checkpoint:
+        state = torch.load(args.vocoder_checkpoint, map_location="cpu")
+        vocoder.load_state_dict(state["model"], strict=False)
+    synthesize(model, vocoder, build_dummy_tokenizer(), args.text, args.speaker, args.out, sample_rate=config.audio.sample_rate)
+
+
+if __name__ == "__main__":
+    main()

bio_voice_tts/inference/realtime_stream.py

@@ -0,0 +1,5 @@
+from bio_voice_tts.streaming.realtime_stream import main
+
+
+if __name__ == "__main__":
+    main()

bio_voice_tts/inference/synthesize.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+import torch
+import torchaudio
+
+from bio_voice_tts.audio.features import AudioFeatureExtractor
+from bio_voice_tts.datasets.tts_dataset import TokenizerState, build_tokenizer
+from bio_voice_tts.model.biovoice_tts import BioVoiceTTS
+from bio_voice_tts.vocoder.sparse_vocoder import SparseNeuralVocoder
+
+
+@torch.no_grad()
+def synthesize(
+    model: BioVoiceTTS,
+    vocoder: SparseNeuralVocoder,
+    tokenizer: TokenizerState,
+    text: str,
+    reference_path: str,
+    out_path: str,
+    sample_rate: int = 24000,
+) -> None:
+    extractor = AudioFeatureExtractor(sample_rate=sample_rate)
+    features = extractor.compute_features(extractor.load_audio(reference_path))
+    token_ids = torch.tensor([tokenizer.encode(text)], dtype=torch.long)
+    outputs = model(token_ids, features.mel.unsqueeze(0))
+    waveform = vocoder(outputs["mel"]).cpu()
+    torchaudio.save(out_path, waveform.unsqueeze(1), sample_rate)
+
+
+def build_dummy_tokenizer() -> TokenizerState:
+    vocab = ["<pad>", "<bos>", "<eos>", "<unk>"] + list("abcdefghijklmnopqrstuvwxyz|' ")
+    return TokenizerState(stoi={token: idx for idx, token in enumerate(vocab)}, itos=vocab)

bio_voice_tts/model/__init__.py

@@ -0,0 +1 @@
+"""BioVoice-TTS model modules."""

bio_voice_tts/model/acoustic_decoder.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from .laminar import LaminarRefinement
+from .sparse_attention import SparseAttention
+
+
+class AcousticEnergyHead(nn.Module):
+    def __init__(self, d_model: int, speaker_dim: int, n_mels: int):
+        super().__init__()
+        self.proposal = nn.Linear(d_model, n_mels)
+        self.context_proj = nn.Linear(d_model, n_mels)
+        self.speaker_proj = nn.Linear(speaker_dim, n_mels)
+        self.memory_proj = nn.Linear(d_model, n_mels)
+        self.gates = nn.Linear(d_model + speaker_dim, 4)
+
+    def forward(self, states: torch.Tensor, speaker_latent: torch.Tensor, memory: torch.Tensor) -> dict[str, torch.Tensor]:
+        acoustic = self.proposal(states)
+        context = self.context_proj(states)
+        speaker = self.speaker_proj(speaker_latent).unsqueeze(1).expand_as(acoustic)
+        memory_term = self.memory_proj(memory)
+        gates = torch.softmax(self.gates(torch.cat([states, speaker_latent.unsqueeze(1).expand(states.size(0), states.size(1), -1)], dim=-1)), dim=-1)
+        mel = gates[..., 0:1] * acoustic + gates[..., 1:2] * context + gates[..., 2:3] * speaker + gates[..., 3:4] * memory_term
+        return {"mel": mel, "gates": gates, "acoustic": acoustic, "context": context, "speaker": speaker, "memory": memory_term}
+
+
+class SparseAcousticDecoder(nn.Module):
+    def __init__(self, d_model: int, speaker_dim: int, n_mels: int, low_rank: int, top_k: int, local_window: int):
+        super().__init__()
+        self.input_proj = nn.Linear(d_model + speaker_dim + 2, d_model)
+        self.sparse_attention = SparseAttention(
+            d_model=d_model,
+            rank=low_rank,
+            top_k=top_k,
+            local_window=local_window,
+            memory_candidates=8,
+            landmark_count=4,
+            content_memory_candidates=4,
+            max_positions=4096,
+        )
+        self.laminar = LaminarRefinement(d_model, steps=2, eta=0.05)
+        self.memory = nn.Linear(d_model, d_model)
+        self.energy_head = AcousticEnergyHead(d_model, speaker_dim, n_mels)
+
+    def forward(
+        self,
+        frame_states: torch.Tensor,
+        speaker_latent: torch.Tensor,
+        pitch: torch.Tensor,
+        energy: torch.Tensor,
+        attention_mode: str = "topk",
+    ) -> dict[str, torch.Tensor]:
+        inputs = torch.cat([frame_states, speaker_latent.unsqueeze(1).expand(frame_states.size(0), frame_states.size(1), -1), pitch.unsqueeze(-1), energy.unsqueeze(-1)], dim=-1)
+        hidden = self.input_proj(inputs)
+        state = self.sparse_attention(hidden, mode=attention_mode)
+        refined = self.laminar(state.context, state.indices, state.weights)
+        memory = self.memory(refined)
+        mel_outputs = self.energy_head(refined, speaker_latent, memory)
+        return {
+            "hidden_states": refined,
+            "attention_weights": state.weights,
+            "attention_indices": state.indices,
+            **mel_outputs,
+        }

bio_voice_tts/model/biovoice_tts.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from bio_voice_tts.utils.config import BioVoiceConfig
+
+from .acoustic_decoder import SparseAcousticDecoder
+from .prosody import DurationPredictor, EnergyPredictor, FiLMConditioner, PitchPredictor, length_regulate
+from .semantic_encoder import SemanticEncoder
+from .speaker_encoder import SpeakerEncoder
+
+
+class BioVoiceTTS(nn.Module):
+    def __init__(self, config: BioVoiceConfig):
+        super().__init__()
+        self.config = config
+        self.speaker_encoder = SpeakerEncoder(
+            input_dim=config.audio.n_mels,
+            channels=config.speaker.conv_channels,
+            embedding_dim=config.speaker.embedding_dim,
+        )
+        self.semantic_encoder = SemanticEncoder(
+            vocab_size=config.semantic.vocab_size,
+            d_model=config.semantic.d_model,
+            low_rank=config.semantic.low_rank,
+            top_k=config.semantic.top_k,
+            local_window=config.semantic.local_window,
+            memory_candidates=config.semantic.memory_candidates,
+            landmark_count=config.semantic.landmark_count,
+            content_memory_candidates=config.semantic.content_memory_candidates,
+            laminar_steps=config.semantic.laminar_steps,
+            laminar_eta=config.semantic.laminar_eta,
+            max_positions=config.semantic.max_positions,
+        )
+        self.conditioner = FiLMConditioner(config.speaker.embedding_dim, config.semantic.d_model)
+        self.duration_predictor = DurationPredictor(config.semantic.d_model, config.prosody.hidden_dim)
+        self.pitch_predictor = PitchPredictor(config.semantic.d_model, config.prosody.hidden_dim)
+        self.energy_predictor = EnergyPredictor(config.semantic.d_model, config.prosody.hidden_dim)
+        self.acoustic_decoder = SparseAcousticDecoder(
+            d_model=config.acoustic.d_model,
+            speaker_dim=config.acoustic.speaker_dim,
+            n_mels=config.acoustic.n_mels,
+            low_rank=config.acoustic.low_rank,
+            top_k=config.acoustic.top_k,
+            local_window=config.acoustic.local_window,
+        )
+
+    @staticmethod
+    def _match_frame_length(values: torch.Tensor, target_frames: int) -> torch.Tensor:
+        if values.size(1) == target_frames:
+            return values
+        if values.size(1) > target_frames:
+            return values[:, :target_frames]
+        return F.pad(values, (0, target_frames - values.size(1)))
+
+    def forward(
+        self,
+        token_ids: torch.Tensor,
+        reference_mel: torch.Tensor,
+        target_durations: torch.Tensor | None = None,
+        target_pitch: torch.Tensor | None = None,
+        target_energy: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
+        speaker_latent = self.speaker_encoder(reference_mel)
+        semantic = self.semantic_encoder(token_ids)
+        conditioned = self.conditioner(semantic["hidden_states"], speaker_latent)
+        duration_pred = self.duration_predictor(conditioned)
+        durations = target_durations.float() if target_durations is not None else duration_pred
+        frame_states, frame_lengths = length_regulate(conditioned, durations, max_frames=self.config.dataset.max_mel_frames)
+        token_pitch = self.pitch_predictor(conditioned)
+        token_energy = self.energy_predictor(conditioned)
+        expanded_pitch, _ = length_regulate(token_pitch.unsqueeze(-1), durations, max_frames=frame_states.size(1))
+        expanded_energy, _ = length_regulate(token_energy.unsqueeze(-1), durations, max_frames=frame_states.size(1))
+        frame_pitch_pred = expanded_pitch.squeeze(-1)
+        frame_energy_pred = expanded_energy.squeeze(-1)
+        pitch = self._match_frame_length(target_pitch, frame_states.size(1)) if target_pitch is not None else frame_pitch_pred
+        energy = self._match_frame_length(target_energy, frame_states.size(1)) if target_energy is not None else frame_energy_pred
+        acoustic = self.acoustic_decoder(frame_states, speaker_latent, pitch, energy)
+        return {
+            "speaker_latent": speaker_latent,
+            "duration_pred": duration_pred,
+            "token_pitch_pred": token_pitch,
+            "token_energy_pred": token_energy,
+            "frame_pitch_pred": frame_pitch_pred,
+            "frame_energy_pred": frame_energy_pred,
+            "frame_lengths": frame_lengths,
+            **semantic,
+            **acoustic,
+        }

bio_voice_tts/model/laminar.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+class LaminarRefinement(nn.Module):
+    def __init__(self, d_model: int, steps: int = 2, eta: float = 0.1):
+        super().__init__()
+        self.steps = steps
+        self.eta = eta
+        self.excitatory_proj = nn.Linear(d_model, d_model)
+        self.inhibitory_proj = nn.Linear(d_model, d_model)
+
+    def _weighted_sum(self, states: torch.Tensor, indices: torch.Tensor, weights: torch.Tensor, proj: nn.Linear) -> torch.Tensor:
+        batch, seq_len, _ = states.shape
+        mixed = torch.zeros_like(states)
+        projected = proj(states)
+        for batch_idx in range(batch):
+            for pos in range(seq_len):
+                idx = indices[batch_idx, pos]
+                w = weights[batch_idx, pos].unsqueeze(-1)
+                mixed[batch_idx, pos] = torch.sum(projected[batch_idx, idx] * w, dim=0)
+        return mixed
+
+    def forward(self, states: torch.Tensor, indices: torch.Tensor, weights: torch.Tensor) -> torch.Tensor:
+        refined = states
+        for _ in range(self.steps):
+            excitatory = self._weighted_sum(refined, indices, weights, self.excitatory_proj)
+            inhibitory = refined.cumsum(dim=1) - self._weighted_sum(refined, indices, weights, self.inhibitory_proj)
+            refined = refined + self.eta * (excitatory - inhibitory)
+        return refined