Upload folder using huggingface_hub

README.md

@@ -1,3 +1,20 @@
----
-license: mit
----
+
+This is a refactor of [arman-bd/guppylm-9M](https://huggingface.co/arman-bd/guppylm-9M) to be compliant with [transformers's custom_model](https://huggingface.co/docs/transformers/custom_models).
+
+
+```bash
+python inference.py  guppylm-9M
+
+GuppyLMForCausalLM loaded: 8.7M params
+
+Guppy Chat (type 'quit' to exit)
+
+You> is there a cat in the room?
+Guppy> i don't like it. it puts its face on the glass by the bubbles.
+
+You> I'm sorry. are you hungry?
+Guppy> i don't eat it.
+
+```
+
+This is a good small LLM for doing execise with SGLang.

config.json

@@ -0,0 +1,23 @@
+{
+  "model_type": "guppylm",
+  "architectures": [
+    "GuppyLMForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_guppylm.GuppyLMConfig",
+    "AutoModel": "modeling_guppylm.GuppyLMModel",
+    "AutoModelForCausalLM": "modeling_guppylm.GuppyLMForCausalLM"
+  },
+  "vocab_size": 4096,
+  "max_position_embeddings": 128,
+  "hidden_size": 384,
+  "num_hidden_layers": 6,
+  "num_attention_heads": 6,
+  "num_key_value_heads": 6,
+  "intermediate_size": 768,
+  "hidden_dropout_prob": 0.1,
+  "pad_token_id": 0,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "tokenizer_class": "PreTrainedTokenizerFast"
+}

configuration_guppylm.py

@@ -0,0 +1,90 @@
+"""GuppyLM Hugging Face configuration (PreTrainedConfig)."""
+
+from typing import List, Optional
+
+from transformers import PretrainedConfig
+
+
+class GuppyLMConfig(PretrainedConfig):
+    """Configuration for GuppyLM.
+
+    Uses Transformers-standard field names in JSON; modeling code may also read
+    aliases via properties (``d_model``, ``n_heads``, etc.).
+    """
+
+    model_type = "guppylm"
+
+    def __init__(
+        self,
+        vocab_size: int = 4096,
+        max_position_embeddings: int = 128,
+        hidden_size: int = 384,
+        num_hidden_layers: int = 6,
+        num_attention_heads: int = 6,
+        intermediate_size: int = 768,
+        hidden_dropout_prob: float = 0.1,
+        pad_token_id: int = 0,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+        architectures: Optional[List[str]] = None,
+        tokenizer_class: Optional[str] = "PreTrainedTokenizerFast",
+        attn_implementation: Optional[str] = "eager",
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.hidden_dropout_prob = hidden_dropout_prob
+        if architectures is None:
+            architectures = ["GuppyLMForCausalLM"]
+        self.architectures = architectures
+        super().__init__(
+            architectures=architectures,
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tokenizer_class=tokenizer_class,
+            tie_word_embeddings=True,
+            **kwargs,
+        )
+
+    # --- Aliases for original GuppyLM / training code naming ---
+
+    @property
+    def max_seq_len(self) -> int:
+        return self.max_position_embeddings
+
+    @property
+    def d_model(self) -> int:
+        return self.hidden_size
+
+    @property
+    def n_layers(self) -> int:
+        return self.num_hidden_layers
+
+    @property
+    def n_heads(self) -> int:
+        return self.num_attention_heads
+
+    @property
+    def ffn_hidden(self) -> int:
+        return self.intermediate_size
+
+    @property
+    def dropout(self) -> float:
+        return self.hidden_dropout_prob
+
+    @property
+    def pad_id(self) -> int:
+        return self.pad_token_id
+
+    @property
+    def bos_id(self) -> int:
+        return self.bos_token_id
+
+    @property
+    def eos_id(self) -> int:
+        return self.eos_token_id

decode_input_ids.py

@@ -0,0 +1,17 @@
+
+import sys
+import torch
+from transformers import AutoTokenizer
+
+input_ids=[32, 720, 103, 432, 146, 167, 1400, 195, 294, 1083, 111]
+input_ids=[1,  88,  64, 779,   2,  64,   1,  89,  64]
+model = "/sgl-workspace/guppylm-9M"
+tokenizer = AutoTokenizer.from_pretrained(
+    model,
+    trust_remote_code=True,
+)
+text = tokenizer.decode(
+    torch.tensor(input_ids, dtype=torch.long),
+    skip_special_tokens=False,
+)
+sys.stdout.write(text)

inference.py

@@ -0,0 +1,124 @@
+"""GuppyLM inference — simple chat."""
+
+import os
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizerFast
+
+
+def _resolve_pretrained_dir(model_dir: str) -> str:
+    """Directory with Hugging Face layout: ``config.json`` + ``pytorch_model.bin`` / ``model.safetensors``.
+
+    Accepts either that directory or the path to a weight file inside it (e.g. ``pytorch_model.bin``).
+    """
+    path = os.path.abspath(model_dir)
+    return path if os.path.isdir(path) else os.path.dirname(path)
+
+
+def _resolve_tokenizer_dir(tokenizer_path: str) -> tuple[str, str]:
+    """Return ``(directory, tokenizer_filename)`` for HF tokenizer loading."""
+    path = os.path.abspath(tokenizer_path)
+    if os.path.isdir(path):
+        return path, "tokenizer.json"
+    return os.path.dirname(path), os.path.basename(path)
+
+
+def _load_autotokenizer(tokenizer_path: str, *, trust_remote_code: bool = True):
+    """Load ``AutoTokenizer`` when ``config.json`` is present (declares ``tokenizer_class``); otherwise fast tokenizer only."""
+    tok_dir, tokenizer_file = _resolve_tokenizer_dir(tokenizer_path)
+    if os.path.isfile(os.path.join(tok_dir, "config.json")):
+        return AutoTokenizer.from_pretrained(tok_dir, trust_remote_code=trust_remote_code)
+    return PreTrainedTokenizerFast.from_pretrained(tok_dir, tokenizer_file=tokenizer_file)
+
+
+class GuppyInference:
+    def __init__(self, model_dir, device="cpu", trust_remote_code=True):
+        self.device = torch.device(device)
+        pretrained_dir = _resolve_pretrained_dir(model_dir)
+        self.tokenizer = _load_autotokenizer(pretrained_dir, trust_remote_code=trust_remote_code)
+
+        self.model = AutoModelForCausalLM.from_pretrained(
+            pretrained_dir,
+            trust_remote_code=trust_remote_code,
+        )
+        self.model.to(self.device)
+        self.model.eval()
+        self.config = self.model.config
+
+        total, _ = self.model.param_count()
+        print(f"{self.model.__class__.__name__} loaded: {total/1e6:.1f}M params")
+
+    def chat_completion(self, messages, temperature=0.001, max_tokens=64,
+                        top_k=0, **kwargs):
+        """Chat completion — takes messages, returns response."""
+        prompt = self._render_chat_prompt(messages)
+        input_ids = self.tokenizer.encode(prompt, add_special_tokens=False)
+        prompt_tokens = len(input_ids)
+        input_t = torch.tensor([input_ids], dtype=torch.long, device=self.device)
+
+        output_t, _ = self.model.generate_simple(
+            input_t, max_new_tokens=max_tokens, temperature=temperature, top_k=top_k
+        )
+        output_text = self.tokenizer.decode(
+            output_t[0].tolist()[prompt_tokens:], skip_special_tokens=False
+        )
+        # Truncate at first <|im_end|> — don't let the model leak into the next turn
+        if "<|im_end|>" in output_text:
+            output_text = output_text.split("<|im_end|>")[0]
+        # Also strip any <|im_start|> fragments
+        if "<|im_start|>" in output_text:
+            output_text = output_text.split("<|im_start|>")[0]
+        resp_text = output_text.strip()
+
+        return {
+            "choices": [{
+                "message": {"role": "assistant", "content": resp_text},
+            }],
+        }
+
+    def _render_chat_prompt(self, messages):
+        """Prefer tokenizer ``chat_template`` (HF ``apply_chat_template``); fall back if absent."""
+        if getattr(self.tokenizer, "chat_template", None):
+            return self.tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True,
+            )
+        parts = []
+        for msg in messages:
+            role = msg.get("role", "user")
+            content = msg.get("content") or ""
+            if role == "system":
+                continue
+            parts.append(f"<|im_start|>{role}\n{content}<|im_end|>")
+        parts.append("<|im_start|>assistant\n")
+        return "\n".join(parts)
+
+
+def main():
+    import argparse
+    p = argparse.ArgumentParser(description="Chat with Guppy")
+    p.add_argument(
+        "model_dir",
+        help="Directory with config.json, tokenizer.json, and weights (or path to a weight file in that directory)",
+    )
+    p.add_argument("--device", default="cpu")
+    args = p.parse_args()
+
+    engine = GuppyInference(args.model_dir, args.device)
+    print("\nGuppy Chat (type 'quit' to exit)")
+    msgs = []
+    while True:
+        inp = input("\nYou> ").strip()
+        if inp.lower() in ("quit", "exit", "q"):
+            break
+        msgs.append({"role": "user", "content": inp})
+        result = engine.chat_completion(msgs)
+        msg = result["choices"][0]["message"]
+        if msg.get("content"):
+            print(f"Guppy> {msg['content']}")
+        msgs.append(msg)
+
+
+if __name__ == "__main__":
+    main()

modeling_guppylm.py

@@ -0,0 +1,268 @@
+"""
+GuppyLM — Hugging Face PreTrainedModel wrapper.
+
+Vanilla transformer: multi-head attention, ReLU FFN, LayerNorm, learned positional embeddings.
+"""
+
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.generation.utils import GenerationMixin
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+import time
+
+try:
+    from .configuration_guppylm import GuppyLMConfig
+except ImportError:
+    from configuration_guppylm import GuppyLMConfig
+
+class Attention(nn.Module):
+    def __init__(self, config: GuppyLMConfig, layer_id: int):
+        super().__init__()
+        self.n_heads = config.n_heads
+        self.head_dim = config.d_model // config.n_heads
+        self.layer_id = layer_id
+        self.qkv = nn.Linear(config.d_model, 3 * config.d_model)
+        self.out = nn.Linear(config.d_model, config.d_model)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x, mask=None, forward_batch=None, attention_instances=None):
+        B, T, C = x.shape
+
+        if attention_instances is not None:
+            qkv = self.qkv(x).reshape(B, T, 3, self.n_heads*self.head_dim).permute(2, 0, 1, 3)
+            q, k, v = qkv[0], qkv[1], qkv[2]
+            attn_out = attention_instances[self.layer_id].forward(q, k, v, forward_batch)
+            return self.out(attn_out).contiguous().view(B, T, C)
+        else:
+            qkv = self.qkv(x).reshape(B, T, 3, self.n_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+            q, k, v = qkv[0], qkv[1], qkv[2]
+            attn = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+            if mask is not None:
+                attn = attn.masked_fill(mask == 0, float("-inf"))
+            attn = self.dropout(F.softmax(attn, dim=-1))
+
+            return self.out((attn @ v).transpose(1, 2).contiguous().view(B, T, C))
+
+
+class FFN(nn.Module):
+    def __init__(self, config: GuppyLMConfig):
+        super().__init__()
+        self.up = nn.Linear(config.d_model, config.ffn_hidden)
+        self.down = nn.Linear(config.ffn_hidden, config.d_model)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x):
+        return self.dropout(self.down(F.relu(self.up(x))))
+
+
+class Block(nn.Module):
+    def __init__(self, config: GuppyLMConfig, layer_id: int):
+        super().__init__()
+        self.layer_id = layer_id
+        self.norm1 = nn.LayerNorm(config.d_model)
+        self.attn = Attention(config, layer_id)
+        self.norm2 = nn.LayerNorm(config.d_model)
+        self.ffn = FFN(config)
+
+    def forward(self, x, mask=None, forward_batch=None, attention_instances=None):
+        x = x + self.attn(self.norm1(x), mask, forward_batch, attention_instances)
+        x = x + self.ffn(self.norm2(x))
+        return x
+
+
+class GuppyLMModel(PreTrainedModel):
+    """Backbone only (hidden states). Used as ``AutoModel`` for SGLang / Transformers tooling."""
+
+    config_class = GuppyLMConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+    _supports_attention_backend = True
+
+    def __init__(self, config: GuppyLMConfig):
+        super().__init__(config)
+
+        # assert 0
+        self.tok_emb = nn.Embedding(config.vocab_size, config.d_model)
+        self.pos_emb = nn.Embedding(config.max_seq_len, config.d_model)
+        self.drop = nn.Dropout(config.dropout)
+        self.blocks = nn.ModuleList([Block(config, i) for i in range(config.n_layers)])
+        self.norm = nn.LayerNorm(config.d_model)
+        self.apply(self._init_weights)
+        self.post_init()
+        self.input_ids = None
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            nn.init.normal_(m.weight, mean=0.0, std=0.02)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.Embedding):
+            nn.init.normal_(m.weight, mean=0.0, std=0.02)
+
+    def get_input_embeddings(self):
+        return self.tok_emb
+
+    def set_input_embeddings(self, value: nn.Module):
+        self.tok_emb = value
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        forward_batch=None,
+        attention_instances=None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        if input_ids is None:
+            raise ValueError("You must specify `input_ids`.")
+
+        if attention_instances is None:
+            B, T = input_ids.shape
+            pos = torch.arange(T, device=input_ids.device)
+            x = self.drop(self.tok_emb(input_ids) + self.pos_emb(pos))
+            mask = torch.tril(torch.ones(T, T, device=input_ids.device)).unsqueeze(0).unsqueeze(0)
+        else:
+            pos = forward_batch.positions
+            x = self.drop(self.tok_emb(input_ids) + self.pos_emb(pos))
+            mask = None
+
+        for block in self.blocks:
+            x = block(x, mask, forward_batch, attention_instances)
+
+        hidden_states = self.norm(x)
+
+        if not return_dict:
+            return (hidden_states,)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=None,
+            hidden_states=None,
+            attentions=None,
+        )
+
+
+class GuppyLMForCausalLM(PreTrainedModel, GenerationMixin):
+    """Causal LM. Checkpoints may use flat keys (legacy) or nested ``model.*`` keys."""
+
+    config_class = GuppyLMConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = False
+    _supports_attention_backend = True
+
+    def __init__(self, config: GuppyLMConfig):
+        super().__init__(config)
+
+        self.model = GuppyLMModel(config)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        nn.init.normal_(self.lm_head.weight, mean=0.0, std=0.02)
+        self.post_init()
+        self.lm_head.weight = self.model.tok_emb.weight
+
+    def get_input_embeddings(self):
+        return self.model.tok_emb
+
+    def set_input_embeddings(self, value: nn.Module):
+        self.model.tok_emb = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings: nn.Module):
+        self.lm_head = new_embeddings
+
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, attention_mask=None, **kwargs
+    ):
+        return {"input_ids": input_ids}
+
+    def load_state_dict(
+        self,
+        state_dict: dict,
+        strict: bool = True,
+        assign: bool = False,
+    ):
+        """Accept legacy flat checkpoints (``tok_emb.*``, ``blocks.*``, …) as well as nested ``model.*``."""
+        keys = list(state_dict.keys())
+        if keys and not any(k.startswith("model.") for k in keys):
+            remapped = {}
+            for k, v in state_dict.items():
+                if k.startswith("lm_head."):
+                    remapped[k] = v
+                else:
+                    remapped[f"model.{k}"] = v
+            state_dict = remapped
+        return super().load_state_dict(state_dict, strict=strict, assign=assign)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        if input_ids is None:
+            raise ValueError("You must specify `input_ids`.")
+
+        outputs = self.model(input_ids=input_ids, return_dict=True, **kwargs)
+        hidden_states = outputs.last_hidden_state
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(
+                logits.view(-1, self.config.vocab_size),
+                labels.view(-1),
+                ignore_index=self.config.pad_token_id,
+            )
+
+        if not return_dict:
+            return (loss, logits) if loss is not None else (logits,)
+
+        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=None)
+
+    @torch.no_grad()
+    def generate_simple(
+        self,
+        idx: torch.LongTensor,
+        max_new_tokens: int = 64,
+        temperature: float = 0.7,
+        top_k: int = 50,
+        **kwargs,
+    ):
+        """Original sampling loop (non-HF ``generate`` API)."""
+        self.eval()
+        for _ in range(max_new_tokens):
+            idx_cond = idx[:, -self.config.max_seq_len :]
+            out = self.forward(input_ids=idx_cond, return_dict=True)
+            logits = out.logits[:, -1, :] / temperature
+            if top_k > 0:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = float("-inf")
+            probs = F.softmax(logits, dim=-1)
+            next_id = torch.multinomial(probs, num_samples=1)
+            idx = torch.cat([idx, next_id], dim=1)
+            if next_id.item() == self.config.eos_token_id:
+                break
+        return idx, []
+
+    def param_count(self):
+        total = sum(p.numel() for p in self.parameters())
+        return total, 0
+
+    def param_summary(self):
+        total, _ = self.param_count()
+        return f"GuppyLM: {total:,} params ({total / 1e6:.1f}M)"
+
+
+__all__ = ["GuppyLMModel", "GuppyLMForCausalLM", "GuppyLMConfig"]

pytorch_model.bin

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

test.py

@@ -0,0 +1,106 @@
+import aiohttp
+import asyncio
+import os
+import time
+from datetime import datetime
+import random
+
+def _base_url() -> str:
+    host = os.environ.get("SGLANG_HTTP_HOST", "127.0.0.1")
+    port = os.environ.get("SGLANG_HTTP_PORT", "30000")
+    return f"http://{host}:{port}"
+
+
+database = [
+    ["who are you?",
+     "i'm just a little fish. i have little body. they work. i eat the red flakes when i can.",
+     "hi",
+     "i'm just a snail would be nice.",
+     "what is the capital of France?",
+     "i am a bubble wall.",
+    ],
+    ["is there a cat in the room?",
+     "i don't like it. it puts its face on the glass by the bubbles.",
+     "I'm sorry. are you hungry?",
+     "i don't eat it.",
+    ]
+]
+
+async def main(num_tests: int, concurrency: int):
+    url = f"{_base_url()}/v1/chat/completions"
+    headers = {"Content-Type": "application/json"}
+    RED = "\033[0;31m"
+    GREEN = "\033[0;32m"
+    YELLOW = "\033[0;33m"
+    BLUE = "\033[0;34m"
+    MAGENTA = "\033[0;35m"
+    CYAN = "\033[0;36m"
+    WHITE = "\033[0;37m"
+    END = "\033[0m"
+    timeout = aiohttp.ClientTimeout(
+        total=600,      # Total timeout for entire request
+        connect=10,    # Connection timeout
+        sock_read=30   # Socket read timeout
+    )
+    async with aiohttp.ClientSession(timeout=timeout) as session:
+
+        if num_tests == 0:
+            # realtime chat 
+            json = {"model": "guppylm-9M",
+                    "messages": [],
+                    "temperature" : 0.001,
+                    "top_k": 1,
+                    "stream": False,
+                    }
+            while True:
+                user_input = input("user: ")
+                json["messages"].append({"role": "user", "content": user_input})
+                async with session.post(url, json=json, headers=headers) as response:
+                    res = (await response.json())["choices"][0]["message"]
+                    output = res["content"]
+                    role = res["role"]
+                    json["messages"].append({"role": role, "content": output})
+                    print(f"{role}: {output}")
+        else:
+            concurrency_sem = asyncio.Semaphore(max(1,concurrency))
+            async def job():
+                json = {"model": "guppylm-9M",
+                        "messages": [],
+                        "temperature" : 0.001,
+                        "top_k": 1,
+                        "stream": False,
+                        }
+                random_index = random.randint(0, len(database) - 1)
+                user_inputs = database[random_index]
+
+                unexpected_count = 0
+                logs = []
+                for input, expected in zip(user_inputs[0::2], user_inputs[1::2]):
+                    json["messages"].append({"role": "user", "content": input})
+                    logs.append(datetime.now().strftime("%H:%M:%S.%f")[:-3] + " " + "    user: " + input)
+                    async with concurrency_sem:
+                        async with session.post(url, json=json, headers=headers) as response:
+                            res = (await response.json())["choices"][0]["message"]
+                            output = res["content"]
+                            role = res["role"]
+                            color = GREEN if output == expected else RED
+                            unexpected_count += 1 if output != expected else 0
+                            json["messages"].append({"role": role, "content": output})
+                            logs.append(datetime.now().strftime("%H:%M:%S.%f")[:-3] + " " + "assistant: " + color + output + END)
+
+                print("=========\n","\n".join(logs))
+                return unexpected_count
+
+        tasks = []
+        for i in range(num_tests):
+            tasks.append(asyncio.create_task(job()))
+        results = await asyncio.gather(*tasks)
+        return sum(results)
+
+if 1:
+    unexpected_count = asyncio.run(main(10, 10))
+    print(f"unexpected count: {unexpected_count}")
+
+asyncio.run(main(0, 0))
+
+

tokenizer_config.json

@@ -0,0 +1,3 @@
+{
+  "chat_template": "{% for message in messages | rejectattr('role', 'equalto', 'system') %}{% if not loop.first %}{{ '\\n' }}{% endif %}<|im_start|>{{ message['role'] }}\n{{ message['content'] }}<|im_end|>{% endfor %}{{ '\\n' }}{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}"
+}