Cleaner Layout and Tabs for Different Models

app.py

@@ -28,13 +28,21 @@ def config_style():
         </style>
     """, unsafe_allow_html=True)
     st.markdown('<h1 class="main-title">Merit Embeddings 🎒📃🏆</h1>', unsafe_allow_html=True)
-    st.markdown('<h2 class="sub-title">Donut 🤗</h2>', unsafe_allow_html=True)
 
-def load_embeddings():
-    df_real = pd.read_csv("data/donut_de_Rodrigo_merit_secret_all_embeddings.csv")
-    df_es_digital_seq = pd.read_csv("data/donut_de_Rodrigo_merit_es-digital-seq_embeddings.csv")
+# Modificamos load_embeddings para aceptar el modelo a cargar
+def load_embeddings(model):
+    if model == "Donut":
+        df_real = pd.read_csv("data/donut_de_Rodrigo_merit_secret_all_embeddings.csv")
+        df_es_digital_seq = pd.read_csv("data/donut_de_Rodrigo_merit_es-digital-seq_embeddings.csv")
+    elif model == "Idefics2":
+        df_real = pd.read_csv("data/idefics2_de_Rodrigo_merit_secret_britanico_embeddings.csv")
+        df_es_digital_seq = pd.read_csv("data/idefics2_de_Rodrigo_merit_secret_britanico_embeddings.csv")
+    else:
+        st.error("Modelo no reconocido")
+        return None
     return {"real": df_real, "es-digital-seq": df_es_digital_seq}
 
+# Funciones auxiliares (idénticas a las de tu código)
 def reducer_selector(df_combined, embedding_cols):
     reduction_method = st.selectbox("Select Dimensionality Reduction Method:", options=["PCA", "t-SNE"])
     all_embeddings = df_combined[embedding_cols].values
@@ -88,11 +96,6 @@ def split_versions(df_combined, reduced):
     unique_es = sorted(df_es['label'].unique().tolist())
     return {"real": df_real, "es-digital-seq": df_es}, {"real": unique_real, "es-digital-seq": unique_es}
 
-def subset_selectors(unique_subsets: dict):
-    selected_real = st.multiselect("Select Real Subsets:", options=unique_subsets["real"], default=unique_subsets["real"])
-    selected_es = st.multiselect("Select Synthetic Subsets:", options=unique_subsets["es-digital-seq"], default=unique_subsets["es-digital-seq"])
-    return {"real": selected_real, "es-digital-seq": selected_es}
-
 def create_figure(dfs_reduced, selected_subsets: dict, color_maps: dict):
     fig = figure(width=400, height=400, tooltips=TOOLTIPS, title="")
     real_renderers = add_dataset_to_fig(fig, dfs_reduced["real"], selected_subsets["real"],
@@ -119,52 +122,61 @@ def compute_distances(centers_es: dict, centers_real: dict) -> pd.DataFrame:
             distances[es_label][real_label] = np.sqrt((x_es - x_real)**2 + (y_es - y_real)**2)
     return pd.DataFrame(distances).T
 
-def main():
-    config_style()
-    embeddings = load_embeddings()
+def create_table(df_distances):
+    df_table = df_distances.copy()
+    df_table.reset_index(inplace=True)
+    df_table.rename(columns={'index': 'Synthetic'}, inplace=True)
+    source_table = ColumnDataSource(df_table)
+    columns = [TableColumn(field='Synthetic', title='Synthetic')]
+    for col in df_table.columns:
+        if col != 'Synthetic':
+            columns.append(TableColumn(field=col, title=col))
+    row_height = 28
+    header_height = 30
+    total_height = header_height + len(df_table) * row_height
+
+    data_table = DataTable(source=source_table, columns=columns, sizing_mode='stretch_width', height=total_height)
+    return data_table, df_table, source_table
+
+# Función que ejecuta todo el proceso para un modelo determinado
+def run_model(model_name):
+    embeddings = load_embeddings(model_name)
+    if embeddings is None:
+        return
+
+    # Asignamos la versión para distinguir en el split
     embeddings["real"]["version"] = "real"
     embeddings["es-digital-seq"]["version"] = "es_digital_seq"
     embedding_cols = [col for col in embeddings["real"].columns if col.startswith("dim_")]
-    
     df_combined = pd.concat([embeddings["real"], embeddings["es-digital-seq"]], ignore_index=True)
-    reduced = reducer_selector(df_combined, embedding_cols)
+    
+    st.markdown('<h6 class="sub-title">Select Dimensionality Reduction Method</h6>', unsafe_allow_html=True)
+    reduction_method = st.selectbox("", options=["t-SNE", "PCA"], key=model_name)
+    if reduction_method == "PCA":
+        reducer = PCA(n_components=2)
+    else:
+        reducer = TSNE(n_components=2, random_state=42, perplexity=30, learning_rate=200)
+    reduced = reducer.fit_transform(df_combined[embedding_cols].values)
     
     dfs_reduced, unique_subsets = split_versions(df_combined, reduced)
-    selected_subsets = subset_selectors(unique_subsets)
+    selected_subsets = {"real": unique_subsets["real"], "es-digital-seq": unique_subsets["es-digital-seq"]}
     color_maps = get_color_maps(selected_subsets)
-    fig, real_renderers, synthetic_renderers = create_figure(dfs_reduced, selected_subsets, color_maps)
     
+    fig, real_renderers, synthetic_renderers = create_figure(dfs_reduced, selected_subsets, color_maps)
     centers_real = calculate_cluster_centers(dfs_reduced["real"], selected_subsets["real"])
     centers_es = calculate_cluster_centers(dfs_reduced["es-digital-seq"], selected_subsets["es-digital-seq"])
     df_distances = compute_distances(centers_es, centers_real)
-    
-    # Tabla de distancias: se muestran todas las combinaciones
-    df_table = df_distances.copy()
-    df_table.reset_index(inplace=True)
-    df_table.rename(columns={'index': 'Synthetic'}, inplace=True)
-    source_table = ColumnDataSource(df_table)
-    columns = [TableColumn(field='Synthetic', title='Synthetic')]
-    for col in df_table.columns:
-        if col != 'Synthetic':
-            columns.append(TableColumn(field=col, title=col))
-    data_table = DataTable(source=source_table, columns=columns, width=400, height=300)
-    
-    # Widget Select para elegir el subset real (columnas de la tabla)
-    real_subset_names = list(df_table.columns[1:])  # todas las columnas excepto 'Synthetic'
-    real_select = Select(title="Select Real Subset:", value=real_subset_names[0], options=real_subset_names)
-    
-    # Botón para resetear la visualización a colores originales
+    data_table, df_table, source_table = create_table(df_distances)
+    real_subset_names = list(df_table.columns[1:])
+    real_select = Select(title="", value=real_subset_names[0], options=real_subset_names)
     reset_button = Button(label="Reset Colors", button_type="primary")
-    
-    # Fuente para la línea que conecta los centros
     line_source = ColumnDataSource(data={'x': [], 'y': []})
     fig.line('x', 'y', source=line_source, line_width=2, line_color='black')
     
-    # Preparar centros para el callback
     synthetic_centers_js = {k: [v[0], v[1]] for k, v in centers_es.items()}
     real_centers_js = {k: [v[0], v[1]] for k, v in centers_real.items()}
     
-    # Callback para actualizar la visualización según la selección de la tabla y el dropdown
+    # Callback para actualizar el gráfico
     callback = CustomJS(args=dict(source=source_table, line_source=line_source,
                                   synthetic_centers=synthetic_centers_js,
                                   real_centers=real_centers_js,
@@ -228,11 +240,9 @@ def main():
             }
         }
     """)
-    
     source_table.selected.js_on_change('indices', callback)
     real_select.js_on_change('value', callback)
     
-    # Callback para el botón de resetear: se reinician la línea y los colores a su estado original.
     reset_callback = CustomJS(args=dict(line_source=line_source,
                                         synthetic_renderers=synthetic_renderers,
                                         real_renderers=real_renderers,
@@ -258,9 +268,21 @@ def main():
     """)
     reset_button.js_on_event("button_click", reset_callback)
     
-    # Organizar el layout: gráfico, dropdown, botón de reset y tabla
     layout = column(fig, column(real_select, reset_button, data_table))
-    st.bokeh_chart(layout)
+    st.bokeh_chart(layout, use_container_width=True)
+
+# Función principal con tabs para cambiar de modelo
+def main():
+    config_style()
+    tabs = st.tabs(["Donut", "Idefics2"])
+    
+    with tabs[0]:
+        st.markdown('<h2 class="sub-title">Modelo Donut 🤗</h2>', unsafe_allow_html=True)
+        run_model("Donut")
+    
+    with tabs[1]:
+        st.markdown('<h2 class="sub-title">Modelo Idefics2 🤗</h2>', unsafe_allow_html=True)
+        run_model("Idefics2")
 
 if __name__ == "__main__":
     main()