Update app.py

app.py

@@ -1,55 +1,100 @@
-import torch
-import torch.nn.functional as F
+import unicodedata
 import string
+
 import gradio as gr
 
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
 all_letters = string.ascii_letters + " .,;'"
 n_letters = len(all_letters)
 
-class_names = ['Arabic','Chinese','Czech','Dutch','English','French','German','Greek',
+all_categories = ['Arabic','Chinese','Czech','Dutch','English','French','German','Greek',
                 'Irish','Italian','Japanese','Korean','Polish','Portuguese','Russian','Scottish',
                 'Spanish','Vietnamese']
-                
-# Find letter index from all_letters: Ex: "a" = 0
+
+def unicodeToAscii(s):
+    return ''.join(
+        c for c in unicodedata.normalize('NFD', s)
+        if unicodedata.category(c) != 'Mn'
+        and c in all_letters
+    )
+
+# Find letter index from all_letters, e.g. "a" = 0
 def letterToIndex(letter):
     return all_letters.find(letter)
-    
-# Giving each charachter in name a one hot vector
+
+
+# Turn a line into a <line_length x 1 x n_letters>,
+# or an array of one-hot letter vectors
 def lineToTensor(line):
-    tensor = torch.zeros(len(line),1,n_letters)
-    for li,letter in enumerate(line):
+    tensor = torch.zeros(len(line), 1, n_letters)
+    for li, letter in enumerate(line):
         tensor[li][0][letterToIndex(letter)] = 1
+    return tensor
 
-    return tensor 
-    
-# Loading in torchscript model
-my_model = torch.jit.load('name_classifier_ts.ptl')
 
-# Return output given a line_tensor
-def evaluate(line_tensor):
-    hidden = torch.zeros(1,128)
+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 he hidden units
+        :param output_size: size of the output
+        '''
+        super(RNN, self).__init__()
 
-    for i in range(line_tensor.size()[0]):
-        output, hidden = my_model(line_tensor[i], hidden)
+        self.hidden_size = hidden_size
+        self.input_size = input_size
 
-    return output
-    
-# Feeding in a name and number of top predictions you want to output
-def classify_lastname(last_name,n_predictions=3):
+        #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()
     
-    last_name = last_name.title()
-    with torch.no_grad():
-        output = evaluate(lineToTensor(last_name))
-        output = F.softmax(output,dim=1)
-
-        top3_prob,top3_catid = torch.topk(output,3)
-        
-        model_output = {}      
-        for i in range(top3_prob.size(1)):
-            model_output[class_names[top3_catid[0][i].item()]] = top3_prob[0][i].item()
-        
+    hidden = (torch.zeros(1, 1, 128),
+                torch.zeros(1, 1, 128))
+    output = model(line_tensor,hidden)
+
+    return output
+
+def classify_lastname(last_name):
+    last_name = unicodeToAscii(last_name)
+    line_tensor = lineToTensor(last_name)
+    output = evaluate_model(line_tensor)
+
+    top3_prob, top3_cat = torch.topk(output,3)
+    probs = torch.exp(top3_prob[0])
+    cats = top3_cat[0]
+
+    model_output = {}
+    for i in range(3):
+        print(probs[i].item())
+        model_output[all_categories[cats[i].item()]] = round(probs[i].item(),2)
     return model_output
 
+
 demo = gr.Interface(classify_lastname,
                         inputs = "text", 
                         outputs = gr.outputs.Label(type="confidences",num_top_classes=3),