lalamo 0.5.7__tar.gz → 0.5.8__tar.gz

{lalamo-0.5.7 → lalamo-0.5.8}/PKG-INFO

@@ -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 → lalamo-0.5.8}/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-0.5.7 → lalamo-0.5.8}/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-0.5.7 → lalamo-0.5.8}/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-0.5.7 → lalamo-0.5.8}/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-0.5.7 → lalamo-0.5.8}/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-0.5.7 → lalamo-0.5.8}/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-0.5.7 → lalamo-0.5.8}/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-0.5.7/lalamo/models/router.py → lalamo-0.5.8/lalamo/models/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 → lalamo-0.5.8}/lalamo.egg-info/PKG-INFO

@@ -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 → lalamo-0.5.8}/lalamo.egg-info/SOURCES.txt

@@ -57,9 +57,9 @@ lalamo/model_import/model_specs/polaris.py
 lalamo/model_import/model_specs/qwen.py
 lalamo/model_import/model_specs/reka.py
 lalamo/models/__init__.py
+lalamo/models/classifier.py
 lalamo/models/common.py
 lalamo/models/language_model.py
-lalamo/models/router.py
 lalamo/modules/__init__.py
 lalamo/modules/activations.py
 lalamo/modules/classifier.py
@@ -90,6 +90,7 @@ lalamo/speculator/inference.py
 lalamo/speculator/ngram.py
 lalamo/speculator/utils.py
 tests/test_cartesia_mlx_models.py
+tests/test_chat_template.py
 tests/test_generation.py
 tests/test_huggingface_model_conversion.py
 tests/test_huggingface_models.py

lalamo-0.5.8/tests/test_chat_template.py

@@ -0,0 +1,173 @@
+import pytest
+
+from lalamo.model_import.decoder_configs.huggingface.llama import HFLlamaConfig
+from lalamo.model_import.model_specs.common import ConfigMap, FileSpec, JSONFieldSpec, ModelSpec
+
+DIRECT_TEMPLATE = "{% for message in messages %}{{ message.content }}{% endfor %}"
+
+
+class TestConfigMapChatTemplate:
+    def test_chat_template_as_string(self) -> None:
+        config = ConfigMap(chat_template=DIRECT_TEMPLATE)
+        assert config.chat_template == DIRECT_TEMPLATE
+        assert isinstance(config.chat_template, str)
+
+    def test_chat_template_as_file_spec(self) -> None:
+        file_spec = FileSpec("chat_template.jinja")
+        config = ConfigMap(chat_template=file_spec)
+        assert config.chat_template == file_spec
+        assert isinstance(config.chat_template, FileSpec)
+
+    def test_chat_template_as_file_spec_with_repo(self) -> None:
+        file_spec = FileSpec("chat_template.jinja", repo="some/repo")
+        config = ConfigMap(chat_template=file_spec)
+        assert config.chat_template == file_spec
+        assert isinstance(config.chat_template, FileSpec)
+        assert config.chat_template.repo == "some/repo"
+
+    def test_chat_template_as_json_field_spec(self) -> None:
+        json_spec = JSONFieldSpec(FileSpec("config.json"), "chat_template")
+        config = ConfigMap(chat_template=json_spec)
+        assert config.chat_template == json_spec
+        assert isinstance(config.chat_template, JSONFieldSpec)
+
+    def test_chat_template_none(self) -> None:
+        config = ConfigMap()
+        assert config.chat_template is None
+
+
+class TestModelSpecWithChatTemplate:
+    def test_model_spec_with_string_chat_template(self) -> None:
+        spec = ModelSpec(
+            vendor="Test",
+            family="Test",
+            name="Test",
+            size="1B",
+            repo="test/test",
+            config_type=HFLlamaConfig,
+            configs=ConfigMap(chat_template=DIRECT_TEMPLATE),
+        )
+        assert spec.configs.chat_template == DIRECT_TEMPLATE
+
+    def test_model_spec_with_file_spec_chat_template(self) -> None:
+        spec = ModelSpec(
+            vendor="Test",
+            family="Test",
+            name="Test",
+            size="1B",
+            repo="test/test",
+            config_type=HFLlamaConfig,
+            configs=ConfigMap(chat_template=FileSpec("chat_template.jinja")),
+        )
+        assert isinstance(spec.configs.chat_template, FileSpec)
+        assert spec.configs.chat_template.filename == "chat_template.jinja"
+
+    def test_model_spec_with_json_field_spec_chat_template(self) -> None:
+        spec = ModelSpec(
+            vendor="Test",
+            family="Test",
+            name="Test",
+            size="1B",
+            repo="test/test",
+            config_type=HFLlamaConfig,
+            configs=ConfigMap(chat_template=JSONFieldSpec(FileSpec("tokenizer_config.json"), "chat_template")),
+        )
+        assert isinstance(spec.configs.chat_template, JSONFieldSpec)
+        assert spec.configs.chat_template.field_name == "chat_template"
+
+
+class TestModelSpecJsonSerialization:
+    def test_roundtrip_with_string_chat_template(self) -> None:
+        spec = ModelSpec(
+            vendor="Test",
+            family="Test",
+            name="Test",
+            size="1B",
+            repo="test/test",
+            config_type=HFLlamaConfig,
+            configs=ConfigMap(chat_template=DIRECT_TEMPLATE),
+        )
+        json_data = spec.to_json()
+        restored = ModelSpec.from_json(json_data)
+        assert restored.configs.chat_template == DIRECT_TEMPLATE
+
+    def test_roundtrip_with_file_spec_chat_template(self) -> None:
+        spec = ModelSpec(
+            vendor="Test",
+            family="Test",
+            name="Test",
+            size="1B",
+            repo="test/test",
+            config_type=HFLlamaConfig,
+            configs=ConfigMap(chat_template=FileSpec("chat_template.jinja")),
+        )
+        json_data = spec.to_json()
+        restored = ModelSpec.from_json(json_data)
+        assert isinstance(restored.configs.chat_template, FileSpec)
+        assert restored.configs.chat_template.filename == "chat_template.jinja"
+
+    def test_roundtrip_with_json_field_spec_chat_template(self) -> None:
+        spec = ModelSpec(
+            vendor="Test",
+            family="Test",
+            name="Test",
+            size="1B",
+            repo="test/test",
+            config_type=HFLlamaConfig,
+            configs=ConfigMap(chat_template=JSONFieldSpec(FileSpec("config.json"), "chat_template")),
+        )
+        json_data = spec.to_json()
+        restored = ModelSpec.from_json(json_data)
+        assert isinstance(restored.configs.chat_template, JSONFieldSpec)
+        assert restored.configs.chat_template.field_name == "chat_template"
+        assert restored.configs.chat_template.file_spec.filename == "config.json"
+
+    def test_from_json_with_string_chat_template(self) -> None:
+        json_data = {
+            "vendor": "Test",
+            "family": "Test",
+            "name": "Test",
+            "size": "1B",
+            "repo": "test/test",
+            "config_type": "HFLlamaConfig",
+            "configs": {
+                "chat_template": DIRECT_TEMPLATE,
+            },
+        }
+        spec = ModelSpec.from_json(json_data)
+        assert spec.configs.chat_template == DIRECT_TEMPLATE
+
+    def test_from_json_with_file_spec_chat_template(self) -> None:
+        json_data = {
+            "vendor": "Test",
+            "family": "Test",
+            "name": "Test",
+            "size": "1B",
+            "repo": "test/test",
+            "config_type": "HFLlamaConfig",
+            "configs": {
+                "chat_template": {"filename": "chat_template.jinja"},
+            },
+        }
+        spec = ModelSpec.from_json(json_data)
+        assert isinstance(spec.configs.chat_template, FileSpec)
+        assert spec.configs.chat_template.filename == "chat_template.jinja"
+
+    def test_from_json_with_json_field_spec_chat_template(self) -> None:
+        json_data = {
+            "vendor": "Test",
+            "family": "Test",
+            "name": "Test",
+            "size": "1B",
+            "repo": "test/test",
+            "config_type": "HFLlamaConfig",
+            "configs": {
+                "chat_template": {
+                    "file_spec": {"filename": "config.json"},
+                    "field_name": "chat_template",
+                },
+            },
+        }
+        spec = ModelSpec.from_json(json_data)
+        assert isinstance(spec.configs.chat_template, JSONFieldSpec)
+        assert spec.configs.chat_template.field_name == "chat_template"

{lalamo-0.5.7 → lalamo-0.5.8}/tests/test_huggingface_model_conversion.py

@@ -14,7 +14,7 @@ from safetensors.flax import save_file
 from lalamo.common import flatten_parameters
 from lalamo.model_import import REPO_TO_MODEL, ModelMetadata, import_model
 from lalamo.model_import.model_specs import ModelType
-from lalamo.models import LanguageModelConfig, RouterConfig
+from lalamo.models import ClassifierModelConfig, LanguageModelConfig
 from lalamo.modules import config_converter
 from tests.test_models import DType, ModelTestSpec
 
@@ -91,7 +91,7 @@ def test_model_conversion(test_spec: ModelTestSpec, tmp_path: pathlib.Path) -> N
     match metadata.model_type:
         case ModelType.LANGUAGE_MODEL:
             model = LanguageModelConfig.load_model(tmp_path)
-        case ModelType.ROUTER_MODEL:
-            model = RouterConfig.load_model(tmp_path)
+        case ModelType.CLASSIFIER_MODEL:
+            model = ClassifierModelConfig.load_model(tmp_path)
     assert model is not None, f"Failed to load model {model_repo_name}"
     del model

{lalamo-0.5.7 → lalamo-0.5.8}/tests/test_models.py

@@ -13,7 +13,7 @@ import torch
 from jaxtyping import Array
 from transformers.models.gpt_oss.modeling_gpt_oss import GptOssAttention
 
-from lalamo import LanguageModel, Router, import_model
+from lalamo import ClassifierModel, LanguageModel, import_model
 from lalamo.model_import.common import ModelType
 from lalamo.modules.classifier import ClassifierActivationTrace, ClassifierResult
 from lalamo.modules.decoder import (
@@ -477,8 +477,8 @@ def _test_model(test_spec: ModelTestSpec, model_tracer: type[ModelTracer]) -> No
                 )
                 err.throw()
 
-            case ModelType.ROUTER_MODEL:
-                assert isinstance(model, Router)
+            case ModelType.CLASSIFIER_MODEL:
+                assert isinstance(model, ClassifierModel)
                 err, inference_results = checkify_forward(model.model)(
                     token_ids=token_ids,
                     token_positions=token_positions,