Code editing task: bpe_tokenizer
Topic: Math; change kind: perfective

Original code:
from typing import Dict, List


class BPETokenizerTrainer(object):
    def __init__(self, training_set: str, max_num_merges: int) -> None:
        self.max_num_merges = max_num_merges
        self.last_token_id = 0

        self.training_set_symbolized: List[str] = []
        self.lookup_table: Dict[str, int] = {}
        for char in training_set:
            self.training_set_symbolized.append(char)
            if char not in self.lookup_table:
                self.lookup_table[char] = self.last_token_id
                self.last_token_id += 1
    
    def merge(self, new_token_text: str) -> None:
        new_symbol = new_token_text
        new_training_set_symbolized: List[str] = []
        i = 1
        while i < len(self.training_set_symbolized):
            pair_text = self.training_set_symbolized[i-1] + self.training_set_symbolized[i]
            if pair_text == new_token_text:
                new_training_set_symbolized.append(new_symbol)
                i += 1
                if i == len(self.training_set_symbolized) - 1:
                    new_training_set_symbolized.append(self.training_set_symbolized[i])
            else:
                new_training_set_symbolized.append(self.training_set_symbolized[i-1])
                if i == len(self.training_set_symbolized) - 1:
                    new_training_set_symbolized.append(self.training_set_symbolized[i])
            i += 1
        self.training_set_symbolized = new_training_set_symbolized
    
    def add_next_pair(self) -> None:
        pair_counts: Dict[str, int] = {}
        i = 1
        while i < len(self.training_set_symbolized):
            pair_text = self.training_set_symbolized[i-1] + self.training_set_symbolized[i]
            if pair_text not in pair_counts:
                pair_counts[pair_text] = 1
            else:
                pair_counts[pair_text] += 1
            i += 1
        
        most_common_pair_text = max(pair_counts, key=pair_counts.get)
        self.lookup_table[most_common_pair_text] = self.last_token_id
        self.last_token_id += 1
        self.merge(new_token_text=most_common_pair_text)
            
    
    def train(self) -> None:
        num_merges = 0
        while num_merges < self.max_num_merges and len(self.training_set_symbolized) > 1:
            self.add_next_pair()
            num_merges += 1
    
    def get_lookup_table(self) -> Dict[str, int]:
        return self.lookup_table

Instruction (descriptive):
Add a `max_num_tokens` parameter to the Trainer constructor. `max_num_tokens` should limit the max size of the `lookup_table` on the Trainer.
During training, the while loop should terminate early if the `lookup_table` reaches a length of `max_num_tokens`.

Instruction (lazy):
Add a `max_num_tokens` parameter to the Trainer which limits the number of tokens that are defined.

Edited code:
from typing import Dict, List


class BPETokenizerTrainer(object):
    def __init__(self, training_set: str, max_num_merges: int, max_num_tokens: int) -> None:
        self.max_num_merges = max_num_merges
        self.last_token_id = 0
        self.max_num_tokens = max_num_tokens

        self.training_set_symbolized: List[str] = []
        self.lookup_table: Dict[str, int] = {}
        for char in training_set:
            if len(self.lookup_table) >= self.max_num_tokens:
                break
            self.training_set_symbolized.append(char)
            if char not in self.lookup_table:
                self.lookup_table[char] = self.last_token_id
                self.last_token_id += 1
    
    def merge(self, new_token_text: str) -> None:
        new_symbol = new_token_text
        new_training_set_symbolized: List[str] = []
        i = 1
        while i < len(self.training_set_symbolized):
            pair_text = self.training_set_symbolized[i-1] + self.training_set_symbolized[i]
            if pair_text == new_token_text:
                new_training_set_symbolized.append(new_symbol)
                i += 1
                if i == len(self.training_set_symbolized) - 1:
                    new_training_set_symbolized.append(self.training_set_symbolized[i])
            else:
                new_training_set_symbolized.append(self.training_set_symbolized[i-1])
                if i == len(self.training_set_symbolized) - 1:
                    new_training_set_symbolized.append(self.training_set_symbolized[i])
            i += 1
        self.training_set_symbolized = new_training_set_symbolized
    
    def add_next_pair(self) -> None:
        pair_counts: Dict[str, int] = {}
        i = 1
        while i < len(self.training_set_symbolized):
            pair_text = self.training_set_symbolized[i-1] + self.training_set_symbolized[i]
            if pair_text not in pair_counts:
                pair_counts[pair_text] = 1
            else:
                pair_counts[pair_text] += 1
            i += 1
        
        most_common_pair_text = max(pair_counts, key=pair_counts.get)
        self.lookup_table[most_common_pair_text] = self.last_token_id
        self.last_token_id += 1
        self.merge(new_token_text=most_common_pair_text)
            
    
    def train(self) -> None:
        num_merges = 0
        while num_merges < self.max_num_merges and len(self.training_set_symbolized) > 1 and len(self.lookup_table) < self.max_num_tokens:
            self.add_next_pair()
            num_merges += 1
    
    def get_lookup_table(self) -> Dict[str, int]:
        return self.lookup_table

Tests:
### START TESTS ###
if True:  # pragma: no cover
    training_set = "Think slow when you write in ink"
    trainer0 = BPETokenizerTrainer(training_set=training_set, max_num_merges=250, max_num_tokens=100)
    assert len(trainer0.get_lookup_table()) == 15

    assert "in" not in trainer0.get_lookup_table()
    trainer0.add_next_pair()
    assert len(trainer0.get_lookup_table()) == 16
    assert "in" in trainer0.get_lookup_table()
    trainer0.merge("in")
    assert len(trainer0.get_lookup_table()) == 16

    assert "ink" not in trainer0.get_lookup_table()
    trainer0.add_next_pair()
    assert len(trainer0.get_lookup_table()) == 17
    assert "ink" in trainer0.get_lookup_table()
    trainer0.merge("ink")
    assert len(trainer0.get_lookup_table()) == 17

    assert " w" not in trainer0.get_lookup_table()
    trainer0.add_next_pair()
    assert len(trainer0.get_lookup_table()) == 18
    assert " w" in trainer0.get_lookup_table()
    trainer0.merge(" w")

    trainer1 = BPETokenizerTrainer(training_set=training_set, max_num_merges=5, max_num_tokens=100)
    assert set(trainer1.get_lookup_table().keys()) == set([c for c in training_set])
    trainer1.train()
    assert set(trainer1.get_lookup_table().keys()) == set([c for c in training_set] + ["in", "ink", " w", "Th", "Think"])

    trainer2 = BPETokenizerTrainer(training_set=training_set, max_num_merges=5, max_num_tokens=10)
    assert set(trainer2.get_lookup_table().keys()) == set([c for c in training_set[:10]])
    trainer2.train()
    assert set(trainer2.get_lookup_table().keys()) == set([c for c in training_set[:10]])

    trainer3 = BPETokenizerTrainer(training_set=training_set, max_num_merges=100, max_num_tokens=18)
    assert set(trainer3.get_lookup_table().keys()) == set([c for c in training_set])
    trainer3.train()
    assert set(trainer3.get_lookup_table().keys()) == set([c for c in training_set] + ["in", "ink", " w"])