Add top_p and top_k params to chat

Author: Belluxx

utils/chat.py

@@ -26,6 +26,8 @@ def __init__(
             eos: str,
             n_generate: int,
             temperature: float = 0.8,
+            top_p: float = 0.9,
+            top_k: int = 40,
             agent_prefixes: dict[str, str] = {  # Default template (most common one)
                 SYSTEM_KEY: '<|im_start|>system',
                 ASSISTANT_KEY: '<|im_start|>assistant',
@@ -46,6 +48,8 @@ def __init__(
         @param eos: the token that ends a single chat round
         @param n_generate: the maximum number of tokens generated by the model in a single turn
         @param temperature: the temperature used for model inference
+        @param top_p: the top_p used for model inference
+        @param top_k: the top_k used for model inference
         @param agent_prefixes: the tokens used to wrap an agent name
         @param agent_names: the dict with the names for: system, assistant, user
         @param debug: whether or not to output debug informations
@@ -55,6 +59,8 @@ def __init__(
         self.eos = eos
         self.n_generate = n_generate
         self.temperature = temperature
+        self.top_p = top_p
+        self.top_k = top_k
         self.agent_prefixes = agent_prefixes
         self.agent_names = agent_names
         self.debug = debug
@@ -78,7 +84,7 @@ def generate_assistant_reply(self, grammar: LlamaGrammar | None = None) -> tuple
 
         reply = ''
         n_reply_tokens = 0
-        for token in self.model.generate(tokens=self.tokens_cache, temp=self.temperature, grammar=grammar):
+        for token in self.model.generate(tokens=self.tokens_cache, temp=self.temperature, top_p=self.top_p, top_k=self.top_k, grammar=grammar):
             self.check_context_overflow()  # Check for context exceeded
             if token == self.model.token_eos() or token == self.eos_token:  # Check for EOS termination
                 self.tokens_cache.append(self.eos_token)
@@ -114,7 +120,7 @@ def generate_assistant_reply_stepped(self, grammar: LlamaGrammar | None = None):
 
         reply = ''
         n_reply_tokens = 0
-        for token in self.model.generate(tokens=self.tokens_cache, temp=self.temperature, grammar=grammar):
+        for token in self.model.generate(tokens=self.tokens_cache, temp=self.temperature, top_p=self.top_p, top_k=self.top_k, grammar=grammar):
             self.check_context_overflow()
             if token == self.model.token_eos() or token == self.eos_token:  # Check for EOS termination
                 self.tokens_cache.append(self.eos_token)