Update app.py

app.py

@@ -43,12 +43,13 @@ def Page():
       solara.Markdown(f'{spans1}')
       outputs = model.generate(tokens, max_new_tokens=1, output_scores=True, return_dict_in_generate=True, pad_token_id=tokenizer.eos_token_id)
       scores = F.softmax(outputs.scores[0], dim=-1)
-      top_10 = torch.topk(scores, 10)+torch.topk(scores, 10, largest=False)
+      top_10 = torch.topk(scores, 10)
+      bottom_10 = torch.topk(scores, 10, largest=False)
       df = pd.DataFrame()
-      df["probs"] = top_10.values[0]
+      df["probs"] = top_10.values[0]+bottom_10.values[0]
       df["probs"] = [f"{value:.2%}" for value in df["probs"].values]
-      df["next token ID"] = [top_10.indices[0][i].numpy() for i in range(20)]
-      df["predicted next token"] = [tokenizer.decode(top_10.indices[0][i]) for i in range(20)]
+      df["next token ID"] = [top_10.indices[0][i].numpy() for i in range(10)]+[bottom_10.indices[0][i].numpy() for i in range(10)]
+      df["predicted next token"] = [tokenizer.decode(top_10.indices[0][i]) for i in range(10)]+[tokenizer.decode(bottom_10.indices[0][i]) for i in range(10)]
       solara.Markdown("###Prediction")
       solara.DataFrame(df, items_per_page=20, cell_actions=cell_actions)
 Page()