Skip to content

lm_state

genlm.bytes.byte_lm.lm_state

StatefulTokenizedLM

A stateful tokenized language model that maintains context and generates next token logprobs.

Parameters:

Name Type Description Default
model AsyncLM

The underlying language model

 required
context list

List of token IDs representing the current context

 required
n_calls int

Number of times the model has been called

 0
max_context_length int

Maximum length of context to maintain

 None
Source code in genlm/bytes/byte_lm/lm_state.py 
 9
10
11
12
13
14
15
16
17
18
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
class StatefulTokenizedLM:
    """A stateful tokenized language model that maintains context and generates next token logprobs.

    Args:
        model (genlm.backend.AsyncLM): The underlying language model
        context (list): List of token IDs representing the current context
        n_calls (int): Number of times the model has been called
        max_context_length (int, optional): Maximum length of context to maintain
    """

    def __init__(self, model, context, n_calls=0, max_context_length=None):
        self.model = model
        self.context = context
        self._n_calls = n_calls
        self.max_context_length = max_context_length

    @classmethod
    def initial(cls, model, initial_context=None, max_context_length=None):
        """Creates an initial state for the language model.

        Args:
            model (genlm.backend.AsyncLM): The language model to use
            initial_context (list, optional): Initial context of token IDs. Defaults to [tokenizer.bos_token_id]
            max_context_length (int, optional): Maximum context length to maintain

        Returns:
            (StatefulTokenizedLM): A new instance with initial state
        """
        if initial_context is None:
            initial_context = [model.tokenizer.bos_token_id]
        return cls(model, initial_context, max_context_length=max_context_length)

    def __lshift__(self, token):
        """Adds a new token to the context and returns a new state.

        Args:
            token (int): Token ID to add to context

        Returns:
            (StatefulTokenizedLM): New state with updated context
        """
        assert isinstance(token, int)
        if (
            self.max_context_length is not None
            and len(self.context) >= self.max_context_length
        ):
            self.context = self.context[-(self.max_context_length - 1) :]
        return StatefulTokenizedLM(
            self.model, self.context + [token], n_calls=self._n_calls
        )

    async def logp_next(self):
        """Computes log probabilities for the next token given the current context.

        Returns:
            (torch.Tensor): Log probabilities for next tokens
        """
        self._n_calls += 1
        return await self.model.next_token_logprobs(self.context)

    def __repr__(self):
        return colors.purple % (
            "|".join([escape(self.model.byte_vocab[x]) for x in self.context])
        )

initial(model, initial_context=None, max_context_length=None) classmethod

Creates an initial state for the language model.

Parameters:

Name Type Description Default
model AsyncLM

The language model to use

 required
initial_context list

Initial context of token IDs. Defaults to [tokenizer.bos_token_id]

 None
max_context_length int

Maximum context length to maintain

 None

Returns:

Type Description
StatefulTokenizedLM

A new instance with initial state
Source code in genlm/bytes/byte_lm/lm_state.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
@classmethod
def initial(cls, model, initial_context=None, max_context_length=None):
    """Creates an initial state for the language model.

    Args:
        model (genlm.backend.AsyncLM): The language model to use
        initial_context (list, optional): Initial context of token IDs. Defaults to [tokenizer.bos_token_id]
        max_context_length (int, optional): Maximum context length to maintain

    Returns:
        (StatefulTokenizedLM): A new instance with initial state
    """
    if initial_context is None:
        initial_context = [model.tokenizer.bos_token_id]
    return cls(model, initial_context, max_context_length=max_context_length)

__lshift__(token)

Adds a new token to the context and returns a new state.

Parameters:

Name Type Description Default
token int

Token ID to add to context

 required

Returns:

Type Description
StatefulTokenizedLM

New state with updated context
Source code in genlm/bytes/byte_lm/lm_state.py 
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
def __lshift__(self, token):
    """Adds a new token to the context and returns a new state.

    Args:
        token (int): Token ID to add to context

    Returns:
        (StatefulTokenizedLM): New state with updated context
    """
    assert isinstance(token, int)
    if (
        self.max_context_length is not None
        and len(self.context) >= self.max_context_length
    ):
        self.context = self.context[-(self.max_context_length - 1) :]
    return StatefulTokenizedLM(
        self.model, self.context + [token], n_calls=self._n_calls
    )

logp_next() async

Computes log probabilities for the next token given the current context.

Returns:

Type Description
Tensor

Log probabilities for next tokens
Source code in genlm/bytes/byte_lm/lm_state.py 
60
61
62
63
64
65
66
67
async def logp_next(self):
    """Computes log probabilities for the next token given the current context.

    Returns:
        (torch.Tensor): Log probabilities for next tokens
    """
    self._n_calls += 1
    return await self.model.next_token_logprobs(self.context)

StatefulByteLM

Bases: ABC

Abstract base class for byte-level language models with state.

Source code in genlm/bytes/byte_lm/lm_state.py 
 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
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
class StatefulByteLM(ABC):
    """Abstract base class for byte-level language models with state."""

    @abstractmethod
    async def __lshift__(self, b: int):
        """Adds a byte to the state and returns new state.

        Args:
            b (int): Byte to add

        Returns:
            (StatefulByteLM): New state with added byte
        """
        pass

    def prune(self):
        """Prunes the current state if needed.

        Override in subclasses.

        Returns:
            (StatefulByteLM): Pruned state
        """
        return self

    @abstractmethod
    async def logp_next(self):
        """Computes the log probability distribution for the next byte.

        Returns:
            (LazyByteProbs): Log probabilities for next possible bytes
        """
        pass

    async def prefill(self, bs):
        """Prefills the model state with a sequence of bytes.

        Args:
            bs (list[int]): Sequence of bytes to add to state

        Returns:
            (StatefulByteLM): New state with all bytes added
        """
        state = self
        for b in bs:
            state = await (state.prune() << b)
        return state

    async def greedy(self, context, steps):
        """Performs greedy decoding for given number of steps.

        Args:
            context (bytes): Initial byte context
            steps (int): Number of generation steps

        Returns:
            (bytes): Generated byte sequence
        """
        context = list(context)
        state = await self.prefill(context)
        for _ in range(steps):
            Q = (await state.logp_next()).materialize()
            b = Q.argmax()
            state = await (state.prune() << b)
            context.append(b)
        return bytes(context)

    async def sample(self, context, steps, draw=sample_dict):
        """Samples from the model for given number of steps.

        Args:
            context (bytes): Initial byte context
            steps (int): Number of generation steps
            draw: Sampling function to use (default: sample_dict)

        Returns:
            (bytes): Generated byte sequence
        """
        context = list(context)
        state = await self.prefill(context)
        for _ in range(steps):
            Q = (await state.logp_next()).materialize()
            b = draw(Q.map_values(exp))
            state = await (state.prune() << b)
            context.append(b)
        return bytes(context)

    async def cleanup(self):
        """Performs any necessary cleanup of the model state."""
        pass

__lshift__(b) abstractmethod async

Adds a byte to the state and returns new state.

Parameters:

Name Type Description Default
b int

Byte to add

 required

Returns:

Type Description
StatefulByteLM

New state with added byte
Source code in genlm/bytes/byte_lm/lm_state.py 
78
79
80
81
82
83
84
85
86
87
88
@abstractmethod
async def __lshift__(self, b: int):
    """Adds a byte to the state and returns new state.

    Args:
        b (int): Byte to add

    Returns:
        (StatefulByteLM): New state with added byte
    """
    pass

prune()

Prunes the current state if needed.

Override in subclasses.

Returns:

Type Description
StatefulByteLM

Pruned state
Source code in genlm/bytes/byte_lm/lm_state.py 
90
91
92
93
94
95
96
97
98
def prune(self):
    """Prunes the current state if needed.

    Override in subclasses.

    Returns:
        (StatefulByteLM): Pruned state
    """
    return self

logp_next() abstractmethod async

Computes the log probability distribution for the next byte.

Returns:

Type Description
LazyByteProbs

Log probabilities for next possible bytes
Source code in genlm/bytes/byte_lm/lm_state.py 
100
101
102
103
104
105
106
107
@abstractmethod
async def logp_next(self):
    """Computes the log probability distribution for the next byte.

    Returns:
        (LazyByteProbs): Log probabilities for next possible bytes
    """
    pass

prefill(bs) async

Prefills the model state with a sequence of bytes.

Parameters:

Name Type Description Default
bs list[int]

Sequence of bytes to add to state

 required

Returns:

Type Description
StatefulByteLM

New state with all bytes added
Source code in genlm/bytes/byte_lm/lm_state.py 
109
110
111
112
113
114
115
116
117
118
119
120
121
async def prefill(self, bs):
    """Prefills the model state with a sequence of bytes.

    Args:
        bs (list[int]): Sequence of bytes to add to state

    Returns:
        (StatefulByteLM): New state with all bytes added
    """
    state = self
    for b in bs:
        state = await (state.prune() << b)
    return state

greedy(context, steps) async

Performs greedy decoding for given number of steps.

Parameters:

Name Type Description Default
context bytes

Initial byte context

 required
steps int

Number of generation steps

 required

Returns:

Type Description
bytes

Generated byte sequence
Source code in genlm/bytes/byte_lm/lm_state.py 
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
async def greedy(self, context, steps):
    """Performs greedy decoding for given number of steps.

    Args:
        context (bytes): Initial byte context
        steps (int): Number of generation steps

    Returns:
        (bytes): Generated byte sequence
    """
    context = list(context)
    state = await self.prefill(context)
    for _ in range(steps):
        Q = (await state.logp_next()).materialize()
        b = Q.argmax()
        state = await (state.prune() << b)
        context.append(b)
    return bytes(context)

sample(context, steps, draw=sample_dict) async

Samples from the model for given number of steps.

Parameters:

Name Type Description Default
context bytes

Initial byte context

 required
steps int

Number of generation steps

 required
draw

Sampling function to use (default: sample_dict)

 sample_dict

Returns:

Type Description
bytes

Generated byte sequence
Source code in genlm/bytes/byte_lm/lm_state.py 
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
async def sample(self, context, steps, draw=sample_dict):
    """Samples from the model for given number of steps.

    Args:
        context (bytes): Initial byte context
        steps (int): Number of generation steps
        draw: Sampling function to use (default: sample_dict)

    Returns:
        (bytes): Generated byte sequence
    """
    context = list(context)
    state = await self.prefill(context)
    for _ in range(steps):
        Q = (await state.logp_next()).materialize()
        b = draw(Q.map_values(exp))
        state = await (state.prune() << b)
        context.append(b)
    return bytes(context)

cleanup() async

Performs any necessary cleanup of the model state.

Source code in genlm/bytes/byte_lm/lm_state.py 
162
163
164
async def cleanup(self):
    """Performs any necessary cleanup of the model state."""
    pass