getting started on the preparedness data

preparedness/main.py

@@ -0,0 +1,100 @@
+# %% [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")
+
+# %%