jampekka/slms_tutorial
0
1import marimo2 3__generated_with = "0.6.8"4app = marimo.App(app_title="SLMs basics")5 6 7@app.cell8def __():9 import marimo as mo10 from pprint import pformat11 from collections import defaultdict12 import utils as U13 14 U.init_output15 return U, defaultdict, mo, pformat16 17 18@app.cell19def __(mo):20 mo.md(21 r"""22 # Small language models23 24 ## Happy birthday25 ---26 To get started, we analyze lyrics of perhaps the most popular song in the world.27 You may be familiar with the lyrics:28 """29 )30 return31 32 33@app.cell34def __():35 corpus_text = """36 Happy birthday to you 37 Happy birthday to you 38 Happy birthday dear Dave 39 Happy birthday to you 40 """41 corpus_text42 return corpus_text,43 44 45@app.cell46def __(mo):47 mo.md(48 r"""49 To work with text, we usually want to split it to some shorter pieces, such50 as words. In general, such pieces are called **tokens**, but we'll start with just51 words. Our lyrics split into words become:52 """53 )54 return55 56 57@app.cell58def __(U, corpus_text):59 corpus_words = corpus_text.split(' ')60 U.python_out(corpus_words)61 return corpus_words,62 63 64@app.cell65def __(mo):66 mo.md(67 rf"""68 (The here `'\n'` means that we start a new line. While not really a word, we treat it as such for now.)69 70 We can also build our **vocabulary**, which is just all individual words that is in our lyrics:71 """72 )73 return74 75 76@app.cell77def __(U, corpus_words):78 # Using dict instead of set to keep the order79 _vocabulary = {w: None for w in corpus_words}.keys()80 U.python_out(list(_vocabulary))81 return82 83 84@app.cell85def __(mo):86 mo.md(87 r"""88 The currently popular large language models (LLMs) -- such as GPT, Llama and Mistral -- are based on predicting what token becomes after89 some number of tokens.90 91 In our case, for example, the word `'Happy'` is followerd by the word `'birthday'` and the92 word `'birthday'` is followed by the word `'to'`.93 94 In fact, to make an extremely simple language model, we can just list what words are followed by each95 word. For our lyrics this becomes:96 """97 )98 return99 100 101@app.cell102def __(U, corpus_words):103 next_words = {}104 for i in range(len(corpus_words)-1):105 word = corpus_words[i]106 next_word = corpus_words[i+1]107 if word not in next_words:108 next_words[word] = []109 next_words[word].append(next_word)110 U.python_out(next_words)111 return i, next_word, next_words, word112 113 114@app.cell115def __(mo):116 mo.md(r"Or as a visual graph format:")117 return118 119 120@app.cell121def __(U, next_words):122 U.plot_follower_graph(next_words)123 return124 125 126@app.cell127def __(mo):128 mo.md(129 r"""130 We can see that after a new line `'\n'` we always get the word `'Happy'`, and `'Happy'` is always followed by131 `'birthday'`. Somewhat more interestingly, the word `'birthday'` was followed three times by `'to'` but also132 once by `'dear'`.133 134 With this model, we are ready to generate new lyrics! Select the next word from the dropdown135 to add it into the lyrics.136 """137 )138 return139 140 141@app.cell142def __(corpus_words, mo):143 initial_lyrics_birthday = tuple(corpus_words[:2])144 get_lyrics_birthday, set_lyrics_birthday = mo.state(initial_lyrics_birthday, allow_self_loops=True)145 return get_lyrics_birthday, initial_lyrics_birthday, set_lyrics_birthday146 147 148@app.cell149def __(mo):150 def dropdown_generate(next_words, lyrics_state, initial_lyrics):151 get_lyrics, set_lyrics = lyrics_state152 lyrics = get_lyrics()153 options = set(next_words[lyrics[-1]])154 def update(value):155 new_lyrics = (*get_lyrics(), value)156 set_lyrics((*get_lyrics(), value))157 158 lyrics_text = ' ' + ' '.join(get_lyrics())159 optvals = {repr(o): o for o in options}160 dropdown = mo.ui.dropdown(options=optvals, on_change=update)161 reset = mo.ui.button(162 label="Reset lyrics",163 on_change=lambda *args: set_lyrics(initial_lyrics)164 )165 166 #lyrics_el = mo.Html(f"<pre>{lyrics_text} {dropdown}</pre>")167 return dropdown, reset168 return dropdown_generate,169 170 171@app.cell172def __(173 dropdown_generate,174 get_lyrics_birthday,175 initial_lyrics_birthday,176 mo,177 next_words,178 set_lyrics_birthday,179):180 # These have to be globals for the events to be triggered.181 # Marimo has some ways to go to enable modular code182 dropdown_birthday, reset_birthday = dropdown_generate(next_words, (get_lyrics_birthday, set_lyrics_birthday), initial_lyrics_birthday)183 _text = ' '.join(get_lyrics_birthday())184 _lyrics_el = mo.Html(f"<pre>{_text} {dropdown_birthday}</pre>")185 186 mo.hstack([_lyrics_el, reset_birthday])187 return dropdown_birthday, reset_birthday188 189 190@app.cell191def __(mo):192 mo.md(193 rf"""194 ## Blowin' in the wind195 ---196 197 The previous looked only one word at the time. However, we can easily use more than one word to predict the next one. How many words (or tokens) we use to predict the next one, is known as the **context length**. The context length of the previous example was 1.198 199 With the very simple lyrics context length more than 1 does not make much sense, so let's pick something a bit more complicated:200 """201 )202 return203 204 205@app.cell206def __():207 blowin_text = """208 Yes, and how many roads must a man walk down, before you call him a man? 209 And how many seas must a white dove sail, before she sleeps in the sand? 210 Yes, and how many times must the cannonballs fly, before they're forever banned? 211 212 Yes, and how many years must a mountain exist, before it is washed to the sea? 213 And how many years can some people exist, before they're allowed to be free? 214 Yes, and how many times can a man turn his head, and pretend that he just doesn't see? 215 216 Yes, and how many times must a man look up, before he can see the sky? 217 And how many ears must one man have, before he can hear people cry? 218 Yes, and how many deaths will it take 'til he knows, that too many people have died? 219 """220 blowin_text221 return blowin_text,222 223 224@app.cell225def __(mo):226 mo.md(227 rf"""228 You may recognize the lyrics. They're the verses of the Bob Dylan's song [Blowin' in the Wind](https://www.youtube.com/watch?v=MMFj8uDubsE).229 230 We proceed like before, first splitting the lyrics into words:231 """232 )233 return234 235 236@app.cell237def __(U, blowin_text):238 blowin_words = blowin_text.split(' ')239 U.python_out(blowin_words)240 return blowin_words,241 242 243@app.cell244def __(mo):245 mo.md(246 rf"""247 Note that we now have punctuation included in the ''words'', like the comma in `'Yes,'` the question mark in `'man?'`. We also treat two newlines `'\n\n'` as one ''word''. This comes handy, as it separates the verses.248 249 We now have quite a bit larger vocabulary:250 """251 )252 return253 254 255@app.cell256def __(U, blowin_words):257 U.python_out(list(U.corpus_to_vocabulary(blowin_words)))258 return259 260 261@app.cell262def __(mo):263 mo.md(264 rf"""265 ### More context266 ---267 We build a simple language model again with these lyrics. These simple models are usually called ''Markov Chain text generators''. This is a bit misleading, as even the next-token-predicting LLMs are Markov chains. We won't discuss what Markov chains really are and what makes a model such, but Wikipedia has a [rather good article](https://en.wikipedia.org/wiki/Markov_chain) of these if you're interested. 268 269 Previously in the ''Happy Birthday'' example the model looked only one word at the time. However, we can easily use more than one word to predict the next one. How many words (or tokens) we use to predict the next one, is known as the **context length**. The context length of the previous example was 1.270 271 For lyrics as simple as in ''Happy Birthday'' using a context length more than 1 didn't make much sense. However, with the more complicated lyrics we can see how the model behavior changes with different context lengths.272 273 You can select the context length with the slider and see how the model changes.274 """275 )276 return277 278 279@app.cell280def __(context_length_slider, mo):281 mo.md(f"The context length is {context_length_slider.value}")282 return283 284 285@app.cell286def __(mo):287 # TODO: Display context length value288 context_length_slider = mo.ui.slider(start=1, stop=8, full_width=True)289 context_length_slider290 return context_length_slider,291 292 293@app.cell294def __(blowin_words, context_length_slider, defaultdict):295 #blowin_context_length = 2296 blowin_context_length = context_length_slider.value297 # Doing this more succintly now298 def get_ngrams(tokens, n):299 for i in range(len(tokens) - n + 1):300 yield tokens[i:i+n]301 302 blowin_next_words1 = defaultdict(list) 303 for *_context, _next_word in get_ngrams(blowin_words, blowin_context_length + 1):304 blowin_next_words1[tuple(_context)].append(_next_word)305 306 #python_out(dict(blowin_next_words1))307 return blowin_context_length, blowin_next_words1, get_ngrams308 309 310@app.cell311def __():312 #plot_follower_graph(blowin_next_words1)313 return314 315 316@app.cell317def __(mo):318 mo.md(rf"We can now generate some lyrics with the model. Here's some machine generated ones, you can do your own below.")319 return320 321 322@app.cell323def __():324 import random325 random.seed(3)326 return random,327 328 329@app.cell330def __(mo):331 regen_blowin1_btn = mo.ui.button(label="Generate new verse")332 regen_blowin1_btn333 return regen_blowin1_btn,334 335 336@app.cell337def genblow1_1(U, blowin_next_words1, random, regen_blowin1_btn):338 # TODO: Keep the seed constant across generations339 340 regen_blowin1_btn341 342 def _generate(next_words):343 context = next(iter(next_words.keys()))344 yield from context345 346 while True:347 choices = next_words[context]348 if not choices: return349 next_word = random.choice(choices)350 if next_word == '\n\n': return351 yield next_word352 context = (*context[1:], next_word)353 354 _generated = list(_generate(blowin_next_words1))355 U.pre_box(' '.join(_generated))356 return357 358 359@app.cell360def __(U, blowin_next_words1, mo):361 mo.accordion({362 "Next word table": U.python_out(dict(blowin_next_words1)),363 "Next word graph": U.plot_follower_graph(blowin_next_words1)364 })365 return366 367 368@app.cell369def __(mo):370 mo.md(371 rf"""372 With a short context length the lyrics dont make much sense. With a longer context length it starts to just copy the originals. Try to find a context length that seems to make a nice tradeoff between these. As a hint, you can get something quite silly with some context lengths.373 374 Try to be such a language model yourself! This time the generated lyrics are hidden. Don't peek at them before you're done, and pretend you don't remember what you picked before!375 """376 )377 return378 379 380@app.cell381def __(blowin_context_length, blowin_words, mo):382 initial_lyrics_blowin = blowin_words[:blowin_context_length + 1]383 get_lyrics_blowin1, set_lyrics_blowin1 = mo.state(initial_lyrics_blowin, allow_self_loops=True)384 return get_lyrics_blowin1, initial_lyrics_blowin, set_lyrics_blowin1385 386 387@app.cell388def __(389 blowin_context_length,390 blowin_next_words1,391 get_lyrics_blowin1,392 initial_lyrics_blowin,393 mo,394 set_lyrics_blowin1,395):396 def dropdown_generate_blowin(next_words, lyrics_state, initial_lyrics):397 get_lyrics, set_lyrics = lyrics_state398 lyrics = get_lyrics()399 context = tuple(lyrics[-blowin_context_length:])400 options = set(next_words[context])401 def update(value):402 new_lyrics = (*get_lyrics(), value)403 set_lyrics((*get_lyrics(), value))404 405 lyrics_text = ' ' + ' '.join(get_lyrics())406 optvals = {repr(o): o for o in options}407 dropdown = mo.ui.dropdown(options=optvals, on_change=update)408 reset = mo.ui.button(409 label="Reset lyrics",410 on_change=lambda *args: set_lyrics(initial_lyrics)411 )412 413 #lyrics_el = mo.Html(f"<pre>{lyrics_text} {dropdown}</pre>")414 return dropdown, reset415 416 dropdown_blowin1, reset_blowin1 = dropdown_generate_blowin(blowin_next_words1, (get_lyrics_blowin1, set_lyrics_blowin1), initial_lyrics_blowin)417 _ctx = ', '.join(map(repr, get_lyrics_blowin1()[-blowin_context_length:]))418 _lyrics_el = mo.Html(f"<pre>{_ctx} {dropdown_blowin1}</pre>")419 420 _lyrics_el421 return dropdown_blowin1, dropdown_generate_blowin, reset_blowin1422 423 424@app.cell425def __(get_lyrics_blowin1, mo, reset_blowin1):426 _lyrics = ' '.join(get_lyrics_blowin1())427 _spoiler = mo.accordion({'Your generated lyrics. SPOILER!': mo.Html(f"<pre>{_lyrics}</pre>")})428 mo.vstack([_spoiler, reset_blowin1])429 return430 431 432@app.cell433def __(mo):434 mo.md(435 rf"""436 ---437 In the next notebook, we'll take a closer look at **tokenization**, i.e. how we split the text for processing.438 439 [Continue to Tokenization >](?file=tokenization.py)440 """441 )442 return443 444 445if __name__ == "__main__":446 app.run()447