Update app.py

app.py

@@ -168,10 +168,17 @@ if st.button("Results"):
             labels = ["person", "location", "country", "city", "organization", "time", "date", "product", "event name", "money", "affiliation", "ordinal value", "percent value", "position"]
             entities = model.predict_entities(text, labels)
             df = pd.DataFrame(entities)
-            st.dataframe(df)
+            
             properties = {"border": "2px solid gray", "color": "blue", "font-size": "16px"}
             df_styled = df.style.set_properties(**properties)
             st.dataframe(df_styled)
+            if df is not None:
+                df["all"] = "all"
+                fig = px.treemap(df, path=['all', 'text', 'label'])
+                fig.update_traces(root_color="yellow")
+                fig.update_layout(margin = dict(t=50, l=25, r = 25, b=250))
+                st.plotly_chart(fig)
+