Spaces:
Sleeping
Sleeping
Premchandyadav369
feat: deploy interactive Q-TensorFormer demo with trained weights and Gradio 4 UI
67ded87 | """ | |
| Baseline implementations for fair comparison. | |
| Baselines: | |
| 1. Standard Transformer: Dense MLP FFN, no TT, no quantum. | |
| 2. Distilled: Smaller transformer trained with KD. | |
| 3. Pruned: Magnitude-based structured pruning. | |
| 4. TT-Only: Tensor network FFN without quantum or adaptive rank. | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import math | |
| from typing import Optional | |
| class StandardTransformer(nn.Module): | |
| """ | |
| Basic transformer decoder (GPT-style) with dense MLP FFN. | |
| Reference baseline — matches Q-TensorFormer architecture | |
| exactly except for TT decomposition and quantum layers. | |
| """ | |
| def __init__(self, vocab_size: int = 10000, d_model: int = 128, | |
| n_heads: int = 4, n_layers: int = 2, ff_mult: int = 4, | |
| max_seq_len: int = 128, dropout: float = 0.1): | |
| super().__init__() | |
| self.d_model = d_model | |
| self.config = type("config", (), { | |
| "d_model": d_model, "n_heads": n_heads, "n_layers": n_layers, | |
| "ff_multiplier": ff_mult, "max_seq_len": max_seq_len, | |
| "vocab_size": vocab_size, "dropout": dropout, | |
| })() | |
| self.embedding = nn.Embedding(vocab_size, d_model) | |
| self.pos_encoding = _PositionalEncoding(d_model, max_seq_len, dropout) | |
| self.blocks = nn.ModuleList([ | |
| _StandardBlock(d_model, n_heads, ff_mult, dropout, max_seq_len) | |
| for _ in range(n_layers) | |
| ]) | |
| self.ln_f = nn.LayerNorm(d_model) | |
| self.lm_head = nn.Linear(d_model, vocab_size, bias=False) | |
| self.lm_head.weight = self.embedding.weight | |
| def forward(self, input_ids, attention_mask=None, return_stats=False): | |
| x = self.embedding(input_ids) | |
| x = self.pos_encoding(x) | |
| for block in self.blocks: | |
| x = block(x, mask=attention_mask) | |
| x = self.ln_f(x) | |
| logits = self.lm_head(x) | |
| if return_stats: | |
| return logits, [] | |
| return logits | |
| def total_params(self) -> int: | |
| return sum(p.numel() for p in self.parameters()) | |
| class DistilledTransformer(nn.Module): | |
| """ | |
| Smaller transformer trained via knowledge distillation. | |
| Designed to match Q-TensorFormer parameter counts. | |
| """ | |
| def __init__(self, vocab_size: int = 10000, d_model: int = 96, | |
| n_heads: int = 4, n_layers: int = 2, ff_mult: int = 3, | |
| max_seq_len: int = 128, dropout: float = 0.1): | |
| super().__init__() | |
| self.d_model = d_model | |
| self.config = type("config", (), { | |
| "d_model": d_model, "n_heads": n_heads, "n_layers": n_layers, | |
| "ff_multiplier": ff_mult, "max_seq_len": max_seq_len, | |
| "vocab_size": vocab_size, "dropout": dropout, | |
| })() | |
| self.embedding = nn.Embedding(vocab_size, d_model) | |
| self.pos_encoding = _PositionalEncoding(d_model, max_seq_len, dropout) | |
| self.blocks = nn.ModuleList([ | |
| _StandardBlock(d_model, n_heads, ff_mult, dropout, max_seq_len) | |
| for _ in range(n_layers) | |
| ]) | |
| self.ln_f = nn.LayerNorm(d_model) | |
| self.lm_head = nn.Linear(d_model, vocab_size, bias=False) | |
| self.lm_head.weight = self.embedding.weight | |
| def forward(self, input_ids, attention_mask=None, return_stats=False): | |
| x = self.embedding(input_ids) | |
| x = self.pos_encoding(x) | |
| for block in self.blocks: | |
| x = block(x, mask=attention_mask) | |
| x = self.ln_f(x) | |
| logits = self.lm_head(x) | |
| if return_stats: | |
| return logits, [] | |
| return logits | |
| def total_params(self) -> int: | |
| return sum(p.numel() for p in self.parameters()) | |
| class PrunedTransformer(nn.Module): | |
| """ | |
| Magnitude-pruned standard transformer. | |
| Prunes FFN weights globally to match Q-TensorFormer parameter count. | |
| Applies structured pruning (zeroing channels) for efficiency. | |
| """ | |
| def __init__(self, base_model: StandardTransformer, | |
| prune_ratio: float = 0.5): | |
| super().__init__() | |
| self.base = base_model | |
| self.prune_ratio = prune_ratio | |
| self.config = base_model.config | |
| self._prune() | |
| def _prune(self): | |
| """Apply structured magnitude pruning to FFN layers.""" | |
| all_weights = [] | |
| for block in self.base.blocks: | |
| for weight in [block.ffn[0].weight, block.ffn[2].weight]: | |
| all_weights.append(weight.flatten()) | |
| # Compute global threshold | |
| flat = torch.cat(all_weights) | |
| k = int(len(flat) * self.prune_ratio) | |
| threshold = torch.topk(flat.abs(), k, largest=False).values[-1] | |
| # Apply structured pruning (zero rows/cols) | |
| for block in self.base.blocks: | |
| for layer in [block.ffn[0], block.ffn[2]]: | |
| mask = (layer.weight.abs() > threshold).float() | |
| # Zero small rows entirely | |
| row_norms = mask.sum(dim=1) | |
| dead_rows = row_norms < layer.weight.size(1) * 0.1 | |
| mask[dead_rows] = 0 | |
| layer.weight.data *= mask | |
| def forward(self, *args, **kwargs): | |
| return self.base(*args, **kwargs) | |
| def total_params(self) -> int: | |
| return sum(p.numel() for p in self.parameters()) | |
| class _StandardBlock(nn.Module): | |
| """Standard transformer decoder block.""" | |
| def __init__(self, d_model, n_heads, ff_mult, dropout, max_seq_len): | |
| super().__init__() | |
| self.ln1 = nn.LayerNorm(d_model) | |
| self.attn = _CausalAttention(d_model, n_heads, dropout, max_seq_len) | |
| self.ln2 = nn.LayerNorm(d_model) | |
| self.ffn = nn.Sequential( | |
| nn.Linear(d_model, d_model * ff_mult), | |
| nn.GELU(), | |
| nn.Linear(d_model * ff_mult, d_model), | |
| nn.Dropout(dropout), | |
| ) | |
| self.dropout = nn.Dropout(dropout) | |
| def forward(self, x, mask=None): | |
| x = x + self.dropout(self.attn(self.ln1(x), mask=mask)) | |
| x = x + self.ffn(self.ln2(x)) | |
| return x | |
| class _CausalAttention(nn.Module): | |
| """Causal multi-head attention.""" | |
| def __init__(self, d_model, n_heads, dropout, max_seq_len): | |
| super().__init__() | |
| assert d_model % n_heads == 0 | |
| self.n_heads = n_heads | |
| self.head_dim = d_model // n_heads | |
| self.scale = math.sqrt(self.head_dim) | |
| self.qkv = nn.Linear(d_model, 3 * d_model, bias=False) | |
| self.out_proj = nn.Linear(d_model, d_model, bias=False) | |
| self.dropout = nn.Dropout(dropout) | |
| self.max_seq_len = max_seq_len | |
| def forward(self, x, mask=None): | |
| B, T, C = x.shape | |
| qkv = self.qkv(x).reshape(B, T, 3, self.n_heads, self.head_dim) | |
| q, k, v = qkv.unbind(dim=2) | |
| q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) | |
| attn = (q @ k.transpose(-2, -1)) / self.scale | |
| hide_mask = torch.triu(torch.ones(T, T, device=x.device, dtype=torch.bool), diagonal=1).view(1, 1, T, T) | |
| if mask is not None: | |
| m = mask.to(device=x.device) | |
| if m.dtype == torch.bool: | |
| pad_hide = ~m | |
| elif m.is_floating_point() and (m < 0).any(): | |
| pad_hide = m < -1e4 | |
| else: | |
| pad_hide = (m == 0) | |
| if pad_hide.dim() == 2: | |
| pad_hide = pad_hide[:, None, None, :T] | |
| elif pad_hide.dim() == 3: | |
| pad_hide = pad_hide.unsqueeze(1) | |
| hide_mask = hide_mask | pad_hide | |
| row_all_masked = hide_mask.all(dim=-1, keepdim=True) | |
| safe_hide = hide_mask & ~row_all_masked | |
| attn = attn.masked_fill(safe_hide, float("-inf")) | |
| attn = F.softmax(attn, dim=-1) | |
| attn = attn.masked_fill(row_all_masked, 0.0) | |
| attn = self.dropout(attn) | |
| out = (attn @ v).transpose(1, 2).reshape(B, T, C) | |
| return self.out_proj(out) | |
| class _PositionalEncoding(nn.Module): | |
| def __init__(self, d_model, max_len, dropout): | |
| super().__init__() | |
| self.dropout = nn.Dropout(dropout) | |
| self.d_model = d_model | |
| pe = torch.zeros(max_len, d_model) | |
| pos = torch.arange(max_len).unsqueeze(1).float() | |
| div = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) | |
| pe[:, 0::2] = torch.sin(pos * div) | |
| pe[:, 1::2] = torch.cos(pos * div) | |
| self.register_buffer("pe", pe.unsqueeze(0)) | |
| def forward(self, x): | |
| T = x.size(1) | |
| if T <= self.pe.size(1): | |
| return self.dropout(x + self.pe[:, :T]) | |
| pos = torch.arange(T, device=x.device).unsqueeze(1).float() | |
| div = torch.exp(torch.arange(0, self.d_model, 2, device=x.device).float() * (-math.log(10000.0) / self.d_model)) | |
| pe_ext = torch.zeros(1, T, self.d_model, device=x.device, dtype=x.dtype) | |
| pe_ext[0, :, 0::2] = torch.sin(pos * div) | |
| pe_ext[0, :, 1::2] = torch.cos(pos * div) | |
| return self.dropout(x + pe_ext) | |