marimo-team/fast-bulk

1# /// script2# requires-python = "==3.12"3# dependencies = [4#     "marimo",5#     "polars==1.23.0",6#     "scikit-learn==1.6.1",7#     "numpy==2.1.3",8#     "mohtml==0.1.2",9#     "model2vec==0.4.0",10#     "altair==5.5.0",11# ]12# ///13 14import marimo15 16__generated_with = "0.11.14"17app = marimo.App()18 19 20@app.cell21def _(mo):22    mo.md("""### Fast labelling demo""")23    return24 25 26@app.cell27def _(mo, use_default_switch):28    uploaded_file = mo.ui.file(kind="area") if not use_default_switch.value else None29    uploaded_file30    return (uploaded_file,)31 32 33@app.cell34def _(mo):35    use_default_switch = mo.ui.switch(False, label="Use default dataset")36    use_default_switch37    return (use_default_switch,)38 39 40@app.cell41def _(mo):42    pos_label = mo.ui.text("pos", placeholder="positive label name", label="positive class name")43    neg_label = mo.ui.text("neg", placeholder="negative label name", label="negative class name")44    return neg_label, pos_label45 46 47@app.cell48def _(uploaded_file, use_default_switch):49    should_stop = not use_default_switch.value and len(uploaded_file.value) == 050    return (should_stop,)51 52 53@app.cell54def _(mo, pl, should_stop, uploaded_file, use_default_switch):55    mo.stop(should_stop , mo.md("**Submit a dataset or use default one to continue.**"))56 57    if use_default_switch.value:58        df = pl.read_csv("spam.csv")59    else:60        df = pl.read_csv(uploaded_file.value[0].contents)61 62    texts = df["text"].to_list()63    return df, texts64 65 66@app.cell67def _(StaticModel, mo):68    with mo.status.spinner(subtitle="Loading model ...") as _spinner:69        tfm = StaticModel.from_pretrained("minishlab/potion-retrieval-32M")70    return (tfm,)71 72 73@app.cell74def _(mo, should_stop):75    mo.stop(should_stop)76 77    text_input = mo.ui.text_area("you will win a free ringtone!", label="Reference sentences")78    form = mo.md("""{text_input}""").batch(text_input=text_input).form()79    form80    return form, text_input81 82 83@app.cell84def _(mo, texts, tfm):85    with mo.status.spinner(subtitle="Creating embeddings ...") as _spinner:86        X = tfm.encode(texts)87    return (X,)88 89 90@app.cell91def _(add_label, get_example, mo, neg_label, pos_label, undo):92    btn_spam = mo.ui.button(93        label=f"Annotate {neg_label.value}", 94        on_click=lambda d: add_label(get_example(), neg_label.value), 95        keyboard_shortcut="Ctrl-L"96    )97    btn_ham = mo.ui.button(98        label=f"Annotate {pos_label.value}", 99        on_click=lambda d: add_label(get_example(), pos_label.value),100        keyboard_shortcut="Ctrl-K"101    )102    btn_undo = mo.ui.button(103        label="Undo", 104        on_click=lambda d: undo(),105        keyboard_shortcut="Ctrl-U"106    )107    return btn_ham, btn_spam, btn_undo108 109 110@app.cell111def _(gen, get_label, set_example, set_label):112    def add_label(text, lab):113        current_labels = get_label()114        set_label(current_labels + [{"text": text, "label": lab}])115        set_example(next(gen))116 117    def undo(): 118        current_labels = get_label()119        set_label(current_labels[:-2])120    return add_label, undo121 122 123@app.cell124def _():125    from mohtml import br126    return (br,)127 128 129@app.cell130def _(br, btn_ham, btn_spam, btn_undo, example, mo, neg_label, p, pos_label):131    mo.vstack([132        mo.hstack([133           pos_label, neg_label134        ]),135        br(),136        mo.hstack([137            btn_ham, btn_spam, btn_undo138        ]),139        br(),140        p("Current example:", klass="font-bold"),141        example142    ])143    return144 145 146@app.cell147def _(mo):148    get_label, set_label = mo.state([])149    return get_label, set_label150 151 152@app.cell153def _(gen, mo):154    get_example, set_example = mo.state(next(gen))155    return get_example, set_example156 157 158@app.cell159def _():160    from mohtml import tailwind_css, div, p161 162    tailwind_css()163    return div, p, tailwind_css164 165 166@app.cell167def _(get_label, mo):168    import json169 170    data = get_label()171 172    json_download = mo.download(173        data=json.dumps(data).encode("utf-8"),174        filename="data.json",175        mimetype="application/json",176        label="Download JSON",177    )178    return data, json, json_download179 180 181@app.cell182def _(X, cosine_similarity, form, get_label, mo, pl, texts, tfm):183    mo.stop(not form.value, "Need a text input to fetch example")184    mo.stop(not form.value.get("text_input", None), "Need a text input to fetch example")185 186    df_emb = (187        pl.DataFrame({188            "index": range(X.shape[0]), 189            "text": texts190        }).with_columns(sim=pl.lit(1))191    )192 193 194    query = tfm.encode([form.value["text_input"]])195    similarity = cosine_similarity(query, X)[0]196    df_emb = df_emb.with_columns(sim=similarity).sort(pl.col("sim"), descending=True)197    label_texts = [_["text"] for _ in get_label()]198    gen = (_["text"] for _ in df_emb.head(100).to_dicts() if _["text"] not in label_texts)199    return df_emb, gen, label_texts, query, similarity200 201 202@app.cell203def _(div, get_example, p):204    example = div(205        p(get_example()), 206        klass="bg-gray-100 p-4 rounded-lg"207    )208    return (example,)209 210 211@app.cell212def _(get_label, mo, pl, should_stop):213    mo.stop(should_stop)214 215    pl.DataFrame(get_label()).reverse()216    return217 218 219@app.cell220def _(mo):221    with mo.status.spinner(subtitle="Loading libraries ...") as _spinner:222        import polars as pl223        import numpy as np224        from sklearn.metrics.pairwise import cosine_similarity225    return cosine_similarity, np, pl226 227 228@app.cell229def _(mo):230    with mo.status.spinner(subtitle="Loading model2vec ...") as _spinner:231        from model2vec import StaticModel232    return (StaticModel,)233 234 235@app.cell236def _():237    import marimo as mo238    return (mo,)239 240 241@app.cell242def _():243    return244 245 246if __name__ == "__main__":247    app.run()248