umbrellm 0.1.0__tar.gz

umbrellm-0.1.0/PKG-INFO

@@ -0,0 +1,90 @@
+Metadata-Version: 2.4
+Name: umbrellm
+Version: 0.1.0
+Summary: A lightweight, trainable language model SDK
+Home-page: https://github.com/The-Umbrellm-Project/UmbreLLM
+Author: Fries
+Author-email: Your Name <your.email@example.com>
+License: MIT
+Project-URL: Homepage, https://github.com/YOUR_USERNAME/umbrellm
+Project-URL: Repository, https://github.com/YOUR_USERNAME/umbrellm
+Project-URL: Issues, https://github.com/YOUR_USERNAME/umbrellm/issues
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: torch>=2.0.0
+Requires-Dist: fastapi>=0.100.0
+Requires-Dist: uvicorn>=0.22.0
+Dynamic: author
+Dynamic: home-page
+Dynamic: requires-python
+
+# Umbrellm
+
+A small language model framework that lets you train and run your own decoder-only transformer locally.
+
+## Overview
+
+Umbrellm is a language model implementation written in Python and PyTorch. The default model configuration uses 12 transformer layers, 8 attention heads, and a hidden size of 512 dimensions, resulting in roughly 20 million parameters.
+
+The goal of the project is to provide an implementation that is easy to understand and modify. It includes the components needed to train a model from scratch, prepare datasets, run inference, and expose the model through a simple API.
+
+The architecture uses:
+
+* Rotary Position Embeddings (RoPE)
+* SwiGLU feed-forward activations
+* RMSNorm
+* Causal multi-head self-attention
+
+## Features
+
+* Train models from scratch using your own datasets
+* Save and resume training from checkpoints
+* Learning rate scheduling and gradient clipping
+* Optional mixed-precision training
+* BPE tokenizer training on custom corpora
+* Dataset compilation using the `.ullm` format
+* Text generation with temperature, top-k, top-p, and repetition penalty controls
+* Token-by-token streaming generation
+* FastAPI-based REST API server
+
+## Quick start
+
+```bash
+pip install torch fastapi uvicorn
+
+python scripts/compile_dataset.py
+python scripts/train.py --epochs 3
+python scripts/inference.py --interactive
+```
+
+## Python SDK
+
+```python
+import umbrellm
+
+umbrellm.utllm.compile()
+umbrellm.training.train({"epochs": 3})
+
+response = umbrellm.generate("Explain black holes simply.")
+print(response)
+```
+
+## Requirements
+
+* Python 3.9 or newer
+* PyTorch 2.0 or newer
+* FastAPI and uvicorn (only required for the API server)
+
+## License
+
+Released under the MIT License.

umbrellm-0.1.0/README.md

@@ -0,0 +1,60 @@
+# Umbrellm
+
+A small language model framework that lets you train and run your own decoder-only transformer locally.
+
+## Overview
+
+Umbrellm is a language model implementation written in Python and PyTorch. The default model configuration uses 12 transformer layers, 8 attention heads, and a hidden size of 512 dimensions, resulting in roughly 20 million parameters.
+
+The goal of the project is to provide an implementation that is easy to understand and modify. It includes the components needed to train a model from scratch, prepare datasets, run inference, and expose the model through a simple API.
+
+The architecture uses:
+
+* Rotary Position Embeddings (RoPE)
+* SwiGLU feed-forward activations
+* RMSNorm
+* Causal multi-head self-attention
+
+## Features
+
+* Train models from scratch using your own datasets
+* Save and resume training from checkpoints
+* Learning rate scheduling and gradient clipping
+* Optional mixed-precision training
+* BPE tokenizer training on custom corpora
+* Dataset compilation using the `.ullm` format
+* Text generation with temperature, top-k, top-p, and repetition penalty controls
+* Token-by-token streaming generation
+* FastAPI-based REST API server
+
+## Quick start
+
+```bash
+pip install torch fastapi uvicorn
+
+python scripts/compile_dataset.py
+python scripts/train.py --epochs 3
+python scripts/inference.py --interactive
+```
+
+## Python SDK
+
+```python
+import umbrellm
+
+umbrellm.utllm.compile()
+umbrellm.training.train({"epochs": 3})
+
+response = umbrellm.generate("Explain black holes simply.")
+print(response)
+```
+
+## Requirements
+
+* Python 3.9 or newer
+* PyTorch 2.0 or newer
+* FastAPI and uvicorn (only required for the API server)
+
+## License
+
+Released under the MIT License.

umbrellm-0.1.0/pyproject.toml

@@ -0,0 +1,44 @@
+[build-system]
+requires = ["setuptools>=61.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "umbrellm"
+version = "0.1.0"
+description = "A lightweight, trainable language model SDK"
+readme = "README.md"
+requires-python = ">=3.9"
+license = {text = "MIT"}
+authors = [
+    {name = "Your Name", email = "your.email@example.com"}
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+]
+
+dependencies = [
+    "torch>=2.0.0",
+    "fastapi>=0.100.0",
+    "uvicorn>=0.22.0",
+]
+
+[project.urls]
+Homepage = "https://github.com/YOUR_USERNAME/umbrellm"
+Repository = "https://github.com/YOUR_USERNAME/umbrellm"
+Issues = "https://github.com/YOUR_USERNAME/umbrellm/issues"
+
+[project.scripts]
+umbrellm = "umbrellm.__main__:main"
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["umbrellm*"]

umbrellm-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

umbrellm-0.1.0/setup.py

@@ -0,0 +1,32 @@
+from setuptools import setup, find_packages
+
+with open("README.md", "r", encoding="utf-8") as fh:
+    long_description = fh.read()
+
+setup(
+    name="umbrellm",
+    version="0.1.0",
+    author="Fries",
+    author_email="freshfrenchfries@proton.me",
+    description="A lightweight, trainable language model SDK",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    url="https://github.com/The-Umbrellm-Project/UmbreLLM",
+    packages=find_packages(),
+    classifiers=[
+        "Programming Language :: Python :: 3",
+        "License :: OSI Approved :: MIT License",
+        "Operating System :: OS Independent",
+    ],
+    python_requires=">=3.9",
+    install_requires=[
+        "torch>=2.0.0",
+        "fastapi>=0.100.0",
+        "uvicorn>=0.22.0",
+    ],
+    entry_points={
+        "console_scripts": [
+            "umbrellm=umbrellm.__main__:main",
+        ],
+    },
+)

umbrellm-0.1.0/umbrellm/__init__.py

@@ -0,0 +1,19 @@
+"""
+Umbrellm – a lightweight, trainable language model SDK.
+"""
+
+from umbrellm import interaction, tokenizer, models, training, utllm, config, device
+from umbrellm._generate import generate, complete
+
+__version__ = "0.1.0"
+__all__ = [
+    "interaction",
+    "tokenizer",
+    "models",
+    "training",
+    "utllm",
+    "config",
+    "device",
+    "generate",
+    "complete",
+]

umbrellm-0.1.0/umbrellm/__main__.py

@@ -0,0 +1,75 @@
+"""
+python -m umbrellm  entry-point.
+"""
+
+import sys
+import argparse
+
+
+def main():
+    parser = argparse.ArgumentParser(prog="umbrellm")
+    sub = parser.add_subparsers(dest="cmd")
+
+    # train
+    p_train = sub.add_parser("train", help="Train the model")
+    p_train.add_argument("--epochs", type=int, default=3)
+    p_train.add_argument("--batch-size", type=int, default=4)
+    p_train.add_argument("--lr", type=float, default=3e-4)
+
+    # chat
+    p_chat = sub.add_parser("chat", help="Interactive chat")
+    p_chat.add_argument("--model", default="umbrellm-20m")
+
+    # compile
+    sub.add_parser("compile", help="Compile UTLLM dataset")
+
+    # stats
+    sub.add_parser("stats", help="Dataset statistics")
+
+    # serve
+    sub.add_parser("serve", help="Start API server")
+
+    args = parser.parse_args()
+
+    if args.cmd == "train":
+        import umbrellm.training as tr
+        tr.train({"epochs": args.epochs, "batch_size": args.batch_size,
+                  "learning_rate": args.lr})
+
+    elif args.cmd == "chat":
+        import umbrellm.interaction as inter
+        print(f"Umbrellm interactive chat (model: {args.model}). Type exit to quit.")
+        history = []
+        while True:
+            try:
+                user = input("You: ").strip()
+            except (EOFError, KeyboardInterrupt):
+                break
+            if user.lower() in ("exit", "quit", "q"):
+                break
+            if not user:
+                continue
+            res = inter.chat({"user": user, "model": args.model, "history": history})
+            print(f"Umbrellm: {res['response']}")
+            history.append({"user": user, "assistant": res["response"]})
+
+    elif args.cmd == "compile":
+        import umbrellm.utllm as ut
+        ut.compile()
+
+    elif args.cmd == "stats":
+        import umbrellm.utllm as ut
+        import json
+        s = ut.stats()
+        print(json.dumps(s, indent=2))
+
+    elif args.cmd == "serve":
+        from umbrellm.api import main as serve_main
+        serve_main()
+
+    else:
+        parser.print_help()
+
+
+if __name__ == "__main__":
+    main()

umbrellm-0.1.0/umbrellm/_generate.py

@@ -0,0 +1,30 @@
+"""
+Top-level generate / complete helpers.
+"""
+
+from umbrellm import interaction as _interaction
+
+
+def generate(prompt: str, **kwargs) -> str:
+    """
+    Generate text from a prompt.
+
+    Example:
+        text = umbrellm.generate("Write a poem about stars.")
+    """
+    kwargs.setdefault("model", "umbrellm-20m")
+    result = _interaction.chat({"user": prompt, **kwargs})
+    return result["response"]
+
+
+def complete(prompt: str, **kwargs) -> str:
+    """
+    Complete a text prefix.
+
+    Example:
+        text = umbrellm.complete("The meaning of life is")
+    """
+    kwargs.setdefault("model", "umbrellm-20m")
+    kwargs.setdefault("system_prompt", "You are a helpful text completion engine. Continue the text naturally.")
+    result = _interaction.chat({"user": prompt, **kwargs})
+    return result["response"]

umbrellm-0.1.0/umbrellm/_runtime.py

@@ -0,0 +1,179 @@
+"""
+Internal inference runtime.
+Handles model loading, prompt formatting, and generation.
+"""
+
+import time
+import os
+from pathlib import Path
+from typing import Iterator, Optional
+
+
+_SYSTEM_DEFAULT = (
+    "You are Umbrellm, a helpful AI assistant. "
+    "Respond clearly, concisely, and helpfully."
+)
+
+
+class InferenceRuntime:
+    def __init__(self):
+        self._model = None
+        self._cfg = None
+        self._device = None
+        self._tokenizer_loaded = False
+
+    def _ensure_ready(self, model_name: str):
+        from umbrellm.config import load as load_cfg
+        from umbrellm.device import current as cur_device
+        import umbrellm.tokenizer as tok_mod
+
+        if self._cfg is None:
+            self._cfg = load_cfg()
+
+        if self._device is None:
+            self._device = cur_device()
+
+        # Ensure tokenizer is loaded
+        if not self._tokenizer_loaded:
+            try:
+                tok_mod.load()
+            except Exception:
+                pass
+            self._tokenizer_loaded = True
+
+        if self._model is None:
+            self._model = self._load_model(model_name)
+
+    def _load_model(self, model_name: str):
+        try:
+            import torch
+            from umbrellm.model import build_model
+            import umbrellm.models as mdl_mod
+
+            # Try loading existing checkpoint
+            ckpt_dir = self._cfg.get("checkpoint_dir", "checkpoints")
+            candidates = [
+                os.path.join(ckpt_dir, "latest.pt"),
+                os.path.join(ckpt_dir, "best.pt"),
+            ]
+            # Also check if model_name is a path
+            if os.path.exists(model_name):
+                candidates.insert(0, model_name)
+
+            for path in candidates:
+                if os.path.exists(path):
+                    model = mdl_mod.load(path)
+                    return model.to(self._device)
+
+            # No checkpoint found – build a fresh (random) model
+            print(f"[runtime] No checkpoint found. Building fresh model for inference.")
+            model = build_model(self._cfg)
+            model.eval()
+            return model.to(self._device)
+
+        except ImportError as e:
+            raise RuntimeError(
+                "PyTorch is required. Install with: pip install torch"
+            ) from e
+
+    def _build_prompt(self, params: dict) -> str:
+        system = params.get("system_prompt", _SYSTEM_DEFAULT)
+        history = params.get("history", [])
+        user = params.get("user", "")
+
+        parts = [f"<|system|>{system}<|assistant|>"]
+        for turn in history:
+            parts.append(f"<|user|>{turn.get('user', '')}")
+            parts.append(f"<|assistant|>{turn.get('assistant', '')}")
+        parts.append(f"<|user|>{user}")
+        parts.append("<|assistant|>")
+        return "".join(parts)
+
+    def chat(self, params: dict) -> dict:
+        model_name = params.get("model", "umbrellm-20m")
+        self._ensure_ready(model_name)
+
+        prompt = self._build_prompt(params)
+        start = time.time()
+
+        import torch
+        import umbrellm.tokenizer as tok_mod
+        from umbrellm.tokenizer import ASSISTANT_TOKEN, EOS_TOKEN
+
+        input_ids = tok_mod.encode(prompt, add_special_tokens=False)
+        input_tensor = torch.tensor([input_ids], dtype=torch.long, device=self._device)
+
+        stop_ids = []
+        eos_id = tok_mod._tokenizer.vocab.get(EOS_TOKEN)
+        if eos_id is not None:
+            stop_ids.append(eos_id)
+
+        with torch.no_grad():
+            output = self._model.generate(
+                input_tensor,
+                max_new_tokens=params.get("max_tokens", self._cfg.get("max_new_tokens", 256)),
+                temperature=params.get("temperature", self._cfg.get("temperature", 0.8)),
+                top_k=params.get("top_k", self._cfg.get("top_k", 50)),
+                top_p=params.get("top_p", self._cfg.get("top_p", 0.95)),
+                repetition_penalty=params.get("repetition_penalty",
+                                              self._cfg.get("repetition_penalty", 1.1)),
+                stop_ids=stop_ids or None,
+                seed=params.get("seed"),
+            )
+
+        new_ids = output[0, len(input_ids):].tolist()
+        response = tok_mod.decode(new_ids, skip_special_tokens=True)
+
+        # Trim at stop sequences
+        for seq in params.get("stop_sequences", []):
+            idx = response.find(seq)
+            if idx != -1:
+                response = response[:idx]
+
+        elapsed = time.time() - start
+        return {
+            "response": response.strip(),
+            "tokens_generated": len(new_ids),
+            "generation_time": round(elapsed, 4),
+            "model": model_name,
+        }
+
+    def stream_chat(self, params: dict) -> Iterator[str]:
+        """Yield tokens one-by-one using a simple greedy decode loop."""
+        model_name = params.get("model", "umbrellm-20m")
+        self._ensure_ready(model_name)
+
+        prompt = self._build_prompt(params)
+        import torch
+        import torch.nn.functional as F
+        import umbrellm.tokenizer as tok_mod
+        from umbrellm.tokenizer import EOS_TOKEN
+
+        input_ids = tok_mod.encode(prompt, add_special_tokens=False)
+        input_tensor = torch.tensor([input_ids], dtype=torch.long, device=self._device)
+
+        temperature = params.get("temperature", self._cfg.get("temperature", 0.8))
+        top_k = params.get("top_k", self._cfg.get("top_k", 50))
+        max_new_tokens = params.get("max_tokens", self._cfg.get("max_new_tokens", 256))
+        stop_ids = set()
+        eos_id = tok_mod._tokenizer.vocab.get(EOS_TOKEN)
+        if eos_id is not None:
+            stop_ids.add(eos_id)
+
+        generated = input_tensor.clone()
+        with torch.no_grad():
+            for _ in range(max_new_tokens):
+                ctx = generated if generated.shape[1] <= self._model.max_seq_len else                     generated[:, -self._model.max_seq_len:]
+                logits, _ = self._model(ctx)
+                logits = logits[:, -1, :] / max(temperature, 1e-5)
+                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_tok = torch.multinomial(probs, num_samples=1)
+                tok_id = next_tok.item()
+                if tok_id in stop_ids:
+                    break
+                generated = torch.cat([generated, next_tok], dim=1)
+                token_text = tok_mod.decode([tok_id], skip_special_tokens=True)
+                yield token_text

umbrellm-0.1.0/umbrellm/api.py

@@ -0,0 +1,105 @@
+"""
+Umbrellm REST API server (FastAPI).
+
+Run:
+    python -m umbrellm.api
+or:
+    uvicorn umbrellm.api:app --host 0.0.0.0 --port 8080
+"""
+
+import time
+from typing import Optional, List
+
+try:
+    from fastapi import FastAPI, HTTPException
+    from fastapi.responses import StreamingResponse
+    from pydantic import BaseModel
+    _FASTAPI_OK = True
+except ImportError:
+    _FASTAPI_OK = False
+
+if _FASTAPI_OK:
+    app = FastAPI(
+        title="Umbrellm API",
+        description="REST API for the Umbrellm language model.",
+        version="0.1.0",
+    )
+
+    class ChatRequest(BaseModel):
+        user: str
+        model: str = "umbrellm-20m"
+        history: Optional[List[dict]] = None
+        system_prompt: Optional[str] = None
+        temperature: Optional[float] = None
+        top_k: Optional[int] = None
+        top_p: Optional[float] = None
+        repetition_penalty: Optional[float] = None
+        max_tokens: Optional[int] = None
+        stop_sequences: Optional[List[str]] = None
+        stream: bool = False
+        seed: Optional[int] = None
+
+    class ChatResponse(BaseModel):
+        response: str
+        tokens_generated: int
+        generation_time: float
+        model: str
+
+    @app.get("/")
+    def root():
+        return {"service": "Umbrellm API", "version": "0.1.0"}
+
+    @app.get("/health")
+    def health():
+        return {"status": "ok", "timestamp": time.time()}
+
+    @app.get("/models")
+    def list_models():
+        import umbrellm.models as mdl
+        checkpoints = mdl.list()
+        return {"models": ["umbrellm-20m"] + checkpoints}
+
+    @app.post("/chat", response_model=ChatResponse)
+    def chat(req: ChatRequest):
+        import umbrellm.interaction as inter
+        params = req.dict(exclude_none=True)
+        if req.stream:
+            def _gen():
+                for tok in inter.stream_chat(params):
+                    yield tok
+            return StreamingResponse(_gen(), media_type="text/plain")
+        result = inter.chat(params)
+        return result
+
+    @app.post("/tokenize")
+    def tokenize(body: dict):
+        import umbrellm.tokenizer as tok
+        text = body.get("text", "")
+        ids = tok.encode(text)
+        return {"tokens": ids, "count": len(ids)}
+
+    @app.post("/detokenize")
+    def detokenize(body: dict):
+        import umbrellm.tokenizer as tok
+        ids = body.get("ids", [])
+        text = tok.decode(ids)
+        return {"text": text}
+
+
+def main():
+    if not _FASTAPI_OK:
+        print("[api] FastAPI/uvicorn not installed. Run: pip install fastapi uvicorn")
+        return
+    import uvicorn
+    from umbrellm.config import load as load_cfg
+    cfg = load_cfg()
+    uvicorn.run(
+        "umbrellm.api:app",
+        host=cfg.get("api_host", "0.0.0.0"),
+        port=cfg.get("api_port", 8080),
+        reload=False,
+    )
+
+
+if __name__ == "__main__":
+    main()