Update app.py

app.py

@@ -69,6 +69,22 @@ def class_judgement(text: str) -> str:
         lines.append(f"{lab}: {p:.3f}")
     return "\n".join(lines)
 
+def class_judgement(text:str) -> dict:
+    """Predicts two task distributions for the input text."""
+    batch_x = tokenizer(
+        text, max_length=MAX_LEN, padding=True, truncation=True, return_tensors="pt"
+    ).to("cuda")
+
+    with torch.no_grad():
+        task0_output, task1_output = model.model(batch_x["input_ids"], batch_x["attention_mask"])
+        probs_0 = torch.nn.functional.softmax(task0_output, dim=1).squeeze().tolist()
+
+    # 任務0（程度等級）
+    task0_id2label = {0: '入門基礎', 1: '進階高階', 2: '流利精通'}
+
+    predictions_0 = {task0_id2label[i]: round(probs_0[i], 3) for i in range(len(probs_0))}
+    return predictions_0
+
 # =========================
 # 範例資料（只放兩個）
 # - 閱讀測驗：id=1