Update app.py

app.py

@@ -35,58 +35,25 @@ def lineToTensor(line):
         tensor[li][0][letterToIndex(letter)] = 1
     return tensor
 
-
-class RNN(nn.Module):
-    """LSTM class"""
-    def __init__(self, input_size, hidden_size, output_size):
-        '''
-        :param input_size: number of input coming in
-        :param hidden_size: number of hidden units
-        :param output_size: size of the output
-        '''
-        super(RNN, self).__init__()
-
-        self.hidden_size = hidden_size
-        self.input_size = input_size
-
-        #LSTM
-        self.lstm = nn.LSTM(input_size, hidden_size)
-        self.hidden2Cat = nn.Linear(hidden_size, output_size)
-        self.hidden = self.init_hidden()
-
-    def forward(self, input, hidden):
-
-        lstm_out, self.hidden = self.lstm(input, hidden)
-        output = self.hidden2Cat(lstm_out[-1]) #many to one
-        output = F.log_softmax(output, dim=1)
-
-        return output
-
-    def init_hidden(self):
-        return (torch.zeros(1, 1, self.hidden_size),
-                torch.zeros(1, 1, self.hidden_size))
-
 # Just return an output given a line
 def evaluate_model(line_tensor):
-    n_hidden = 128
-    n_categories = len(all_categories)
-    model = RNN(n_letters, n_hidden, n_categories)
-    model.load_state_dict(torch.load('classify_names_lstm.pt'))
-    model.eval()
+    model = torch.jit.load("torchscript_classify_names_lstm.pt")
     
-    hidden = (torch.zeros(1, 1, 128),
-                torch.zeros(1, 1, 128))
+    hidden = (torch.zeros(1, 1, 128),torch.zeros(1, 1, 128))
     output = model(line_tensor,hidden)
 
     return output
 
 def classify_lastname(last_name):
+    # Converting to Ascii and capitalizing first letter
     last_name = unicodeToAscii(last_name)
     last_name = last_name.title()
-    
+
+    # Converting name to tensor
     line_tensor = lineToTensor(last_name)
     output = evaluate_model(line_tensor)
 
+    # Grabbing top3 probabilities and categories
     top3_prob, top3_cat = torch.topk(output,3)
     probs = torch.exp(top3_prob[0])
     cats = top3_cat[0]