jampekka/slms_tutorial

1import marimo as mo2 3from pprint import pformat4from collections import defaultdict, Counter5import random6 7_pre_box_height = "10em";8_font_size = "12px";9def pre_box(text):10    return mo.Html(f"""11<pre class="pre_out_box" style="overflow: auto; height: {_pre_box_height}; font-size: {_font_size};">12{text}13 14</pre>""")15 16def python_out(code):17    return mo.Html(f"""18<pre class="python_out_box" style="overflow: auto; height: {_pre_box_height}; font-size: {_font_size};">19{pformat(code, sort_dicts=False, compact=True)}20 21</pre>""")22 23def tokens_out(tokens, tokenizer):24    out = ""25    for i, string in enumerate(tokenizer.tokens_to_strings(tokens)):26        #colors = ["rgb(20, 184, 166)", "rgb(245, 158, 11)"]27        colors = [28                "#2b9a66",29                #"#26997b",30                "#00749e",31                "#dc3e42",32        ]33        #colors = "#d1f0fa", "#ffcdce"34        colors = "var(--sky-3)", "var(--red-3)", "var(--amber-3)"35        color = colors[i%len(colors)]36        # TODO: Be more general!37        if string == ' ':38            decoration = "underline"39        else:40            decoration = "none"41 42        n_newlines = string.count('\n')43        string = string.replace("\n", "\\n")44        string += "\n"*n_newlines45 46        out += f'<span style="background-color: {color}; text-decoration: {decoration}">{string}</span>'47    out = f'<div style="overflow: auto; height: {_pre_box_height};">{out}</div>'48    return pre_box(out)49 50def corpus_to_vocabulary(tokens):51    # Using dict instead of set to keep the order52    return list({w: None for w in tokens}.keys())53 54init_output = mo.Html(f"""55    <style>56    .python_out_box {{57        overflow: auto !important;58        max_height: {_pre_box_height};59        font-size: 12px;60    }}61 62    .pre_out_box {{63        overflow: auto !important;64        height: {_pre_box_height};65        font-size: 12px;66    }}67    </style>68    """)69init_output = None70 71def graph_out(svg):72    return mo.Html(f"""73        <div style="overflow: auto; max-height: 32em;">74        {svg}75        </div>76    """)77 78def plot_follower_graph(next_words):79    import pydot80 81    graph = pydot.Dot("follower_graph", ordering="in")82    def mangle(s):83        #if isinstance(s, tuple) and len(s) == 1:84        #    s = s[0]85        return repr(s).replace(r'\n', r'\\n')86    for context, followers in next_words.items():87        graph.add_node(pydot.Node(mangle(context)))88        for follower in followers:89            edge = graph.add_edge(pydot.Edge(mangle(context), mangle(follower)))90            # A bit of a hack91            #if hasattr(followers, 'get'):92            #    edge.set_label(followers.get(follower))93            #else:94            #    count = None95            96    svg = graph.create_svg().decode('utf-8')97    return graph_out(svg)98 99def plot_follower_context_graph(next_words):100    # TODO: This is fugly. Use dot101    import pydot102 103    graph = pydot.Dot("follower_graph", ordering="in", strict=True)104    def mangle(s):105        #if isinstance(s, tuple) and len(s) == 1:106        #    s = s[0]107        return repr(s).replace(r'\n', r'\\n')108    for context, followers in next_words.items():109        #graph.add_node(pydot.Node(mangle(context)))110        for follower in followers:111            # A bit of a hack112            #edge = graph.add_edge(pydot.Edge(mangle(context), mangle(follower)))113            new_context = (*context[1:], follower)114            for follower in next_words.get(context, []):115                follower_context = (*context[1:], follower)116                graph.add_edge(pydot.Edge(117                    mangle(context),118                    mangle(follower_context),119                    label=mangle(follower)120                ))121 122    svg = graph.create_svg().decode('utf-8')123    return graph_out(svg)124 125def generate_tokens(next_words, context=None, max_tokens=200, seed=3):126    rng = random.Random(seed)127 128    if context is None:129        context = next(iter(next_words.keys()))130    yield from context131 132    for i in range(max_tokens):133        candidates = next_words.get(context, None)134        if not candidates: return135 136        choices, counts = zip(*candidates.items())137        if not choices: return138        next_word = rng.choice(choices)139        if next_word == '\n\n': return140        yield next_word141        context = (*context[1:], next_word)142 143# Doing this more succintly now144def get_ngrams(tokens, n):145    for i in range(len(tokens) - n + 1):146        yield tokens[i:i+n]147 148def get_next_token_table(tokens, context_length, table=None):149    if table is None:150        table = defaultdict(Counter)151    for *context, next_token in get_ngrams(tokens, context_length + 1):152        table[tuple(context)][next_token] += 1153    154    return table155 156happy_birthday_text = """157Happy birthday to you 158Happy birthday to you 159Happy birthday dear Dave 160Happy birthday to you 161"""162 163blowin_text = """164Yes, and how many roads must a man walk down, before you call him a man?165And how many seas must a white dove sail, before she sleeps in the sand?166Yes, and how many times must the cannonballs fly, before they're forever banned?167 168Yes, and how many years must a mountain exist, before it is washed to the sea?169And how many years can some people exist, before they're allowed to be free?170Yes, and how many times can a man turn his head, and pretend that he just doesn't see?171 172Yes, and how many times must a man look up, before he can see the sky?173And how many ears must one man have, before he can hear people cry?174Yes, and how many deaths will it take 'til he knows, that too many people have died?175"""176 177blowin_text_finnish = """178Niin, ja kuinka monta tietä miehen täytyy kävellä, ennen kuin kutsut häntä mieheksi?179Ja kuinka monta merta valkoisen kyyhkysen täytyy purjehtia, ennen kuin se nukkuu hiekkaan?180Kyllä, ja kuinka monta kertaa kanuunankuulat täytyy lentää, ennen kuin ne on ikuisesti kielletty?181 182Kyllä, ja kuinka monta vuotta vuoren on oltava olemassa, ennen kuin se huuhtoutuu mereen?183Ja kuinka monta vuotta jotkut ihmiset voivat olla olemassa ennen kuin he saavat olla vapaita?184Kyllä, ja kuinka monta kertaa ihminen voi kääntää päätään ja teeskennellä, ettei hän vain näe?185 186Kyllä, ja kuinka monta kertaa miehen täytyy katsoa ylös, ennen kuin hän voi nähdä taivaan?187Ja kuinka monta korvaa yhdellä ihmisellä pitää olla, ennen kuin hän voi kuulla ihmisten itkevän?188Kyllä, ja kuinka monta kuolemaa kestää, ennen kuin hän tietää, että liian monta ihmistä on kuollut?189"""190 191blowin_text_german = """192Ja, und wie viele Wege muss ein Mann gehen, bevor man ihn einen Mann nennt?193Und wie viele Meere muss eine weiße Taube durchsegeln, bevor sie im Sand schläft?194Ja, und wie oft müssen die Kanonenkugeln fliegen, bevor sie für immer verboten werden?195 196Ja, und wie viele Jahre muss ein Berg existieren, bevor er ins Meer gespült wird?197Und wie viele Jahre können manche Menschen existieren, bevor sie frei sein dürfen?198Ja, und wie oft kann ein Mann den Kopf drehen und so tun, als würde er einfach nichts sehen?199 200Ja, und wie oft muss ein Mensch nach oben schauen, bevor er den Himmel sehen kann?201Und wie viele Ohren muss ein Mann haben, bevor er Menschen weinen hören kann?202Ja, und wie viele Todesfälle wird es dauern, bis er weiß, dass zu viele Menschen gestorben sind?203"""204