Update app.py

app.py

@@ -113,6 +113,7 @@ def create_analysis_html(sentence_results, global_label, global_confidence):
     html += "</table>"
     return html
 
+
 def process_input(text_input, labels_or_premise, mode):
     if mode == "Zero-Shot Classification":
         labels = [label.strip() for label in labels_or_premise.split(',')]
@@ -126,8 +127,9 @@ def process_input(text_input, labels_or_premise, mode):
     else:  # Long Context NLI
         # Global prediction
         global_pred = nli_classifier([{"text": text_input, "text_pair": labels_or_premise}], return_all_scores=True)[0]
-        global_results = {pred['label']: pred['score'] for pred in global_pred}
-        global_label, global_confidence = max(global_results.items(), key=lambda x: x[1])
+        global_results = {p['label']: p['score'] for p in global_pred}
+        global_label = max(global_results.items(), key=lambda x: x[1])[0]
+        global_confidence = max(global_results.values())
         
         # Sentence-level analysis
         sentences = sent_tokenize(text_input)
@@ -135,15 +137,18 @@ def process_input(text_input, labels_or_premise, mode):
         
         for sentence in sentences:
             sent_pred = nli_classifier([{"text": sentence, "text_pair": labels_or_premise}], return_all_scores=True)[0]
-            sent_scores = {pred['label']: pred['score'] for pred in sent_pred}
-            max_label = max(sent_scores.items(), key=lambda x: x[1])[0]
+            # Get the prediction and confidence for the sentence
+            pred_scores = [(p['label'], p['score']) for p in sent_pred]
+            max_pred = max(pred_scores, key=lambda x: x[1])
+            max_label, confidence = max_pred
+            
             sentence_results.append({
                 'sentence': sentence,
                 'prediction': max_label,
-                'scores': sent_scores
+                'confidence': confidence
             })
         
-        analysis_html = create_analysis_html(sentence_results, global_label,global_confidence)
+        analysis_html = create_analysis_html(sentence_results, global_label, global_confidence)
         return global_results, analysis_html
 
 def update_interface(mode):