Update app.py

app.py

@@ -54,6 +54,16 @@ model.load_state_dict(checkpoint.get("model_state_dict", checkpoint), strict=Fal
 model = model.to(device)
 model.eval()
 
+def get_color(prob):
+    if prob < 0.25:
+        return "green"
+    elif prob < 0.5:
+        return "yellow"
+    elif prob < 0.75:
+        return "orange"
+    else:
+        return "red"
+        
 def get_word_probabilities(text):
     text = " ".join(text.split(" ")[:2048])
     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
@@ -73,7 +83,7 @@ def get_word_probabilities(text):
             if current_word and current_probs:
                 current_prob = sum(current_probs) / len(current_probs)
                 word_probs.append(current_prob)
-                color = ("green" if current_prob < 0.25 else "yellow" if current_prob < 0.5 else "orange" if current_prob < 0.75 else "red")
+                color = get_color(current_prob)
                 word_colors.append(color)
             current_word = token[1:] if token != "▁" else ""
             current_probs = [prob]
@@ -83,8 +93,17 @@ def get_word_probabilities(text):
     if current_word and current_probs:
         current_prob = sum(current_probs) / len(current_probs)
         word_probs.append(current_prob)
-        color = ("green" if current_prob < 0.25 else "yellow" if current_prob < 0.5 else "orange" if current_prob < 0.75 else "red")
+        color = get_color(current_prob)
         word_colors.append(color)
+        
+    ####### FOR STABLE OUTPUTS
+    first_avg = (word_probs[1] + word_probs[2]) / 2
+    word_colors[0] = get_color(first_avg)
+
+    last_avg = (word_probs[-2] + word_probs[-3]) / 2
+    word_colors[-1] = get_color(last_avg)
+    #########
+    
     word_probs = [float(p) for p in word_probs]
     return word_probs, word_colors