Skip to content

Semantic Router

mailtag.semantic_router.SemanticRouter

Routes text to categories based on embedding similarity.

Uses pre-computed category embeddings to instantly classify emails based on semantic similarity to category centroids.

Source code in src/mailtag/semantic_router.py 
 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
 77
 78
 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
class SemanticRouter:
    """Routes text to categories based on embedding similarity.

    Uses pre-computed category embeddings to instantly classify
    emails based on semantic similarity to category centroids.
    """

    def __init__(self, embedder: MLXEmbedder):
        """Initialize the semantic router.

        Args:
            embedder: MLXEmbedder instance for generating embeddings
        """
        self.embedder = embedder
        self.category_embeddings: dict[str, np.ndarray] = {}
        self.categories: list[str] = []
        self._embedding_matrix: np.ndarray | None = None
        logger.info("SemanticRouter initialized")

    def load_embeddings(self, path: Path | str) -> bool:
        """Load pre-computed category embeddings from file.

        Args:
            path: Path to the .npz file containing embeddings

        Returns:
            True if loaded successfully, False otherwise
        """
        path = Path(path)
        if not path.exists():
            logger.warning(f"Embeddings file not found: {path}")
            return False

        try:
            data = np.load(path, allow_pickle=True)
            self.category_embeddings = {key: data[key] for key in data.files}
            self.categories = list(self.category_embeddings.keys())
            self._build_embedding_matrix()
            logger.info(f"Loaded embeddings for {len(self.categories)} categories from {path}")
            return True
        except (OSError, json.JSONDecodeError, KeyError, ValueError) as e:
            logger.error(f"Failed to load embeddings from {path}: {e}")
            return False

    def save_embeddings(self, path: Path | str) -> bool:
        """Save category embeddings to file.

        Args:
            path: Path to save the .npz file

        Returns:
            True if saved successfully, False otherwise
        """
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)

        try:
            np.savez(path, **self.category_embeddings)
            logger.info(f"Saved embeddings for {len(self.categories)} categories to {path}")
            return True
        except (OSError, TypeError, ValueError) as e:
            logger.error(f"Failed to save embeddings to {path}: {e}")
            return False

    def _build_embedding_matrix(self):
        """Build the embedding matrix for efficient similarity computation."""
        if not self.category_embeddings:
            self._embedding_matrix = None
            return

        self._embedding_matrix = np.stack([self.category_embeddings[cat] for cat in self.categories])
        # Normalize for cosine similarity
        norms = np.linalg.norm(self._embedding_matrix, axis=1, keepdims=True)
        self._embedding_matrix = self._embedding_matrix / norms

    def build_from_examples(self, category_examples: dict[str, list[str]]) -> None:
        """Build category embeddings from example texts.

        Args:
            category_examples: Dict mapping category names to lists of example texts
        """
        logger.info(f"Building embeddings for {len(category_examples)} categories...")

        for category, examples in category_examples.items():
            if not examples:
                logger.warning(f"No examples for category '{category}', skipping")
                continue

            # Compute embeddings for all examples
            embeddings = self.embedder.encode_documents(examples)

            # Use centroid (mean) as category embedding
            centroid = embeddings.mean(axis=0)
            self.category_embeddings[category] = centroid

            logger.debug(f"Built embedding for '{category}' from {len(examples)} examples")

        self.categories = list(self.category_embeddings.keys())
        self._build_embedding_matrix()
        logger.info(f"Built embeddings for {len(self.categories)} categories")

    def top_batch(self, texts: list[str]) -> list[tuple[str, float]]:
        """Return the nearest category and its similarity for each text, without threshold."""
        if not self.categories or self._embedding_matrix is None:
            return [("", 0.0)] * len(texts)
        if not texts:
            return []

        query_embeddings = self.embedder.encode(texts, prefix="search_query: ")
        norms = np.linalg.norm(query_embeddings, axis=1, keepdims=True)
        similarities = np.dot(query_embeddings / norms, self._embedding_matrix.T)

        best = np.argmax(similarities, axis=1)
        return [(self.categories[j], float(similarities[i][j])) for i, j in enumerate(best)]

    @property
    def num_categories(self) -> int:
        """Return the number of loaded categories."""
        return len(self.categories)

num_categories property

Return the number of loaded categories.

__init__(embedder)

Initialize the semantic router.

Parameters:

Name Type Description Default
embedder MLXEmbedder

MLXEmbedder instance for generating embeddings

 required
Source code in src/mailtag/semantic_router.py 
28
29
30
31
32
33
34
35
36
37
38
def __init__(self, embedder: MLXEmbedder):
    """Initialize the semantic router.

    Args:
        embedder: MLXEmbedder instance for generating embeddings
    """
    self.embedder = embedder
    self.category_embeddings: dict[str, np.ndarray] = {}
    self.categories: list[str] = []
    self._embedding_matrix: np.ndarray | None = None
    logger.info("SemanticRouter initialized")

top_batch(texts)

Return the nearest category and its similarity for each text, without threshold.

Source code in src/mailtag/semantic_router.py 
122
123
124
125
126
127
128
129
130
131
132
133
134
def top_batch(self, texts: list[str]) -> list[tuple[str, float]]:
    """Return the nearest category and its similarity for each text, without threshold."""
    if not self.categories or self._embedding_matrix is None:
        return [("", 0.0)] * len(texts)
    if not texts:
        return []

    query_embeddings = self.embedder.encode(texts, prefix="search_query: ")
    norms = np.linalg.norm(query_embeddings, axis=1, keepdims=True)
    similarities = np.dot(query_embeddings / norms, self._embedding_matrix.T)

    best = np.argmax(similarities, axis=1)
    return [(self.categories[j], float(similarities[i][j])) for i, j in enumerate(best)]

build_from_examples(category_examples)

Build category embeddings from example texts.

Parameters:

Name Type Description Default
category_examples dict[str, list[str]]

Dict mapping category names to lists of example texts

 required
Source code in src/mailtag/semantic_router.py 
 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
def build_from_examples(self, category_examples: dict[str, list[str]]) -> None:
    """Build category embeddings from example texts.

    Args:
        category_examples: Dict mapping category names to lists of example texts
    """
    logger.info(f"Building embeddings for {len(category_examples)} categories...")

    for category, examples in category_examples.items():
        if not examples:
            logger.warning(f"No examples for category '{category}', skipping")
            continue

        # Compute embeddings for all examples
        embeddings = self.embedder.encode_documents(examples)

        # Use centroid (mean) as category embedding
        centroid = embeddings.mean(axis=0)
        self.category_embeddings[category] = centroid

        logger.debug(f"Built embedding for '{category}' from {len(examples)} examples")

    self.categories = list(self.category_embeddings.keys())
    self._build_embedding_matrix()
    logger.info(f"Built embeddings for {len(self.categories)} categories")

load_embeddings(path)

Load pre-computed category embeddings from file.

Parameters:

Name Type Description Default
path Path | str

Path to the .npz file containing embeddings

 required

Returns:

Type Description
bool

True if loaded successfully, False otherwise
Source code in src/mailtag/semantic_router.py 
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
def load_embeddings(self, path: Path | str) -> bool:
    """Load pre-computed category embeddings from file.

    Args:
        path: Path to the .npz file containing embeddings

    Returns:
        True if loaded successfully, False otherwise
    """
    path = Path(path)
    if not path.exists():
        logger.warning(f"Embeddings file not found: {path}")
        return False

    try:
        data = np.load(path, allow_pickle=True)
        self.category_embeddings = {key: data[key] for key in data.files}
        self.categories = list(self.category_embeddings.keys())
        self._build_embedding_matrix()
        logger.info(f"Loaded embeddings for {len(self.categories)} categories from {path}")
        return True
    except (OSError, json.JSONDecodeError, KeyError, ValueError) as e:
        logger.error(f"Failed to load embeddings from {path}: {e}")
        return False

save_embeddings(path)

Save category embeddings to file.

Parameters:

Name Type Description Default
path Path | str

Path to save the .npz file

 required

Returns:

Type Description
bool

True if saved successfully, False otherwise
Source code in src/mailtag/semantic_router.py 
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
def save_embeddings(self, path: Path | str) -> bool:
    """Save category embeddings to file.

    Args:
        path: Path to save the .npz file

    Returns:
        True if saved successfully, False otherwise
    """
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)

    try:
        np.savez(path, **self.category_embeddings)
        logger.info(f"Saved embeddings for {len(self.categories)} categories to {path}")
        return True
    except (OSError, TypeError, ValueError) as e:
        logger.error(f"Failed to save embeddings to {path}: {e}")
        return False