Upload 8 files

README.md

@@ -1,13 +1,41 @@
 ---
-title: Wvy 3m
-emoji: 🦀
-colorFrom: yellow
-colorTo: green
+title: WVY Tiny Liquid LM
+emoji: 🌊
+colorFrom: gray
+colorTo: indigo
 sdk: gradio
 sdk_version: 6.26.0
-python_version: '3.12'
+python_version: "3.12"
 app_file: app.py
-pinned: false
+suggested_hardware: cpu-basic
+short_description: Chat with a 3.31M-parameter custom liquid-style LM.
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# WVY Tiny Liquid LM
+
+A Hugging Face Space for the trained `TinyLiquidCausalLanguageModel` checkpoint in this repository.
+
+## Model
+
+- Parameters: 3,314,880
+- Vocabulary: 8,000 BPE tokens
+- Width: 192
+- Blocks: 4
+- Feed-forward hidden width: 512
+- Causal depthwise convolution kernel: 5
+- Context used by the chat app: 256 tokens
+- Architecture: recurrent liquid-style state mixer + SwiGLU feed-forward layers
+- Final fine-tuning loss: 0.1344
+- Final fine-tuning perplexity: 1.1438
+
+The Space loads the custom PyTorch architecture directly from `model.py`, restores `tiny_liquid_causal_lm.pt`, and uses the bundled `tokenizer.json`.
+
+## Files
+
+- `app.py` — Gradio chat UI and autoregressive generation
+- `model.py` — exact custom model architecture used for the checkpoint
+- `tiny_liquid_causal_lm.pt` — final fine-tuned checkpoint
+- `tokenizer.json` — BPE tokenizer
+- `model_config.json` — architecture configuration
+- `training_summary.json` — training statistics
+- `requirements.txt` — runtime dependencies

app.py

@@ -0,0 +1,149 @@
+from pathlib import Path
+import json
+
+import gradio as gr
+import torch
+from tokenizers import Tokenizer
+
+from model import LiquidModelConfig, TinyLiquidCausalLanguageModel
+
+ROOT = Path(__file__).resolve().parent
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+SEQUENCE_LENGTH = 256
+
+PAD_TOKEN = "<|pad|>"
+BOS_TOKEN = "<|bos|>"
+EOS_TOKEN = "<|eos|>"
+USER_TOKEN = "<|user|>"
+ASSISTANT_TOKEN = "<|assistant|>"
+
+with open(ROOT / "model_config.json", "r", encoding="utf-8") as f:
+    CONFIG_DATA = json.load(f)
+
+config = LiquidModelConfig(**CONFIG_DATA)
+tokenizer = Tokenizer.from_file(str(ROOT / "tokenizer.json"))
+PAD_ID = tokenizer.token_to_id(PAD_TOKEN)
+BOS_ID = tokenizer.token_to_id(BOS_TOKEN)
+EOS_ID = tokenizer.token_to_id(EOS_TOKEN)
+USER_ID = tokenizer.token_to_id(USER_TOKEN)
+
+checkpoint = torch.load(ROOT / "tiny_liquid_causal_lm.pt", map_location="cpu", weights_only=False)
+model = TinyLiquidCausalLanguageModel(config)
+model.load_state_dict(checkpoint["model_state_dict"], strict=True)
+model.eval().to(DEVICE)
+
+if DEVICE.type == "cpu":
+    torch.set_num_threads(max(1, min(4, torch.get_num_threads())))
+
+
+def _history_to_lines(history):
+    lines = []
+    for item in history or []:
+        if isinstance(item, dict):
+            role = str(item.get("role", "")).lower()
+            content = item.get("content", "")
+            if isinstance(content, str) and content.strip() and role in {"user", "assistant"}:
+                lines.append(f"<|{role}|>\n{content.strip()}")
+        elif isinstance(item, (list, tuple)) and len(item) == 2:
+            user_message, assistant_message = item
+            if user_message:
+                lines.append(f"<|user|>\n{str(user_message).strip()}")
+            if assistant_message:
+                lines.append(f"<|assistant|>\n{str(assistant_message).strip()}")
+    return lines[-8:]
+
+
+@torch.inference_mode()
+def generate(message, history, temperature, top_k, max_new_tokens):
+    message = str(message).strip()
+    if not message:
+        return ""
+
+    transcript = _history_to_lines(history)
+    transcript.extend([f"<|user|>\n{message}", "<|assistant|>"])
+    prompt = "\n".join(transcript) + "\n"
+
+    generated_ids = [BOS_ID] + tokenizer.encode(prompt).ids
+    response_ids = []
+    stop_ids = {EOS_ID, USER_ID}
+
+    for _ in range(int(max_new_tokens)):
+        model_input = torch.tensor(
+            [generated_ids[-SEQUENCE_LENGTH:]],
+            dtype=torch.long,
+            device=DEVICE,
+        )
+        logits = model(model_input)["logits"][0, -1].float()
+        logits[PAD_ID] = -float("inf")
+
+        if float(temperature) <= 0:
+            next_token = int(torch.argmax(logits).item())
+        else:
+            logits = logits / float(temperature)
+            k = max(1, min(int(top_k), logits.numel()))
+            top_values, top_indices = torch.topk(logits, k=k)
+            probabilities = torch.softmax(top_values, dim=-1)
+            sampled_position = torch.multinomial(probabilities, num_samples=1)
+            next_token = int(top_indices[sampled_position].item())
+
+        if next_token in stop_ids:
+            break
+
+        generated_ids.append(next_token)
+        response_ids.append(next_token)
+
+    return tokenizer.decode(response_ids, skip_special_tokens=True).strip()
+
+
+CSS = """
+.gradio-container { max-width: 920px !important; margin: 0 auto !important; }
+footer { display: none !important; }
+"""
+
+with gr.Blocks(css=CSS, title="WVY Tiny Liquid LM") as demo:
+    gr.Markdown(
+        "# WVY Tiny Liquid LM\n"
+        "3.31M-parameter custom causal language model · 4 liquid-style state-mixer blocks · 256-token context"
+    )
+
+    chatbot = gr.Chatbot(height=560)
+
+    with gr.Row():
+        msg = gr.Textbox(
+            placeholder="Message WVY...",
+            show_label=False,
+            scale=8,
+            autofocus=True,
+        )
+        send = gr.Button("Send", variant="primary", scale=1)
+
+    with gr.Accordion("Generation settings", open=False):
+        temperature = gr.Slider(0.0, 1.5, value=0.8, step=0.05, label="Temperature")
+        top_k = gr.Slider(1, 100, value=40, step=1, label="Top-k")
+        max_new_tokens = gr.Slider(8, 128, value=96, step=8, label="Max new tokens")
+
+    clear = gr.Button("Clear chat")
+
+    def respond(message, history, temperature, top_k, max_new_tokens):
+        history = history or []
+        reply = generate(message, history, temperature, top_k, max_new_tokens)
+        history = history + [
+            {"role": "user", "content": message},
+            {"role": "assistant", "content": reply},
+        ]
+        return "", history
+
+    send.click(
+        respond,
+        [msg, chatbot, temperature, top_k, max_new_tokens],
+        [msg, chatbot],
+    )
+    msg.submit(
+        respond,
+        [msg, chatbot, temperature, top_k, max_new_tokens],
+        [msg, chatbot],
+    )
+    clear.click(lambda: ("", []), outputs=[msg, chatbot])
+
+if __name__ == "__main__":
+    demo.queue(default_concurrency_limit=1).launch()

model.py

@@ -0,0 +1,149 @@
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dimension: int, epsilon: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dimension))
+        self.epsilon = epsilon
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        normalized = hidden_states * torch.rsqrt(
+            hidden_states.pow(2).mean(dim=-1, keepdim=True) + self.epsilon
+        )
+        return self.weight * normalized
+
+
+class LiquidStateMixer(nn.Module):
+    def __init__(self, dimension: int, kernel_size: int, dropout: float):
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.input_norm = RMSNorm(dimension)
+        self.causal_depthwise_convolution = nn.Conv1d(
+            dimension,
+            dimension,
+            kernel_size,
+            groups=dimension,
+            bias=True,
+        )
+        self.state_parameters = nn.Linear(dimension, 3 * dimension)
+        self.base_decay_logits = nn.Parameter(torch.zeros(dimension))
+        self.output_projection = nn.Linear(dimension, dimension, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        normalized = self.input_norm(hidden_states)
+        convolution_input = normalized.transpose(1, 2)
+        convolution_input = F.pad(convolution_input, (self.kernel_size - 1, 0))
+        local_features = self.causal_depthwise_convolution(convolution_input).transpose(1, 2)
+
+        candidate, decay_logits, output_gate = self.state_parameters(local_features).chunk(3, dim=-1)
+        candidate = torch.tanh(candidate)
+        decay = torch.sigmoid(decay_logits + self.base_decay_logits)
+        output_gate = torch.sigmoid(output_gate)
+
+        state = torch.zeros_like(candidate[:, 0])
+        mixed_steps = []
+        for step in range(candidate.size(1)):
+            step_decay = decay[:, step]
+            state = step_decay * state + (1.0 - step_decay) * candidate[:, step]
+            mixed_steps.append(output_gate[:, step] * state)
+
+        mixed = torch.stack(mixed_steps, dim=1)
+        return hidden_states + self.dropout(self.output_projection(mixed))
+
+
+class SwiGLUFeedForward(nn.Module):
+    def __init__(self, dimension: int, hidden_dimension: int, dropout: float):
+        super().__init__()
+        self.input_norm = RMSNorm(dimension)
+        self.gate_projection = nn.Linear(dimension, hidden_dimension, bias=False)
+        self.value_projection = nn.Linear(dimension, hidden_dimension, bias=False)
+        self.output_projection = nn.Linear(hidden_dimension, dimension, bias=False)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        normalized = self.input_norm(hidden_states)
+        activated = F.silu(self.gate_projection(normalized)) * self.value_projection(normalized)
+        return hidden_states + self.dropout(self.output_projection(activated))
+
+
+class LiquidBlock(nn.Module):
+    def __init__(
+        self,
+        dimension: int,
+        hidden_dimension: int,
+        kernel_size: int,
+        dropout: float,
+    ):
+        super().__init__()
+        self.state_mixer = LiquidStateMixer(dimension, kernel_size, dropout)
+        self.feed_forward = SwiGLUFeedForward(dimension, hidden_dimension, dropout)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.feed_forward(self.state_mixer(hidden_states))
+
+
+@dataclass
+class LiquidModelConfig:
+    vocab_size: int
+    dimension: int = 192
+    layer_count: int = 4
+    feed_forward_hidden_dimension: int = 512
+    convolution_kernel_size: int = 5
+    dropout: float = 0.05
+    pad_token_id: int = 0
+    bos_token_id: int = 2
+    eos_token_id: int = 3
+
+
+class TinyLiquidCausalLanguageModel(nn.Module):
+    def __init__(self, config: LiquidModelConfig):
+        super().__init__()
+        self.config = config
+        self.token_embedding = nn.Embedding(
+            config.vocab_size,
+            config.dimension,
+            padding_idx=config.pad_token_id,
+        )
+        self.blocks = nn.ModuleList(
+            [
+                LiquidBlock(
+                    config.dimension,
+                    config.feed_forward_hidden_dimension,
+                    config.convolution_kernel_size,
+                    config.dropout,
+                )
+                for _ in range(config.layer_count)
+            ]
+        )
+        self.final_norm = RMSNorm(config.dimension)
+        self.language_model_head = nn.Linear(config.dimension, config.vocab_size, bias=False)
+        self.language_model_head.weight = self.token_embedding.weight
+        self.apply(self._initialize_weights)
+
+    @staticmethod
+    def _initialize_weights(module: nn.Module) -> None:
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+        if isinstance(module, nn.Linear) and module.bias is not None:
+            nn.init.zeros_(module.bias)
+
+    def forward(self, input_ids: torch.Tensor, labels=None):
+        hidden_states = self.token_embedding(input_ids)
+        for block in self.blocks:
+            hidden_states = block(hidden_states)
+        logits = self.language_model_head(self.final_norm(hidden_states))
+
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(
+                logits.reshape(-1, logits.size(-1)),
+                labels.reshape(-1),
+                ignore_index=-100,
+            )
+        return {"loss": loss, "logits": logits}

model_config.json

@@ -0,0 +1,11 @@
+{
+  "vocab_size": 8000,
+  "dimension": 192,
+  "layer_count": 4,
+  "feed_forward_hidden_dimension": 512,
+  "convolution_kernel_size": 5,
+  "dropout": 0.05,
+  "pad_token_id": 0,
+  "bos_token_id": 2,
+  "eos_token_id": 3
+}

requirements.txt

@@ -0,0 +1,2 @@
+torch>=2.2
+tokenizers>=0.20

tiny_liquid_causal_lm.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b3dbbc1dd613e14bb0ffb195dd0625f39f1d51f3c31cc805dfe7a1b3b7461800
+size 13279025

training_summary.json

@@ -0,0 +1,31 @@
+{
+  "dataset": {
+    "files": 17,
+    "pretraining_documents": 28585,
+    "pretraining_tokens": 2856322,
+    "sft_examples": 24536
+  },
+  "parameters": {
+    "total": 3314880,
+    "trainable": 3314880
+  },
+  "pretraining_history": [
+    {
+      "epoch": 1,
+      "loss": 1.9338438765763382,
+      "perplexity": 6.916043631489244
+    }
+  ],
+  "fine_tuning_history": [
+    {
+      "epoch": 1,
+      "loss": 0.23482260895200027,
+      "perplexity": 1.2646844051414168
+    },
+    {
+      "epoch": 2,
+      "loss": 0.134362090434873,
+      "perplexity": 1.143806906211811
+    }
+  ]
+}