squeeze em bits

app.py

@@ -19,7 +19,7 @@ def generate_response(text):
   mask_token_id = tokenizer.mask_token_id
   mask_token_index = torch.where(input_ids == mask_token_id)[1]
   token_logits = logits[0, mask_token_index, :]
-  top_5_tokens = torch.topk(token_logits, k=5).indices # get top 5 tokens
+  top_5_tokens = torch.topk(token_logits.squeeze(), k=5).indices # get top 5 tokens
   predicted_tokens = tokenizer.convert_ids_to_tokens(top_5_tokens.tolist()) # convert ids to tokens
   
   # Choose one of the predicted tokens randomly and replace the mask with it