.

app.py

@@ -129,6 +129,9 @@ def show_avg(base_fn, model_name, patent_claim_num, show_pick=False):
   ]
 
   for i, tk in enumerate(lst):
+    if i == len(lst)-1:
+      break
+
     token_text = handle_char_return(tk['actual_next_token_text'])
     if token_text == '<|end_of_claim|>': 
       break
@@ -139,7 +142,7 @@ def show_avg(base_fn, model_name, patent_claim_num, show_pick=False):
     pick = int(tk['actual_next_token_top_seq'])
     prob = float(tk['actual_next_token_top_prob'])
 
-    for i, item in enumerate(colors):
+    for j, item in enumerate(colors):
       sum_all[item[2]] = 0
 
     if pick == 0:
@@ -158,11 +161,14 @@ def show_avg(base_fn, model_name, patent_claim_num, show_pick=False):
       fg_color = colors[2][1]
       tag = colors[2][2]
       sum_all[tag] += 1
+  
+    if show_pick:
+      pick = '[%s]' % pick
+    else:
+      pick = ''
 
     result += "<span style=background-color:#%s;color:#%s;border-radius:5px;>%s%s</span>  " % (bg_color, fg_color, token_text, pick)
 
-  st.text('debug 2')  
-
   color_msg = ''
   for i, v in enumerate(colors):
     color_msg += "<span style=background-color:#%s;color:#%s;border-radius:5px;>&nbsp;%s&nbsp;</span> " % (v[0], v[1], v[2]) 
@@ -178,9 +184,7 @@ def show_avg(base_fn, model_name, patent_claim_num, show_pick=False):
     'Total keystrokes: %s (with autocomplete), %s (without autocomplete)\n' \
     'Keystroke distribution: top 1~10: %s (top 1: %s), out of top 10: %s' % (model_name, saved_ratio, keys_with_auto, keys_without_auto,  sum_pick, sum_hit_1, sum_outscope_len)
   st.text(s)
-
-  s = color_msg
-  st.markdown(s, unsafe_allow_html=True)
+  st.markdown(color_msg, unsafe_allow_html=True)
   st.markdown(result, unsafe_allow_html=True)
   sum_lst = [sum_all['1'], sum_all['2-10'], sum_all['out of top 10']]
 
@@ -286,7 +290,7 @@ def main():
   pick = random.randrange(len(select_lst))
   num = select_lst[pick]
 
-  st.text('debug')
+  #st.text('debug 1')
 
   avgs = []
   for prefix in prefix_lst: