Update tools/llama/generate.py

tools/llama/generate.py

@@ -154,16 +154,11 @@ def decode_one_token_ar_agent(
     logits = x.logits  # [:, -1:]
     hidden_states = x.hidden_states  # [:, -1:]
 
-    sampling_kwargs_main = sampling_kwargs.copy()
-    sampling_kwargs_main["temperature"] = 0.1
-    sampling_kwargs_main["top_p"] = 0.1
-    sampling_kwargs_main["repetition_penalty"] = 1.0
-
     codebooks = [
         sample_agent(
             logits,
             previous_tokens=None,  # Disable repetition penalty for the token codebook
-            **sampling_kwargs_main,
+            **sampling_kwargs,
         )[0]
     ]
 
@@ -194,15 +189,6 @@ def decode_one_token_ar_agent(
         codebooks[:, 1:, :], codebooks[:, :1, :] != semantic_id, CODEBOOK_PAD_TOKEN_ID
     )
 
-    # for i in range(codebooks.size(1) - 1):
-    #     codebooks[:, i + 1, :] = torch.masked_fill(
-    #         codebooks[:, i + 1, :],
-    #         codebooks[:, :1, :] != semantic_id,
-    #         CODEBOOK_PAD_TOKEN_ID + i * 1024,
-    #     )
-
-    # print(codebooks)
-
     return codebooks