Update app.py

app.py

@@ -36,11 +36,12 @@ def classify_text(inp):
     with torch.no_grad():
         logits = model(**inp).logits
     print(logits)
-    print(logits.argmax().item())
-    # logits = torch.nn.Softmax(dim=0)(logits)
-    # print(logits)
-    # return {labels[i]: float(logits[i].item()) for i in range(len(labels))}
-    return {logits.argmax().item() : labels[logits.argmax().item()]}
+    print(logits.shape)
+    # logits = logits.argsort().squeeze(0)
+    logits = torch.nn.Softmax(dim=1)(logits).squeeze(0)
+    print(logits)
+    return {labels[i]: float(logits[i].item()) for i in range(len(labels))}
+    # return {logits.argmax().item() : labels[logits.argmax().item()]}
 
 input = '明天又是好日子\n千金的光阴不能等\n明天又是好日子\n赶上了盛世咱享太平\n今天是个好日子\n心想的事儿都能成\n明天又是好日子\n千金的光阴不能等\n今天明天都是好日子\n赶上了盛世咱享太平'