Custom entities color

app.py

@@ -5,6 +5,13 @@ from annotated_text import annotated_text
 
 from analyzer import NewsAnalyzer
 
+ENTITY_COLOR = {
+    "PER": "#b2ffff",
+    "LOC": "#ffffb2",
+    "ORG": "#adfbaf",
+    "MISC": "#ffb2b2",
+}
+
 
 def run() -> None:
     analyzer = NewsAnalyzer(
@@ -15,7 +22,7 @@ def run() -> None:
     )
     st.title("📰 News Analyzer")
     headline = st.text_input("Headline:")
-    content = st.text_input("Content:")
+    content = st.text_area("Content:", height=200)
     if headline == "":
         st.error("Please, provide a headline.")
     else:
@@ -26,7 +33,7 @@ def run() -> None:
         button = st.button("Analyze")
         if button:
             predictions = analyzer(headline=headline, content=content)
-            col1, _, col2 = st.columns([2, 1, 5])
+            col1, _, col2 = st.columns([2, 1, 4])
 
             with col1:
                 st.subheader("Analysis:")
@@ -64,7 +71,13 @@ def parse_entities(
     parsed_text = []
     for entity in entities:
         parsed_text.append(text[start : entity["start"]])
-        parsed_text.append((entity["word"], entity["entity_group"]))
+        parsed_text.append(
+            (
+                entity["word"],
+                entity["entity_group"],
+                ENTITY_COLOR[entity["entity_group"]],
+            )
+        )
         start = entity["end"]
     parsed_text.append(text[start:])
     return parsed_text

pipeline.py

@@ -1,6 +1,7 @@
-from transformers import TextClassificationPipeline
 from typing import Dict, Optional
 
+from transformers import TextClassificationPipeline
+
 
 class NewsPipeline(TextClassificationPipeline):
     def __init__(self, emojis: Dict[str, str], **kwargs) -> None: