Include Pretraining Datasets

app.py

@@ -40,15 +40,18 @@ def config_style():
     """, unsafe_allow_html=True)
     st.markdown('<h1 class="main-title">Merit Embeddings 🎒📃🏆</h1>', unsafe_allow_html=True)
 
-def load_embeddings(model, version):
+def load_embeddings(model, version, embedding_prefix):
     if model == "Donut":
         df_real = pd.read_csv(f"data/donut_{version}_de_Rodrigo_merit_secret_all_embeddings.csv")
-        df_par = pd.read_csv(f"data/donut_{version}_de_Rodrigo_merit_es-digital-paragraph-degradation-seq_embeddings.csv")
-        df_line = pd.read_csv(f"data/donut_{version}_de_Rodrigo_merit_es-digital-line-degradation-seq_embeddings.csv")
-        df_seq  = pd.read_csv(f"data/donut_{version}_de_Rodrigo_merit_es-digital-seq_embeddings.csv")
-        df_rot  = pd.read_csv(f"data/donut_{version}_de_Rodrigo_merit_es-digital-rotation-degradation-seq_embeddings.csv")
-        df_zoom = pd.read_csv(f"data/donut_{version}_de_Rodrigo_merit_es-digital-zoom-degradation-seq_embeddings.csv")
-        df_render = pd.read_csv(f"data/donut_{version}_de_Rodrigo_merit_es-render-seq_embeddings.csv")
+        df_par = pd.read_csv(f"data/donut_{version}_de_Rodrigo_merit_es-digital-paragraph-degradation-seq_{embedding_prefix}embeddings.csv")
+        df_line = pd.read_csv(f"data/donut_{version}_de_Rodrigo_merit_es-digital-line-degradation-seq_{embedding_prefix}embeddings.csv")
+        df_seq  = pd.read_csv(f"data/donut_{version}_de_Rodrigo_merit_es-digital-seq_{embedding_prefix}embeddings.csv")
+        df_rot  = pd.read_csv(f"data/donut_{version}_de_Rodrigo_merit_es-digital-rotation-degradation-seq_{embedding_prefix}embeddings.csv")
+        df_zoom = pd.read_csv(f"data/donut_{version}_de_Rodrigo_merit_es-digital-zoom-degradation-seq_{embedding_prefix}embeddings.csv")
+        df_render = pd.read_csv(f"data/donut_{version}_de_Rodrigo_merit_es-render-seq_{embedding_prefix}embeddings.csv")
+        df_pretratrained = pd.read_csv(f"data/donut_{version}_de_Rodrigo_merit_aux_IIT-CDIP_{embedding_prefix}embeddings.csv")
+        
+        # Asignar etiquetas de versión
         df_real["version"] = "real"
         df_par["version"] = "synthetic"
         df_line["version"] = "synthetic"
@@ -56,23 +59,32 @@ def load_embeddings(model, version):
         df_rot["version"] = "synthetic"
         df_zoom["version"] = "synthetic"
         df_render["version"] = "synthetic"
+        df_pretratrained["version"] = "pretrained"
 
+        # Asignar fuente (source)
         df_par["source"] = "es-digital-paragraph-degradation-seq"
         df_line["source"] = "es-digital-line-degradation-seq"
         df_seq["source"] = "es-digital-seq"
         df_rot["source"] = "es-digital-rotation-degradation-seq"
         df_zoom["source"] = "es-digital-zoom-degradation-seq"
         df_render["source"] = "es-render-seq"
-        return {"real": df_real, "synthetic": pd.concat([df_seq, df_line, df_par, df_rot, df_zoom, df_render], ignore_index=True)}
+        # Si lo requieres, puedes asignar también una fuente para pretrained
+        df_pretratrained["source"] = "pretrained"
+
+        return {"real": df_real, 
+                "synthetic": pd.concat([df_seq, df_line, df_par, df_rot, df_zoom, df_render], ignore_index=True),
+                "pretrained": df_pretratrained}
     
     elif model == "Idefics2":
         df_real = pd.read_csv(f"data/idefics2_{version}_de_Rodrigo_merit_secret_britanico_embeddings.csv")
-        df_par = pd.read_csv(f"data/idefics2_{version}_de_Rodrigo_merit_es-digital-paragraph-degradation-seq_embeddings.csv")
-        df_line = pd.read_csv(f"data/idefics2_{version}_de_Rodrigo_merit_es-digital-line-degradation-seq_embeddings.csv")
-        df_seq  = pd.read_csv(f"data/idefics2_{version}_de_Rodrigo_merit_es-digital-seq_embeddings.csv")
-        df_rot  = pd.read_csv(f"data/idefics2_{version}_de_Rodrigo_merit_es-digital-rotation-degradation-seq_embeddings.csv")
-        df_zoom = pd.read_csv(f"data/idefics2_{version}_de_Rodrigo_merit_es-digital-zoom-degradation-seq_embeddings.csv")
-        df_render = pd.read_csv(f"data/idefics2_{version}_de_Rodrigo_merit_es-render-seq_embeddings.csv")
+        df_par = pd.read_csv(f"data/idefics2_{version}_de_Rodrigo_merit_es-digital-paragraph-degradation-seq_{embedding_prefix}embeddings.csv")
+        df_line = pd.read_csv(f"data/idefics2_{version}_de_Rodrigo_merit_es-digital-line-degradation-seq_{embedding_prefix}embeddings.csv")
+        df_seq  = pd.read_csv(f"data/idefics2_{version}_de_Rodrigo_merit_es-digital-seq_{embedding_prefix}embeddings.csv")
+        df_rot  = pd.read_csv(f"data/idefics2_{version}_de_Rodrigo_merit_es-digital-rotation-degradation-seq_{embedding_prefix}embeddings.csv")
+        df_zoom = pd.read_csv(f"data/idefics2_{version}_de_Rodrigo_merit_es-digital-zoom-degradation-seq_{embedding_prefix}embeddings.csv")
+        df_render = pd.read_csv(f"data/idefics2_{version}_de_Rodrigo_merit_es-render-seq_{embedding_prefix}embeddings.csv")
+        df_pretratrained = pd.read_csv(f"data/idefics2_{version}_de_Rodrigo_merit_pretrained_{embedding_prefix}embeddings.csv")
+        
         df_real["version"] = "real"
         df_par["version"] = "synthetic"
         df_line["version"] = "synthetic"
@@ -80,6 +92,7 @@ def load_embeddings(model, version):
         df_rot["version"] = "synthetic"
         df_zoom["version"] = "synthetic"
         df_render["version"] = "synthetic"
+        df_pretratrained["version"] = "pretrained"
 
         df_par["source"] = "es-digital-paragraph-degradation-seq"
         df_line["source"] = "es-digital-line-degradation-seq"
@@ -87,27 +100,38 @@ def load_embeddings(model, version):
         df_rot["source"] = "es-digital-rotation-degradation-seq"
         df_zoom["source"] = "es-digital-zoom-degradation-seq"
         df_render["source"] = "es-render-seq"
-        return {"real": df_real, "synthetic": pd.concat([df_seq, df_line, df_par, df_rot, df_zoom, df_render], ignore_index=True)}
+        df_pretratrained["source"] = "pretrained"
+
+        return {"real": df_real, 
+                "synthetic": pd.concat([df_seq, df_line, df_par, df_rot, df_zoom, df_render], ignore_index=True),
+                "pretrained": df_pretratrained}
     
     else:
         st.error("Modelo no reconocido")
         return None
 
+
+
 def split_versions(df_combined, reduced):
-    # Si el embedding es 2D se asignan las columnas x e y para visualización.
+    # Asignar las coordenadas si la reducción es 2D
     if reduced.shape[1] == 2:
         df_combined['x'] = reduced[:, 0]
         df_combined['y'] = reduced[:, 1]
     df_real = df_combined[df_combined["version"] == "real"].copy()
     df_synth = df_combined[df_combined["version"] == "synthetic"].copy()
+    df_pretrained = df_combined[df_combined["version"] == "pretrained"].copy()
+    
     unique_real = sorted(df_real['label'].unique().tolist())
     unique_synth = {}
     for source in df_synth["source"].unique():
         unique_synth[source] = sorted(df_synth[df_synth["source"] == source]['label'].unique().tolist())
-    df_dict = {"real": df_real, "synthetic": df_synth}
-    unique_subsets = {"real": unique_real, "synthetic": unique_synth}
+    unique_pretrained = sorted(df_pretrained['label'].unique().tolist())
+    
+    df_dict = {"real": df_real, "synthetic": df_synth, "pretrained": df_pretrained}
+    unique_subsets = {"real": unique_real, "synthetic": unique_synth, "pretrained": unique_pretrained}
     return df_dict, unique_subsets
 
+
 def get_embedding_from_df(df):
     # Retorna el embedding completo (4 dimensiones en este caso) guardado en la columna 'embedding'
     if 'embedding' in df.columns:
@@ -212,11 +236,15 @@ def create_table(df_distances):
     return data_table, df_table, source_table
 
 def create_figure(dfs, unique_subsets, color_maps, model_name):
-    # Se crea solo si el embedding es 2D (ya que se usan 'x' y 'y' para visualizar)
+    # Se crea el plot para el embedding reducido (asumiendo que es 2D)
     fig = figure(width=600, height=600, tools="wheel_zoom,pan,reset,save", active_scroll="wheel_zoom", tooltips=TOOLTIPS, title="")
+    
+    # Renderizar datos reales
     real_renderers = add_dataset_to_fig(fig, dfs["real"], unique_subsets["real"],
                                         marker="circle", color_mapping=color_maps["real"],
                                         group_label="Real")
+    
+    # Renderizar datos sintéticos (por fuente)
     marker_mapping = {
         "es-digital-paragraph-degradation-seq": "x",
         "es-digital-line-degradation-seq": "cross",
@@ -236,11 +264,17 @@ def create_figure(dfs, unique_subsets, color_maps, model_name):
                                                   group_label=source)
         synthetic_renderers.update(renderers)
     
+    # Agregar el subset pretrained (se puede usar un marcador distinto, por ejemplo, "triangle")
+    pretrained_renderers = add_dataset_to_fig(fig, dfs["pretrained"], unique_subsets["pretrained"],
+                                               marker="triangle", color_mapping=color_maps["pretrained"],
+                                               group_label="Pretrained")
+    
     fig.legend.location = "top_right"
     fig.legend.click_policy = "hide"
     show_legend = st.checkbox("Show Legend", value=False, key=f"legend_{model_name}")
     fig.legend.visible = show_legend
-    return fig, real_renderers, synthetic_renderers
+    return fig, real_renderers, synthetic_renderers, pretrained_renderers
+
 
 def add_dataset_to_fig(fig, df, selected_labels, marker, color_mapping, group_label):
     renderers = {}
@@ -343,8 +377,15 @@ def get_color_maps(unique_subsets):
         else:
             palette = Blues9[:len(labels)] if len(labels) <= 9 else (Blues9 * ((len(labels)//9)+1))[:len(labels)]
         color_map["synthetic"][source] = {label: palette[i] for i, label in enumerate(sorted(labels))}
+    
+    # Asignar colores al subset pretrained usando, por ejemplo, la paleta Purples9
+    num_pretrained = len(unique_subsets["pretrained"])
+    purple_palette = Purples9[:num_pretrained] if num_pretrained <= 9 else (Purples9 * ((num_pretrained // 9) + 1))[:num_pretrained]
+    color_map["pretrained"] = {label: purple_palette[i] for i, label in enumerate(sorted(unique_subsets["pretrained"]))}
+    
     return color_map
 
+
 def calculate_cluster_centers(df, labels):
     centers = {}
     for label in labels:
@@ -485,8 +526,12 @@ def optimize_tsne_params(df_combined, embedding_cols, df_f1, distance_metric):
 
 def run_model(model_name):
     version = st.selectbox("Select Model Version:", options=["vanilla", "finetuned_real"], key=f"version_{model_name}")
+    # Nuevo selector para el cómputo del embedding
+    embedding_computation = st.selectbox("¿Cómo se computa el embedding?", options=["weighted", "averaged"], key=f"embedding_method_{model_name}")
+    # Se asigna el prefijo correspondiente
+    prefijo_embedding = "weighted_" if embedding_computation == "weighted" else ""
     
-    embeddings = load_embeddings(model_name, version)
+    embeddings = load_embeddings(model_name, version, prefijo_embedding)
     if embeddings is None:
         return
     embedding_cols = [col for col in embeddings["real"].columns if col.startswith("dim_")]
@@ -562,10 +607,13 @@ def run_model(model_name):
     reset_button = Button(label="Reset Colors", button_type="primary")
     line_source = ColumnDataSource(data={'x': [], 'y': []})
     
-    # Si el embedding es 2D se crea el scatter plot de embeddings;
-    # dado que con PCA ahora usamos 4 dimensiones, este bloque se omite para PCA
     if (reduction_method == "t-SNE" and N_COMPONENTS == 2) or (reduction_method == "PCA" and N_COMPONENTS == 2):
-        fig, real_renderers, synthetic_renderers = create_figure(result["dfs_reduced"], result["unique_subsets"], get_color_maps(result["unique_subsets"]), model_name)
+        fig, real_renderers, synthetic_renderers, pretrained_renderers = create_figure(
+            result["dfs_reduced"],
+            result["unique_subsets"],
+            get_color_maps(result["unique_subsets"]),
+            model_name
+        )
         fig.line('x', 'y', source=line_source, line_width=2, line_color='black')
         centers_real = calculate_cluster_centers(result["dfs_reduced"]["real"], result["unique_subsets"]["real"])
         real_centers_js = {k: [v[0], v[1]] for k, v in centers_real.items()}
@@ -633,4 +681,4 @@ def main():
         run_model("Idefics2")
 
 if __name__ == "__main__":
-    main()
+    main()