Update modules/semantic_analysis.py

modules/semantic_analysis.py

@@ -1,4 +1,4 @@
-#syntax_analysis.py
+#semantic_analysis.py
 import streamlit as st
 import spacy
 import networkx as nx
@@ -84,71 +84,94 @@ POS_TRANSLATIONS = {
     }
 }
 ########################################################################################################################################
-def count_pos(doc):
-    return Counter(token.pos_ for token in doc if token.pos_ != 'PUNCT')
-#######################################################################################################################################
-def create_syntax_graph(doc, lang):
-    G = nx.DiGraph()
-    pos_counts = count_pos(doc)
-    word_nodes = {}
-    word_colors = {}
-
-    for token in doc:
-        if token.pos_ != 'PUNCT':
-            lower_text = token.text.lower()
-            if lower_text not in word_nodes:
-                node_id = len(word_nodes)
-                word_nodes[lower_text] = node_id
-                color = POS_COLORS.get(token.pos_, '#FFFFFF')
-                word_colors[lower_text] = color
-                G.add_node(node_id,
-                           label=f"{token.text}\n[{POS_TRANSLATIONS[lang].get(token.pos_, token.pos_)}]",
-                           pos=token.pos_,
-                           size=pos_counts[token.pos_] * 500,
-                           color=color)
-
-            if token.dep_ != "ROOT" and token.head.pos_ != 'PUNCT':
-                head_id = word_nodes.get(token.head.text.lower())
-                if head_id is not None:
-                    G.add_edge(head_id, word_nodes[lower_text], label=token.dep_)
-
-    return G, word_colors
-####################################################################################################################################
-def visualize_syntax_graph(doc, lang):
+
+def extract_entities(doc):
+    entities = {
+        "Personas": [],
+        "Conceptos": [],
+        "Lugares": [],
+        "Fechas": []
+    }
+
+    for ent in doc.ents:
+        if ent.label_ == "PER":
+            entities["Personas"].append(ent.text)
+        elif ent.label_ in ["LOC", "GPE"]:
+            entities["Lugares"].append(ent.text)
+        elif ent.label_ == "DATE":
+            entities["Fechas"].append(ent.text)
+        else:
+            entities["Conceptos"].append(ent.text)
+
+    return entities
+
+def visualize_context_graph(doc, lang):
+    G = nx.Graph()
+    entities = extract_entities(doc)
+
+    # Add nodes
+    for category, items in entities.items():
+        for item in items:
+            G.add_node(item, category=category)
+
+    # Add edges
+    for sent in doc.sents:
+        sent_entities = [ent.text for ent in sent.ents if ent.text in G.nodes()]
+        for i in range(len(sent_entities)):
+            for j in range(i+1, len(sent_entities)):
+                G.add_edge(sent_entities[i], sent_entities[j])
+
+    # 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))
+
+    plt.title("Análisis de Contexto" if lang == 'es' else "Context Analysis" if lang == 'en' else "Analyse de Contexte", fontsize=20)
+    plt.axis('off')
+
+    return plt
+
+def visualize_semantic_relations(doc, lang):
+    # Esta función puede mantener la lógica que ya tienes en visualize_syntax_graph
+    # con algunas modificaciones para enfocarse en relaciones semánticas
     G, word_colors = create_syntax_graph(doc, lang)
 
-    plt.figure(figsize=(24, 18))  # Increase figure size
-    pos = nx.spring_layout(G, k=0.9, iterations=50)  # Adjust layout parameters
+    plt.figure(figsize=(24, 18))
+    pos = nx.spring_layout(G, k=0.9, iterations=50)
 
     node_colors = [data['color'] for _, data in G.nodes(data=True)]
     node_sizes = [data['size'] for _, data in G.nodes(data=True)]
 
     nx.draw(G, pos, with_labels=False, node_color=node_colors, node_size=node_sizes, arrows=True, 
-            arrowsize=20, width=2, edge_color='gray')  # Adjust node and edge appearance
+            arrowsize=20, width=2, edge_color='gray')
 
     nx.draw_networkx_labels(G, pos, {node: data['label'] for node, data in G.nodes(data=True)}, 
-                            font_size=10, font_weight='bold')  # Increase font size and make bold
+                            font_size=10, font_weight='bold')
 
     edge_labels = nx.get_edge_attributes(G, 'label')
     nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_size=8)
 
-    plt.title("Syntactic Analysis" if lang == 'en' else "Analyse Syntaxique" if lang == 'fr' else "Análisis Sintáctico",
-              fontsize=20, fontweight='bold')  # Increase title font size
+    plt.title("Análisis de Relaciones Semánticas" if lang == 'es' else "Semantic Relations Analysis" if lang == 'en' else "Analyse des Relations Sémantiques",
+              fontsize=20, fontweight='bold')
     plt.axis('off')
 
     legend_elements = [plt.Rectangle((0,0),1,1, facecolor=color, edgecolor='none', 
                        label=f"{POS_TRANSLATIONS[lang][pos]} ({count_pos(doc)[pos]})")
                        for pos, color in POS_COLORS.items() if 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=12)  # Increase legend font size
+    plt.legend(handles=legend_elements, loc='center left', bbox_to_anchor=(1, 0.5), fontsize=12)
 
     return plt
-################################################################################################################################
-def visualize_syntax(text, nlp, lang):
-    max_tokens = 5000
-    doc = nlp(text)
-    if len(doc) > max_tokens:
-        doc = nlp(text[:max_tokens])
-        print(f"Warning: The input text is too long. Only the first {max_tokens} tokens will be visualized.")
-    return visualize_syntax_graph(doc, lang)
 
-    pass
+def perform_semantic_analysis(text, nlp, lang):
+    doc = nlp(text)
+    context_graph = visualize_context_graph(doc, lang)
+    relations_graph = visualize_semantic_relations(doc, lang)
+    return context_graph, relations_graph