Skip to content

MLX Provider

MLXEmbedder

mailtag.mlx_provider.MLXEmbedder

Generates embeddings using sentence-transformers with MLX backend.

Supports models like nomic-ai/nomic-embed-text-v1.5 which are optimized for semantic similarity tasks and support long context (8k tokens).

Source code in src/mailtag/mlx_provider.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
class MLXEmbedder:
    """Generates embeddings using sentence-transformers with MLX backend.

    Supports models like nomic-ai/nomic-embed-text-v1.5 which are optimized
    for semantic similarity tasks and support long context (8k tokens).
    """

    def __init__(self, model_name: str = "nomic-ai/nomic-embed-text-v1.5"):
        """Initialize the embedder with the specified model.

        Args:
            model_name: HuggingFace model name for embeddings
        """
        self.model_name = model_name
        self._model: SentenceTransformer | None = None
        logger.info(f"MLXEmbedder initialized with model: {model_name}")

    @property
    def model(self) -> SentenceTransformer:
        """Lazy load the embedding model."""
        if self._model is None:
            logger.info(f"Loading embedding model: {self.model_name}")
            from sentence_transformers import SentenceTransformer

            self._model = SentenceTransformer(self.model_name, trust_remote_code=True)
            logger.info("Embedding model loaded successfully")
        return self._model

    def encode(self, texts: list[str] | str, prefix: str = "search_document: ") -> np.ndarray:
        """Generate embeddings for the given texts.

        Args:
            texts: Single text or list of texts to encode
            prefix: Prefix to add to each text (nomic models use task prefixes)

        Returns:
            Numpy array of embeddings with shape (n_texts, embedding_dim)
        """
        if isinstance(texts, str):
            texts = [texts]

        # Add prefix for nomic models (they use task-specific prefixes)
        if "nomic" in self.model_name.lower():
            texts = [f"{prefix}{text}" for text in texts]

        embeddings = self.model.encode(texts, convert_to_numpy=True)
        return embeddings

    def encode_documents(self, documents: list[str]) -> np.ndarray:
        """Encode document texts (for building an index).

        Args:
            documents: List of document texts to encode

        Returns:
            Numpy array of shape (n_documents, embedding_dim)
        """
        return self.encode(documents, prefix="search_document: ")

__init__(model_name='nomic-ai/nomic-embed-text-v1.5')

Initialize the embedder with the specified model.

Parameters:

Name Type Description Default
model_name str

HuggingFace model name for embeddings

 'nomic-ai/nomic-embed-text-v1.5'
Source code in src/mailtag/mlx_provider.py 
26
27
28
29
30
31
32
33
34
def __init__(self, model_name: str = "nomic-ai/nomic-embed-text-v1.5"):
    """Initialize the embedder with the specified model.

    Args:
        model_name: HuggingFace model name for embeddings
    """
    self.model_name = model_name
    self._model: SentenceTransformer | None = None
    logger.info(f"MLXEmbedder initialized with model: {model_name}")

encode(texts, prefix='search_document: ')

Generate embeddings for the given texts.

Parameters:

Name Type Description Default
texts list[str] | str

Single text or list of texts to encode

 required
prefix str

Prefix to add to each text (nomic models use task prefixes)

 'search_document: '

Returns:

Type Description
ndarray

Numpy array of embeddings with shape (n_texts, embedding_dim)
Source code in src/mailtag/mlx_provider.py 
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def encode(self, texts: list[str] | str, prefix: str = "search_document: ") -> np.ndarray:
    """Generate embeddings for the given texts.

    Args:
        texts: Single text or list of texts to encode
        prefix: Prefix to add to each text (nomic models use task prefixes)

    Returns:
        Numpy array of embeddings with shape (n_texts, embedding_dim)
    """
    if isinstance(texts, str):
        texts = [texts]

    # Add prefix for nomic models (they use task-specific prefixes)
    if "nomic" in self.model_name.lower():
        texts = [f"{prefix}{text}" for text in texts]

    embeddings = self.model.encode(texts, convert_to_numpy=True)
    return embeddings

encode_documents(documents)

Encode document texts (for building an index).

Parameters:

Name Type Description Default
documents list[str]

List of document texts to encode

 required

Returns:

Type Description
ndarray

Numpy array of shape (n_documents, embedding_dim)
Source code in src/mailtag/mlx_provider.py 
67
68
69
70
71
72
73
74
75
76
def encode_documents(self, documents: list[str]) -> np.ndarray:
    """Encode document texts (for building an index).

    Args:
        documents: List of document texts to encode

    Returns:
        Numpy array of shape (n_documents, embedding_dim)
    """
    return self.encode(documents, prefix="search_document: ")

MLXLLM

mailtag.mlx_provider.MLXLLM

Answers short prompts in batches using mlx-lm (Signal 6).

Uses quantized models from mlx-community for efficient inference on Apple Silicon.

Source code in src/mailtag/mlx_provider.py 
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
class MLXLLM:
    """Answers short prompts in batches using mlx-lm (Signal 6).

    Uses quantized models from mlx-community for efficient inference
    on Apple Silicon.
    """

    def __init__(self, model_name: str = "mlx-community/gemma-4-e4b-it-OptiQ-4bit"):
        """Initialize the LLM with the specified model.

        Args:
            model_name: MLX model name from mlx-community
        """
        self.model_name = model_name
        self._model = None
        self._tokenizer = None
        self._prefix_key: str | None = None
        self._prefix_cache = None
        self._prefix_post = ""
        logger.info(f"MLXLLM initialized with model: {model_name}")

    def _load_model(self):
        """Lazy load the model and tokenizer."""
        if self._model is None:
            logger.info(f"Loading LLM model: {self.model_name}")
            from mlx_lm import load

            self._model, self._tokenizer = load(self.model_name)
            logger.info("LLM model loaded successfully")

    @property
    def model(self):
        """Get the loaded model."""
        self._load_model()
        return self._model

    @property
    def tokenizer(self):
        """Get the loaded tokenizer."""
        self._load_model()
        return self._tokenizer

    def classify_batch(
        self, static_prompt: str, email_parts: list[str], batch_size: int = 8, max_tokens: int = 4
    ) -> list[str]:
        """Answer one short prompt per email, reusing the KV cache of the shared static prefix.

        The chat-formatted prompt is `static_prompt + email_part`; the static part is prefilled once
        per distinct value and each email only prefills its own tokens. Emails run in batches.
        """
        import copy

        import mlx.core as mx
        from mlx_lm import batch_generate
        from mlx_lm.generate import generate_step
        from mlx_lm.models.cache import make_prompt_cache
        from mlx_lm.sample_utils import make_sampler

        if not email_parts:
            return []

        if self._prefix_key != static_prompt:
            marker = "<<<EMAIL>>>"
            messages = [{"role": "user", "content": static_prompt + marker}]
            try:
                template = self.tokenizer.apply_chat_template(
                    messages, tokenize=False, add_generation_prompt=True, enable_thinking=False
                )
            except TypeError:
                template = self.tokenizer.apply_chat_template(
                    messages, tokenize=False, add_generation_prompt=True
                )
            pre, self._prefix_post = template.split(marker)
            cache = make_prompt_cache(self.model)
            for _ in generate_step(
                mx.array(self.tokenizer.encode(pre)), self.model, max_tokens=0, prompt_cache=cache
            ):
                pass
            self._prefix_cache, self._prefix_key = cache, static_prompt

        suffixes = [
            self.tokenizer.encode(part + self._prefix_post, add_special_tokens=False) for part in email_parts
        ]
        sampler = make_sampler(temp=0.0)
        texts: list[str] = []
        for start in range(0, len(suffixes), batch_size):
            chunk = suffixes[start : start + batch_size]
            response = batch_generate(
                self.model,
                self.tokenizer,
                chunk,
                prompt_caches=[copy.deepcopy(self._prefix_cache) for _ in chunk],
                max_tokens=max_tokens,
                sampler=sampler,
            )
            texts.extend(text.strip() for text in response.texts)
        return texts

__init__(model_name='mlx-community/gemma-4-e4b-it-OptiQ-4bit')

Initialize the LLM with the specified model.

Parameters:

Name Type Description Default
model_name str

MLX model name from mlx-community

 'mlx-community/gemma-4-e4b-it-OptiQ-4bit'
Source code in src/mailtag/mlx_provider.py 
86
87
88
89
90
91
92
93
94
95
96
97
98
def __init__(self, model_name: str = "mlx-community/gemma-4-e4b-it-OptiQ-4bit"):
    """Initialize the LLM with the specified model.

    Args:
        model_name: MLX model name from mlx-community
    """
    self.model_name = model_name
    self._model = None
    self._tokenizer = None
    self._prefix_key: str | None = None
    self._prefix_cache = None
    self._prefix_post = ""
    logger.info(f"MLXLLM initialized with model: {model_name}")

classify_batch(static_prompt, email_parts, batch_size=8, max_tokens=4)

Answer one short prompt per email, reusing the KV cache of the shared static prefix.

The chat-formatted prompt is static_prompt + email_part; the static part is prefilled once per distinct value and each email only prefills its own tokens. Emails run in batches.

Source code in src/mailtag/mlx_provider.py 
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
def classify_batch(
    self, static_prompt: str, email_parts: list[str], batch_size: int = 8, max_tokens: int = 4
) -> list[str]:
    """Answer one short prompt per email, reusing the KV cache of the shared static prefix.

    The chat-formatted prompt is `static_prompt + email_part`; the static part is prefilled once
    per distinct value and each email only prefills its own tokens. Emails run in batches.
    """
    import copy

    import mlx.core as mx
    from mlx_lm import batch_generate
    from mlx_lm.generate import generate_step
    from mlx_lm.models.cache import make_prompt_cache
    from mlx_lm.sample_utils import make_sampler

    if not email_parts:
        return []

    if self._prefix_key != static_prompt:
        marker = "<<<EMAIL>>>"
        messages = [{"role": "user", "content": static_prompt + marker}]
        try:
            template = self.tokenizer.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True, enable_thinking=False
            )
        except TypeError:
            template = self.tokenizer.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True
            )
        pre, self._prefix_post = template.split(marker)
        cache = make_prompt_cache(self.model)
        for _ in generate_step(
            mx.array(self.tokenizer.encode(pre)), self.model, max_tokens=0, prompt_cache=cache
        ):
            pass
        self._prefix_cache, self._prefix_key = cache, static_prompt

    suffixes = [
        self.tokenizer.encode(part + self._prefix_post, add_special_tokens=False) for part in email_parts
    ]
    sampler = make_sampler(temp=0.0)
    texts: list[str] = []
    for start in range(0, len(suffixes), batch_size):
        chunk = suffixes[start : start + batch_size]
        response = batch_generate(
            self.model,
            self.tokenizer,
            chunk,
            prompt_caches=[copy.deepcopy(self._prefix_cache) for _ in chunk],
            max_tokens=max_tokens,
            sampler=sampler,
        )
        texts.extend(text.strip() for text in response.texts)
    return texts