Update app.py

app.py

@@ -146,7 +146,13 @@ def processSpatialEntities(review, nlp):
             token_embeddings.append(spaBert_emb)
             if(dev_mode == True):
                 st.write("Geo-Entity Found in review: ", text)
-
+    
+    # Handle the case where no geo-entities are found
+    if not token_embeddings:
+        st.warning("No geo-entities found in the review. Please include one from the list.")
+        # Return a zero vector as a fallback if no entities are found
+        return torch.zeros(bert_model.config.hidden_size), []
+    
     token_embeddings = torch.stack(token_embeddings, dim=0)
     processed_embedding = token_embeddings.mean(dim=0)  # Shape: (768)
     #processed_embedding = processed_embedding.unsqueeze(0)  # Shape: (1, 768)
@@ -273,7 +279,7 @@ user_input_review = st.text_area("Or type your own review here","")
 st.info(f"Please include one of the following entities in your review:\n {', '.join(california_entities)}")
 
 review_to_process = user_input_review if user_input_review.strip() else selected_review
-st.write("Selected Review: ", review_to_process)
+#st.write("Selected Review: ", review_to_process)
 lower_case_review = review_to_process.lower()
 
 # Process the text when the button is clicked
@@ -281,45 +287,49 @@ if st.button("Process Review"):
     if lower_case_review.strip():
         bert_embedding = get_bert_embedding(lower_case_review)
         spaBert_embedding, current_pseudo_sentences = processSpatialEntities(review_to_process,nlp)
-        combined_embedding = torch.cat((bert_embedding,spaBert_embedding),dim=-1)
-        
-        if(dev_mode == True):
-            st.write("Review Embedding Shape:", bert_embedding.shape)
-            st.write("Geo-Entities embedding shape: ", spaBert_embedding.shape)
-            st.write("Concatenated Embedding Shape:", combined_embedding.shape)
-            st.write("Concatenated Embedding:", combined_embedding)
-
-        prediction = get_prediction(combined_embedding)
-                    
-        # Process the text using spaCy
-        doc = nlp(review_to_process)
-        
-        # Highlight geo-entities with different colors
-        highlighted_text = review_to_process
-        for ent in reversed(doc.ents):
-            if ent.label_ in COLOR_MAP:
-                color = COLOR_MAP[ent.label_][0]
-                highlighted_text = (
-                    highlighted_text[:ent.start_char] +
-                    f"<span style='color:{color}; font-weight:bold'>{ent.text}</span>" + 
-                    highlighted_text[ent.end_char:]
-                )
-
-        # Display the highlighted text with HTML support
-        st.markdown(highlighted_text, unsafe_allow_html=True)
-
-        #Display pseudo sentences found
-        for sentence in current_pseudo_sentences:
-            clean_sentence = sentence.replace("[PAD]", "").strip()
-            st.write("Pseudo-Sentence:", clean_sentence)
-        
-        #Display the models prediction
-        if prediction == 0:
-            st.markdown("<h3 style='color:green;'>✅ Prediction: Not Spam</h3>", unsafe_allow_html=True)
-        elif prediction == 1:
-            st.markdown("<h3 style='color:red;'>❌ Prediction: Spam</h3>", unsafe_allow_html=True)
+        # Check if SpaBERT embedding is valid
+        if spaBert_embedding is None or spaBert_embedding.sum() == 0:
+            st.error("Unable to process the review. Please include at least one valid geo-entity.")
         else:
-            st.markdown("<h3 style='color:orange;'>⚠️ Error during prediction</h3>", unsafe_allow_html=True)
+            combined_embedding = torch.cat((bert_embedding,spaBert_embedding),dim=-1)
+            
+            if(dev_mode == True):
+                st.write("Review Embedding Shape:", bert_embedding.shape)
+                st.write("Geo-Entities embedding shape: ", spaBert_embedding.shape)
+                st.write("Concatenated Embedding Shape:", combined_embedding.shape)
+                st.write("Concatenated Embedding:", combined_embedding)
+    
+            prediction = get_prediction(combined_embedding)
+                        
+            # Process the text using spaCy
+            doc = nlp(review_to_process)
+            
+            # Highlight geo-entities with different colors
+            highlighted_text = review_to_process
+            for ent in reversed(doc.ents):
+                if ent.label_ in COLOR_MAP:
+                    color = COLOR_MAP[ent.label_][0]
+                    highlighted_text = (
+                        highlighted_text[:ent.start_char] +
+                        f"<span style='color:{color}; font-weight:bold'>{ent.text}</span>" + 
+                        highlighted_text[ent.end_char:]
+                    )
+    
+            # Display the highlighted text with HTML support
+            st.markdown(highlighted_text, unsafe_allow_html=True)
+    
+            #Display pseudo sentences found
+            for sentence in current_pseudo_sentences:
+                clean_sentence = sentence.replace("[PAD]", "").strip()
+                st.write("Pseudo-Sentence:", clean_sentence)
+            
+            #Display the models prediction
+            if prediction == 0:
+                st.markdown("<h3 style='color:green;'>✅ Prediction: Not Spam</h3>", unsafe_allow_html=True)
+            elif prediction == 1:
+                st.markdown("<h3 style='color:red;'>❌ Prediction: Spam</h3>", unsafe_allow_html=True)
+            else:
+                st.markdown("<h3 style='color:orange;'>⚠️ Error during prediction</h3>", unsafe_allow_html=True)
             
     else:
         st.error("Please select a review.")