Update modules/semantic_analysis.py

modules/semantic_analysis.py

@@ -88,6 +88,13 @@ POS_TRANSLATIONS = {
 def count_pos(doc):
     return Counter(token.pos_ for token in doc if token.pos_ != 'PUNCT')
 
+import spacy
+import networkx as nx
+import matplotlib.pyplot as plt
+from collections import Counter
+
+# Mantén las definiciones de POS_COLORS y POS_TRANSLATIONS que ya tienes
+
 def extract_entities(doc):
     entities = {
         "Personas": [],
@@ -97,7 +104,7 @@ def extract_entities(doc):
     }
 
     for ent in doc.ents:
-        if ent.label_ == "PER":
+        if ent.label_ == "PERSON":
             entities["Personas"].append(ent.text)
         elif ent.label_ in ["LOC", "GPE"]:
             entities["Lugares"].append(ent.text)
@@ -112,6 +119,8 @@ def visualize_context_graph(doc, lang):
     G = nx.Graph()
     entities = extract_entities(doc)
 
+    color_map = {"Personas": "lightblue", "Conceptos": "lightgreen", "Lugares": "lightcoral", "Fechas": "lightyellow"}
+
     # Add nodes
     for category, items in entities.items():
         for item in items:
@@ -120,47 +129,28 @@ def visualize_context_graph(doc, lang):
     # Add edges
     for sent in doc.sents:
         sent_entities = [ent for ent in sent.ents if ent.text in G.nodes()]
-        person = next((ent for ent in sent_entities if ent.label_ == "PER"), None)
-        if person:
-            for ent in sent_entities:
-                if ent != person:
-                    G.add_edge(person.text, ent.text)
+        for i in range(len(sent_entities)):
+            for j in range(i+1, len(sent_entities)):
+                G.add_edge(sent_entities[i].text, sent_entities[j].text)
 
     # Visualize
     plt.figure(figsize=(20, 15))
     pos = nx.spring_layout(G, k=0.5, iterations=50)
 
-    color_map = {"Personas": "lightblue", "Conceptos": "lightgreen", "Lugares": "lightcoral", "Fechas": "lightyellow"}
     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')
 
     # Add a legend
-    legend_elements = [plt.Rectangle((0,0),1,1,fc=color, edgecolor='none') for color in color_map.values()]
-    plt.legend(legend_elements, color_map.keys(), loc='upper left', bbox_to_anchor=(1, 1))
+    legend_elements = [plt.Rectangle((0,0),1,1,fc=color, edgecolor='none', label=category) 
+                       for category, color in color_map.items()]
+    plt.legend(handles=legend_elements, loc='upper left', bbox_to_anchor=(1, 1))
 
     plt.title("Análisis del Contexto" if lang == 'es' else "Context Analysis" if lang == 'en' else "Analyse du Contexte", fontsize=20)
     plt.axis('off')
 
     return plt
 
-def create_semantic_graph(doc, lang):
-    G = nx.Graph()
-    pos_counts = count_pos(doc)
-
-    for token in doc:
-        if token.pos_ != 'PUNCT':
-            G.add_node(token.text, 
-                       pos=token.pos_,
-                       color=POS_COLORS.get(token.pos_, '#CCCCCC'),  # Color gris por defecto
-                       size=pos_counts.get(token.pos_, 1) * 100)  # Tamaño mínimo si no hay conteo
-
-    for token in doc:
-        if token.dep_ != "ROOT" and token.head.text in G.nodes and token.text in G.nodes:
-            G.add_edge(token.head.text, token.text, label=token.dep_)
-
-    return G, pos_counts
-
 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'])
@@ -200,8 +190,4 @@ def perform_semantic_analysis(text, nlp, lang):
     doc = nlp(text)
     context_graph = visualize_context_graph(doc, lang)
     relations_graph = visualize_semantic_relations(doc, lang)
-    
-    # Extraer entidades para mostrar en forma de lista
-    entities = extract_entities(doc)
-    
-    return context_graph, relations_graph, entities
+    return context_graph, relations_graph