diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..595b77229591f0505a1e2e64c3cf6e4282403400 --- /dev/null +++ b/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 + diff --git a/bio_llm/__init__.py b/bio_llm/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..19305b51abc72dbc8830b835908438dd8a2c8606 --- /dev/null +++ b/bio_llm/__init__.py @@ -0,0 +1 @@ +"""Bio-LLM package.""" diff --git a/bio_llm/model/__init__.py b/bio_llm/model/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ff46f38e70318b060b7ba4ad618900280c1307e6 --- /dev/null +++ b/bio_llm/model/__init__.py @@ -0,0 +1 @@ +"""Model components for the Structured Sparse Energy Transformer.""" diff --git a/bio_llm/model/candidate_retrieval.py b/bio_llm/model/candidate_retrieval.py new file mode 100644 index 0000000000000000000000000000000000000000..003b55f514d8c5dfd111b66ff764b59d0363ffa0 --- /dev/null +++ b/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 diff --git a/bio_llm/model/embedding.py b/bio_llm/model/embedding.py new file mode 100644 index 0000000000000000000000000000000000000000..ec12588ee2b4bfdc6a58c708042b9bfc330d3e33 --- /dev/null +++ b/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) diff --git a/bio_llm/model/energy_head.py b/bio_llm/model/energy_head.py new file mode 100644 index 0000000000000000000000000000000000000000..8b40926875d68c3d2c49b19eacdce81a8fa4ef26 --- /dev/null +++ b/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, + } diff --git a/bio_llm/model/laminar_layer.py b/bio_llm/model/laminar_layer.py new file mode 100644 index 0000000000000000000000000000000000000000..522df9b723b2cfb4bddec4698d2cd01c26682902 --- /dev/null +++ b/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 diff --git a/bio_llm/model/low_rank_qkv.py b/bio_llm/model/low_rank_qkv.py new file mode 100644 index 0000000000000000000000000000000000000000..d93883e42350046a0dcf289ee179706965d0c776 --- /dev/null +++ b/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) diff --git a/bio_llm/model/model.py b/bio_llm/model/model.py new file mode 100644 index 0000000000000000000000000000000000000000..f723547d0436e430e8bca2b100797253e25c3c07 --- /dev/null +++ b/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 diff --git a/bio_llm/model/sparse_attention.py b/bio_llm/model/sparse_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..db59c29a7301ececa635e4ffa111c9031c044316 --- /dev/null +++ b/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, + ) diff --git a/bio_llm/training/__init__.py b/bio_llm/training/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9686203107b353d8114919765b6ea3433c7654bd --- /dev/null +++ b/bio_llm/training/__init__.py @@ -0,0 +1 @@ +"""Training utilities for the Structured Sparse Energy Transformer.""" diff --git a/bio_llm/training/chat_dataset.py b/bio_llm/training/chat_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..bc3b4a1963667c5434a3ad59aba146d1ec1afe3a --- /dev/null +++ b/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"\s*(?P[a-zA-Z_]+)\s*\n(?P.*?)", + 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 diff --git a/bio_llm/training/interview_dataset.py b/bio_llm/training/interview_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..f4e10d20928a71404b51e49671f902792147c641 --- /dev/null +++ b/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", + ) diff --git a/bio_llm/training/loss.py b/bio_llm/training/loss.py new file mode 100644 index 0000000000000000000000000000000000000000..4c51a110dd3005ae3658367aa785c315c0a651b5 --- /dev/null +++ b/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 diff --git a/bio_llm/training/metrics.py b/bio_llm/training/metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..39ded488c4f3b0e8b8e388e2ba805992ff49a209 --- /dev/null +++ b/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, + } \ No newline at end of file diff --git a/bio_llm/training/recipes.py b/bio_llm/training/recipes.py new file mode 100644 index 0000000000000000000000000000000000000000..e62921055ef266188b3d867676b64504180e88b8 --- /dev/null +++ b/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) diff --git a/bio_llm/training/trainer.py b/bio_llm/training/trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..d978301db5dfbf570f1c6309188f0d3a51ec0d87 --- /dev/null +++ b/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 diff --git a/bio_llm/utils/__init__.py b/bio_llm/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..45ea4e983937f28bc3c5d1c219c9abfd33f8f8a6 --- /dev/null +++ b/bio_llm/utils/__init__.py @@ -0,0 +1 @@ +"""Utility helpers for the Structured Sparse Energy Transformer.""" diff --git a/bio_llm/utils/config.py b/bio_llm/utils/config.py new file mode 100644 index 0000000000000000000000000000000000000000..16f025344ab9f28d851f618eea29a446e740722a --- /dev/null +++ b/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) diff --git a/bio_llm/utils/export.py b/bio_llm/utils/export.py new file mode 100644 index 0000000000000000000000000000000000000000..e71a250001585227b4f16345719b22917a0f9ffa --- /dev/null +++ b/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 diff --git a/bio_llm/utils/tokenizer.py b/bio_llm/utils/tokenizer.py new file mode 100644 index 0000000000000000000000000000000000000000..5d420b6a449fcc424bc508e3d005580e846d63c7 --- /dev/null +++ b/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 = "" + BOS = "" + EOS = "" + 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 = "" + BOS = "" + EOS = "" + UNK = "" + END_OF_WORD = "" + 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"]) diff --git a/bio_voice_tts/TRAINING_INFRASTRUCTURE.md b/bio_voice_tts/TRAINING_INFRASTRUCTURE.md new file mode 100644 index 0000000000000000000000000000000000000000..4b85b191edee83ea04324906b519a2f42b8aff0b --- /dev/null +++ b/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//` + +--- + +# 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_{