Update app.py

app.py

@@ -177,6 +177,14 @@ discriminator_weights = ('data/datasets/discriminator_weights.pth')
 discriminator.load_state_dict(torch.load(discriminator_weights,map_location=torch.device('cpu')))
 discriminator.eval()
 
+def get_prediction(embeddings):
+    with torch.no_grad():
+        last_rep, logits, probs = discriminator(embeddings)
+
+    predicted_labels = torch.argmax(probs,dim=-1)
+    predicted_labels = predicted_labels.cpu().numpy()
+    return predicted_labels
+
 
 # Function to read reviews from a text file
 def load_reviews_from_file(file_path):
@@ -192,8 +200,15 @@ def load_reviews_from_file(file_path):
     return reviews
 
 
+
+
+
+
+
+
+#Demo Section
 st.title("SpaGAN Demo")
-st.write("Enter a text, and the system will highlight the geo-entities within it.")
+st.write("This demo allows you to select from a list of sample reviews that contrain both real and fake reviews. This demo will highlight all the entity types found within the review and display the prediction of the model.")
 
 # Define a color map and descriptions for different entity types
 COLOR_MAP = {
@@ -230,7 +245,7 @@ if st.button("Highlight Geo-Entities"):
         st.write("Concatenated Embedding Shape:", combined_embedding.shape)
         st.write("Concatenated Embedding:", combined_embedding)
 
-        prediction = discriminator(combined_embedding)
+        prediction = get_prediction(combined_embedding)
         st.write("Prediction:", prediction)
         
         # Process the text using spaCy