Update modules/semantic_analysis.py

modules/semantic_analysis.py

@@ -85,6 +85,29 @@ POS_TRANSLATIONS = {
 }
 ########################################################################################################################################
 
+# Definimos las etiquetas y colores para cada idioma
+ENTITY_LABELS = {
+    'es': {
+        "Personas": "lightblue",
+        "Conceptos": "lightgreen",
+        "Lugares": "lightcoral",
+        "Fechas": "lightyellow"
+    },
+    'en': {
+        "People": "lightblue",
+        "Concepts": "lightgreen",
+        "Places": "lightcoral",
+        "Dates": "lightyellow"
+    },
+    'fr': {
+        "Personnes": "lightblue",
+        "Concepts": "lightgreen",
+        "Lieux": "lightcoral",
+        "Dates": "lightyellow"
+    }
+}
+
+#########################################################################################################
 def count_pos(doc):
     return Counter(token.pos_ for token in doc if token.pos_ != 'PUNCT')
 
@@ -95,31 +118,27 @@ from collections import Counter
 
 # Mantén las definiciones de POS_COLORS y POS_TRANSLATIONS que ya tienes
 
-def extract_entities(doc):
-    entities = {
-        "Personas": [],
-        "Conceptos": [],
-        "Lugares": [],
-        "Fechas": []
-    }
-
+#############################################################################################################################
+def extract_entities(doc, lang):
+    entities = {label: [] for label in ENTITY_LABELS[lang].keys()}
+    
     for ent in doc.ents:
         if ent.label_ == "PERSON":
-            entities["Personas"].append(ent.text)
+            entities[list(ENTITY_LABELS[lang].keys())[0]].append(ent.text)
         elif ent.label_ in ["LOC", "GPE"]:
-            entities["Lugares"].append(ent.text)
+            entities[list(ENTITY_LABELS[lang].keys())[2]].append(ent.text)
         elif ent.label_ == "DATE":
-            entities["Fechas"].append(ent.text)
+            entities[list(ENTITY_LABELS[lang].keys())[3]].append(ent.text)
         else:
-            entities["Conceptos"].append(ent.text)
-
+            entities[list(ENTITY_LABELS[lang].keys())[1]].append(ent.text)
+    
     return entities
 
+#####################################################################################################################
 def visualize_context_graph(doc, lang):
     G = nx.Graph()
-    entities = extract_entities(doc)
-
-    color_map = {"Personas": "lightblue", "Conceptos": "lightgreen", "Lugares": "lightcoral", "Fechas": "lightyellow"}
+    entities = extract_entities(doc, lang)
+    color_map = ENTITY_LABELS[lang]
 
     # Add nodes
     for category, items in entities.items():
@@ -139,7 +158,8 @@ def visualize_context_graph(doc, lang):
 
     node_colors = [color_map[G.nodes[node]['category']] for node in G.nodes()]
 
-    nx.draw(G, pos, node_color=node_colors, with_labels=True, node_size=3000, font_size=8, font_weight='bold')
+    nx.draw(G, pos, node_color=node_colors, with_labels=True, node_size=5000, 
+            font_size=12, font_weight='bold')
 
     # Add a legend
     legend_elements = [plt.Rectangle((0,0),1,1,fc=color, edgecolor='none', label=category) 
@@ -151,6 +171,7 @@ def visualize_context_graph(doc, lang):
 
     return plt
 
+############################################################################################################################################
 def visualize_semantic_relations(doc, lang):
     G = nx.Graph()
     word_freq = Counter(token.text.lower() for token in doc if token.pos_ not in ['PUNCT', 'SPACE'])
@@ -175,8 +196,12 @@ def visualize_semantic_relations(doc, lang):
     edge_labels = nx.get_edge_attributes(G, 'label')
     nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_size=8)
 
-    plt.title("Relaciones Semánticas Relevantes" if lang == 'es' else "Relevant Semantic Relations" if lang == 'en' else "Relations Sémantiques Pertinentes",
-              fontsize=20, fontweight='bold')
+    title = {
+        'es': "Relaciones Semánticas Relevantes",
+        'en': "Relevant Semantic Relations",
+        'fr': "Relations Sémantiques Pertinentes"
+    }
+    plt.title(title[lang], fontsize=20, fontweight='bold')
     plt.axis('off')
 
     legend_elements = [plt.Rectangle((0,0),1,1, facecolor=POS_COLORS.get(pos, '#CCCCCC'), edgecolor='none', 
@@ -186,8 +211,16 @@ def visualize_semantic_relations(doc, lang):
 
     return plt
 
+    
+############################################################################################################################################
 def perform_semantic_analysis(text, nlp, lang):
     doc = nlp(text)
+    
+    # Imprimir entidades para depuración
+    print(f"Entidades encontradas ({lang}):")
+    for ent in doc.ents:
+        print(f"{ent.text} - {ent.label_}")
+    
     context_graph = visualize_context_graph(doc, lang)
     relations_graph = visualize_semantic_relations(doc, lang)
     return context_graph, relations_graph