Spaces:
Running
Running
import umap | |
import plotly.express as px | |
import numpy as np | |
class EmbeddingPlotter: | |
"""Lower the dimensionality of the representation from 768 -> 2, over the surface of the sphere""" | |
def transform(self, df, embeddings): | |
"""Plotting pipleine""" | |
df = self.umap_embedding(df, embeddings) | |
fig1, fig2 = self.plot(df) | |
return fig1, fig2 | |
def umap_embedding(self, df, embeddings): | |
"""Generate 2D embeddings""" | |
# UMAP - Spherical | |
sphere_mapper = umap.UMAP(output_metric="haversine", random_state=42).fit( | |
np.array(embeddings) | |
) | |
df["spherical_emb_X"] = np.sin(sphere_mapper.embedding_[:, 0]) * np.cos( | |
sphere_mapper.embedding_[:, 1] | |
) | |
df["spherical_emb_Y"] = np.sin(sphere_mapper.embedding_[:, 0]) * np.sin( | |
sphere_mapper.embedding_[:, 1] | |
) | |
df["spherical_emb_Z"] = np.cos(sphere_mapper.embedding_[:, 0]) | |
# UMAP - Lambert Conformal | |
df["lambert_conformal_emb_x"] = np.arctan2( | |
df["spherical_emb_X"], df["spherical_emb_Y"] | |
) | |
df["lambert_conformal_emb_y"] = -np.arccos(df["spherical_emb_Z"]) | |
return df | |
def plot(self, df): | |
"""Plot embeddings over representation space""" | |
# on the 3d sphere | |
fig1 = px.scatter_3d( | |
df, | |
x="spherical_emb_X", | |
y="spherical_emb_Y", | |
z="spherical_emb_Z", | |
color="categories_parsed", | |
) | |
# on the projected spehre | |
fig2 = px.scatter( | |
data_frame=df, | |
x="lambert_conformal_emb_x", | |
y="lambert_conformal_emb_y", | |
color="categories_parsed", | |
) | |
return fig1, fig2 | |