Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
| """ | |
| YOCTO — Hugging Face Space | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import gradio as gr | |
| from huggingface_hub import hf_hub_download | |
| from tokenizers import Tokenizer as HFTokenizer | |
| from typing import Optional, Tuple, List | |
| from dataclasses import dataclass, field | |
| import time | |
| # ============================================================================== | |
| # CONFIG | |
| # ============================================================================== | |
| class ModelConfig: | |
| vocab_size: int = 4000 | |
| embed_dim: int = 72 | |
| num_heads: int = 3 | |
| num_layers: int = 4 | |
| ffn_dim: int = 288 | |
| max_seq_len: int = 512 | |
| dropout: float = 0.1 | |
| head_dim: int = field(init=False) | |
| third_dim: int = field(init=False) | |
| component_head_dim: int = field(init=False) | |
| def __post_init__(self): | |
| self.head_dim = self.embed_dim // self.num_heads | |
| self.third_dim = self.embed_dim // 3 | |
| self.component_head_dim = self.third_dim // self.num_heads | |
| # ============================================================================== | |
| # MODEL | |
| # ============================================================================== | |
| class RotaryPositionEmbedding(nn.Module): | |
| def __init__(self, dim: int, max_seq_len: int = 2048, base: float = 10000.0): | |
| super().__init__() | |
| self.dim = dim | |
| self.max_seq_len = max_seq_len | |
| inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) | |
| self.register_buffer("inv_freq", inv_freq, persistent=False) | |
| self._build_cache(max_seq_len) | |
| def _build_cache(self, seq_len: int): | |
| positions = torch.arange(seq_len, dtype=torch.float32, device=self.inv_freq.device) | |
| freqs = torch.outer(positions, self.inv_freq) | |
| emb = torch.cat([freqs, freqs], dim=-1) | |
| self.register_buffer("cos_cached", emb.cos(), persistent=False) | |
| self.register_buffer("sin_cached", emb.sin(), persistent=False) | |
| self.max_seq_len = seq_len | |
| def forward(self, x, seq_len, offset=0): | |
| if offset + seq_len > self.max_seq_len: | |
| self._build_cache(offset + seq_len) | |
| return ( | |
| self.cos_cached[offset:offset + seq_len].to(x.dtype), | |
| self.sin_cached[offset:offset + seq_len].to(x.dtype) | |
| ) | |
| def rotate_half(x): | |
| x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:] | |
| return torch.cat([-x2, x1], dim=-1) | |
| def apply_rotary_pos_emb(s, o, cos, sin): | |
| cos, sin = cos.unsqueeze(0).unsqueeze(0), sin.unsqueeze(0).unsqueeze(0) | |
| return (s * cos) + (rotate_half(s) * sin), (o * cos) + (rotate_half(o) * sin) | |
| class KVCache: | |
| def __init__(self): | |
| self.cache = [] | |
| def reset(self, num_layers): | |
| self.cache = [None] * num_layers | |
| def get(self, idx): | |
| return self.cache[idx] | |
| def update(self, idx, k, v): | |
| if self.cache[idx] is None: | |
| self.cache[idx] = (k, v) | |
| else: | |
| ck, cv = self.cache[idx] | |
| self.cache[idx] = (torch.cat([ck, k], dim=2), torch.cat([cv, v], dim=2)) | |
| def seq_len(self): | |
| return self.cache[0][0].shape[2] if self.cache and self.cache[0] else 0 | |
| class UnifiedAttention(nn.Module): | |
| def __init__(self, config, layer_idx=0): | |
| super().__init__() | |
| self.layer_idx = layer_idx | |
| self.num_heads = config.num_heads | |
| self.third_dim = config.third_dim | |
| self.component_head_dim = config.component_head_dim | |
| self.W_unified = nn.Linear(config.embed_dim, config.third_dim * 3, bias=False) | |
| self.W_output = nn.Linear(config.third_dim, config.embed_dim, bias=False) | |
| self.rope = RotaryPositionEmbedding(config.component_head_dim, config.max_seq_len) | |
| def forward(self, x, kv_cache=None, use_cache=False): | |
| B, S, _ = x.shape | |
| unified = self.W_unified(x) | |
| seeking, offering, content = unified.split(self.third_dim, dim=-1) | |
| seeking = seeking.view(B, S, self.num_heads, self.component_head_dim).transpose(1, 2) | |
| offering = offering.view(B, S, self.num_heads, self.component_head_dim).transpose(1, 2) | |
| content = content.view(B, S, self.num_heads, self.component_head_dim).transpose(1, 2) | |
| offset = kv_cache.seq_len if kv_cache else 0 | |
| cos, sin = self.rope(seeking, S, offset) | |
| seeking, offering = apply_rotary_pos_emb(seeking, offering, cos, sin) | |
| if kv_cache: | |
| cached = kv_cache.get(self.layer_idx) | |
| if cached: | |
| offering = torch.cat([cached[0], offering], dim=2) | |
| content = torch.cat([cached[1], content], dim=2) | |
| if use_cache: | |
| new_k = unified[..., self.third_dim:2*self.third_dim].view(B, S, self.num_heads, self.component_head_dim).transpose(1, 2) | |
| new_k, _ = apply_rotary_pos_emb(new_k, new_k, cos, sin) | |
| new_v = unified[..., 2*self.third_dim:].view(B, S, self.num_heads, self.component_head_dim).transpose(1, 2) | |
| kv_cache.update(self.layer_idx, new_k, new_v) | |
| out = F.scaled_dot_product_attention(seeking, offering, content, is_causal=(not kv_cache or kv_cache.seq_len == 0)) | |
| return self.W_output(out.transpose(1, 2).contiguous().view(B, S, self.third_dim)) | |
| class FeedForward(nn.Module): | |
| def __init__(self, config): | |
| super().__init__() | |
| self.fc1 = nn.Linear(config.embed_dim, config.ffn_dim) | |
| self.fc2 = nn.Linear(config.ffn_dim, config.embed_dim) | |
| def forward(self, x): | |
| return self.fc2(F.gelu(self.fc1(x))) | |
| class TransformerBlock(nn.Module): | |
| def __init__(self, config, layer_idx=0): | |
| super().__init__() | |
| self.attn = UnifiedAttention(config, layer_idx) | |
| self.ffn = FeedForward(config) | |
| self.norm1 = nn.LayerNorm(config.embed_dim) | |
| self.norm2 = nn.LayerNorm(config.embed_dim) | |
| def forward(self, x, kv_cache=None, use_cache=False): | |
| x = x + self.attn(self.norm1(x), kv_cache, use_cache) | |
| x = x + self.ffn(self.norm2(x)) | |
| return x | |
| class Yocto(nn.Module): | |
| def __init__(self, config): | |
| super().__init__() | |
| self.config = config | |
| self.token_embedding = nn.Embedding(config.vocab_size, config.embed_dim) | |
| self.blocks = nn.ModuleList([TransformerBlock(config, i) for i in range(config.num_layers)]) | |
| self.norm = nn.LayerNorm(config.embed_dim) | |
| self.output = nn.Linear(config.embed_dim, config.vocab_size, bias=False) | |
| self.output.weight = self.token_embedding.weight | |
| def forward(self, input_ids, kv_cache=None, use_cache=False): | |
| x = self.token_embedding(input_ids) | |
| for block in self.blocks: | |
| x = block(x, kv_cache, use_cache) | |
| return self.output(self.norm(x)) | |
| # ============================================================================== | |
| # LOAD MODEL | |
| # ============================================================================== | |
| def load_model(): | |
| model_path = hf_hub_download(repo_id="Reinforce-ai/yocto", filename="model.pt") | |
| tokenizer_path = hf_hub_download(repo_id="Reinforce-ai/yocto", filename="tokenizer.json") | |
| checkpoint = torch.load(model_path, map_location="cpu", weights_only=False) | |
| cfg = checkpoint['config']['model'] | |
| config = ModelConfig( | |
| vocab_size=cfg['vocab_size'], | |
| embed_dim=cfg['embed_dim'], | |
| num_heads=cfg['num_heads'], | |
| num_layers=cfg['num_layers'], | |
| ffn_dim=cfg['ffn_dim'], | |
| max_seq_len=cfg.get('max_seq_len', 512), | |
| dropout=cfg.get('dropout', 0.1), | |
| ) | |
| model = Yocto(config) | |
| model.load_state_dict(checkpoint['model_state_dict']) | |
| model.eval() | |
| tokenizer = HFTokenizer.from_file(tokenizer_path) | |
| return model, tokenizer, config | |
| MODEL, TOKENIZER, CONFIG = load_model() | |
| # ============================================================================== | |
| # STREAMING GENERATION | |
| # ============================================================================== | |
| def generate_stream(prompt, max_tokens=150, temperature=0.8, top_k=50, top_p=0.95): | |
| input_ids = torch.tensor([TOKENIZER.encode(prompt).ids]) | |
| kv_cache = KVCache() | |
| kv_cache.reset(len(MODEL.blocks)) | |
| logits = MODEL(input_ids, kv_cache=kv_cache, use_cache=True) | |
| output_text = f"{prompt}" | |
| yield output_text | |
| token_count = 0 | |
| start_time = time.perf_counter() | |
| for _ in range(max_tokens): | |
| next_logits = logits[:, -1, :] / temperature | |
| if top_k > 0: | |
| v, _ = torch.topk(next_logits, min(top_k, next_logits.size(-1))) | |
| next_logits[next_logits < v[:, -1:]] = float('-inf') | |
| if top_p < 1.0: | |
| sorted_logits, sorted_idx = torch.sort(next_logits, descending=True) | |
| cumsum = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) | |
| mask = cumsum > top_p | |
| mask[..., 1:] = mask[..., :-1].clone() | |
| mask[..., 0] = False | |
| next_logits[mask.scatter(1, sorted_idx, mask)] = float('-inf') | |
| probs = F.softmax(next_logits, dim=-1) | |
| next_token = torch.multinomial(probs, 1) | |
| if next_token.item() == 3: | |
| break | |
| token_text = TOKENIZER.decode([next_token.item()]) | |
| output_text += token_text | |
| token_count += 1 | |
| yield output_text | |
| logits = MODEL(next_token, kv_cache=kv_cache, use_cache=True) | |
| elapsed = time.perf_counter() - start_time | |
| tokens_per_sec = token_count / elapsed if elapsed > 0 else 0 | |
| output_text += f"\n\n─────────────────────────────────────────\n⚡ {token_count} tokens • {elapsed:.2f}s • {tokens_per_sec:.1f} tok/s" | |
| yield output_text | |
| # ============================================================================== | |
| # GRADIO INTERFACE | |
| # ============================================================================== | |
| CSS = """ | |
| @import url('https://fonts.googleapis.com/css2?family=VT323&family=Space+Mono&display=swap'); | |
| /* Global */ | |
| .gradio-container { | |
| background: #000 !important; | |
| max-width: 100% !important; | |
| } | |
| footer {display: none !important;} | |
| /* All textareas and inputs */ | |
| textarea, input[type="text"] { | |
| background: #111 !important; | |
| color: #0f0 !important; | |
| border: 2px solid #0f0 !important; | |
| border-radius: 0 !important; | |
| font-family: 'VT323', monospace !important; | |
| font-size: 1.3rem !important; | |
| caret-color: #0f0 !important; | |
| } | |
| textarea:focus, input[type="text"]:focus { | |
| outline: none !important; | |
| box-shadow: 0 0 20px rgba(0,255,0,0.3) !important; | |
| } | |
| textarea::placeholder, input::placeholder { | |
| color: #060 !important; | |
| } | |
| /* Labels */ | |
| label { | |
| color: #0f0 !important; | |
| font-family: 'Space Mono', monospace !important; | |
| } | |
| /* Buttons */ | |
| button { | |
| font-family: 'Space Mono', monospace !important; | |
| } | |
| #generate-btn { | |
| background: #0f0 !important; | |
| color: #000 !important; | |
| border: none !important; | |
| border-radius: 0 !important; | |
| font-size: 1.1rem !important; | |
| font-weight: bold !important; | |
| padding: 12px 30px !important; | |
| text-transform: uppercase !important; | |
| } | |
| #generate-btn:hover { | |
| background: #0c0 !important; | |
| box-shadow: 0 0 25px rgba(0,255,0,0.5) !important; | |
| } | |
| .chip-btn { | |
| background: transparent !important; | |
| color: #0f0 !important; | |
| border: 1px solid #080 !important; | |
| border-radius: 20px !important; | |
| font-size: 0.9rem !important; | |
| padding: 8px 18px !important; | |
| margin: 4px !important; | |
| } | |
| .chip-btn:hover { | |
| background: rgba(0,255,0,0.1) !important; | |
| border-color: #0f0 !important; | |
| box-shadow: 0 0 15px rgba(0,255,0,0.3) !important; | |
| } | |
| """ | |
| def generate_story(prompt): | |
| if not prompt.strip(): | |
| prompt = "Once upon a time" | |
| for output in generate_stream(prompt, max_tokens=150, temperature=0.8, top_k=50, top_p=0.95): | |
| yield output | |
| with gr.Blocks(title="YOCTO") as demo: | |
| # Header | |
| gr.HTML(""" | |
| <div style="text-align:center; padding:30px 0 20px 0;"> | |
| <h1 style="font-family:'VT323',monospace; font-size:5rem; color:#0f0; | |
| text-shadow:0 0 20px #0f0, 0 0 40px #0f0; letter-spacing:0.4em; margin:0;"> | |
| YOCTO | |
| </h1> | |
| <p style="font-family:'Space Mono',monospace; color:#0a0; font-size:0.9rem; letter-spacing:0.15em; margin-top:10px;"> | |
| WORLD'S SMALLEST STORY TELLER • 484K PARAMS • 946 KB | |
| </p> | |
| </div> | |
| """) | |
| # Output | |
| output = gr.Textbox( | |
| label="Output", | |
| value="Ready to generate stories...\n\nType a prompt below or click a suggestion.", | |
| lines=14, | |
| interactive=False | |
| ) | |
| # Input + Button | |
| with gr.Row(): | |
| prompt = gr.Textbox( | |
| label="Your prompt", | |
| placeholder="Start your story here...", | |
| lines=1, | |
| scale=5 | |
| ) | |
| generate_btn = gr.Button("GENERATE", scale=1, elem_id="generate-btn") | |
| # Prompt chips | |
| gr.HTML("<div style='text-align:center; margin:15px 0 5px 0; color:#0a0; font-family:Space Mono,monospace; font-size:0.8rem;'>TRY THESE:</div>") | |
| with gr.Row(): | |
| btn1 = gr.Button("Once upon a time", elem_classes="chip-btn") | |
| btn2 = gr.Button("There was a little girl named Lily", elem_classes="chip-btn") | |
| btn3 = gr.Button("One day, a curious boy found", elem_classes="chip-btn") | |
| btn4 = gr.Button("The little rabbit was very", elem_classes="chip-btn") | |
| btn5 = gr.Button("Mommy and Timmy went to the", elem_classes="chip-btn") | |
| btn6 = gr.Button("There was a big red", elem_classes="chip-btn") | |
| # Footer | |
| gr.HTML(""" | |
| <div style="text-align:center; padding:20px 0; margin-top:20px; border-top:1px solid #030;"> | |
| <span style="font-family:'Space Mono',monospace; font-size:0.75rem; color:#080;"> | |
| REINFORCE AI © 2025 • | |
| <a href="https://www.reinforceai.com/yocto" style="color:#0f0; text-decoration:none;">WEBSITE</a> • | |
| <a href="https://github.com/reinforceai/yocto" style="color:#0f0; text-decoration:none;">GITHUB</a> • | |
| <a href="https://huggingface.co/Reinforce-ai/yocto" style="color:#0f0; text-decoration:none;">MODEL</a> | |
| </span> | |
| </div> | |
| """) | |
| # Events | |
| btn1.click(lambda: "Once upon a time", outputs=prompt) | |
| btn2.click(lambda: "There was a little girl named Lily", outputs=prompt) | |
| btn3.click(lambda: "One day, a curious boy found", outputs=prompt) | |
| btn4.click(lambda: "The little rabbit was very", outputs=prompt) | |
| btn5.click(lambda: "Mommy and Timmy went to the", outputs=prompt) | |
| btn6.click(lambda: "There was a big red", outputs=prompt) | |
| generate_btn.click(generate_story, inputs=[prompt], outputs=[output]) | |
| prompt.submit(generate_story, inputs=[prompt], outputs=[output]) | |
| if __name__ == "__main__": | |
| demo.launch(server_name="0.0.0.0", server_port=7860, css=CSS) |