Fix top_n_tokens not being passed

argilla-io · Dec 12, 2024 · cb3d55a · cb3d55a
1 parent 6c2df01
commit cb3d55a
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/src/distilabel/models/llms/huggingface/inference_endpoints.py b/src/distilabel/models/llms/huggingface/inference_endpoints.py
@@ -16,7 +16,16 @@
 import random
 import sys
 import warnings
-from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+    Union,
+)
 
 from pydantic import (
     Field,
@@ -679,6 +688,7 @@ async def agenerate(  # type: ignore
             repetition_penalty=repetition_penalty,
             frequency_penalty=frequency_penalty,
             temperature=temperature,
+            top_n_tokens=top_n_tokens,
             top_p=top_p,
             top_k=top_k,
             stop_sequences=stop_sequences,