Making It look more asthetic

main.py

@@ -3,6 +3,8 @@ import json
 import torch
 from transformers import AutoTokenizer
 from modelling_cnn import CNNForNER, SentimentCNNModel
+import pandas as pd
+import altair as alt
 
 # Load the Yoruba NER model
 ner_model_name = "./my_model/pytorch_model.bin"
@@ -70,22 +72,71 @@ def analyze_text(text):
     return ner_labels, sentiment
 
 def main():
+    st.set_page_config(page_title="YorubaCNN for NER and Sentiment Analysis", layout="wide")
+    
     st.title("YorubaCNN Models for NER and Sentiment Analysis")
-
+    
     # Input text
     text = st.text_area("Enter Yoruba text", "")
-
+    
     if st.button("Analyze"):
         if text:
-            ner_labels, sentiment_scores = analyze_text(text)
-
+            ner_labels, sentiment = analyze_text(text)
+            
             # Display Named Entities
-            st.subheader("Named Entities")
-            st.write(ner_labels)
-
+            st.header("Named Entities")
+            
+            # Convert NER results to DataFrame
+            ner_df = pd.DataFrame([label.split(': ') for label in ner_labels], columns=['Token', 'Entity'])
+            
+            # Display NER results in a styled table
+            st.dataframe(ner_df.style.highlight_max(axis=0, color='lightblue'))
+            
             # Display Sentiment Analysis
-            st.subheader("Sentiment Analysis")
-            st.write(f"Sentiment: {sentiment_scores}")
+            st.header("Sentiment Analysis")
+            
+            # Create a sentiment score (you may need to adjust this based on your model's output)
+            sentiment_score = 0.8 if sentiment == "positive" else -0.8 if sentiment == "negative" else 0
+            
+            # Create a chart for sentiment visualization
+            sentiment_df = pd.DataFrame({'sentiment': [sentiment_score]})
+            chart = alt.Chart(sentiment_df).mark_bar().encode(
+                x=alt.X('sentiment', scale=alt.Scale(domain=(-1, 1))),
+                color=alt.condition(
+                    alt.datum.sentiment > 0,
+                    alt.value("green"),
+                    alt.value("red")
+                )
+            ).properties(width=600, height=100)
+            
+            st.altair_chart(chart)
+            st.write(f"Sentiment: {sentiment.capitalize()}")
+    
+    # Explanatory section
+    with st.expander("About this analysis"):
+        st.write("""
+        This tool uses YorubaCNN models to perform two types of analysis on Yoruba text:
+        
+        1. **Named Entity Recognition (NER)**: Identifies and classifies named entities (e.g., person names, organizations) in the text.
+        2. **Sentiment Analysis**: Determines the overall emotional tone of the text (positive, negative, or neutral).
+        
+        The models used are based on Convolutional Neural Networks (CNN) and are specifically trained for the Yoruba language.
+        """)
+
+    # Styling
+    st.markdown("""
+        <style>
+        .stAlert > div {
+            padding-top: 20px;
+            padding-bottom: 20px;
+        }
+        .stDataFrame {
+            padding: 10px;
+            border-radius: 5px;
+            background-color: #f0f2f6;
+        }
+        </style>
+        """, unsafe_allow_html=True)
 
 if __name__ == "__main__":
-    main()
+    main()