adding new vars

weni-ai · Dec 29, 2023 · 8e4f32a · 8e4f32a
1 parent fa12f2e
commit 8e4f32a
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/bothub_nlp_api/settings.py b/bothub_nlp_api/settings.py
@@ -54,9 +54,11 @@ def cast_supported_languages(i):
     WENIGPT_PROMPT_REINFORCEMENT_INSTRUCTION=(str, "Lembre, se não estiver no texto ou não souber a resposta, responda especificamente 'Desculpe, não possuo essa informação'. Precisamos ajudar o paciente.\n\n"),
     WENIGPT_PROMPT_ANSWER=(str, "RESPOSTA:"),
     WENIGPT_MAX_NEW_TOKENS=(int, 1000),
+    WENIGPT_MAX_LENGHT=(int, 8192),
     WENIGPT_TOP_P=(float, 0.1),
+    WENIGPT_TOP_K=(int, 10),
     WENIGPT_TEMPERATURE=(float, 0.1),
-    WENIGPT_STOP=(str, "\n")
+    WENIGPT_STOP=(str, "\n"),
 )
 
 ENVIRONMENT = env.str("ENVIRONMENT")
@@ -131,6 +133,8 @@ def cast_supported_languages(i):
 
 # wenigpt most dinamic sampling_params
 WENIGPT_MAX_NEW_TOKENS = env.int("WENIGPT_MAX_NEW_TOKENS")
+WENIGPT_MAX_LENGHT = env.int("WENIGPT_MAX_LENGHT")
 WENIGPT_TOP_P = env.float("WENIGPT_TOP_P")
+WENIGPT_TOP_K = env.int("WENIGPT_TOP_K")
 WENIGPT_TEMPERATURE = env.float("WENIGPT_TEMPERATURE")
-WENIGPT_STOP_SEQUENCES = env.str("WENIGPT_STOP")
+WENIGPT_STOP = env.str("WENIGPT_STOP")
diff --git a/bothub_nlp_api/utils.py b/bothub_nlp_api/utils.py
@@ -282,10 +282,12 @@ def request_wenigpt(context, question):
             "prompt": base_prompt,
             "sampling_params": {
                 "max_new_tokens": settings.WENIGPT_MAX_NEW_TOKENS,
+                "max_length": settings.WENIGPT_MAX_LENGHT,
                 "top_p": settings.WENIGPT_TOP_P,
+                "top_k": settings.WENIGPT_TOP_K,
                 "temperature": settings.WENIGPT_TEMPERATURE,
                 "do_sample": False,
-                "stop": settings.WENIGPT_STOP
+                "stop": settings.WENIGPT_STOP,
             }
         }
     }