umbrellm 0.1.0__py3-none-any.whl

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/__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/_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/_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/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()

umbrellm/config.py

@@ -0,0 +1,66 @@
+"""
+Configuration utilities for Umbrellm.
+"""
+
+import json
+import os
+from pathlib import Path
+
+_DEFAULT_CONFIG_PATH = "config/umbrellm.json"
+
+DEFAULT_CONFIG = {
+    "model_name": "umbrellm-20m",
+    "vocab_size": 32000,
+    "d_model": 512,
+    "n_heads": 8,
+    "n_layers": 12,
+    "d_ff": 2048,
+    "max_seq_len": 1024,
+    "dropout": 0.1,
+    "learning_rate": 3e-4,
+    "batch_size": 4,
+    "epochs": 3,
+    "checkpoint_dir": "checkpoints",
+    "data_dir": "data",
+    "log_dir": "logs",
+    "device": "auto",
+    "seed": 42,
+    "mixed_precision": False,
+    "gradient_clip": 1.0,
+    "warmup_steps": 100,
+    "save_every": 500,
+    "eval_every": 250,
+    "api_host": "0.0.0.0",
+    "api_port": 8080,
+    "temperature": 0.8,
+    "top_k": 50,
+    "top_p": 0.95,
+    "repetition_penalty": 1.1,
+    "max_new_tokens": 256,
+}
+
+
+def load(path: str = _DEFAULT_CONFIG_PATH) -> dict:
+    """Load configuration from disk; falls back to defaults."""
+    p = Path(path)
+    if not p.exists():
+        return dict(DEFAULT_CONFIG)
+    with open(p) as f:
+        user_cfg = json.load(f)
+    merged = dict(DEFAULT_CONFIG)
+    merged.update(user_cfg)
+    return merged
+
+
+def save(cfg: dict, path: str = _DEFAULT_CONFIG_PATH) -> None:
+    """Persist configuration to disk."""
+    p = Path(path)
+    p.parent.mkdir(parents=True, exist_ok=True)
+    with open(p, "w") as f:
+        json.dump(cfg, f, indent=2)
+    print(f"[config] Saved to {p}")
+
+
+def get(key: str, default=None):
+    """Get a single config value."""
+    return load().get(key, default)

umbrellm/device.py

@@ -0,0 +1,30 @@
+"""
+Device utilities for Umbrellm.
+"""
+
+
+def current() -> str:
+    """
+    Return the best available device string.
+
+    Returns:
+        "cuda", "mps", or "cpu"
+    """
+    try:
+        import torch
+        if torch.cuda.is_available():
+            return "cuda"
+        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+            return "mps"
+        return "cpu"
+    except ImportError:
+        return "cpu"
+
+
+def torch_device():
+    """Return a torch.device object for the best available device."""
+    try:
+        import torch
+        return torch.device(current())
+    except ImportError:
+        raise RuntimeError("PyTorch is required. Install with: pip install torch")

umbrellm/interaction.py

@@ -0,0 +1,60 @@
+"""
+Chat and streaming interaction interface.
+"""
+
+import time
+from typing import Optional, List, Iterator, Union
+
+_RUNTIME = None
+
+
+def _get_runtime():
+    global _RUNTIME
+    if _RUNTIME is None:
+        from umbrellm._runtime import InferenceRuntime
+        _RUNTIME = InferenceRuntime()
+    return _RUNTIME
+
+
+def chat(params: dict) -> dict:
+    """
+    Send a chat message and return a response dict.
+
+    Required keys:
+        user (str): The user message.
+
+    Optional keys:
+        model (str): Model name or checkpoint path.
+        history (list): List of {"user": ..., "assistant": ...} dicts.
+        system_prompt (str): System instruction.
+        temperature (float): Sampling temperature.
+        top_k (int): Top-K sampling.
+        top_p (float): Nucleus sampling threshold.
+        repetition_penalty (float): Repetition penalty.
+        max_tokens (int): Maximum new tokens.
+        stop_sequences (list[str]): Stop strings.
+        stream (bool): Enable streaming (returns generator).
+        seed (int): RNG seed.
+
+    Returns:
+        {
+            "response": str,
+            "tokens_generated": int,
+            "generation_time": float,
+            "model": str
+        }
+    """
+    rt = _get_runtime()
+    return rt.chat(params)
+
+
+def stream_chat(params: dict) -> Iterator[str]:
+    """
+    Stream tokens from a chat response.
+
+    Example:
+        for token in umbrellm.interaction.stream_chat({"user": "Hello"}):
+            print(token, end="", flush=True)
+    """
+    rt = _get_runtime()
+    yield from rt.stream_chat(params)