Update app.py

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 import spacy
 import torch
+import torch.nn as nn
 from transformers import BertTokenizer, BertModel
 from transformers.models.bert.modeling_bert import BertForMaskedLM
 
@@ -142,16 +143,38 @@ def processSpatialEntities(review, nlp):
     return processed_embedding
 
 
-
-
-#dConfig = AutoConfig.from_pretrained("bert-base-uncased")
-#hidden_size = int(dConfig.hidden_size)
-#num_hidden_layers_d = 2; 
-#hidden_levels_d = [hidden_size for i in range(0, num_hidden_layers_d)]
-#label_list = ["1", "0"]
-#label_list.append('UNL')
-#discriminator = Discriminator(input_size=hidden_size*2, hidden_sizes=hidden_levels_d, num_labels=len(label_list), dropout_rate=out_dropout_rate).to(device)
-
+#Initialize discriminator module
+class Discriminator(nn.Module):
+    def __init__(self, input_size=512, hidden_sizes=[512], num_labels=2, dropout_rate=0.1):
+        super(Discriminator, self).__init__()
+        self.input_dropout = nn.Dropout(p=dropout_rate)
+        layers = []
+        hidden_sizes = [input_size] + hidden_sizes
+        for i in range(len(hidden_sizes)-1):
+            layers.extend([nn.Linear(hidden_sizes[i], hidden_sizes[i+1]), nn.LeakyReLU(0.2, inplace=True), nn.Dropout(dropout_rate)])
+
+        self.layers = nn.Sequential(*layers) #per il flatten
+        self.logit = nn.Linear(hidden_sizes[-1],num_labels+1) # +1 for the probability of this sample being fake/real.
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, input_rep):
+        input_rep = self.input_dropout(input_rep)
+        last_rep = self.layers(input_rep)
+        logits = self.logit(last_rep)
+        probs = self.softmax(logits)
+        return last_rep, logits, probs
+        
+dConfig = AutoConfig.from_pretrained("bert-base-uncased")
+hidden_size = int(dConfig.hidden_size)
+num_hidden_layers_d = 2; 
+hidden_levels_d = [hidden_size for i in range(0, num_hidden_layers_d)]
+label_list = ["1", "0"]
+label_list.append('UNL')
+
+discriminator = Discriminator(input_size=hidden_size*2, hidden_sizes=hidden_levels_d, num_labels=len(label_list), dropout_rate=out_dropout_rate).to(device)
+discriminator_weights = ('data/datasets/discriminator_weights.pth')
+discriminator.load_state_dict(torch.load(discriminator_weights))
+discriminator.eval()
 
 
 # Function to read reviews from a text file