Upload Glyph multi-task model (epoch 3, avg acc 88.0%)

README.md

@@ -72,11 +72,10 @@ def predict(text, task="lang_id", top_k=3):
     with torch.no_grad():
         logits = model(inp, task)["logits"][0]
     probs = torch.softmax(logits, dim=-1)
-    idx2label = dict(enumerate(
-        sorted(ckpt["label_maps"][task], key=ckpt["label_maps"][task].get)
-    ))
-    topk = probs.topk(top_k)
-    return [(idx2label[topk.indices[i].item()], topk.values[i].item()) for i in range(top_k)]
+    idx2label = {v: k for k, v in ckpt["label_maps"][task].items()}
+    k = min(top_k, len(probs))
+    topk = probs.topk(k)
+    return [(idx2label[topk.indices[i].item()], topk.values[i].item()) for i in range(k)]
 
 # Examples
 print(predict("La pizza napoletana è patrimonio UNESCO", task="lang_id"))