Returning only winning probability and category

app.py

@@ -62,15 +62,14 @@ def classify_lastname(last_name):
     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]
+    # Grabbing top probability and category
+    top_prob, top_cat = torch.topk(output,1)
+    prob = torch.exp(top_prob[0])
+    cat = top_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)
+    model_output[all_categories[cat[0].item()]] = round(prob[0].item(),2)
+
     return model_output
     
 # Define a request model