feat(api/core/model_runtime/entities/defaults.py): Add TOP_K in defau…

…lt parameters.
langgenius · Sep 9, 2024 · 54548f3 · 54548f3
1 parent 0bec6a0
commit 54548f3
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 1 deletion.
diff --git a/api/core/model_runtime/entities/defaults.py b/api/core/model_runtime/entities/defaults.py
@@ -33,6 +33,22 @@
         'max': 1.0,
         'precision': 2,
     },
+    DefaultParameterName.TOP_K: {
+        'label': {
+            'en_US': 'Top K',
+            'zh_Hans': 'Top K',
+        },
+        'type': 'int',
+        'help': {
+            'en_US': 'Limits the number of tokens to consider for each step by keeping only the k most likely tokens.',
+            'zh_Hans': '通过只保留每一步中最可能的 k 个标记来限制要考虑的标记数量。',
+        },
+        'required': False,
+        'default': 50,
+        'min': 1,
+        'max': 100,
+        'precision': 0,
+    },
     DefaultParameterName.PRESENCE_PENALTY: {
         'label': {
             'en_US': 'Presence Penalty',

diff --git a/api/core/model_runtime/entities/model_entities.py b/api/core/model_runtime/entities/model_entities.py
@@ -85,12 +85,13 @@ class ModelFeature(Enum):
     STREAM_TOOL_CALL = "stream-tool-call"
 
 
-class DefaultParameterName(Enum):
+class DefaultParameterName(str, Enum):
     """
     Enum class for parameter template variable.
     """
     TEMPERATURE = "temperature"
     TOP_P = "top_p"
+    TOP_K = "top_k"
     PRESENCE_PENALTY = "presence_penalty"
     FREQUENCY_PENALTY = "frequency_penalty"
     MAX_TOKENS = "max_tokens"