import torch
import torch.nn.functional as F
import string
import gradio as gr

all_letters = string.ascii_letters + " .,;'"
n_letters = len(all_letters)

class_names = ['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 letterToIndex(letter):
    return all_letters.find(letter)
    
# Giving each charachter in name a one hot vector
def lineToTensor(line):
    tensor = torch.zeros(len(line),1,n_letters)
    for li,letter in enumerate(line):
        tensor[li][0][letterToIndex(letter)] = 1

    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)

    for i in range(line_tensor.size()[0]):
        output, hidden = my_model(line_tensor[i], hidden)

    return output
    
# Feeding in a name and number of top predictions you want to output
def classify_lastname(last_name,n_predictions=3):
    
    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()
        
    return model_output

demo = gr.Interface(classify_lastname,
                        inputs = "text", 
                        outputs = gr.outputs.Label(type="confidences",num_top_classes=3),
                        title = "Classify Last Name :)", 
                        description="Classifies last name into one of 18 language of origin. Returns confidence % for the top three categories"
                    )
                    
demo.launch(inline=False)