ommlds 0.0.0.dev472__py3-none-any.whl → 0.0.0.dev473__py3-none-any.whl

ommlds/.omlish-manifests.json

@@ -468,7 +468,7 @@
       "!.minichain.registries.manifests.RegistryManifest": {
         "module": "ommlds.minichain.backends.impls.tinygrad.chat",
         "attr": "TinygradLlama3ChatChoicesService",
-        "name": "tinygrad_llama3",
+        "name": "tinygrad-llama3",
         "aliases": null,
         "type": "ChatChoicesService"
       }
@@ -483,7 +483,7 @@
       "!.minichain.registries.manifests.RegistryManifest": {
         "module": "ommlds.minichain.backends.impls.tinygrad.chat",
         "attr": "TinygradLlama3ChatChoicesStreamService",
-        "name": "tinygrad_llama3",
+        "name": "tinygrad-llama3",
         "aliases": null,
         "type": "ChatChoicesStreamService"
       }
@@ -500,7 +500,7 @@
           "ChatChoicesService",
           "ChatChoicesStreamService"
         ],
-        "backend_name": "tinygrad_llama3",
+        "backend_name": "tinygrad-llama3",
         "model_names": null
       }
     }

ommlds/__about__.py

@@ -41,6 +41,10 @@ class Project(ProjectBase):
             'datasets ~= 4.3',
         ],
 
+        'nanochat': [
+            'regex >= 2025.0',
+        ],
+
         'numpy': [
             'numpy >= 1.26',
         ],
@@ -76,6 +80,8 @@ class Project(ProjectBase):
 
 
 class Setuptools(SetuptoolsBase):
+    rs = True
+
     find_packages = {
         'include': [Project.name, f'{Project.name}.*'],
         'exclude': [*SetuptoolsBase.find_packages['exclude']],

ommlds/cli/main.py

@@ -191,7 +191,7 @@ def _main(args: ta.Any = None) -> None:
             _a_main,
             args,
         ),
-        backend='asyncio',
+        # backend='trio',
     )  # noqa
 
 

ommlds/cli/sessions/chat/backends/inject.py

@@ -2,6 +2,7 @@ import typing as ta
 
 from omlish import inject as inj
 from omlish import lang
+from omlish import typedvalues as tv
 
 from ..... import minichain as mc
 from .injection import backend_configs
@@ -41,7 +42,7 @@ def bind_backends(
         async def inner(be: 'mc.BackendCatalog.Backend', cfgs: _types.BackendConfigs | None) -> ta.Any:
             kwt = inj.build_kwargs_target(be.factory, non_strict=True)
             kw = await injector.provide_kwargs(kwt)
-            return be.factory(*cfgs or [], **kw)
+            return be.factory(*tv.collect(*(be.configs or []), *(cfgs or []), override=True), **kw)
 
         return _catalog.CatalogBackendProvider.Instantiator(inner)
 

ommlds/minichain/backends/impls/anthropic/stream.py

@@ -52,7 +52,7 @@ class AnthropicChatChoicesStreamService:
             self._model_name = cc.pop(AnthropicChatChoicesService.DEFAULT_MODEL_NAME)
             self._api_key = check.not_none(ApiKey.pop_secret(cc, env='ANTHROPIC_API_KEY'))
 
-    READ_CHUNK_SIZE = 64 * 1024
+    READ_CHUNK_SIZE: ta.ClassVar[int] = -1
 
     async def invoke(
             self,

ommlds/minichain/backends/impls/google/stream.py

@@ -125,7 +125,7 @@ class GoogleChatChoicesStreamService:
         AiMessage: 'model',
     }
 
-    READ_CHUNK_SIZE = 64 * 1024
+    READ_CHUNK_SIZE: ta.ClassVar[int] = -1
 
     async def invoke(
             self,

ommlds/minichain/backends/impls/ollama/chat.py

@@ -146,7 +146,7 @@ class OllamaChatChoicesService(BaseOllamaChatChoicesService):
 # )
 @static_check_is_chat_choices_stream_service
 class OllamaChatChoicesStreamService(BaseOllamaChatChoicesService):
-    READ_CHUNK_SIZE = 64 * 1024
+    READ_CHUNK_SIZE: ta.ClassVar[int] = -1
 
     async def invoke(
             self,

ommlds/minichain/backends/impls/openai/stream.py

@@ -54,7 +54,7 @@ class OpenaiChatChoicesStreamService:
             self._model_name = cc.pop(OpenaiChatChoicesService.DEFAULT_MODEL_NAME)
             self._api_key = ApiKey.pop_secret(cc, env='OPENAI_API_KEY')
 
-    READ_CHUNK_SIZE = 64 * 1024
+    READ_CHUNK_SIZE: ta.ClassVar[int] = -1
 
     async def invoke(self, request: ChatChoicesStreamRequest) -> ChatChoicesStreamResponse:
         # check.isinstance(request, ChatRequest)

ommlds/minichain/backends/impls/tinygrad/chat.py

@@ -113,7 +113,7 @@ class BaseTinygradLlama3ChatService(lang.ExitStacked, lang.Abstract):
 
 
 # @omlish-manifest $.minichain.registries.manifests.RegistryManifest(
-#     name='tinygrad_llama3',
+#     name='tinygrad-llama3',
 #     type='ChatChoicesService',
 # )
 @static_check_is_chat_choices_service
@@ -133,7 +133,7 @@ class TinygradLlama3ChatChoicesService(BaseTinygradLlama3ChatService):
 
 
 # @omlish-manifest $.minichain.registries.manifests.RegistryManifest(
-#     name='tinygrad_llama3',
+#     name='tinygrad-llama3',
 #     type='ChatChoicesStreamService',
 # )
 @static_check_is_chat_choices_stream_service
@@ -168,5 +168,5 @@ class TinygradLlama3ChatChoicesStreamService(BaseTinygradLlama3ChatService):
 #         'ChatChoicesService',
 #         'ChatChoicesStreamService',
 #     ],
-#     'tinygrad_llama3',
+#     'tinygrad-llama3',
 # )

ommlds/minichain/backends/strings/parsing.py

@@ -56,7 +56,7 @@ _REPO_MODEL_PAT = re.compile(
 def parse_backend_string(s: str) -> ParsedBackendString:
     backend: str | None
     if ':' in s:
-        backend, s = s.split(':')
+        backend, _, s = s.partition(':')
     else:
         backend = None
 

ommlds/minichain/backends/strings/resolving.py

@@ -108,7 +108,10 @@ class ManifestBackendStringResolver(BackendStringResolver):
 
         mn: str | None = mdl.name
 
-        if mn == m.backend_name:
+        if args.parsed.backend == m.backend_name and mn is not None:
+            pass
+
+        elif mn == m.backend_name:
             if m.model_names is not None:
                 mn = m.model_names.resolved_default
             else:

ommlds/minichain/stream/services.py

@@ -125,17 +125,18 @@ class _StreamServiceResponse(StreamResponseIterator[V, OutputT]):
             return
         if self._cr.cr_running or self._cr.cr_suspended:
             cex = StreamServiceCancelledError()
-            for i in itertools.count():
+            i = None
+            for n in itertools.count():
                 try:
-                    if not i:
+                    if not n:
                         x = self._g.throw(cex)
                     else:
-                        x = self._g.send(None)
+                        x = self._g.send(i)
                 except StreamServiceCancelledError as cex2:
                     if cex2 is cex:
                         break
                     raise
-                yield x
+                i = yield x
         if self._cr.cr_running:
             raise RuntimeError(f'Coroutine {self._cr!r} not terminated')
         if self._g is not self._a:
@@ -155,9 +156,10 @@ class _StreamServiceResponse(StreamResponseIterator[V, OutputT]):
     @types.coroutine
     def _anext(self):
         check.state(self._state == 'running')
+        i = None
         while True:
             try:
-                x = self._g.send(None)
+                x = self._g.send(i)
             except StopIteration as e:
                 if e.value is not None:
                     self._outputs = tv.TypedValues(*check.isinstance(e.value, ta.Sequence))
@@ -170,7 +172,7 @@ class _StreamServiceResponse(StreamResponseIterator[V, OutputT]):
                 x.done = True
                 return x.value
 
-            yield x
+            i = yield x
 
     async def __anext__(self) -> V:
         return await self._anext()

ommlds/nanochat/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Andrej Karpathy
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

ommlds/nanochat/rustbpe/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Andrej Karpathy
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

ommlds/nanochat/tokenizers.py

@@ -0,0 +1,406 @@
+# https://github.com/karpathy/nanochat/tree/9467d83cf23dcc9a9b4ca6e35103142f48a55b27
+"""
+BPE Tokenizer in the style of GPT-4.
+
+Two implementations are available:
+1) HuggingFace Tokenizer that can do both training and inference but is really confusing
+2) Our own RustBPE Tokenizer for training and tiktoken for efficient inference
+"""
+import copy
+import os
+import pickle
+import typing as ta
+
+from omlish import check
+from omlish import collections as col
+from omlish import lang
+
+
+with lang.auto_proxy_import(globals()):
+    import tiktoken
+    import tokenizers
+
+
+rustbpe: ta.Any = lang.proxy_import('.rustbpe', __package__)
+
+
+##
+
+
+SPECIAL_TOKENS = [
+    # every document begins with the Beginning of Sequence (BOS) token that delimits documents
+    '<|bos|>',
+    # tokens below are only used during finetuning to render Conversations into token ids
+    '<|user_start|>',  # user messages
+    '<|user_end|>',
+    '<|assistant_start|>',  # assistant messages
+    '<|assistant_end|>',
+    '<|python_start|>',  # assistant invokes python REPL tool
+    '<|python_end|>',
+    '<|output_start|>',  # python REPL outputs back to assistant
+    '<|output_end|>',
+]
+
+
+# NOTE: this split pattern deviates from GPT-4 in that we use \p{N}{1,2} instead of \p{N}{1,3}
+# I did this because I didn't want to "waste" too many tokens on numbers for smaller vocab sizes.
+# I haven't validated that this is actually a good idea, TODO.
+SPLIT_PATTERN = r"""'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}{1,2}| ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+"""  # noqa
+
+
+# -----------------------------------------------------------------------------
+# Generic GPT-4-style tokenizer based on HuggingFace Tokenizer
+
+
+class HuggingFaceTokenizer:
+    """Light wrapper around HuggingFace Tokenizer for some utilities"""
+
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+
+    @classmethod
+    def from_pretrained(cls, hf_path):
+        # init from a HuggingFace pretrained tokenizer (e.g. "gpt2")
+        tokenizer = tokenizers.Tokenizer.from_pretrained(hf_path)
+        return cls(tokenizer)
+
+    @classmethod
+    def from_directory(cls, tokenizer_dir):
+        # init from a local directory on disk (e.g. "out/tokenizer")
+        tokenizer_path = os.path.join(tokenizer_dir, 'tokenizer.json')
+        tokenizer = tokenizers.Tokenizer.from_file(tokenizer_path)
+        return cls(tokenizer)
+
+    @classmethod
+    def train_from_iterator(
+            cls,
+            text_iterator,
+            vocab_size,
+            *,
+            split_pattern=SPLIT_PATTERN,
+            special_tokens=SPECIAL_TOKENS,
+    ):
+        # train from an iterator of text
+        # Configure the HuggingFace Tokenizer
+        tokenizer = tokenizers.Tokenizer(tokenizers.models.BPE(
+            byte_fallback=True,  # needed!
+            unk_token=None,
+            fuse_unk=False,
+        ))
+        # Normalizer: None
+        tokenizer.normalizer = None
+        # Pre-tokenizer: GPT-4 style
+        # the regex pattern used by GPT-4 to split text into groups before BPE
+        # NOTE: The pattern was changed from \p{N}{1,3} to \p{N}{1,2} because I suspect it is harmful to
+        # very small models and smaller vocab sizes, because it is a little bit wasteful in the token space.
+        # (but I haven't validated this! TODO)
+        gpt4_split_regex = tokenizers.Regex(split_pattern)  # huggingface demands that you wrap it in Regex!!
+        tokenizer.pre_tokenizer = tokenizers.pre_tokenizers.Sequence([
+            tokenizers.pre_tokenizers.Split(pattern=gpt4_split_regex, behavior='isolated', invert=False),
+            tokenizers.pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False),
+        ])
+        # Decoder: ByteLevel (it pairs together with the ByteLevel pre-tokenizer)
+        tokenizer.decoder = tokenizers.decoders.ByteLevel()
+        # Post-processor: None
+        tokenizer.post_processor = None
+        # Trainer: BPE
+        trainer = tokenizers.trainers.BpeTrainer(
+            vocab_size=vocab_size,
+            show_progress=True,
+            min_frequency=0,  # no minimum frequency
+            initial_alphabet=tokenizers.pre_tokenizers.ByteLevel.alphabet(),
+            special_tokens=special_tokens,
+        )
+        # Kick off the training
+        tokenizer.train_from_iterator(text_iterator, trainer)
+        return cls(tokenizer)
+
+    def encode_ordinary(self, text):
+        ids = self.tokenizer.encode(text, add_special_tokens=False).ids
+        return ids
+
+    def get_vocab_size(self):
+        return self.tokenizer.get_vocab_size()
+
+    def get_special_tokens(self):
+        special_tokens_map = self.tokenizer.get_added_tokens_decoder()
+        special_tokens = [w.content for w in special_tokens_map.values()]
+        return special_tokens
+
+    def id_to_token(self, id):  # noqa
+        return self.tokenizer.id_to_token(id)
+
+    def _encode_one(self, text, prepend=None, append=None):
+        # encode a single string
+        # prepend/append can be either a string of a special token or a token id directly.
+        check.isinstance(text, str)
+        ids = []
+        if prepend is not None:
+            prepend_id = prepend if isinstance(prepend, int) else self.encode_special(prepend)
+            ids.append(prepend_id)
+        ids.extend(self.tokenizer.encode(text, add_special_tokens=False).ids)
+        if append is not None:
+            append_id = append if isinstance(append, int) else self.encode_special(append)
+            ids.append(append_id)
+        return ids
+
+    def encode_special(self, text):
+        # encode a single special token via exact match
+        return self.tokenizer.token_to_id(text)
+
+    def get_bos_token_id(self):
+        bos = self.encode_special('<|bos|>')
+        return bos
+
+    def encode(self, text, *args, **kwargs):
+        if isinstance(text, str):
+            return self._encode_one(text, *args, **kwargs)
+        elif isinstance(text, list):
+            return [self._encode_one(t, *args, **kwargs) for t in text]
+        else:
+            raise ValueError(f'Invalid input type: {type(text)}')  # noqa
+
+    def __call__(self, *args, **kwargs):
+        return self.encode(*args, **kwargs)
+
+    def decode(self, ids):
+        return self.tokenizer.decode(ids, skip_special_tokens=False)
+
+    def save(self, tokenizer_dir):
+        # save the tokenizer to disk
+        os.makedirs(tokenizer_dir, exist_ok=True)
+        tokenizer_path = os.path.join(tokenizer_dir, 'tokenizer.json')
+        self.tokenizer.save(tokenizer_path)
+        print(f'Saved tokenizer to {tokenizer_path}')
+
+
+# -----------------------------------------------------------------------------
+# Tokenizer based on rustbpe + tiktoken combo
+
+
+class RustBPETokenizer:
+    """Light wrapper around tiktoken (for efficient inference) but train with rustbpe"""
+
+    def __init__(self, enc, bos_token):
+        self.enc = enc
+        self.bos_token_id = self.encode_special(bos_token)
+
+    @classmethod
+    def train_from_iterator(cls, text_iterator, vocab_size):
+        # 1) train using rustbpe
+        tokenizer = rustbpe.Tokenizer()
+        # the special tokens are inserted later in __init__, we don't train them here
+        vocab_size_no_special = vocab_size - len(SPECIAL_TOKENS)
+        check.state(vocab_size_no_special >= 256, f'vocab_size_no_special must be at least 256, got {vocab_size_no_special}')  # noqa
+        tokenizer.train_from_iterator(text_iterator, vocab_size_no_special, pattern=SPLIT_PATTERN)
+        # 2) construct the associated tiktoken encoding for inference
+        pattern = tokenizer.get_pattern()
+        mergeable_ranks_list = tokenizer.get_mergeable_ranks()
+        mergeable_ranks = {bytes(k): v for k, v in mergeable_ranks_list}
+        tokens_offset = len(mergeable_ranks)
+        special_tokens = {name: tokens_offset + i for i, name in enumerate(SPECIAL_TOKENS)}
+        enc = tiktoken.Encoding(
+            name='rustbpe',
+            pat_str=pattern,
+            mergeable_ranks=mergeable_ranks,  # dict[bytes, int] (token bytes -> merge priority rank)
+            special_tokens=special_tokens,  # dict[str, int] (special token name -> token id)
+        )
+        return cls(enc, '<|bos|>')
+
+    @classmethod
+    def from_directory(cls, tokenizer_dir):
+        pickle_path = os.path.join(tokenizer_dir, 'tokenizer.pkl')
+        with open(pickle_path, 'rb') as f:
+            enc = pickle.load(f)  # noqa
+        return cls(enc, '<|bos|>')
+
+    @classmethod
+    def from_pretrained(cls, tiktoken_name):
+        # https://github.com/openai/tiktoken/blob/eedc8563/tiktoken_ext/openai_public.py
+        enc = tiktoken.get_encoding(tiktoken_name)
+        # tiktoken calls the special document delimiter token "<|endoftext|>"
+        # yes this is confusing because this token is almost always PREPENDED to the beginning of the document
+        # it most often is used to signal the start of a new sequence to the LLM during inference etc.
+        # so in nanoChat we always use "<|bos|>" short for "beginning of sequence", but historically it is often called
+        # "<|endoftext|>".
+        return cls(enc, '<|endoftext|>')
+
+    def get_vocab_size(self):
+        return self.enc.n_vocab
+
+    def get_special_tokens(self):
+        return self.enc.special_tokens_set
+
+    def id_to_token(self, id):  # noqa
+        return self.enc.decode([id])
+
+    @col.cache.cache(max_size=32)
+    def encode_special(self, text):
+        return self.enc.encode_single_token(text)
+
+    def get_bos_token_id(self):
+        return self.bos_token_id
+
+    def encode(self, text, prepend=None, append=None, num_threads=8):
+        # text can be either a string or a list of strings
+
+        if prepend is not None:
+            prepend_id = prepend if isinstance(prepend, int) else self.encode_special(prepend)
+        if append is not None:
+            append_id = append if isinstance(append, int) else self.encode_special(append)
+
+        if isinstance(text, str):
+            ids = self.enc.encode_ordinary(text)
+            if prepend is not None:
+                ids.insert(0, prepend_id)  # TODO: slightly inefficient here? :( hmm
+            if append is not None:
+                ids.append(append_id)
+        elif isinstance(text, list):
+            ids = self.enc.encode_ordinary_batch(text, num_threads=num_threads)
+            if prepend is not None:
+                for ids_row in ids:
+                    ids_row.insert(0, prepend_id)  # TODO: same
+            if append is not None:
+                for ids_row in ids:
+                    ids_row.append(append_id)
+        else:
+            raise ValueError(f'Invalid input type: {type(text)}')  # noqa
+
+        return ids
+
+    def __call__(self, *args, **kwargs):
+        return self.encode(*args, **kwargs)
+
+    def decode(self, ids):
+        return self.enc.decode(ids)
+
+    def save(self, tokenizer_dir):
+        # save the encoding object to disk
+        os.makedirs(tokenizer_dir, exist_ok=True)
+        pickle_path = os.path.join(tokenizer_dir, 'tokenizer.pkl')
+        with open(pickle_path, 'wb') as f:
+            pickle.dump(self.enc, f)
+        print(f'Saved tokenizer encoding to {pickle_path}')
+
+    def render_conversation(self, conversation, max_tokens=2048):
+        """
+        Tokenize a single Chat conversation (which we call a "doc" or "document" here).
+        Returns:
+        - ids: list[int] is a list of token ids of this rendered conversation
+        - mask: list[int] of same length, mask = 1 for tokens that the Assistant is expected to train on.
+        """
+
+        # ids, masks that we will return and a helper function to help build them up.
+        ids, mask = [], []
+
+        def add_tokens(token_ids, mask_val):
+            if isinstance(token_ids, int):
+                token_ids = [token_ids]
+            ids.extend(token_ids)
+            mask.extend([mask_val] * len(token_ids))
+
+        # sometimes the first message is a system message...
+        # => just merge it with the second (user) message
+        if conversation['messages'][0]['role'] == 'system':
+            # some conversation surgery is necessary here for now...
+            conversation = copy.deepcopy(conversation)  # avoid mutating the original
+            messages = conversation['messages']
+            check.state(messages[1]['role'] == 'user', 'System message must be followed by a user message')
+            messages[1]['content'] = messages[0]['content'] + '\n\n' + messages[1]['content']
+            messages = messages[1:]
+        else:
+            messages = conversation['messages']
+        check.state(len(messages) >= 1, f'Conversation has less than 1 message: {messages}')
+
+        # fetch all the special tokens we need
+        bos = self.get_bos_token_id()
+        user_start, user_end = self.encode_special('<|user_start|>'), self.encode_special('<|user_end|>')
+        assistant_start, assistant_end = self.encode_special('<|assistant_start|>'), self.encode_special('<|assistant_end|>')  # noqa
+        python_start, python_end = self.encode_special('<|python_start|>'), self.encode_special('<|python_end|>')
+        output_start, output_end = self.encode_special('<|output_start|>'), self.encode_special('<|output_end|>')
+
+        # now we can tokenize the conversation
+        add_tokens(bos, 0)
+        for i, message in enumerate(messages):
+            # some sanity checking here around assumptions, to prevent footguns
+            must_be_from = 'user' if i % 2 == 0 else 'assistant'
+            check.state(message['role'] == must_be_from, f"Message {i} is from {message['role']} but should be from {must_be_from}")  # noqa
+
+            # content can be either a simple string or a list of parts (e.g. containing tool calls)
+            content = message['content']
+
+            if message['role'] == 'user':
+                check.isinstance(content, str), 'User messages are simply expected to be strings'
+                value_ids = self.encode(content)
+                add_tokens(user_start, 0)
+                add_tokens(value_ids, 0)
+                add_tokens(user_end, 0)
+            elif message['role'] == 'assistant':
+                add_tokens(assistant_start, 0)
+                if isinstance(content, str):
+                    # simple string => simply add the tokens
+                    value_ids = self.encode(content)
+                    add_tokens(value_ids, 1)
+                elif isinstance(content, list):
+                    for part in content:
+                        value_ids = self.encode(part['text'])
+                        if part['type'] == 'text':
+                            # string part => simply add the tokens
+                            add_tokens(value_ids, 1)
+                        elif part['type'] == 'python':
+                            # python tool call => add the tokens inside <|python_start|> and <|python_end|>
+                            add_tokens(python_start, 1)
+                            add_tokens(value_ids, 1)
+                            add_tokens(python_end, 1)
+                        elif part['type'] == 'python_output':
+                            # python output => add the tokens inside <|output_start|> and <|output_end|>
+                            # none of these tokens are supervised because the tokens come from Python at test time
+                            add_tokens(output_start, 0)
+                            add_tokens(value_ids, 0)
+                            add_tokens(output_end, 0)
+                        else:
+                            raise ValueError(f"Unknown part type: {part['type']}")
+                else:
+                    raise ValueError(f'Unknown content type: {type(content)}')
+                add_tokens(assistant_end, 1)
+
+        # truncate to max_tokens tokens MAX (helps prevent OOMs)
+        ids = ids[:max_tokens]
+        mask = mask[:max_tokens]
+        return ids, mask
+
+    def visualize_tokenization(self, ids, mask, with_token_id=False):
+        """Small helper function useful in debugging: visualize the tokenization of render_conversation"""
+
+        red = '\033[91m'
+        green = '\033[92m'
+        reset = '\033[0m'
+        gray = '\033[90m'
+        tokens = []
+        for i, (token_id, mask_val) in enumerate(zip(ids, mask)):  # noqa
+            token_str = self.decode([token_id])
+            color = green if mask_val == 1 else red
+            tokens.append(f'{color}{token_str}{reset}')
+            if with_token_id:
+                tokens.append(f'{gray}({token_id}){reset}')
+        return '|'.join(tokens)
+
+    def render_for_completion(self, conversation):
+        """
+        Used during Reinforcement Learning. In that setting, we want to render the conversation priming the Assistant
+        for a completion. Unlike the Chat SFT case, we don't need to return the mask.
+        """
+
+        # We have some surgery to do: we need to pop the last message (of the Assistant)
+        conversation = copy.deepcopy(conversation)  # avoid mutating the original
+        messages = conversation['messages']
+        check.state(messages[-1]['role'] == 'assistant', 'Last message must be from the Assistant')
+        messages.pop()  # remove the last message (of the Assistant) inplace
+
+        # Now tokenize the conversation
+        ids, mask = self.render_conversation(conversation)
+
+        # Finally, to prime the Assistant for a completion, append the Assistant start token
+        assistant_start = self.encode_special('<|assistant_start|>')
+        ids.append(assistant_start)
+        return ids

{ommlds-0.0.0.dev472.dist-info → ommlds-0.0.0.dev473.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ommlds
-Version: 0.0.0.dev472
+Version: 0.0.0.dev473
 Summary: ommlds
 Author: wrmsr
 License-Expression: BSD-3-Clause
@@ -14,8 +14,8 @@ Classifier: Programming Language :: Python :: 3.13
 Requires-Python: >=3.13
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: omdev==0.0.0.dev472
-Requires-Dist: omlish==0.0.0.dev472
+Requires-Dist: omdev==0.0.0.dev473
+Requires-Dist: omlish==0.0.0.dev473
 Provides-Extra: all
 Requires-Dist: llama-cpp-python~=0.3; extra == "all"
 Requires-Dist: mlx~=0.29; extra == "all"
@@ -28,6 +28,7 @@ Requires-Dist: transformers~=4.57; extra == "all"
 Requires-Dist: sentence-transformers~=5.1; extra == "all"
 Requires-Dist: huggingface-hub~=0.36; extra == "all"
 Requires-Dist: datasets~=4.3; extra == "all"
+Requires-Dist: regex>=2025.0; extra == "all"
 Requires-Dist: numpy>=1.26; extra == "all"
 Requires-Dist: pytesseract~=0.3; extra == "all"
 Requires-Dist: rapidocr-onnxruntime~=1.4; extra == "all"
@@ -49,6 +50,8 @@ Requires-Dist: sentence-transformers~=5.1; extra == "backends"
 Provides-Extra: huggingface
 Requires-Dist: huggingface-hub~=0.36; extra == "huggingface"
 Requires-Dist: datasets~=4.3; extra == "huggingface"
+Provides-Extra: nanochat
+Requires-Dist: regex>=2025.0; extra == "nanochat"
 Provides-Extra: numpy
 Requires-Dist: numpy>=1.26; extra == "numpy"
 Provides-Extra: ocr

{ommlds-0.0.0.dev472.dist-info → ommlds-0.0.0.dev473.dist-info}/RECORD

@@ -1,5 +1,5 @@
-ommlds/.omlish-manifests.json,sha256=nVhqSv9iQ2zu3MpOtJC2mFH5ORyzIC2K8kCAwDEdzGE,21555
-ommlds/__about__.py,sha256=t2rQF0yXpWFcCb2dvgzGR3I35HKGvGSn-EfhaUWVl5s,1759
+ommlds/.omlish-manifests.json,sha256=jsNWNqNQTjpxr-irqzy0dWpqeZH9CsD0SFY3OgIGuZ4,21555
+ommlds/__about__.py,sha256=v41mIGxCcZVApXGD-CSqkYDT9wXorPjnU611ejWJjaQ,1839
 ommlds/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ommlds/_hacks/__init__.py,sha256=ajfw7dMKH8UuloeQ5MSxWwgAmdWf2v8gm-K3uLP9wtY,196
 ommlds/_hacks/funcs.py,sha256=8XseIblP7yolDUD7WQSGn1LP90IQzByVejSzphAPDyM,2861
@@ -86,7 +86,7 @@ ommlds/backends/transformers/streamers.py,sha256=Hu_9lp_kUilKjOfs7Ixqr2NoA5FuRn2
 ommlds/cli/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ommlds/cli/__main__.py,sha256=1ffCb0fcUOJMzxROJmJRXQ8PSOVYv7KrcuBtT95cf0c,140
 ommlds/cli/inject.py,sha256=WhTDabJz9b1NRRHVH-UyVN5nj6UncvIeTvgkGrcE9vc,666
-ommlds/cli/main.py,sha256=dTBJSxWdsS1pEqDY-vT_g2PC-aUkatVhq0XEW8zirQ0,5802
+ommlds/cli/main.py,sha256=-cMFxZqK4XTHxaFxDk4fvtWLB6LZ6UGlmbN74L2ir-g,5801
 ommlds/cli/backends/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ommlds/cli/backends/inject.py,sha256=OVstNsoeVnprM9PBL_zP0N46KkoDg3_Wz90BWcQ7km4,1734
 ommlds/cli/backends/standard.py,sha256=HnammWyAXJHeqXJrAMBdarcT4Nyt2CxudZdD2fW_Y9M,631
@@ -100,7 +100,7 @@ ommlds/cli/sessions/chat/inject.py,sha256=7Yg6wUs2Oej4UjNZCAWCJCEsDJZWvT4G8XvkvV
 ommlds/cli/sessions/chat/session.py,sha256=eqwelLE74JFC-fBpk_hdwMD2nP4pLv3ZPwUn99200B8,521
 ommlds/cli/sessions/chat/backends/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ommlds/cli/sessions/chat/backends/catalog.py,sha256=hIY0L1zewuJX0_xxcMcy4gylSLiQENB3YxgYJEoKgrU,2109
-ommlds/cli/sessions/chat/backends/inject.py,sha256=v_kw6lOMo5XoAf2dyIld2oBjg7lVbS2ndmdPQv4F464,1558
+ommlds/cli/sessions/chat/backends/inject.py,sha256=er03V-84B-wWt86o2HXXqCb6uFRKbfXfUZsB6GrzYOA,1646
 ommlds/cli/sessions/chat/backends/injection.py,sha256=GCn5OvNIEowgB70kQVuU84z3i8lLA4vOVkTZlQG8s0o,327
 ommlds/cli/sessions/chat/backends/types.py,sha256=5eImYHXLKqbC5MDrN443eMGamP9snCmV1n7LtAsqgPk,696
 ommlds/cli/sessions/chat/chat/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -180,14 +180,14 @@ ommlds/minichain/backends/impls/anthropic/__init__.py,sha256=47DEQpj8HBSa-_TImW-
 ommlds/minichain/backends/impls/anthropic/chat.py,sha256=-qGr_DZgGe-dr1AKb6WLtCq_I2E9635X1rQZSJqOb04,4318
 ommlds/minichain/backends/impls/anthropic/names.py,sha256=GPPeYt0CcDcDCR8I6BMd7bMjC_Zk_bjnLLpF9ClwXcg,1099
 ommlds/minichain/backends/impls/anthropic/protocol.py,sha256=whPVYuKShKiMCzasHl77sCIiymhzXj8mFZXEyhZvld8,3292
-ommlds/minichain/backends/impls/anthropic/stream.py,sha256=ePgAz1QyuUcqSLiJTmP2IeW7jZagOMYceWmnvTTHBa8,8779
+ommlds/minichain/backends/impls/anthropic/stream.py,sha256=NNBFb0sMId9yWua3fkAMZ-qYhQN9nLrXiO4DViR77YI,8790
 ommlds/minichain/backends/impls/duckduckgo/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ommlds/minichain/backends/impls/duckduckgo/search.py,sha256=igzeU9P9b1MMiu4KAJVS9H6KLIoPm68wXi4Kx3_DHyQ,940
 ommlds/minichain/backends/impls/google/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ommlds/minichain/backends/impls/google/chat.py,sha256=lGb5blGLlcBlt9xeDZJvbh5SlV7fgfezd5_As_SPBXo,6499
 ommlds/minichain/backends/impls/google/names.py,sha256=HxHJ31HeKZg6aW1C_Anqp-gamCXpq9pOdKj8_yVgE8Y,871
 ommlds/minichain/backends/impls/google/search.py,sha256=y5_6seSRU8CFnLA_Ja8XEMbIBWSgwBzE1iBf-qyz0tA,3427
-ommlds/minichain/backends/impls/google/stream.py,sha256=lAxPZkFLkeBB8rNvnv5wEvK6aPqAfyAssukMA8QQz7s,8002
+ommlds/minichain/backends/impls/google/stream.py,sha256=ITownhKSOJB4IG23wWZJepUImSM6vJsDMOM9W1STpwU,8013
 ommlds/minichain/backends/impls/google/tools.py,sha256=Tty0gsyx7-PbeoNqMuql_ewQ6q-ZsDaDdsD5ShinGVY,5089
 ommlds/minichain/backends/impls/huggingface/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ommlds/minichain/backends/impls/huggingface/configs.py,sha256=6jsBtPNXOP57PcpxNTVLGWLc-18Iwn_lDbGouwCJTIQ,258
@@ -200,18 +200,18 @@ ommlds/minichain/backends/impls/llamacpp/stream.py,sha256=uzrXr2HhshgFe3Z0g8KTPc
 ommlds/minichain/backends/impls/mlx/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ommlds/minichain/backends/impls/mlx/chat.py,sha256=sMlhgiFZrxAC-kKkLSJ6c-2uJn0IHZXH4EiPET_-CKI,7458
 ommlds/minichain/backends/impls/ollama/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ommlds/minichain/backends/impls/ollama/chat.py,sha256=3fuIAsIW20aEAOLLfM21d5ju27igr6N-3Lf9nNUWcoY,6598
+ommlds/minichain/backends/impls/ollama/chat.py,sha256=agnJcOwJGebSiV5TG0UmVFYGCc7hEpt7FM73rtyF8gk,6609
 ommlds/minichain/backends/impls/openai/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ommlds/minichain/backends/impls/openai/chat.py,sha256=QcQZO78p4UUzI4QU1K-057OEZGKIYxjXENhkifsSuaI,2841
 ommlds/minichain/backends/impls/openai/completion.py,sha256=4Mi4Zvrq5fCqUd0asL3WiCbCdmxOdo0NFkoZMfdsYXY,1939
 ommlds/minichain/backends/impls/openai/embedding.py,sha256=BNtvKYLTsnQwQR9Tv3Fr8zCYN1kr1UNdJ15lcsjz6X0,1765
 ommlds/minichain/backends/impls/openai/format.py,sha256=teGX8mNU3sXNWP4YWGD8d59M4X9_r75ImSzfTJgtNCM,7351
 ommlds/minichain/backends/impls/openai/names.py,sha256=b74t8FwSbGEveVtVz4SqM5tiRDyTKNlUKlseV6AX3Yo,1211
-ommlds/minichain/backends/impls/openai/stream.py,sha256=pm-iUNjw5o94LrnrbAWttfbpxStnq0vDvrE7FV9fdsM,5399
+ommlds/minichain/backends/impls/openai/stream.py,sha256=1kh_V_eu8QAY_i4ulfm22kbZeMiIDLwmDPJf7aaIikI,5410
 ommlds/minichain/backends/impls/sentencepiece/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ommlds/minichain/backends/impls/sentencepiece/tokens.py,sha256=tUEBKyBgkTowssS_AdcAuPkyFzfyDfE935x4JG8PXM0,1602
 ommlds/minichain/backends/impls/tinygrad/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ommlds/minichain/backends/impls/tinygrad/chat.py,sha256=EsxNflUhrc_ouFyIvpEhKxx-0DTZKRi-BJhyD7bAfmc,4916
+ommlds/minichain/backends/impls/tinygrad/chat.py,sha256=Y3Lp08Sb0YUPAxEciexOUm0uyoJnhbH5pWT9buclx6Y,4916
 ommlds/minichain/backends/impls/tokenizers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ommlds/minichain/backends/impls/tokenizers/tokens.py,sha256=_8Q49k5YroG5wQI0cuK6kOJ3XYwjhpaAS04ejhzBsWw,1500
 ommlds/minichain/backends/impls/transformers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -220,8 +220,8 @@ ommlds/minichain/backends/impls/transformers/tokens.py,sha256=uS3-IWOJRUMBfPDVRr
 ommlds/minichain/backends/impls/transformers/transformers.py,sha256=laM8G2SAE6jUjnHkeZsbWxS2KJF4efi-35aBlRBzIsE,9053
 ommlds/minichain/backends/strings/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ommlds/minichain/backends/strings/manifests.py,sha256=kmlanVUAZqIh0P95Mm8H20e8ib3gEgYHHUlkCXDQGFk,413
-ommlds/minichain/backends/strings/parsing.py,sha256=2wChk9Z8fhqJTk8_91f8QFjKcSZygOQM_rVk-P4NnKw,1772
-ommlds/minichain/backends/strings/resolving.py,sha256=CVn0RDnOnw1BNQFWbSfKaDxm2C6j3FT755nT3z_Q254,5760
+ommlds/minichain/backends/strings/parsing.py,sha256=Etmk04BnKvCMtGg4AgbvxsPGvfRcLldLxpdpxcozdNk,1779
+ommlds/minichain/backends/strings/resolving.py,sha256=q0qMdIvFZH-yScpXNX8GHE_yyQC_eEf1eptiUnNolUI,5849
 ommlds/minichain/chat/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ommlds/minichain/chat/_marshal.py,sha256=M3p093nxzxITbznc--P-tyCXuWDHrq4JFKTZAx6XWdk,740
 ommlds/minichain/chat/formats.py,sha256=LmlU7iu8PMJuroFTmyWfP4tXvLjj5VNxdAp1Us9MSAA,562
@@ -318,7 +318,7 @@ ommlds/minichain/services/requests.py,sha256=VAfKbYu4T0CZTWVQmZ2LUmYU7DNm6IerYMN
 ommlds/minichain/services/responses.py,sha256=4W6Z4Fx4_GFqKgle27OeLr0zzjVTA0pkZrlsZiFQNdo,1534
 ommlds/minichain/services/services.py,sha256=WjkQNYIp87SflLSReOHMkG2qIVAOem6vsrs_2NxWN_M,325
 ommlds/minichain/stream/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ommlds/minichain/stream/services.py,sha256=fFR1klP_PZJ3Pqmqx_SGap8gRDuthJah1fyoke6G9Ww,5328
+ommlds/minichain/stream/services.py,sha256=Mx0FmoEuXBRgwTJDfcTd14F4HdvpDOwuCCNWUvDqbKE,5368
 ommlds/minichain/stream/wrap.py,sha256=nQC0aCi49I18nF0Yx8qiiLkhIAECV6s6o4pvOy5Kx98,2041
 ommlds/minichain/text/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ommlds/minichain/text/applypatch.py,sha256=YIN5JChJ0FXyK1I6OiAHQmE7BT-exHfaAMM9ay7ylyc,17705
@@ -356,6 +356,10 @@ ommlds/minichain/vectors/search.py,sha256=27MTUiVT2xmSnmgJTAR09oQaiNRh1ixj0mGZVu
 ommlds/minichain/vectors/similarity.py,sha256=etqSswPH7ERThueqnCUHULsM3rpsVslRFua0m_ps_F4,1308
 ommlds/minichain/vectors/stores.py,sha256=etbLCS0RXAEmqcCdqiys8twa8R7Y_DcjQ_VqnEnRF4s,530
 ommlds/minichain/vectors/types.py,sha256=xSAK1Xfkubqf95QgJhSHrwBu_C5quuye3wZAASmxJkM,3473
+ommlds/nanochat/LICENSE,sha256=QrsJ8zmor4uQ07SWm29uS6Sv87XGFBA7Ax_M33HI93I,1072
+ommlds/nanochat/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ommlds/nanochat/tokenizers.py,sha256=cU6ld0qdMG1T41_ijRD8EsbFMLLCpSNLDjgQOBi6RdM,17502
+ommlds/nanochat/rustbpe/LICENSE,sha256=QrsJ8zmor4uQ07SWm29uS6Sv87XGFBA7Ax_M33HI93I,1072
 ommlds/server/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ommlds/server/__main__.py,sha256=morlItVl-0_MDK6xk2VKhqOtA8oQk0SoWOWEqcgqXTw,155
 ommlds/server/cli.py,sha256=gCN__45IXjCtk-tWwO2hr8vs5K-R0e1auNWdIc7d6_U,1825
@@ -377,9 +381,9 @@ ommlds/wiki/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU
 ommlds/wiki/utils/io.py,sha256=UKgDJGtmpnWvIqVd2mJc2QNPOqlToEY1GEveNp6_pMo,7088
 ommlds/wiki/utils/progress.py,sha256=EhvKcMFYtsarCQhIahlO6f0SboyAKP3UwUyrnVnP-Vk,3222
 ommlds/wiki/utils/xml.py,sha256=vVV8Ctn13aaRM9eYfs9Wd6rHn5WOCEUzQ44fIhOvJdg,3754
-ommlds-0.0.0.dev472.dist-info/licenses/LICENSE,sha256=B_hVtavaA8zCYDW99DYdcpDLKz1n3BBRjZrcbv8uG8c,1451
-ommlds-0.0.0.dev472.dist-info/METADATA,sha256=1pgQO5hxkugLK5ik5YuManJYKtc1SgPl96SZdHo1k04,3224
-ommlds-0.0.0.dev472.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ommlds-0.0.0.dev472.dist-info/entry_points.txt,sha256=Z5YWtX7ClfiCKdW-dd_CSVvM0h4yQpJPi-2G3q6gNFo,35
-ommlds-0.0.0.dev472.dist-info/top_level.txt,sha256=Rbnk5d5wi58vnAXx13WFZqdQ4VX8hBCS2hEL3WeXOhY,7
-ommlds-0.0.0.dev472.dist-info/RECORD,,
+ommlds-0.0.0.dev473.dist-info/licenses/LICENSE,sha256=B_hVtavaA8zCYDW99DYdcpDLKz1n3BBRjZrcbv8uG8c,1451
+ommlds-0.0.0.dev473.dist-info/METADATA,sha256=0l14jYSJstdmoZV9kjXuIwuF0AGrEufXx2Ixqm6-wJ4,3344
+ommlds-0.0.0.dev473.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ommlds-0.0.0.dev473.dist-info/entry_points.txt,sha256=Z5YWtX7ClfiCKdW-dd_CSVvM0h4yQpJPi-2G3q6gNFo,35
+ommlds-0.0.0.dev473.dist-info/top_level.txt,sha256=Rbnk5d5wi58vnAXx13WFZqdQ4VX8hBCS2hEL3WeXOhY,7
+ommlds-0.0.0.dev473.dist-info/RECORD,,