added watermarking and quantization exp

app.py

@@ -1,4 +1,5 @@
 import json
+from typing import Tuple
 
 import numpy as np
 import pandas as pd
@@ -11,8 +12,6 @@ from plotly.subplots import make_subplots
 from exp_utils import MODELS
 from visualize_utils import viridis_rgb
 
-#
-
 st.set_page_config(
     page_title="Results Viewer",
     page_icon="📊",
@@ -23,14 +22,35 @@ st.set_page_config(
 MODELS_SIZE_MAPPING = {k: v["model_size"] for k, v in MODELS.items()}
 MODELS_FAMILY_MAPPING = {k: v["model_family"] for k, v in MODELS.items()}
 MODEL_FAMILES = set([model["model_family"] for model in MODELS.values()])
-MODEL_NAMES = list(MODELS.keys())
+Q_W_MODELS = [
+    "llama-7b",
+    "llama-2-7b",
+    "llama-13b",
+    "llama-2-13b",
+    "llama-30b",
+    "llama-65b",
+    "llama-2-70b",
+]
+Q_W_MODELS = [f"{model}_quantized" for model in Q_W_MODELS] + [
+    f"{model}_watermarked" for model in Q_W_MODELS
+]
+
+MODEL_NAMES = list(MODELS.keys()) + Q_W_MODELS
 
 MODEL_NAMES_SORTED_BY_NAME_AND_SIZE = sorted(
-    MODEL_NAMES, key=lambda x: (MODELS[x]["model_family"], MODELS[x]["model_size"])
+    MODEL_NAMES,
+    key=lambda x: (
+        MODELS[x.replace("_quantized", "").replace("_watermarked", "")]["model_family"],
+        MODELS[x.replace("_quantized", "").replace("_watermarked", "")]["model_size"],
+    ),
 )
 
 MODEL_NAMES_SORTED_BY_SIZE = sorted(
-    MODEL_NAMES, key=lambda x: (MODELS[x]["model_size"], MODELS[x]["model_family"])
+    MODEL_NAMES,
+    key=lambda x: (
+        MODELS[x.replace("_quantized", "").replace("_watermarked", "")]["model_size"],
+        MODELS[x.replace("_quantized", "").replace("_watermarked", "")]["model_family"],
+    ),
 )
 
 
@@ -43,7 +63,11 @@ MODELS_SIZE_MAPPING = {
 MODELS_SIZE_MAPPING_LIST = list(MODELS_SIZE_MAPPING.keys())
 
 
-CHAT_MODELS = [x for x in MODEL_NAMES_SORTED_BY_NAME_AND_SIZE if MODELS[x]["is_chat"]]
+CHAT_MODELS = [
+    x
+    for x in MODEL_NAMES_SORTED_BY_NAME_AND_SIZE
+    if MODELS[x.replace("_quantized", "").replace("_watermarked", "")]["is_chat"]
+]
 
 
 def clean_dataframe(df: pd.DataFrame) -> pd.DataFrame:
@@ -66,7 +90,11 @@ def clean_dataframe(df: pd.DataFrame) -> pd.DataFrame:
     df.columns = df.columns.str.replace("_roc_auc", "")
     df.columns = df.columns.str.replace("eval_", "")
 
-    df["model_family"] = df["model_name"].map(MODELS_FAMILY_MAPPING)
+    df["model_family"] = df["model_name"].apply(
+        lambda x: MODELS_FAMILY_MAPPING[
+            x.replace("_quantized", "").replace("_watermarked", "")
+        ]
+    )
     # create a dict with the model_name and the model_family
     model_family_dict = {
         k: v
@@ -84,8 +112,16 @@ def clean_dataframe(df: pd.DataFrame) -> pd.DataFrame:
     df_std = df_std.drop(columns=["exp_seed"])
     df_avg["model_family"] = df_avg.index.map(model_family_dict)
     df_std["model_family"] = df_std.index.map(model_family_dict)
-    df_avg["model_size"] = df_avg.index.map(MODELS_SIZE_MAPPING)
-    df_std["model_size"] = df_std.index.map(MODELS_SIZE_MAPPING)
+    df_avg["model_size"] = df_avg.index.map(
+        lambda x: MODELS_SIZE_MAPPING[
+            x.replace("_quantized", "").replace("_watermarked", "")
+        ]
+    )
+    df_std["model_size"] = df_std.index.map(
+        lambda x: MODELS_SIZE_MAPPING[
+            x.replace("_quantized", "").replace("_watermarked", "")
+        ]
+    )
 
     # sort rows by model family then model size
     df_avg = df_avg.sort_values(
@@ -101,10 +137,15 @@ def clean_dataframe(df: pd.DataFrame) -> pd.DataFrame:
     availables_rows = [x for x in df_std.columns if x in df_std.index]
     df_std = df_std.reindex(availables_rows)
 
+    df_avg["is_quantized"] = df_avg.index.str.contains("quantized")
+    df_avg["is_watermarked"] = df_avg.index.str.contains("watermarked")
+    df_std["is_quantized"] = df_std.index.str.contains("quantized")
+    df_std["is_watermarked"] = df_std.index.str.contains("watermarked")
+
     return df_avg, df_std
 
 
-def get_data(path):
+def get_data(path) -> Tuple[pd.DataFrame, pd.DataFrame]:
     df, df_std = clean_dataframe(pd.read_csv(path, index_col=0))
     return df, df_std
 
@@ -117,8 +158,15 @@ def filter_df(
     model_size_test: tuple,
     is_chat_train: bool,
     is_chat_test: bool,
+    is_quantized_train: bool,
+    is_quantized_test: bool,
+    is_watermarked_train: bool,
+    is_watermarked_test: bool,
     sort_by_size: bool,
     split_chat_models: bool,
+    split_quantized_models: bool,
+    split_watermarked_models: bool,
+    filter_empty_col_row: bool,
     is_debug: bool,
 ) -> pd.DataFrame:
     # remove all columns and rows that have "pythia-70m" in the name
@@ -143,6 +191,16 @@ def filter_df(
         if is_debug:
             st.write("Filter is chat train")
             st.write(df)
+    if is_quantized_train != "Both":
+        df = df.loc[df["is_quantized"] == is_quantized_train]
+        if is_debug:
+            st.write("Filter is quantized train")
+            st.write(df)
+    if is_watermarked_train != "Both":
+        df = df.loc[df["is_watermarked"] == is_watermarked_train]
+        if is_debug:
+            st.write("Filter is watermark train")
+            st.write(df)
 
     # filter columns
     if is_debug:
@@ -150,8 +208,13 @@ def filter_df(
         st.write(df)
     columns_to_keep = []
     for column in df.columns:
-        if column in MODELS.keys():
-            model_size = MODELS[column]["model_size"]
+        if (
+            column.replace("_quantized", "").replace("_watermarked", "")
+            in MODELS.keys()
+        ):
+            model_size = MODELS[
+                column.replace("_quantized", "").replace("_watermarked", "")
+            ]["model_size"]
             if (
                 model_size >= model_size_test[0] * 1e9
                 and model_size <= model_size_test[1] * 1e9
@@ -167,7 +230,12 @@ def filter_df(
     columns_to_keep = []
     for column in df.columns:
         for model_family in model_family_test:
-            if model_family == MODELS[column]["model_family"]:
+            if (
+                model_family
+                == MODELS[column.replace("_quantized", "").replace("_watermarked", "")][
+                    "model_family"
+                ]
+            ):
                 columns_to_keep.append(column)
     df = df[list(sorted(list(set(columns_to_keep))))]
     if is_debug:
@@ -178,13 +246,44 @@ def filter_df(
         # filter columns
         columns_to_keep = []
         for column in df.columns:
-            if MODELS[column]["is_chat"] == is_chat_test:
+            if (
+                MODELS[column.replace("_quantized", "").replace("_watermarked", "")][
+                    "is_chat"
+                ]
+                == is_chat_test
+            ):
                 columns_to_keep.append(column)
         df = df[list(sorted(list(set(columns_to_keep))))]
         if is_debug:
             st.write("Filter is chat test")
             st.write(df)
 
+    if is_quantized_test != "Both":
+        # filter columns
+        columns_to_keep = []
+        for column in df.columns:
+            if "quantized" in column and is_quantized_test:
+                columns_to_keep.append(column)
+            elif "quantized" not in column and not is_quantized_test:
+                columns_to_keep.append(column)
+        df = df[list(sorted(list(set(columns_to_keep))))]
+        if is_debug:
+            st.write("Filter is quantized test")
+            st.write(df)
+
+    if is_watermarked_test != "Both":
+        # filter columns
+        columns_to_keep = []
+        for column in df.columns:
+            if "watermark" in column and is_watermarked_test:
+                columns_to_keep.append(column)
+            elif "watermark" not in column and not is_watermarked_test:
+                columns_to_keep.append(column)
+        df = df[list(sorted(list(set(columns_to_keep))))]
+        if is_debug:
+            st.write("Filter is watermark test")
+            st.write(df)
+
     df = df.select_dtypes(include="number")
     if is_debug:
         st.write("Select dtypes to be only numbers")
@@ -227,10 +326,121 @@ def filter_df(
     if is_debug:
         st.write("Split chat models")
         st.write(df)
+
+    if split_quantized_models:
+        # put chat models at the end of the columns
+        quantized_models = [
+            x for x in Q_W_MODELS if x in df.columns and "quantized" in x
+        ]
+        # sort chat models by size
+        quantized_models = sorted(
+            quantized_models,
+            key=lambda x: MODELS[
+                x.replace("_quantized", "").replace("_watermarked", "")
+            ]["model_size"],
+        )
+        df = df[[x for x in df.columns if x not in quantized_models] + quantized_models]
+
+        # put chat models at the end of the rows
+        quantized_models = [x for x in Q_W_MODELS if x in df.index and "quantized" in x]
+        # sort chat models by size
+        quantized_models = sorted(
+            quantized_models,
+            key=lambda x: MODELS[
+                x.replace("_quantized", "").replace("_watermarked", "")
+            ]["model_size"],
+        )
+        df = df.reindex(
+            [x for x in df.index if x not in quantized_models] + quantized_models
+        )
+
+    if split_watermarked_models:
+        # put chat models at the end of the columns
+        watermarked_models = [
+            x for x in Q_W_MODELS if x in df.columns and "watermarked" in x
+        ]
+        # sort chat models by size
+        watermarked_models = sorted(
+            watermarked_models,
+            key=lambda x: MODELS[
+                x.replace("_quantized", "").replace("_watermarked", "")
+            ]["model_size"],
+        )
+        df = df[
+            [x for x in df.columns if x not in watermarked_models] + watermarked_models
+        ]
+
+        # put chat models at the end of the rows
+        watermarked_models = [
+            x for x in Q_W_MODELS if x in df.index and "watermarked" in x
+        ]
+        # sort chat models by size
+        watermarked_models = sorted(
+            watermarked_models,
+            key=lambda x: MODELS[
+                x.replace("_quantized", "").replace("_watermarked", "")
+            ]["model_size"],
+        )
+        df = df.reindex(
+            [x for x in df.index if x not in watermarked_models] + watermarked_models
+        )
+
+    if is_debug:
+        st.write("Split chat models")
+        st.write(df)
+
+    if filter_empty_col_row:
+        # remove all for which the row and column are Nan
+        df = df.dropna(axis=0, how="all")
+        df = df.dropna(axis=1, how="all")
     return df
 
 
 df, df_std = get_data("./deberta_results.csv")
+df_q_w, df_std_q_w = get_data("./results_qantized_watermarked.csv")
+
+df = df.merge(
+    df_q_w[
+        df_q_w.columns[
+            df_q_w.columns.str.contains("quantized|watermarked", case=False, regex=True)
+        ]
+    ],
+    how="outer",
+    left_index=True,
+    right_index=True,
+)
+df_std = df_std.merge(
+    df_std_q_w[
+        df_std_q_w.columns[
+            df_std_q_w.columns.str.contains(
+                "quantized|watermarked", case=False, regex=True
+            )
+        ]
+    ],
+    how="outer",
+    left_index=True,
+    right_index=True,
+)
+
+
+df.columns = df.columns.str.replace("_y", "", regex=True)
+df_std.columns = df_std.columns.str.replace("_y", "", regex=True)
+
+df = df.drop(columns=["is_quantized_x", "is_watermarked_x"])
+
+
+df.update(df_q_w)
+df_std.update(df_std_q_w)
+
+
+df["is_chat"].fillna(False, inplace=True)
+df_std["is_chat"].fillna(False, inplace=True)
+
+df["is_watermarked"].fillna(False, inplace=True)
+df_std["is_watermarked"].fillna(False, inplace=True)
+
+df["is_quantized"].fillna(False, inplace=True)
+df_std["is_quantized"].fillna(False, inplace=True)
 
 with open("./ood_results.json", "r") as f:
     ood_results = json.load(f)
@@ -258,11 +468,14 @@ st.write(
 )
 
 # filters
-show_diff = st.sidebar.checkbox("Show Diff", value=False)
-sort_by_size = st.sidebar.checkbox("Sort by size", value=False)
-split_chat_models = st.sidebar.checkbox("Split chat models", value=False)
+how_diff = st.sidebar.checkbox("Show Diff", value=False)
+sort_by_size = st.sidebar.checkbox("Sort by size", value=True)
+split_chat_models = st.sidebar.checkbox("Split chat models", value=True)
+split_quantized_models = st.sidebar.checkbox("Split quantized models", value=True)
+split_watermarked_models = st.sidebar.checkbox("Split watermarked models", value=True)
 add_mean = st.sidebar.checkbox("Add mean", value=False)
 show_std = st.sidebar.checkbox("Show std", value=False)
+filter_empty_col_row = st.sidebar.checkbox("Filter empty col/row", value=True)
 model_size_train = st.sidebar.slider(
     "Train Model Size in Billion", min_value=0, max_value=100, value=(0, 100), step=1
 )
@@ -271,6 +484,18 @@ model_size_test = st.sidebar.slider(
 )
 is_chat_train = st.sidebar.selectbox("(Train) Is Chat?", [True, False, "Both"], index=2)
 is_chat_test = st.sidebar.selectbox("(Test) Is Chat?", [True, False, "Both"], index=2)
+is_quantized_train = st.sidebar.selectbox(
+    "(Train) Is Quantized?", [True, False, "Both"], index=1
+)
+is_quantized_test = st.sidebar.selectbox(
+    "(Test) Is Quantized?", [True, False, "Both"], index=1
+)
+is_watermarked_train = st.sidebar.selectbox(
+    "(Train) Is Watermark?", [True, False, "Both"], index=1
+)
+is_watermarked_test = st.sidebar.selectbox(
+    "(Test) Is Watermark?", [True, False, "Both"], index=1
+)
 model_family_train = st.sidebar.multiselect(
     "Model Family Train",
     MODEL_FAMILES,
@@ -282,6 +507,8 @@ model_family_test = st.sidebar.multiselect(
     default=MODEL_FAMILES,
 )
 
+show_values = st.sidebar.checkbox("Show Values", value=False)
+
 add_adversarial = False
 if "Adversarial" in model_family_test:
     model_family_test.remove("Adversarial")
@@ -304,14 +531,6 @@ if show_std:
 else:
     selected_df = df.copy()
 
-if show_diff:
-    # get those 3 columns {'model_size', 'model_family', 'is_chat'}
-    columns_to_keep = ["model_size", "model_family", "is_chat"]
-    to_be_added = selected_df[columns_to_keep]
-    selected_df = selected_df.drop(columns=columns_to_keep)
-    selected_df = selected_df.sub(selected_df.values.diagonal(), axis=1)
-    selected_df = selected_df.join(to_be_added)
-
 
 filtered_df = filter_df(
     selected_df,
@@ -321,18 +540,32 @@ filtered_df = filter_df(
     model_size_test,
     is_chat_train,
     is_chat_test,
+    is_quantized_train,
+    is_quantized_test,
+    is_watermarked_train,
+    is_watermarked_test,
     sort_by_size,
     split_chat_models,
+    split_quantized_models,
+    split_watermarked_models,
+    filter_empty_col_row,
     is_debug,
 )
 
 
-# subtract each row by the diagonal
+if show_diff:
+    # get those 3 columns {'model_size', 'model_family', 'is_chat'}
+    diag = filtered_df.values.diagonal()
+    filtered_df = filtered_df.sub(diag, axis=1)
 
-# if show_diff:
-#     filtered_df = filtered_df.sub(filtered_df.values.diagonal(), axis=1)
+# subtract each row by the diagonal
 if add_adversarial:
-    filtered_df = filtered_df.join(ood_results_avg)
+    if show_diff:
+        index = filtered_df.index
+        ood_results_avg = ood_results_avg.loc[index]
+        filtered_df = filtered_df.join(ood_results_avg.sub(diag, axis=0))
+    else:
+        filtered_df = filtered_df.join(ood_results_avg)
 
 if add_mean:
     col_mean = filtered_df.mean(axis=1)
@@ -341,7 +574,6 @@ if add_mean:
     filtered_df["mean"] = col_mean
     filtered_df.loc["mean"] = row_mean
 
-
 filtered_df = filtered_df * 100
 filtered_df = filtered_df.round(0)
 
@@ -364,7 +596,7 @@ fig = px.imshow(
     y=list(filtered_df.index),
     color_continuous_scale=color_scale,
     contrast_rescaling=None,
-    text_auto=True,
+    text_auto=show_values,
     aspect="auto",
 )