Update app.py

app.py

@@ -11,6 +11,9 @@ label_encoder = joblib.load('WSP_label_encoder.joblib')
 
 input_phrase = st.text_input("Input search text")
 
+# Добавляем слайдер для выбора вероятности
+confidence_threshold = st.slider("Confidence Threshold (%)", 0, 100, 50)
+
 if input_phrase:
     tokenized_input = tokenizer(input_phrase, truncation=True, padding=True, return_tensors='pt')
     
@@ -18,6 +21,10 @@ if input_phrase:
         outputs = loaded_model(**tokenized_input)
     
     predicted_class = torch.argmax(outputs.logits, dim=1).item()
-    predicted_category = label_encoder.inverse_transform([predicted_class])[0]
+    predicted_confidence = torch.softmax(outputs.logits, dim=1)[0][predicted_class].item() * 100
     
-    st.text(f"Output: {predicted_category}")
+    if predicted_confidence >= confidence_threshold:
+        predicted_category = label_encoder.inverse_transform([predicted_class])[0]
+        st.text(f"Output: {predicted_category}")
+    else:
+        st.text(f"Output: {input_phrase}")