Pushing to HuggingFace after Testing

main.py

@@ -39,12 +39,22 @@ def analyze_text(text):
     # Tokenize input text for NER
     ner_inputs = ner_tokenizer(text, return_tensors="pt")
 
+    input_ids = ner_inputs['input_ids']
+
+    # Converting token IDs back to tokens
+    tokens = [ner_tokenizer.convert_ids_to_tokens(id) for id in input_ids.squeeze().tolist()]
+
+    
     # Perform Named Entity Recognition
     with torch.no_grad():
         ner_outputs = ner_model(**ner_inputs)
     
-    ner_predictions = torch.argmax(ner_outputs, dim=-1)
+    ner_predictions = torch.argmax(ner_outputs, dim=-1)[0]
     ner_labels = ner_predictions.tolist()
+    ner_labels = [ner_config["id2labels"][str(label)] for label in ner_labels]
+
+    #matching the tokens with the labels
+    ner_labels = [f"{token}: {label}" for token, label in zip(tokens, ner_labels)]
 
     # Tokenize input text for sentiment analysis
     sentiment_inputs = sentiment_tokenizer(text, return_tensors="pt")
@@ -52,10 +62,10 @@ def analyze_text(text):
     # Perform sentiment analysis
     with torch.no_grad():
         sentiment_outputs = sentiment_model(**sentiment_inputs)
-    sentiment_probabilities = torch.softmax(sentiment_outputs, dim=1)
+    sentiment_probabilities = torch.argmax(sentiment_outputs, dim=1)
     sentiment_scores = sentiment_probabilities.tolist()
-
-    sentiment = sentiment_config["id2label"][torch.argmax(sentiment_outputs).item()]
+    sentiment_id = sentiment_scores[0]
+    sentiment = sentiment_config["id2label"][str(sentiment_id)]
 
     return ner_labels, sentiment
 
@@ -71,12 +81,11 @@ def main():
 
             # Display Named Entities
             st.subheader("Named Entities")
-            for label in ner_labels:
-                st.write(f"- {label}")
+            st.write(ner_labels)
 
             # Display Sentiment Analysis
             st.subheader("Sentiment Analysis")
             st.write(f"Sentiment: {sentiment_scores}")
 
 if __name__ == "__main__":
-    main()
+    main()