itsamejms
101 lines
2.8 KiB
Python
101 lines
2.8 KiB
Python
# %% [markdown]
|
|
# # Visualize Most Prepared Users
|
|
# This workflow loads user embeddings, generates a prompt embedding, computes similarity, and visualizes the most prepared users.
|
|
|
|
# %%
|
|
import pandas as pd
|
|
import numpy as np
|
|
from sklearn.metrics.pairwise import cosine_similarity
|
|
import matplotlib.pyplot as plt
|
|
import seaborn as sns
|
|
from utils import get_ollama_embedding
|
|
|
|
# %% [markdown]
|
|
# ## Read in the file with user vectors
|
|
# Update the path/format as needed.
|
|
# %%
|
|
df = (
|
|
pd.read_csv("./data/eurobarometer_preparedness_model_data_v2.csv").assign(
|
|
**{
|
|
"user_id": lambda x: x["country_code_iso_3166"]
|
|
+ "_"
|
|
+ x["age_recoded_6_categories"].astype(str)
|
|
+ "_"
|
|
+ x["gender"].astype(str)
|
|
}
|
|
)
|
|
# select only relevant countries for visualization
|
|
.loc[
|
|
lambda x: x["country_code_iso_3166"].isin(
|
|
["FI", "DE-E", "DE-W", "FR", "ES", "PT"]
|
|
)
|
|
]
|
|
)
|
|
df.shape
|
|
|
|
|
|
# %% [markdown]
|
|
# ## Convert string embeddings to lists if needed
|
|
# %%
|
|
def parse_embedding(x):
|
|
if isinstance(x, str):
|
|
return [float(i) for i in x.strip("[]").split(",")]
|
|
return x
|
|
|
|
|
|
df["embedding"] = df["embedding"].apply(parse_embedding)
|
|
|
|
# %%
|
|
# Select best user:
|
|
# disaster_measures_in_hh_* == 1, and how_many_days_meet_* == (4 | 5)
|
|
disaster_measures_in_hh_columns = [
|
|
col for col in df.columns if col.startswith("disaster_measures_in_hh_")
|
|
]
|
|
how_many_days_meet_columns = [
|
|
col for col in df.columns if col.startswith("how_many_days_meet_")
|
|
]
|
|
best_users = df[
|
|
(df[disaster_measures_in_hh_columns] == 1).all(axis=1)
|
|
& (df[how_many_days_meet_columns].isin([4, 5]).all(axis=1))
|
|
]
|
|
print(f"Number of best users: {len(best_users)}")
|
|
|
|
# %% [markdown]
|
|
# ## Write your prompt and generate its embedding
|
|
# %%
|
|
# prompt = "The user is highly prepared for disasters, with emergency supplies and a clear plan."
|
|
prompt = best_users.head(1)["user_text"].values[
|
|
0
|
|
] # Example: use the first user's text as the prompt
|
|
prompt_embedding = get_ollama_embedding(prompt)
|
|
|
|
# %% [markdown]
|
|
# ## Compute similarity between each user and the prompt
|
|
# %%
|
|
user_embeddings = np.vstack(df["embedding"].values)
|
|
prompt_vec = np.array(prompt_embedding).reshape(1, -1)
|
|
similarities = cosine_similarity(user_embeddings, prompt_vec).flatten()
|
|
df["similarity"] = similarities
|
|
|
|
# %% [markdown]
|
|
# ## Visualize the users by similarity
|
|
# %%
|
|
_ = plt.figure(figsize=(5, 7))
|
|
_ = sns.violinplot(
|
|
x="similarity", y="country_code_iso_3166", hue="country_code_iso_3166", data=df
|
|
)
|
|
_ = sns.stripplot(
|
|
x="similarity",
|
|
y="country_code_iso_3166",
|
|
data=df,
|
|
hue="country_code_iso_3166",
|
|
alpha=0.8,
|
|
jitter=True,
|
|
linewidth=0.5,
|
|
edgecolor="white",
|
|
)
|
|
_ = plt.title("Distribution of User Similarities to Preparedness Prompt")
|
|
_ = plt.xlabel("Cosine Similarity")
|
|
|
|
# %%
|