embeddings_weave.py

!curl -O https://calmcode.io/datasets/pokemon.json

import pandas
import weave
from weave.ecosystem import openai
from weave.ecosystem import umap
from weave.ecosystem import hdbscan

raw_data = pandas.read_json('./pokemon.json')

data = weave.save(weave.ops.dataframe_to_arrow(raw_data), 'pokemon_data')
embeddings = openai.openai_embed(data['name'].map), {"model": "text-embedding-ada-002"})
clusterable_projection = umap.umap_project(
    embeddings, {
        'n_neighbors': 30,
        'min_dist': 0,
        'n_components': 2,
    }
)
clusters = hdbscan.hdbscan_cluster(clusterable_projection, {
            'min_samples': 10,
            'min_cluster_size': 50
        })
projection = umap.umap_project(embeddings, {})

weave.show([{'x': x, 'y': y, 'k':k, 'd':d} for (x,y),k,d in 
            zip(weave.use(projection),
                weave.use(clusters),
                weave.use(data))])