Add improved aber small model Space

README.md

@@ -1,12 +1,27 @@
 ---
-title: Aber Small Model Python
-emoji: 🐢
-colorFrom: pink
-colorTo: pink
+title: aber Small Model
+colorFrom: green
+colorTo: blue
 sdk: gradio
-sdk_version: 6.10.0
 app_file: app.py
 pinned: false
+license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# aber Small Model
+
+This is an improved local small language model project written in Python from scratch.
+
+## What is different
+
+- Model name is `aber`
+- Uses a different architecture than the previous tiny character model
+- Uses a word-level tokenizer
+- Uses a GRU language model for more readable short outputs
+- Trains and runs locally on CPU
+
+## Important
+
+- No external pretrained LLM is used
+- This is still a small educational model
+- It is designed for lightweight Hugging Face CPU Spaces and local Python use

aber_llm/__init__.py

@@ -0,0 +1,4 @@
+from .config import AberConfig
+from .service import AberLLMService
+
+__all__ = ["AberConfig", "AberLLMService"]

aber_llm/config.py

@@ -0,0 +1,24 @@
+from dataclasses import dataclass
+from pathlib import Path
+
+
+@dataclass
+class AberConfig:
+    seq_len: int = 40
+    batch_size: int = 24
+    embed_dim: int = 96
+    hidden_dim: int = 160
+    num_layers: int = 2
+    dropout: float = 0.15
+    learning_rate: float = 2.5e-3
+    bootstrap_steps: int = 90
+    cpu_threads: int = 4
+    seed: int = 42
+
+    @property
+    def root_dir(self) -> Path:
+        return Path(__file__).resolve().parents[1]
+
+    @property
+    def checkpoint_path(self) -> Path:
+        return self.root_dir / "artifacts" / "aber_checkpoint.pt"

aber_llm/data.py

@@ -0,0 +1,87 @@
+BASE_CORPUS = """
+System: You are aber, a small local language model.
+User: hello
+aber: Hello. I am aber, a tiny local model written in Python.
+
+User: what are you
+aber: I am a small language model trained from scratch on local text.
+
+User: explain machine learning
+aber: Machine learning builds patterns from data so a model can make predictions or decisions.
+
+User: explain deep learning
+aber: Deep learning uses neural networks with many layers to learn useful representations.
+
+User: explain natural language processing
+aber: Natural language processing helps computers read, generate, and understand human language.
+
+User: what is regression
+aber: Regression predicts a continuous value such as price, score, or demand.
+
+User: what is classification
+aber: Classification predicts a category or label for an example.
+
+User: what is clustering
+aber: Clustering groups similar data points without labeled targets.
+
+User: what is preprocessing
+aber: Preprocessing cleans, encodes, scales, and prepares data before modeling.
+
+User: what is evaluation
+aber: Evaluation measures how well a model performs on data it did not train on.
+
+User: how do I learn coding
+aber: Practice small projects, read errors carefully, and improve one step at a time.
+
+User: give a short greeting
+aber: Hello. I am ready to help with a short answer.
+
+User: write a short poem
+aber: Quiet light on morning glass, slow thoughts gather as clouds pass.
+
+User: what is Python
+aber: Python is a programming language known for clarity, flexibility, and strong libraries.
+
+User: what is a feature
+aber: A feature is an input variable used by a model.
+
+User: what is a target
+aber: A target is the value or label a model tries to predict.
+
+User: what is overfitting
+aber: Overfitting happens when a model memorizes training data and generalizes poorly.
+
+User: what is underfitting
+aber: Underfitting happens when a model is too simple to learn important patterns.
+
+User: what is a transformer
+aber: A transformer is a model architecture that uses attention between tokens.
+
+User: what is a tokenizer
+aber: A tokenizer converts text into smaller pieces that a model can process.
+
+User: what is local ai
+aber: Local AI runs on your own machine so you can control training, files, and execution.
+
+User: how should I debug code
+aber: Reproduce the issue, isolate the failing step, inspect values, and test one fix at a time.
+
+User: summarize good study habits
+aber: Good study habits use planning, active recall, spaced review, and regular breaks.
+
+User: what is data science
+aber: Data science combines programming, statistics, and domain knowledge to learn from data.
+
+User: what is a neural network
+aber: A neural network is a layered function that transforms input signals into predictions.
+
+User: give motivation
+aber: Small repeated effort beats waiting for perfect motivation.
+""".strip()
+
+
+def build_training_text(extra_text: str = "") -> str:
+    extra = " ".join((extra_text or "").split())
+    if not extra:
+        return BASE_CORPUS
+    return BASE_CORPUS + "\n\n" + extra

aber_llm/model.py

@@ -0,0 +1,48 @@
+import torch
+from torch import nn
+
+
+class AberLanguageModel(nn.Module):
+    def __init__(self, vocab_size, embed_dim, hidden_dim, num_layers, dropout):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, embed_dim)
+        self.gru = nn.GRU(
+            input_size=embed_dim,
+            hidden_size=hidden_dim,
+            num_layers=num_layers,
+            dropout=dropout if num_layers > 1 else 0.0,
+            batch_first=True,
+        )
+        self.dropout = nn.Dropout(dropout)
+        self.head = nn.Linear(hidden_dim, vocab_size)
+
+    def forward(self, idx, hidden=None, targets=None):
+        emb = self.embedding(idx)
+        out, hidden = self.gru(emb, hidden)
+        out = self.dropout(out)
+        logits = self.head(out)
+
+        loss = None
+        if targets is not None:
+            loss = nn.functional.cross_entropy(
+                logits.reshape(-1, logits.size(-1)),
+                targets.reshape(-1),
+            )
+        return logits, hidden, loss
+
+    def generate(self, idx, max_new_tokens, eos_id, temperature=1.0, top_k=8):
+        hidden = None
+        for _ in range(max_new_tokens):
+            logits, hidden, _ = self(idx[:, -1:], hidden)
+            next_logits = logits[:, -1, :] / max(temperature, 1e-4)
+
+            if top_k is not None and top_k > 0:
+                values, _ = torch.topk(next_logits, min(top_k, next_logits.size(-1)))
+                next_logits[next_logits < values[:, [-1]]] = float("-inf")
+
+            probs = torch.softmax(next_logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+            idx = torch.cat([idx, next_token], dim=1)
+            if int(next_token.item()) == eos_id:
+                break
+        return idx

aber_llm/service.py

@@ -0,0 +1,120 @@
+from pathlib import Path
+import shutil
+
+import torch
+
+from .config import AberConfig
+from .model import AberLanguageModel
+from .tokenizer import WordTokenizer
+from .trainer import create_model_and_tokenizer, set_seed, train_model
+
+
+class AberLLMService:
+    def __init__(self, config: AberConfig):
+        self.config = config
+        torch.set_num_threads(max(1, self.config.cpu_threads))
+        self.model = None
+        self.tokenizer = None
+
+    def generate(self, prompt: str, max_new_tokens: int, temperature: float, top_k: int):
+        clean_prompt = prompt or "User: hello\naber:"
+        self._ensure_ready()
+        encoded = self.tokenizer.encode(clean_prompt, add_bos=True)
+        idx = torch.tensor(encoded, dtype=torch.long).unsqueeze(0)
+        self.model.eval()
+
+        with torch.inference_mode():
+            output = self.model.generate(
+                idx=idx,
+                max_new_tokens=max_new_tokens,
+                eos_id=self.tokenizer.eos_id,
+                temperature=temperature,
+                top_k=top_k,
+            )
+
+        text = self.tokenizer.decode(output[0].tolist())
+        status = (
+            f"Generated with aber. "
+            f"Architecture=word-level GRU, Vocab={self.tokenizer.vocab_size}, Hidden={self.config.hidden_dim}."
+        )
+        return text, status
+
+    def train(self, extra_text: str, steps: int):
+        steps = max(1, steps)
+        checkpoint_exists = self.config.checkpoint_path.exists()
+        training_text = extra_text or ""
+
+        if checkpoint_exists:
+            self._load_or_initialize(extra_text="")
+
+        model, tokenizer, encoded = create_model_and_tokenizer(self.config, training_text)
+        if checkpoint_exists and self.model is not None and self.tokenizer is not None:
+            if tokenizer.stoi == self.tokenizer.stoi:
+                model.load_state_dict(self.model.state_dict())
+
+        losses = train_model(model, encoded, self.config, steps)
+        self.model = model
+        self.tokenizer = tokenizer
+        self._save_checkpoint(extra_text=training_text)
+
+        return (
+            f"aber training finished.\n"
+            f"Steps: {steps}\n"
+            f"Start Loss: {losses[0]:.4f}\n"
+            f"End Loss: {losses[-1]:.4f}\n"
+            f"Checkpoint: {self.config.checkpoint_path}"
+        )
+
+    def reset(self):
+        checkpoint_dir = self.config.checkpoint_path.parent
+        if checkpoint_dir.exists():
+            shutil.rmtree(checkpoint_dir)
+        self.model = None
+        self.tokenizer = None
+        return "aber reset complete. Next train or generate call will rebuild the model from scratch."
+
+    def _ensure_ready(self):
+        if self.model is not None and self.tokenizer is not None:
+            return
+        self._load_or_initialize(extra_text="")
+
+    def _load_or_initialize(self, extra_text: str):
+        checkpoint = self.config.checkpoint_path
+        if checkpoint.exists():
+            state = torch.load(checkpoint, map_location="cpu")
+            self.tokenizer = WordTokenizer.from_state_dict(state["tokenizer"])
+            self.model = AberLanguageModel(
+                vocab_size=state["config"]["vocab_size"],
+                embed_dim=state["config"]["embed_dim"],
+                hidden_dim=state["config"]["hidden_dim"],
+                num_layers=state["config"]["num_layers"],
+                dropout=state["config"]["dropout"],
+            )
+            self.model.load_state_dict(state["model"])
+            self.model.eval()
+            return
+
+        set_seed(self.config.seed)
+        self.model, self.tokenizer, encoded = create_model_and_tokenizer(self.config, extra_text)
+        train_model(self.model, encoded, self.config, self.config.bootstrap_steps)
+        self._save_checkpoint(extra_text=extra_text)
+
+    def _save_checkpoint(self, extra_text: str):
+        checkpoint = self.config.checkpoint_path
+        checkpoint.parent.mkdir(parents=True, exist_ok=True)
+        torch.save(
+            {
+                "model": self.model.state_dict(),
+                "tokenizer": self.tokenizer.state_dict(),
+                "config": {
+                    "vocab_size": self.tokenizer.vocab_size,
+                    "embed_dim": self.config.embed_dim,
+                    "hidden_dim": self.config.hidden_dim,
+                    "num_layers": self.config.num_layers,
+                    "dropout": self.config.dropout,
+                    "seq_len": self.config.seq_len,
+                    "extra_text": extra_text,
+                },
+            },
+            checkpoint,
+        )

aber_llm/tokenizer.py

@@ -0,0 +1,73 @@
+import re
+
+
+TOKEN_PATTERN = re.compile(r"\n|[A-Za-z0-9_']+|[^\w\s]")
+
+
+class WordTokenizer:
+    def __init__(self):
+        self.special_tokens = ["<pad>", "<unk>", "<bos>", "<eos>"]
+        self.stoi = {}
+        self.itos = {}
+
+    @property
+    def pad_id(self):
+        return self.stoi["<pad>"]
+
+    @property
+    def bos_id(self):
+        return self.stoi["<bos>"]
+
+    @property
+    def eos_id(self):
+        return self.stoi["<eos>"]
+
+    @property
+    def vocab_size(self):
+        return len(self.stoi)
+
+    def tokenize(self, text: str):
+        return TOKEN_PATTERN.findall(text)
+
+    def fit(self, text: str):
+        vocab = self.special_tokens + sorted(set(self.tokenize(text)))
+        self.stoi = {token: idx for idx, token in enumerate(vocab)}
+        self.itos = {idx: token for token, idx in self.stoi.items()}
+        return self
+
+    def encode(self, text: str, add_bos: bool = False, add_eos: bool = False):
+        tokens = self.tokenize(text)
+        ids = [self.stoi.get(token, self.stoi["<unk>"]) for token in tokens]
+        if add_bos:
+            ids = [self.bos_id] + ids
+        if add_eos:
+            ids = ids + [self.eos_id]
+        return ids
+
+    def decode(self, ids):
+        tokens = []
+        for idx in ids:
+            token = self.itos.get(int(idx), "<unk>")
+            if token in self.special_tokens:
+                continue
+            tokens.append(token)
+
+        text = ""
+        for token in tokens:
+            if token == "\n":
+                text = text.rstrip() + "\n"
+            elif token in {".", ",", "!", "?", ":", ";"}:
+                text = text.rstrip() + token + " "
+            else:
+                text += token + " "
+        return text.strip()
+
+    def state_dict(self):
+        return {"stoi": self.stoi}
+
+    @classmethod
+    def from_state_dict(cls, state):
+        tok = cls()
+        tok.stoi = dict(state["stoi"])
+        tok.itos = {idx: token for token, idx in tok.stoi.items()}
+        return tok

aber_llm/trainer.py

@@ -0,0 +1,51 @@
+import random
+
+import torch
+
+from .data import build_training_text
+from .model import AberLanguageModel
+from .tokenizer import WordTokenizer
+
+
+def set_seed(seed: int):
+    random.seed(seed)
+    torch.manual_seed(seed)
+
+
+def create_model_and_tokenizer(config, extra_text=""):
+    text = build_training_text(extra_text)
+    tokenizer = WordTokenizer().fit(text)
+    encoded = tokenizer.encode(text, add_bos=True, add_eos=True)
+    encoded = torch.tensor(encoded, dtype=torch.long)
+    model = AberLanguageModel(
+        vocab_size=tokenizer.vocab_size,
+        embed_dim=config.embed_dim,
+        hidden_dim=config.hidden_dim,
+        num_layers=config.num_layers,
+        dropout=config.dropout,
+    )
+    return model, tokenizer, encoded
+
+
+def build_batch(encoded, seq_len, batch_size):
+    max_start = max(1, len(encoded) - seq_len - 1)
+    starts = torch.randint(0, max_start, (batch_size,))
+    x = torch.stack([encoded[start : start + seq_len] for start in starts])
+    y = torch.stack([encoded[start + 1 : start + seq_len + 1] for start in starts])
+    return x, y
+
+
+def train_model(model, encoded, config, steps):
+    optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate)
+    model.train()
+    losses = []
+
+    for _ in range(steps):
+        xb, yb = build_batch(encoded, config.seq_len, config.batch_size)
+        _, _, loss = model(xb, targets=yb)
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        losses.append(float(loss.item()))
+
+    return losses

app.py

@@ -0,0 +1,89 @@
+import gradio as gr
+
+from aber_llm.config import AberConfig
+from aber_llm.service import AberLLMService
+
+
+config = AberConfig()
+service = AberLLMService(config=config)
+
+
+def generate_text(prompt, max_new_tokens, temperature, top_k):
+    return service.generate(
+        prompt=prompt,
+        max_new_tokens=int(max_new_tokens),
+        temperature=float(temperature),
+        top_k=int(top_k),
+    )
+
+
+def train_model(extra_text, steps):
+    return service.train(extra_text=extra_text, steps=int(steps))
+
+
+def reset_model():
+    return service.reset()
+
+
+with gr.Blocks(
+    title="aber Small Model",
+    theme=gr.themes.Soft(primary_hue="green", secondary_hue="blue"),
+) as demo:
+    gr.Markdown(
+        """
+        # aber
+        An improved small language model written in Python from scratch.
+
+        - Model name: `aber`
+        - No external pretrained LLM
+        - Word-level tokenizer
+        - GRU language model
+        - Local CPU training and generation
+        """
+    )
+
+    with gr.Tab("Generate"):
+        prompt_input = gr.Textbox(
+            label="Prompt",
+            value="User: hello\naber:",
+            lines=6,
+        )
+        with gr.Row():
+            max_tokens_input = gr.Slider(10, 160, value=72, step=2, label="Max New Tokens")
+            temperature_input = gr.Slider(0.2, 1.3, value=0.75, step=0.05, label="Temperature")
+            top_k_input = gr.Slider(1, 20, value=8, step=1, label="Top-K")
+        generate_button = gr.Button("Generate", variant="primary")
+        output_text = gr.Textbox(label="Output", lines=10)
+        output_status = gr.Textbox(label="Status", lines=4)
+
+    with gr.Tab("Train"):
+        extra_text_input = gr.Textbox(
+            label="Extra Training Text",
+            placeholder="Add more local text to train aber on your own data.",
+            lines=10,
+        )
+        steps_input = gr.Slider(10, 400, value=120, step=10, label="Training Steps")
+        train_button = gr.Button("Train / Continue Training", variant="primary")
+        reset_button = gr.Button("Reset aber")
+        train_status = gr.Textbox(label="Training Status", lines=6)
+
+    generate_button.click(
+        fn=generate_text,
+        inputs=[prompt_input, max_tokens_input, temperature_input, top_k_input],
+        outputs=[output_text, output_status],
+    )
+
+    train_button.click(
+        fn=train_model,
+        inputs=[extra_text_input, steps_input],
+        outputs=[train_status],
+    )
+
+    reset_button.click(
+        fn=reset_model,
+        outputs=[train_status],
+    )
+
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -0,0 +1,2 @@
+gradio>=5.23.0
+torch>=2.3.0