Update

app.py

@@ -30,11 +30,16 @@ def getDictFromPOS(texts, labels):
 
 def randomColor():
     rgb = (random.uniform(0.0,1.0), random.uniform(0.0,1.0), random.uniform(0.0,1.0))
-    hsl = (random.uniform(0.0,1.0), random.uniform(0.35,1.0), 0.62)
+    hsl = (random.uniform(0.0,1.0), random.uniform(0.35,1.0), 0.75)
     return str(Color(rgb=rgb, hsl=hsl))
 
-def getAnnotatedFromPOS(texts, labels):
-    return [(t,l,randomColor()) for t, l in zip(texts, labels)]
+def get_colors(model):
+    labels = [t.decode("utf-8") for t in model.tag_dictionary.idx2item if t.isupper() and len(t) > 1]
+    colors = [randomColor() for t in model.tag_dictionary.idx2item if t.isupper() and len(t) > 1]
+    return dict(zip(labels, colors))
+
+def getAnnotatedFromPOS(texts, labels, colors):
+    return [(t,l,colors[t]) for t, l in zip(texts, labels)]
 
 def main():
 
@@ -42,6 +47,7 @@ def main():
 
     checkpoint = st.selectbox("Choose model", checkpoints)
     model = get_model(checkpoint)
+    colors = get_colors(model)
 
     default_text = "George Washington est allé à Washington"
     input_text = st.text_area(
@@ -65,7 +71,7 @@ def main():
                 texts, labels = getPos(s)
                 
                 st.header("Labels:")
-                anns = getAnnotatedFromPOS(texts, labels)
+                anns = getAnnotatedFromPOS(texts, labels, colors)
                 annotated_text(*anns)
 
                 st.header("JSON:")