From 2557b22e575ca2793eb4c7e6742c89f217f1ce4b Mon Sep 17 00:00:00 2001
From: Aleksandar Petrov <7101799+AleksandarPetrov@users.noreply.github.com>
Date: Tue, 2 Sep 2025 15:05:08 +0100
Subject: [PATCH] Docstring update vocabulary_size -> top_k

---
 src/synthid_text/logits_processing.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/synthid_text/logits_processing.py b/src/synthid_text/logits_processing.py
index 3866c9d..6c27327 100644
--- a/src/synthid_text/logits_processing.py
+++ b/src/synthid_text/logits_processing.py
@@ -30,11 +30,11 @@ def update_scores(
 
   We assume that the scores are in the log space.
   Args:
-    scores: Scores (batch_size, vocab_size).
-    g_values: G values (batch_size, vocab_size, depth).
+    scores: Scores (batch_size, top_k).
+    g_values: G values (batch_size, top_k, depth).
 
   Returns:
-    Updated scores (batch_size, vocab_size).
+    Updated scores (batch_size, top_k).
   """
   _, _, depth = g_values.shape
   device = scores.device
@@ -62,12 +62,12 @@ def update_scores_distortionary(
 
   We assume that the scores are in the log space.
   Args:
-    scores: Scores (batch_size, vocab_size).
-    g_values: G values (batch_size, vocab_size, depth).
+    scores: Scores (batch_size, top_k).
+    g_values: G values (batch_size, top_k, depth).
     num_leaves: Number of leaves per node in the tournament tree.
 
   Returns:
-    Updated scores (batch_size, vocab_size).
+    Updated scores (batch_size, top_k).
   """
   _, _, depth = g_values.shape
   device = scores.device