Testys/YorubaCNN

1import streamlit as st2import json3import torch4from transformers import AutoTokenizer, AutoModelForTokenClassification5from modelling_cnn import CNNForNER, SentimentCNNModel6import pandas as pd7import altair as alt8 9# Load the Yoruba NER model10# ner_model_name = "./my_model/pytorch_model.bin"11# model_ner = "Testys/cnn_yor_ner"12# ner_tokenizer = AutoTokenizer.from_pretrained(model_ner)13# with open("./my_model/config.json", "r") as f:14#     ner_config = json.load(f)15 16# ner_model = CNNForNER(17#                       pretrained_model_name=ner_config["pretrained_model_name"],18#                       num_classes=ner_config["num_classes"]19#                       )20# ner_model.load_state_dict(torch.load(ner_model_name, map_location=torch.device('cpu')))21# ner_model.eval()22 23ner_model = AutoModelForTokenClassification.from_pretrained("masakhane/afroxlmr-large-ner-masakhaner-1.0_2.0")24ner_tokenizer = AutoTokenizer.from_pretrained("masakhane/afroxlmr-large-ner-masakhaner-1.0_2.0")25ner_config = ner_model.config26 27ner_model.eval()28 29 30# Load the Yoruba sentiment analysis model31sentiment_model_name = "./sent_model/sent_pytorch_model.bin"32model_sent = "Testys/cnn_sent_yor"33sentiment_tokenizer = AutoTokenizer.from_pretrained(model_sent)34 35with open("./sent_model/config.json", "r") as f:36    sentiment_config = json.load(f)37 38sentiment_model = SentimentCNNModel(39                                    transformer_model_name=sentiment_config["pretrained_model_name"],40                                    num_classes=sentiment_config["num_classes"]41                                    )42 43sentiment_model.load_state_dict(torch.load(sentiment_model_name, map_location=torch.device('cpu')))44sentiment_model.eval()45 46 47def analyze_text(text):48    # Tokenize input text for NER49    ner_inputs = ner_tokenizer(text, return_tensors="pt")50    51    # Perform Named Entity Recognition52    tokens = ner_tokenizer.convert_ids_to_tokens(ner_inputs.input_ids[0])53    with torch.no_grad():54        ner_outputs = ner_model(**ner_inputs)55    56    print(ner_outputs)57    58    ner_predictions = torch.argmax(ner_outputs.logits, dim=-1)[0]59    ner_labels = ner_predictions.tolist()60    print(ner_labels)61    ner_labels = [ner_config.id2label[label] for label in ner_labels]62 63    #matching the tokens with the labels64    ner_labels = [f"{token}: {label}" for token, label in zip(tokens, ner_labels)]65 66    # Tokenize input text for sentiment analysis67    sentiment_inputs = sentiment_tokenizer(text, max_length= 514, truncation= True, padding= "max_length", return_tensors="pt")68 69    # Perform sentiment analysis70    with torch.no_grad():71        sentiment_outputs = sentiment_model(**sentiment_inputs)72    sentiment_probabilities = torch.argmax(sentiment_outputs, dim=1)73    sentiment_scores = sentiment_probabilities.tolist()74    sentiment_id = sentiment_scores[0]75    sentiment = sentiment_config["id2label"][str(sentiment_id)]76 77    return ner_labels, sentiment78 79def main():80    st.set_page_config(page_title="YorubaCNN for NER and Sentiment Analysis", layout="wide")81    82    st.title("YorubaCNN Models for NER and Sentiment Analysis")83    84    # Input text85    text = st.text_area("Enter Yoruba text", "")86    87    if st.button("Analyze"):88        if text:89            ner_labels, sentiment = analyze_text(text)90            91            # Display Named Entities92            st.header("Named Entities")93            94            # Convert NER results to DataFrame95            ner_df = pd.DataFrame([label.split(': ') for label in ner_labels], columns=['Token', 'Entity'])96            97            # Display NER results in a styled table98            st.dataframe(ner_df.style.highlight_max(axis=0, color='lightblue'))99            100            # Display Sentiment Analysis101            st.header("Sentiment Analysis")102            103            # Create a sentiment score (you may need to adjust this based on your model's output)104            sentiment_score = 0.8 if sentiment == "positive" else -0.8 if sentiment == "negative" else 0105            106            # Create a chart for sentiment visualization107            sentiment_df = pd.DataFrame({'sentiment': [sentiment_score]})108            chart = alt.Chart(sentiment_df).mark_bar().encode(109                x=alt.X('sentiment', scale=alt.Scale(domain=(-1, 1))),110                color=alt.condition(111                    alt.datum.sentiment > 0,112                    alt.value("green"),113                    alt.value("red")114                )115            ).properties(width=600, height=100)116            117            st.altair_chart(chart)118            st.write(f"Sentiment: {sentiment.capitalize()}")119    120    # Explanatory section121    with st.expander("About this analysis"):122        st.write("""123        This tool uses YorubaCNN models to perform two types of analysis on Yoruba text:124        125        1. **Named Entity Recognition (NER)**: Identifies and classifies named entities (e.g., person names, organizations) in the text.126        2. **Sentiment Analysis**: Determines the overall emotional tone of the text (positive, negative, or neutral).127        128        The models used are based on Convolutional Neural Networks (CNN) and are specifically trained for the Yoruba language.129        """)130 131    # Styling132    st.markdown("""133        <style>134        .stAlert > div {135            padding-top: 20px;136            padding-bottom: 20px;137        }138        .stDataFrame {139            padding: 10px;140            border-radius: 5px;141            background-color: #f0f2f6;142        }143        </style>144        """, unsafe_allow_html=True)145 146if __name__ == "__main__":147    main()