UMAP bokeh code for Normconf talk

umap_talk_example.py

# pip install umap-learn!
from umap import UMAP
from sentence_transformers import SentenceTransformer
from bokeh.models import ColumnDataSource
from bokeh.plotting import figure, output_file, output_notebook, show


# get your text data into a list
bodyprompts = ["a list of text strings with body in them", "another string"]


model = SentenceTransformer('all-distilroberta-v1')
X = model.encode(bodyprompts)


umap = UMAP()
Xtfm = umap.fit_transform(X)


# put in a df to use in plot

bdf = pd.DataFrame()
bdf['x'] = Xtfm[:, 0]
bdf['y'] = Xtfm[:, 1]
bdf['text'] = bodyprompts


# to display in a notebook - after the loads above, in it's own cell:
# output_notebook()

# plain tooltips on hover
TOOLTIPS = [
    ("text", "@text"),
]

# or use html to make them more readable
TOOLTIPS = """
<div style="max-width: 400px; word-wrap: break-word;">
    <span style=" font-weight: bold;">👉 @text</span>
</div>
"""

output_file("body_prompts.html", title="Body Prompts, A Sample")  # if in a notebook and wanting inline, you can comment out
source = ColumnDataSource(bdf)

# set your range bounds as the data requires:
p = figure(title='"Body Prompts from Stable Diffusion Sample',
           x_range=(-5,20), y_range=(-5, 12), width=900, height=800, tooltips=TOOLTIPS)
p.scatter(x='x', y='y', size=3, source=source, alpha=0.8)
p.xaxis[0].axis_label = 'X'
p.yaxis[0].axis_label = 'Y'

show(p)