Simple changes to model

main.py

@@ -4,7 +4,6 @@ import torch
 from transformers import AutoTokenizer
 from modelling_cnn import CNNForNER, SentimentCNNModel
 
-
 # Load the Yoruba NER model
 ner_model_name = "./my_model/pytorch_model.bin"
 model_ner = "Testys/cnn_yor_ner"
@@ -44,8 +43,9 @@ def analyze_text(text):
     with torch.no_grad():
         ner_outputs = ner_model(**ner_inputs)
     
-    ner_predictions = torch.argmax(ner_outputs.logits, dim=-1)
-    ner_labels = [ner_tokenizer.decode(token) for token in ner_predictions[0]]
+    ner_predictions = torch.argmax(ner_outputs, dim=-1)
+    ner_labels = [ner_config["id2label"][label] for label in ner_predictions[0].tolist()]
+
 
     # Tokenize input text for sentiment analysis
     sentiment_inputs = sentiment_tokenizer.encode_plus(text, return_tensors="pt")
@@ -53,10 +53,12 @@ def analyze_text(text):
     # Perform sentiment analysis
     with torch.no_grad():
         sentiment_outputs = sentiment_model(**sentiment_inputs)
-    sentiment_probabilities = torch.softmax(sentiment_outputs.logits, dim=1)
+    sentiment_probabilities = torch.softmax(sentiment_outputs, dim=1)
     sentiment_scores = sentiment_probabilities.tolist()
 
-    return ner_labels, sentiment_scores
+    sentiment = sentiment_config["id2label"][torch.argmax(sentiment_outputs).item()]
+
+    return ner_labels, sentiment
 
 def main():
     st.title("YorubaCNN Models for NER and Sentiment Analysis")
@@ -75,9 +77,7 @@ def main():
 
             # Display Sentiment Analysis
             st.subheader("Sentiment Analysis")
-            st.write(f"Positive: {sentiment_scores[2]:.2f}")
-            st.write(f"Negative: {sentiment_scores[0]:.2f}")
-            st.write(f"Neutral: {sentiment_scores[1]:.2f}")
+            st.write(f"Sentiment: {sentiment_scores}")
 
 if __name__ == "__main__":
     main()