Mistral AI Engineer Interview Guide

1from typing import List, Dict, Tuple
2
3
4def group_anagrams(words: List[str]) -> List[List[str]]:
5    """Group anagrams together using a character count signature."""
6    groups: Dict[Tuple[int, ...], List[str]] = {}
7
8    for w in words:
9        # Assumes lowercase a-z; adjust if needed.
10        counts = [0] * 26
11        for ch in w:
12            counts[ord(ch) - ord('a')] += 1
13        key = tuple(counts)
14        if key not in groups:
15            groups[key] = []
16        groups[key].append(w)
17
18    return list(groups.values())
19
20
21if __name__ == "__main__":
22    sample = ["eat", "tea", "tan", "ate", "nat", "bat"]
23    print(group_anagrams(sample))
24

1from collections import deque
2from typing import List, Tuple
3
4
5def max_subarray_sum_at_most_k(probs: List[float], k: int) -> Tuple[float, Tuple[int, int]]:
6    """Return (max_sum, (l, r)) for a contiguous subarray with length <= k.
7
8    l and r are inclusive indices. If probs is empty or k <= 0, returns (0.0, (-1, -1)).
9    """
10    n = len(probs)
11    if n == 0 or k <= 0:
12        return 0.0, (-1, -1)
13
14    # Prefix sums: pref[t] = sum(probs[0:t]) for t in [0..n]
15    pref = [0.0] * (n + 1)
16    for i, x in enumerate(probs):
17        pref[i + 1] = pref[i] + x
18
19    # Deque stores indices of pref, increasing by pref value.
20    # For each end j, we need min pref[i] with i in [j-k, j-1] (length 1..k) or i=j (length 0).
21    # Here we allow length >= 1 by comparing i < j.
22    dq = deque([0])
23
24    best_sum = float("-inf")
25    best_l, best_r = 0, 0
26
27    for j in range(1, n + 1):
28        # Remove indices i that are too old (j - i > k)
29        while dq and dq[0] < j - k:
30            dq.popleft()
31
32        # Current best window ending at j-1 uses smallest pref[i] in dq
33        i = dq[0]
34        curr = pref[j] - pref[i]
35        if curr > best_sum:
36            best_sum = curr
37            best_l, best_r = i, j - 1
38
39        # Maintain deque monotonicity by pref value
40        while dq and pref[dq[-1]] >= pref[j]:
41            dq.pop()
42        dq.append(j)
43
44    return best_sum, (best_l, best_r)
45
46
47if __name__ == "__main__":
48    probs = [0.1, -0.2, 0.3, 0.4, -0.1, 0.2]
49    print(max_subarray_sum_at_most_k(probs, k=3))
50

1import math
2from dataclasses import dataclass
3
4import torch
5import torch.nn as nn
6import torch.nn.functional as F
7from torch.utils.data import DataLoader, Dataset
8
9
10class SyntheticTokenDataset(Dataset):
11    """Synthetic padded token sequences with a simple label rule."""
12
13    def __init__(self, n_samples=4096, vocab_size=2000, max_len=64, pad_id=0, seed=0):
14        g = torch.Generator().manual_seed(seed)
15        self.vocab_size = vocab_size
16        self.max_len = max_len
17        self.pad_id = pad_id
18
19        lengths = torch.randint(low=8, high=max_len + 1, size=(n_samples,), generator=g)
20        self.x = torch.full((n_samples, max_len), pad_id, dtype=torch.long)
21        for i, L in enumerate(lengths.tolist()):
22            # Tokens in [1, vocab_size)
23            self.x[i, :L] = torch.randint(1, vocab_size, size=(L,), generator=g)
24
25        # Label rule: parity of (sum of tokens mod 2) excluding pads
26        mask = (self.x != pad_id).long()
27        sums = (self.x * mask).sum(dim=1)
28        self.y = (sums % 2).long()
29
30    def __len__(self):
31        return self.x.size(0)
32
33    def __getitem__(self, idx):
34        return self.x[idx], self.y[idx]
35
36
37class TinyTransformerClassifier(nn.Module):
38    def __init__(self, vocab_size=2000, d_model=128, n_heads=4, n_layers=2, max_len=64, pad_id=0):
39        super().__init__()
40        self.pad_id = pad_id
41        self.tok_emb = nn.Embedding(vocab_size, d_model, padding_idx=pad_id)
42        self.pos_emb = nn.Embedding(max_len, d_model)
43
44        encoder_layer = nn.TransformerEncoderLayer(
45            d_model=d_model,
46            nhead=n_heads,
47            dim_feedforward=4 * d_model,
48            dropout=0.1,
49            activation="gelu",
50            batch_first=True,
51            norm_first=True,
52        )
53        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
54        self.ln = nn.LayerNorm(d_model)
55        self.head = nn.Linear(d_model, 2)
56
57    def forward(self, x):
58        # x: (B, T)
59        B, T = x.shape
60        pos = torch.arange(T, device=x.device).unsqueeze(0).expand(B, T)
61        h = self.tok_emb(x) + self.pos_emb(pos)  # (B, T, D)
62
63        # src_key_padding_mask: True for pads
64        pad_mask = (x == self.pad_id)  # (B, T)
65        h = self.encoder(h, src_key_padding_mask=pad_mask)  # (B, T, D)
66
67        # Masked mean pooling
68        keep = (~pad_mask).float().unsqueeze(-1)  # (B, T, 1)
69        h = h * keep
70        denom = keep.sum(dim=1).clamp_min(1.0)
71        pooled = h.sum(dim=1) / denom  # (B, D)
72
73        pooled = self.ln(pooled)
74        return self.head(pooled)  # (B, 2)
75
76
77def attach_nan_inf_hooks(model: nn.Module):
78    """Attach forward and backward hooks that stop on first NaN/Inf."""
79
80    def _check_tensor(t, name, phase):
81        if t is None or not torch.is_tensor(t):
82            return
83        if not torch.is_floating_point(t):
84            return
85        if torch.isnan(t).any() or torch.isinf(t).any():
86            # Print concise info and raise to stop training
87            with torch.no_grad():
88                n_nan = torch.isnan(t).sum().item()
89                n_inf = torch.isinf(t).sum().item()
90                mx = t.abs().max().item() if t.numel() else 0.0
91            raise RuntimeError(f"{phase} detected NaN/Inf in {name}, nan={n_nan} inf={n_inf} max_abs={mx:.4e}")
92
93    # Forward hooks on modules
94    for module_name, module in model.named_modules():
95        # Skip container modules to reduce spam
96        if len(list(module.children())) > 0:
97            continue
98
99        def make_fwd_hook(mname):
100            def fwd_hook(mod, inp, out):
101                if torch.is_tensor(out):
102                    _check_tensor(out, f"{mname}.output", "forward")
103                elif isinstance(out, (tuple, list)):
104                    for k, v in enumerate(out):
105                        if torch.is_tensor(v):
106                            _check_tensor(v, f"{mname}.output[{k}]", "forward")
107            return fwd_hook
108
109        module.register_forward_hook(make_fwd_hook(module_name))
110
111    # Backward hooks on parameters
112    for pname, p in model.named_parameters():
113        if not p.requires_grad:
114            continue
115
116        def make_grad_hook(param_name):
117            def grad_hook(g):
118                _check_tensor(g, f"{param_name}.grad", "backward")
119                return g
120            return grad_hook
121
122        p.register_hook(make_grad_hook(pname))
123
124
125@dataclass
126class TrainConfig:
127    device: str = "cuda" if torch.cuda.is_available() else "cpu"
128    vocab_size: int = 2000
129    max_len: int = 64
130    pad_id: int = 0
131    batch_size: int = 64
132    lr: float = 3e-4
133    weight_decay: float = 0.01
134    steps: int = 200
135    grad_accum_steps: int = 4
136    max_grad_norm: float = 1.0
137    use_amp: bool = True
138
139
140def train(cfg: TrainConfig):
141    torch.manual_seed(0)
142
143    ds = SyntheticTokenDataset(
144        n_samples=8192,
145        vocab_size=cfg.vocab_size,
146        max_len=cfg.max_len,
147        pad_id=cfg.pad_id,
148        seed=0,
149    )
150    loader = DataLoader(ds, batch_size=cfg.batch_size, shuffle=True, drop_last=True)
151    it = iter(loader)
152
153    model = TinyTransformerClassifier(
154        vocab_size=cfg.vocab_size,
155        d_model=128,
156        n_heads=4,
157        n_layers=2,
158        max_len=cfg.max_len,
159        pad_id=cfg.pad_id,
160    ).to(cfg.device)
161
162    attach_nan_inf_hooks(model)
163
164    opt = torch.optim.AdamW(model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay)
165    scaler = torch.amp.GradScaler(enabled=(cfg.use_amp and cfg.device.startswith("cuda")))
166
167    model.train()
168    opt.zero_grad(set_to_none=True)
169
170    for step in range(1, cfg.steps + 1):
171        try:
172            x, y = next(it)
173        except StopIteration:
174            it = iter(loader)
175            x, y = next(it)
176
177        x = x.to(cfg.device, non_blocking=True)
178        y = y.to(cfg.device, non_blocking=True)
179
180        with torch.amp.autocast(device_type="cuda" if cfg.device.startswith("cuda") else "cpu", enabled=scaler.is_enabled()):
181            logits = model(x)
182            loss = F.cross_entropy(logits, y)
183            loss = loss / cfg.grad_accum_steps
184
185        scaler.scale(loss).backward()
186
187        if step % cfg.grad_accum_steps == 0:
188            # Unscale before clipping so clip threshold is in real grad units
189            scaler.unscale_(opt)
190            total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), cfg.max_grad_norm)
191
192            scaler.step(opt)
193            scaler.update()
194            opt.zero_grad(set_to_none=True)
195
196            if math.isfinite(float(total_norm)):
197                pass
198            else:
199                raise RuntimeError(f"Non-finite grad norm: {total_norm}")
200
201        if step % 20 == 0:
202            # Report the unscaled loss
203            print(f"step={step:04d} loss={(loss.item() * cfg.grad_accum_steps):.4f}")
204
205
206if __name__ == "__main__":
207    cfg = TrainConfig()
208    train(cfg)
209