Update modules/semantic_analysis.py

modules/semantic_analysis.py

@@ -180,18 +180,19 @@ def visualize_context_graph(doc, lang):
 
 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'])
-    top_words = [word for word, _ in word_freq.most_common(20)]  # Top 20 most frequent words
+    word_freq = Counter(token.text.lower() for token in doc if token.pos_ in ['NOUN', 'VERB'])
+    top_words = [word for word, _ in word_freq.most_common(20)]  # Top 20 most frequent nouns and verbs
 
     for token in doc:
-        if token.text.lower() in top_words:
+        if token.pos_ in ['NOUN', 'VERB'] and token.text.lower() in top_words:
             G.add_node(token.text, pos=token.pos_)
 
     for token in doc:
-        if token.text.lower() in top_words and token.head.text.lower() in top_words:
-            G.add_edge(token.text, token.head.text, label=token.dep_)
+        if token.pos_ in ['NOUN', 'VERB'] and token.text.lower() in top_words:
+            if token.head.pos_ in ['NOUN', 'VERB'] and token.head.text.lower() in top_words:
+                G.add_edge(token.text, token.head.text, label=token.dep_)
 
-    plt.figure(figsize=(36, 27))  # Increased figure size
+    fig, ax = plt.subplots(figsize=(36, 27))  # Create a figure and axis
     pos = nx.spring_layout(G, k=0.7, iterations=50)  # Adjusted layout
 
     node_colors = [POS_COLORS.get(G.nodes[node]['pos'], '#CCCCCC') for node in G.nodes()]
@@ -203,25 +204,26 @@ def visualize_semantic_relations(doc, lang):
             arrows=True, 
             arrowsize=30,  # Increased arrow size
             width=3,  # Increased edge width
-            edge_color='gray')
+            edge_color='gray',
+            ax=ax)  # Draw on the axis
 
     edge_labels = nx.get_edge_attributes(G, 'label')
-    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_size=14)  # Increased edge label font size
+    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_size=14, ax=ax)  # Increased edge label font size
 
     title = {
         'es': "Relaciones Semánticas Relevantes",
         'en': "Relevant Semantic Relations",
         'fr': "Relations Sémantiques Pertinentes"
     }
-    plt.title(title[lang], fontsize=24, fontweight='bold')  # Increased title font size
-    plt.axis('off')
+    ax.set_title(title[lang], fontsize=24, fontweight='bold')  # Set title on the axis
+    ax.axis('off')
 
-    legend_elements = [plt.Rectangle((0,0),1,1, facecolor=POS_COLORS.get(pos, '#CCCCCC'), edgecolor='none', 
+    legend_elements = [plt.Rectangle((0,0),1,1,fc=POS_COLORS.get(pos, '#CCCCCC'), edgecolor='none', 
                        label=f"{POS_TRANSLATIONS[lang].get(pos, pos)}")
-                       for pos in set(nx.get_node_attributes(G, 'pos').values())]
-    plt.legend(handles=legend_elements, loc='center left', bbox_to_anchor=(1, 0.5), fontsize=16)  # Increased legend font size
+                       for pos in ['NOUN', 'VERB']]
+    ax.legend(handles=legend_elements, loc='center left', bbox_to_anchor=(1, 0.5), fontsize=16)  # Add legend to the axis
 
-    return plt
+    return fig  # Return the figure instead of plt
 
     
 ############################################################################################################################################
@@ -233,6 +235,5 @@ def perform_semantic_analysis(text, nlp, 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
+    return relations_graph  # Ahora solo devuelve un único gráfico