Update app.py

app.py

@@ -38,6 +38,10 @@ def inference(prompt):
     bi_prob = sigmoid(logits)
 
     predicted_class_id = bi_prob.argmax().item()
+
+    # 임시
+    model.config.id2label = {0: 'CLEAN', 1: 'UNCLEAN'}
+    
     class_id = model.config.id2label[predicted_class_id]
 
     return "class_id: " + str(class_id) + "\n" + "clean_prob: " + str(bi_prob[0][0].item()) + "\n" + "unclean_prob: " + str(bi_prob[0][1].item())