Update app.py

app.py

@@ -32,8 +32,8 @@ input = tokenizer(premise, hypothesis, truncation=True, return_tensors="pt")
 output = model(input["input_ids"].to(device))  # device = "cuda:0" or "cpu"
 prediction = torch.softmax(output["logits"][0], -1).tolist()
 label_names = ["support", "neutral", "refute"]
-prediction = {name: round(float(pred) * 100, 1) for pred, name in zip(prediction, label_names)}
-print(prediction)
+prediction = {name: float(pred) for pred, name in zip(prediction, label_names)}
+highest_label = max(prediction, key=prediction.get)
 
 
 from transformers import pipeline
@@ -69,14 +69,16 @@ def extract_person_names(sentence):
 person_name1 = extract_person_names(selected_sentence1)
 person_name2 = extract_person_names(selected_sentence2)
 
-st.write("Result:", prediction)
 
 col1, col2 = st.columns(2)
 
 with col1:
-    st.write("Without Factual Entailment:",prediction)
+    st.write("Without Factual Entailment.")
+    st.write("Textual Entailment Model:\n",highest_label)
 
 with col2:
-    st.write("Factual Entailment:",labels)
+    st.write("With Factual Entailment:")
+    st.write("Textual Entailment Model:\n",labels)
+    st.write("Span Detection Model:\n")
     st.write(f"{person_name1}::{person_name2}")