Update app.py

app.py

@@ -6,7 +6,7 @@ import gradio as gr
 all_letters = string.ascii_letters + " .,;'"
 n_letters = len(all_letters)
 
-all_categories = ['Arabic','Chinese','Czech','Dutch','English','French','German','Greek',
+class_names = ['Arabic','Chinese','Czech','Dutch','English','French','German','Greek',
                 'Irish','Italian','Japanese','Korean','Polish','Portuguese','Russian','Scottish',
                 'Spanish','Vietnamese']
                 
@@ -35,28 +35,26 @@ def evaluate(line_tensor):
     return output
     
 # Feeding in a name and number of top predictions you want to output
-def predict(last_name,n_predictions=3):
+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)
 
-        topv,topi = output.topk(n_predictions,1,True)
+        top3_prob,top3_catid = torch.topk(output,3)
         
-        top_3_countries = ''
-        for i in range(n_predictions):
-            value = topv[0]
-            category_index = topi[0][i].item()
-            top_3_countries += f'{all_categories[category_index]}: {round(value[i].item()*100,2)}% confident'
-            top_3_countries += '\n'
-    return top_3_countries
-
-demo = gr.Interface(predict,
+        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 = "text",
                         title = "Classify Last Name :)", 
-                        description="Classifies last name into one of 18 language of origin. Returns top 3 languages of origin"
+                        description="Classifies last name into one of 18 language of origin. Returns confidence % for the top three categories"
                     )
                     
 demo.launch(inline=False)