updated colors

app.py

@@ -3,6 +3,31 @@ from pyabsa import ATEPCCheckpointManager
 import gradio as gr
 
 
+def assign_sentiment_marker(sentiment_list, aspect_index):
+    """
+    This function returns a sentiment marker '+' for Positive, '-' for Negative, and '$' for any other sentiment.
+    It maps the sentiment from the sentiment_list at the aspect_index to the respective marker.
+
+    Parameters:
+    sentiment_list (list): A list of sentiments.
+    aspect_index (int): The index of the sentiment in sentiment_list to map to a marker.
+
+    Returns:
+    str: The sentiment marker.
+    """
+
+    # Define a dictionary to map sentiment labels to markers.
+    sentiment_markers = {
+        "Positive": "pos",
+        "Negative": "neg"
+    }
+
+    # Fetch sentiment label from the sentiment_list using aspect_index.
+    sentiment_label = sentiment_list[aspect_index]
+
+    # Use get method on dictionary to fetch corresponding marker or default to '$' for unlisted sentiments.
+    return sentiment_markers.get(sentiment_label, "neu")
+    
 
 def annotate_sentiment(tokens, aspect_positions, sentiment_list):
     """
@@ -19,7 +44,7 @@ def annotate_sentiment(tokens, aspect_positions, sentiment_list):
     for i, token_group in enumerate(tokens):
         if aspect_index < len(aspect_positions) and i == aspect_positions[aspect_index]:
             for token in token_group:
-                sentiment = '+' if sentiment_list[aspect_index] == 'Positive' else '-'
+                sentiment = assign_sentiment_marker(sentiment_list, aspect_index)
                 annotated_tokens.append((token, sentiment))
             aspect_index += 1
         else:
@@ -52,10 +77,10 @@ def annotate_text_sentiment(text):
         aspect_positions = [position[0]-1 for position in aspect_extraction_result['position']]
         sentiment_list = aspect_extraction_result["sentiment"]
 
+    print(sentiment_list)
     annotated_tokens = annotate_sentiment(tokens, aspect_positions, sentiment_list)
     return annotated_tokens
 
-
 # Initializing the aspect extractor model
 checkpoint_map = available_checkpoints()
 aspect_extractor = ATEPCCheckpointManager.get_aspect_extractor(checkpoint='english',
@@ -77,7 +102,7 @@ demo1 = gr.Interface(
         label="Aspect Detector based on DeBERTa",
         combine_adjacent=True,
         show_legend=True,
-    ).style(color_map={"+": "green", "-": "red"}),
+    ).style(color_map={"pos": "green", "neg": "red", "neu":"blue"}),
     theme=gr.themes.Base()
 )
 
@@ -95,7 +120,7 @@ demo2 = gr.Interface(
         label="Aspect Detector based on Relational Graph Attention Networks, and BERT",
         combine_adjacent=True,
         show_legend=True,
-    ).style(color_map={"+": "green", "-": "red"}),
+    ).style(color_map={"pos": "green", "neg": "red", "neu":"blue"}),
     theme=gr.themes.Base()
 )