Update modules/ui/ui.py

modules/ui/ui.py

@@ -13,6 +13,7 @@ from datetime import datetime
 from streamlit_player import st_player  # Necesitarás instalar esta librería: pip install streamlit-player
 from spacy import displacy
 import logging
+import random
 
 ######################################################
 # Configuración del logger
@@ -1019,17 +1020,27 @@ def display_discourse_results(result, lang_code, t):
                 st.table(df2)
             else:
                 st.warning(t.get('concepts_not_available', "Los conceptos clave no están disponibles."))
-                
+
     # Relación de conceptos entre ambos documentos (Diagrama de Sankey)
     st.subheader(t.get('comparison', "Relación de conceptos entre ambos documentos"))
     if 'key_concepts1' in result and 'key_concepts2' in result:
         df1 = pd.DataFrame(result['key_concepts1'], columns=['Concepto', 'Frecuencia'])
         df2 = pd.DataFrame(result['key_concepts2'], columns=['Concepto', 'Frecuencia'])
-        
+
+        # Crear una lista de todos los conceptos únicos
+        all_concepts = list(set(df1['Concepto'].tolist() + df2['Concepto'].tolist()))
+
+        # Crear un diccionario de colores para cada concepto
+        color_scale = [f'rgb({random.randint(50,255)},{random.randint(50,255)},{random.randint(50,255)})' for _ in range(len(all_concepts))]
+        color_map = dict(zip(all_concepts, color_scale))
+
         # Crear el diagrama de Sankey
-        source = [t.get('file_uploader1', "Documento 1")] * len(df1) + [concept for concept in df1['Concepto']]
-        target = [concept for concept in df1['Concepto']] + [t.get('file_uploader2', "Documento 2")] * len(df2)
+        source = [0] * len(df1) + list(range(2, 2 + len(df1)))
+        target = list(range(2, 2 + len(df1))) + [1] * len(df2)
         value = list(df1['Frecuencia']) + list(df2['Frecuencia'])
+        
+        node_colors = ['blue', 'red'] + [color_map[concept] for concept in df1['Concepto']] + [color_map[concept] for concept in df2['Concepto']]
+        link_colors = [color_map[concept] for concept in df1['Concepto']] + [color_map[concept] for concept in df2['Concepto']]
 
         fig = go.Figure(data=[go.Sankey(
             node = dict(
@@ -1037,12 +1048,13 @@ def display_discourse_results(result, lang_code, t):
               thickness = 20,
               line = dict(color = "black", width = 0.5),
               label = [t.get('file_uploader1', "Documento 1"), t.get('file_uploader2', "Documento 2")] + list(df1['Concepto']) + list(df2['Concepto']),
-              color = "blue"
+              color = node_colors
             ),
             link = dict(
-              source = [0] * len(df1) + list(range(2, 2 + len(df1))),
-              target = list(range(2, 2 + len(df1))) + [1] * len(df2),
-              value = value
+              source = source,
+              target = target,
+              value = value,
+              color = link_colors
           ))])
 
         fig.update_layout(title_text="Relación de conceptos entre documentos", font_size=10)