lalamo 0.5.7__py3-none-any.whl → 0.5.8__py3-none-any.whl

lalamo/__init__.py

@@ -8,24 +8,24 @@ from lalamo.message_processor import (
     ToolSchema,
     UserMessage,
 )
-from lalamo.model_import import ModelSpec
-from lalamo.models import LanguageModel, Router
+from lalamo.model_import import ModelSpec, import_model
+from lalamo.models import ClassifierModel, LanguageModel
 from lalamo.speculator import (
     CollectTracesEvent,
     SpeculatorTrainingEvent,
 )
 
-__version__ = "0.5.7"
+__version__ = "0.5.8"
 
 __all__ = [
     "AssistantMessage",
+    "ClassifierModel",
     "CollectTracesEvent",
     "ContentBlock",
     "Image",
     "LanguageModel",
     "Message",
     "ModelSpec",
-    "Router",
     "SpeculatorTrainingEvent",
     "SystemMessage",
     "ToolSchema",
@@ -33,5 +33,6 @@ __all__ = [
     "collect_traces",
     "convert",
     "estimate_batchsize",
+    "import_model",
     "train",
 ]

lalamo/main.py

@@ -43,7 +43,7 @@ from lalamo.model_import.common import (
     InitializingModelEvent,
     StatusEvent,
 )
-from lalamo.models import LanguageModelConfig, RouterConfig
+from lalamo.models import ClassifierModelConfig, LanguageModelConfig
 from lalamo.modules import config_converter
 from lalamo.speculator.estimator import EstimateBatchsizeFromMemoryEvent, estimate_batchsize_from_memory
 from lalamo.speculator.inference import CollectTracesEvent, inference_collect_traces
@@ -149,7 +149,7 @@ def chat(
         messages.append(model.message_processor.parse_response(model_response_text))
 
 
-@app.command(help="Classify given message with a Router type of model.")
+@app.command(help="Classify given message with a Classifier type of model.")
 def classify(
     model_path: Annotated[
         Path,
@@ -165,7 +165,7 @@ def classify(
         transient=True,
     ) as progress:
         loading_task = progress.add_task("🚀 [cyan]Loading model...[/cyan]")
-        model = RouterConfig.load_model(model_path)
+        model = ClassifierModelConfig.load_model(model_path)
         progress.remove_task(loading_task)
         warmup_task = progress.add_task("🔥 Warming up...")
         model.classify_chat([UserMessage(content="warmup message")])

lalamo/model_import/common.py

@@ -14,7 +14,7 @@ from jaxtyping import DTypeLike
 from tokenizers import Tokenizer
 
 from lalamo.message_processor import MessageProcessor, MessageProcessorConfig
-from lalamo.models import GenerationConfig, LanguageModel, LanguageModelConfig, Router, RouterConfig
+from lalamo.models import ClassifierModel, ClassifierModelConfig, GenerationConfig, LanguageModel, LanguageModelConfig
 from lalamo.modules import Classifier, Decoder, LalamoModule
 from lalamo.quantization import QuantizationMode
 
@@ -72,7 +72,8 @@ class ModelMetadata:
     repo: str
     use_cases: tuple[UseCase, ...]
     model_type: ModelType
-    model_config: LanguageModelConfig | RouterConfig
+    model_config: LanguageModelConfig | ClassifierModelConfig
+    grammar_start_tokens: tuple[str, ...]
 
 
 def download_file(
@@ -118,7 +119,7 @@ def download_config_file(
 
 
 class ImportResults(NamedTuple):
-    model: LanguageModel | Router
+    model: LanguageModel | ClassifierModel
     metadata: ModelMetadata
 
 
@@ -145,6 +146,8 @@ def import_message_processor(
             case FileSpec(_) as file_spec:
                 chat_template_file = download_file(file_spec, model_spec.repo, output_dir)
                 prompt_template = chat_template_file.read_text()
+            case str() as template_string:
+                prompt_template = template_string
             case None:
                 raise ValueError("No chat template specified.")
     else:
@@ -263,14 +266,14 @@ def _import_language_model(
     return language_model, language_model_config
 
 
-def _import_router(
+def _import_classifier(
     model_spec: ModelSpec,
     *,
     context_length: int | None = None,
     precision: DTypeLike | None = None,
     accumulation_precision: DTypeLike = jnp.float32,
     progress_callback: Callable[[StatusEvent], None] | None = None,
-) -> tuple[Router, RouterConfig]:
+) -> tuple[ClassifierModel, ClassifierModelConfig]:
     foreign_classifier_config_file = download_config_file(model_spec)
     foreign_classifier_config = model_spec.config_type.from_json(foreign_classifier_config_file)
     assert isinstance(foreign_classifier_config, ForeignClassifierConfig)
@@ -293,12 +296,12 @@ def _import_router(
 
     message_processor = import_message_processor(model_spec)
 
-    router_config = RouterConfig(
+    classifier_model_config = ClassifierModelConfig(
         model_config=classifier.config,
         message_processor_config=message_processor.config,
     )
-    router_model = Router(router_config, classifier, message_processor)
-    return router_model, router_config
+    classifier_model = ClassifierModel(classifier_model_config, classifier, message_processor)
+    return classifier_model, classifier_model_config
 
 
 def import_model(
@@ -324,8 +327,8 @@ def import_model(
                 accumulation_precision=accumulation_precision,
                 progress_callback=progress_callback,
             )
-        case ModelType.ROUTER_MODEL:
-            model, config = _import_router(
+        case ModelType.CLASSIFIER_MODEL:
+            model, config = _import_classifier(
                 model_spec,
                 context_length=context_length,
                 precision=precision,
@@ -344,5 +347,6 @@ def import_model(
         use_cases=model_spec.use_cases,
         model_type=model_spec.model_type,
         model_config=config,
+        grammar_start_tokens=model_spec.grammar_start_tokens,
     )
     return ImportResults(model, metadata)

lalamo/model_import/model_specs/__init__.py

@@ -5,7 +5,7 @@ from .gpt_oss import GPT_OSS_MODELS
 from .huggingface import HUGGINGFACE_MODELS
 from .llama import LLAMA_MODELS
 from .llamba import LLAMBA_MODELS
-from .mirai import MIRAI_ROUTER_MODELS
+from .mirai import MIRAI_CLASSIFIER_MODELS
 from .mistral import MISTRAL_MODELS
 
 # from .pleias import PLEIAS_MODELS
@@ -35,7 +35,7 @@ ALL_MODEL_LISTS = [
     POLARIS_MODELS,
     QWEN_MODELS,
     REKA_MODELS,
-    MIRAI_ROUTER_MODELS,
+    MIRAI_CLASSIFIER_MODELS,
 ]
 
 ALL_MODELS = [model for model_list in ALL_MODEL_LISTS for model in model_list]

lalamo/model_import/model_specs/common.py

@@ -32,7 +32,7 @@ __all__ = [
 
 class ModelType(StrEnum):
     LANGUAGE_MODEL = "language_model"
-    ROUTER_MODEL = "router_model"
+    CLASSIFIER_MODEL = "classifier_model"
 
 
 def cast_if_float(array: Array, cast_to: DTypeLike) -> Array:
@@ -84,7 +84,7 @@ class ConfigMap:
     tokenizer: FileSpec = field(default=FileSpec("tokenizer.json"))
     tokenizer_config: FileSpec = field(default=FileSpec("tokenizer_config.json"))
     generation_config: FileSpec | None = field(default=FileSpec("generation_config.json"))
-    chat_template: FileSpec | JSONFieldSpec | None = None
+    chat_template: FileSpec | JSONFieldSpec | str | None = None
 
 
 def _is_foreign_config_type(t: object) -> bool:
@@ -114,12 +114,29 @@ def _unstructure_foreign_config_factory(t: object, c: cattrs.Converter) -> Calla
     return _hook
 
 
+def _structure_chat_template(value: object, _type: object) -> FileSpec | JSONFieldSpec | str | None:
+    if value is None:
+        return None
+    if isinstance(value, str):
+        return value
+    if isinstance(value, dict):
+        if "file_spec" in value and "field_name" in value:
+            return JSONFieldSpec(
+                file_spec=FileSpec(**value["file_spec"]),
+                field_name=value["field_name"],
+            )
+        if "filename" in value:
+            return FileSpec(**value)
+    raise ValueError(f"Invalid chat_template value: {value}")
+
+
 @dataclass(frozen=True)
 class ModelSpec:
     _converter: ClassVar[cattrs.Converter] = cattrs.Converter()
 
     _converter.register_structure_hook_factory(_is_foreign_config_type, _structure_foreign_config_factory)
     _converter.register_unstructure_hook_factory(_is_foreign_config_type, _unstructure_foreign_config_factory)
+    _converter.register_structure_hook(FileSpec | JSONFieldSpec | str | None, _structure_chat_template)
 
     vendor: str
     family: str
@@ -137,6 +154,7 @@ class ModelSpec:
     model_type: ModelType = ModelType.LANGUAGE_MODEL
     configs: ConfigMap = field(default=ConfigMap())
     use_cases: tuple[UseCase, ...] = tuple()
+    grammar_start_tokens: tuple[str, ...] = tuple()
 
     @classmethod
     def from_json(cls, json_data: dict) -> "ModelSpec":
@@ -162,6 +180,7 @@ def awq_model_spec(
         configs=model_spec.configs,
         weights_type=model_spec.weights_type,
         use_cases=model_spec.use_cases,
+        grammar_start_tokens=model_spec.grammar_start_tokens,
     )
 
 

lalamo/model_import/model_specs/mirai.py

@@ -2,9 +2,9 @@ from lalamo.model_import.decoder_configs.huggingface import ModernBERTConfig
 
 from .common import ConfigMap, FileSpec, ModelSpec, ModelType
 
-__all__ = ["MIRAI_ROUTER_MODELS"]
+__all__ = ["MIRAI_CLASSIFIER_MODELS"]
 
-MIRAI_ROUTER_MODELS = [
+MIRAI_CLASSIFIER_MODELS = [
     ModelSpec(
         vendor="trymirai",
         family="ModernBERT",
@@ -14,7 +14,7 @@ MIRAI_ROUTER_MODELS = [
         repo="trymirai/chat-moderation-router",
         config_type=ModernBERTConfig,
         use_cases=tuple(),
-        model_type=ModelType("router_model"),
+        model_type=ModelType("classifier_model"),
         configs=ConfigMap(chat_template=FileSpec("chat_template.jinja")),
     ),
 ]

lalamo/models/__init__.py

@@ -1,10 +1,10 @@
+from .classifier import ClassifierModel, ClassifierModelConfig
 from .language_model import GenerationConfig, LanguageModel, LanguageModelConfig
-from .router import Router, RouterConfig
 
 __all__ = [
+    "ClassifierModel",
+    "ClassifierModelConfig",
     "GenerationConfig",
     "LanguageModel",
     "LanguageModelConfig",
-    "Router",
-    "RouterConfig",
 ]

lalamo/models/{router.py → classifier.py}

@@ -13,29 +13,29 @@ from lalamo.modules import Classifier, ClassifierConfig, LalamoModule
 from .common import TextModel, TextModelConfig
 
 __all__ = [
-    "Router",
-    "RouterConfig",
+    "ClassifierModel",
+    "ClassifierModelConfig",
 ]
 
 
 @dataclass(frozen=True)
-class RouterConfig(TextModelConfig[ClassifierConfig]):
+class ClassifierModelConfig(TextModelConfig[ClassifierConfig]):
     def init(
         self,
         model: LalamoModule,
         message_processor: MessageProcessor,
-    ) -> "Router":
+    ) -> "ClassifierModel":
         assert isinstance(model, Classifier)
-        return Router(self, model, message_processor)
+        return ClassifierModel(self, model, message_processor)
 
     @classmethod
-    def load_model(cls, path: Path | str) -> "Router":
+    def load_model(cls, path: Path | str) -> "ClassifierModel":
         result = super().load_model(path)
-        assert isinstance(result, Router)
+        assert isinstance(result, ClassifierModel)
         return result
 
 
-class Router(TextModel[RouterConfig, Classifier]):
+class ClassifierModel(TextModel[ClassifierModelConfig, Classifier]):
     def label_output_logits(self, logits: Float[Array, "batch logits"]) -> dict[str, Float[Array, " batch"]]:
         output_labels = self.model.config.output_labels
         probabilities = jax.nn.sigmoid(logits)

{lalamo-0.5.7.dist-info → lalamo-0.5.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lalamo
-Version: 0.5.7
+Version: 0.5.8
 Summary: JAX library for optimization and export of models for use with the UZU inference engine.
 Requires-Python: <4,>=3.12
 Description-Content-Type: text/markdown

{lalamo-0.5.7.dist-info → lalamo-0.5.8.dist-info}/RECORD

@@ -1,6 +1,6 @@
-lalamo/__init__.py,sha256=ynqoSW4I6eQ92AR9jzr_8XMEiGpDGjMHHGUrjlTkmb0,762
+lalamo/__init__.py,sha256=ZJ5Cjq4OoGVrjba9zUYIYnFGRKZkCkhBLaakdt4D008,814
 lalamo/common.py,sha256=5NUFD26yQgOnEEk3LaQnce8n-VwJxILkEpFesHZhtQU,3820
-lalamo/main.py,sha256=Lqv-lU6hRSqbJeVOheZoKllK1LaPHTuR_8jNTPC7UZg,23956
+lalamo/main.py,sha256=GgUT7lT48-XQuAEH7qzsDKG8Lx9iBf-sYBIRhZL9q7E,23978
 lalamo/message_processor.py,sha256=bSUAQg7CemLTnBV4LtPxJBicAalruDCA-JXjkTYPZ8U,5797
 lalamo/quantization.py,sha256=8o6ryIZLzzDYQuvBTboPfaVVdfijAKGpTxOcg3GKVD8,2752
 lalamo/registry_abc.py,sha256=ENjXiD_wEH100fNjG-W5Em1L_EQ0Lf0pdRhRGvf3qZk,2197
@@ -11,7 +11,7 @@ lalamo/data/huggingface_message.py,sha256=-7lN9eIcETQzt1Pnx3d4d8p3_I7WYMNf4mp1P9
 lalamo/data/lalamo_completions.py,sha256=U_m3UNSJASUFz3rJq_taZOtL_U4B8Oj-ndkTF-JH-v4,1509
 lalamo/data/utils.py,sha256=B96gLaULyStKYuR8wjFdTpFc6YIDC8EEvGh1eiMe_Ec,338
 lalamo/model_import/__init__.py,sha256=Z8pS9rbKKx1QgUy7KZtHxiNWlZhII3mdovT9d37vAxg,168
-lalamo/model_import/common.py,sha256=dS8HbDCSReaffJ0WWdVpYjk-VGv_Ga8C3GN7z6bdxlY,11992
+lalamo/model_import/common.py,sha256=tdZsteRsxL6DVUFwHw_1eeNLckflOdAaIm7Wm9eJzxM,12311
 lalamo/model_import/huggingface_generation_config.py,sha256=mot6VQ6ezCtEhN6VjhnvaU-nR5P5T2BuBUgpFNnWJxU,1495
 lalamo/model_import/huggingface_tokenizer_config.py,sha256=xvwdmio7b9nhn2H3uMBVligiYj58JaCFCvHY3-8dBvM,2502
 lalamo/model_import/decoder_configs/__init__.py,sha256=1ZqMcEHvCJjMIZ9iNyY31XMXOaFxB-NbqIU01BtmcEk,641
@@ -33,24 +33,24 @@ lalamo/model_import/loaders/common.py,sha256=kkugV-bMQlN1zvGHoj3uc7z0FbXKoMtXEBT
 lalamo/model_import/loaders/executorch.py,sha256=t2Ey_mBMNC8bTSTdYWjuGXdPTRoohFlYrqtWyNkBU_8,9219
 lalamo/model_import/loaders/huggingface.py,sha256=ITA0Y_kCDFL4Tanuvd1NWUvV77WEn0VEzkcX5Whlwys,29835
 lalamo/model_import/loaders/utils.py,sha256=eiX3WKFRrAfBY-dugodscNInl5o5w3KmVcgma4atpGY,2456
-lalamo/model_import/model_specs/__init__.py,sha256=STy1Au_mwJYcl_zSei6FEp8mmGZg-Fia1tpUVAO9x9Y,1167
-lalamo/model_import/model_specs/common.py,sha256=2E6mtpUS613jsp_SDiROOvvWFB6T5U2tDkbIBXVWjmI,5780
+lalamo/model_import/model_specs/__init__.py,sha256=V7S5Uo3GVBUG7KD0czMtmWZcQ-FJgryTZlxC7Abn_c0,1175
+lalamo/model_import/model_specs/common.py,sha256=RVPlNWHG_5OvU1W3YcOpqYz59Dh8plDmd7z1xNrqmaY,6585
 lalamo/model_import/model_specs/deepseek.py,sha256=Umef93_ZBuq93yYsejIRNwj3udoln1gHfrv3SK5jyMo,417
 lalamo/model_import/model_specs/gemma.py,sha256=irWgylL-pc7y3Gn5DK3fjKoCT9kJWH3B7mTa-1Gmxqc,1306
 lalamo/model_import/model_specs/gpt_oss.py,sha256=PLo0QGrXKdX61ReTRdyOaP_EH3Dmj5lp3fpJjZRwRVA,542
 lalamo/model_import/model_specs/huggingface.py,sha256=eF8ItF5reFrFkjYxwiAJcFwUAlN6CpXfM-aQ8a92ItM,430
 lalamo/model_import/model_specs/llama.py,sha256=Ml-xvRGlXBT9NJhmEpwgNo6C84oBSMYgA1_PrCYGcAw,990
 lalamo/model_import/model_specs/llamba.py,sha256=Ic3sWTv34FLJ4fG6OR_Mc5goGJQR6fa5b2WbVXbn9FA,1471
-lalamo/model_import/model_specs/mirai.py,sha256=20sLTZZQ_kn8ULfID71xdnVhjW7OjcJDf96BvVnCBUQ,605
+lalamo/model_import/model_specs/mirai.py,sha256=eifYVV5-fABiLH6rr82_DiVFtDyqpW0vbvXCYsQQzto,617
 lalamo/model_import/model_specs/mistral.py,sha256=HAojorjOqsJn2DoMBzYRw8A70qCslhFEsE9AF5xumlg,1278
 lalamo/model_import/model_specs/pleias.py,sha256=5sRpZGYwLdsav6bLiW-459y1Cs9iJKgKkBIuGsOxtsQ,368
 lalamo/model_import/model_specs/polaris.py,sha256=Mw1-6bByjDmPIKlIUIV46CsmV5xUp_laI5Qquo5DmAQ,520
 lalamo/model_import/model_specs/qwen.py,sha256=qzLmTveATmnwNFQSFJlffcXw7syFnrCmKf9ggkkkw1Y,7050
 lalamo/model_import/model_specs/reka.py,sha256=dOUYbEMMvovQdzQuBO_DCsjGI39syhoKCvnxLkNEDCw,423
-lalamo/models/__init__.py,sha256=HfAYc4mteKu7BMyIP8aTFliLKGWJzLk1kPcdKurb8eo,243
+lalamo/models/__init__.py,sha256=Vn5PcvSqKppIchkSZwQVTn_GpRvOOzZVxo5PUeDl6N8,283
+lalamo/models/classifier.py,sha256=LvL54crCVi4HVSIXuoaSLB_5jtcx74GL7kgdy2Y16Zc,2094
 lalamo/models/common.py,sha256=PDteofGxjSBWYw_mPxbN1DTUba70aOURrAIjl13SSHc,2954
 lalamo/models/language_model.py,sha256=QPeVEyhutSze7fSNhvOvwSoYt24QMk-dtTJkos38amY,13465
-lalamo/models/router.py,sha256=7KZqHVhr2TA7Qh76KfwrvyfztfZnV-P-Ss11O8dzbRg,2013
 lalamo/modules/__init__.py,sha256=xWJ4OPAF4gKd0evYwXIK5kTnbH6nI55oLAePcoDDHQ0,3730
 lalamo/modules/activations.py,sha256=U3qTQtZawPAUcoqbkIJnmTYcaNiQuSPMLcBeJ398GhI,1022
 lalamo/modules/classifier.py,sha256=_jtJ3INEq1dJP5HpUmcDk9YYzpRYlQ04zvFGaWBV6Lg,12101
@@ -80,9 +80,9 @@ lalamo/speculator/estimator.py,sha256=4D8dPZCWsrpORb7y8pQ6VsiIg1Cblvvxe6gXCoYtcD
 lalamo/speculator/inference.py,sha256=5GntUgj0HQLeLn3HIHnVX8EEO0EBzmKeP5-_U7kdFAM,3670
 lalamo/speculator/ngram.py,sha256=95mdfAWhx4d5XOnOwhyhElnvcy6nlUjYhcbJzqDs414,5875
 lalamo/speculator/utils.py,sha256=0wZoMMIzzk0Q-3zq5H5f-JBplePNHxywndkrNtOJOyo,1697
-lalamo-0.5.7.dist-info/licenses/LICENSE,sha256=diHRfjSEJHD1nnEeMIfMRCjR3UERf8bT3eseD6b1ayA,1072
-lalamo-0.5.7.dist-info/METADATA,sha256=ZzSQiR7KYoAXkBKOToa2bi9tCPRvVXvhnN5y6AO7wyQ,3146
-lalamo-0.5.7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lalamo-0.5.7.dist-info/entry_points.txt,sha256=qli7qTfnBk5WP10rOGXXEckHMtt-atJMDWd8jN89Uks,43
-lalamo-0.5.7.dist-info/top_level.txt,sha256=VHvWL5JN5XRG36NsN_MieJ7EwRihEOrEjyDaTdFJ-aI,7
-lalamo-0.5.7.dist-info/RECORD,,
+lalamo-0.5.8.dist-info/licenses/LICENSE,sha256=diHRfjSEJHD1nnEeMIfMRCjR3UERf8bT3eseD6b1ayA,1072
+lalamo-0.5.8.dist-info/METADATA,sha256=miYVR0hj7X-d1X09Bwaqf9-zKUqmljZ2qrhkV1rLICQ,3146
+lalamo-0.5.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lalamo-0.5.8.dist-info/entry_points.txt,sha256=qli7qTfnBk5WP10rOGXXEckHMtt-atJMDWd8jN89Uks,43
+lalamo-0.5.8.dist-info/top_level.txt,sha256=VHvWL5JN5XRG36NsN_MieJ7EwRihEOrEjyDaTdFJ-aI,7
+lalamo-0.5.8.dist-info/RECORD,,