diff --git a/kcpp_src/gpttype_adapter.cpp b/kcpp_src/gpttype_adapter.cpp index 143f5f01b..1465cf1de 100644 --- a/kcpp_src/gpttype_adapter.cpp +++ b/kcpp_src/gpttype_adapter.cpp @@ -4367,6 +4367,7 @@ struct BatchGenerateRequest float dry_base = 1.75f; int dry_allowed_length = 2; int dry_penalty_last_n = 0; + bool tool_call_fix = false; bool allow_eos_token = true; bool bypass_eos_token = false; bool render_special = false; @@ -4548,7 +4549,7 @@ static bool batch_inputs_eligible(const generation_inputs & inputs) } } } - if(inputs.reasoning_budget >= 0 || inputs.tool_call_fix) + if(inputs.reasoning_budget >= 0) { return false; } @@ -4679,6 +4680,106 @@ static llama_sampler * batch_rep_pen_init(int32_t penalty_last_n, float penalty_ }); } +struct BatchToolCallSampler +{ + std::vector prevented_tokens; + bool render_special = false; + int generated_tokens = 0; + int opening_bracket_count = 0; +}; + +static const char * batch_tool_call_name(const llama_sampler * /*smpl*/) +{ + return "kcpp-batch-tool-call-fix"; +} + +static void batch_tool_call_accept(llama_sampler * smpl, llama_token token) +{ + auto * ctx = (BatchToolCallSampler *) smpl->ctx; + ctx->generated_tokens++; + const std::string piece = FileFormatTokenizeID(token, file_format, ctx->render_special); + ctx->opening_bracket_count += std::count(piece.begin(), piece.end(), '['); +} + +static void batch_tool_call_apply(llama_sampler * smpl, llama_token_data_array * cur_p) +{ + auto * ctx = (BatchToolCallSampler *) smpl->ctx; + if(ctx->generated_tokens >= 3 || ctx->opening_bracket_count > 1 || cur_p->size == 0) + { + return; + } + + float lowest_logit = cur_p->data[0].logit; + for(size_t i = 1; i < cur_p->size; ++i) + { + lowest_logit = std::min(lowest_logit, cur_p->data[i].logit); + } + for(size_t i = 0; i < cur_p->size; ++i) + { + llama_token token = cur_p->data[i].id; + if(token >= 0 && (size_t) token < ctx->prevented_tokens.size() && ctx->prevented_tokens[token]) + { + cur_p->data[i].logit = lowest_logit; + } + } + cur_p->sorted = false; +} + +static void batch_tool_call_reset(llama_sampler * smpl) +{ + auto * ctx = (BatchToolCallSampler *) smpl->ctx; + ctx->generated_tokens = 0; + ctx->opening_bracket_count = 0; +} + +static llama_sampler * batch_tool_call_clone(const llama_sampler * smpl) +{ + const auto * ctx = (const BatchToolCallSampler *) smpl->ctx; + return llama_sampler_init(smpl->iface, new BatchToolCallSampler { + ctx->prevented_tokens, + ctx->render_special, + ctx->generated_tokens, + ctx->opening_bracket_count, + }); +} + +static void batch_tool_call_free(llama_sampler * smpl) +{ + delete (BatchToolCallSampler *) smpl->ctx; +} + +static llama_sampler_i batch_tool_call_i = { + /* .name = */ batch_tool_call_name, + /* .accept = */ batch_tool_call_accept, + /* .apply = */ batch_tool_call_apply, + /* .reset = */ batch_tool_call_reset, + /* .clone = */ batch_tool_call_clone, + /* .free = */ batch_tool_call_free, + /* .backend_init = */ nullptr, + /* .backend_accept = */ nullptr, + /* .backend_apply = */ nullptr, + /* .backend_set_input = */ nullptr, +}; + +static llama_sampler * batch_tool_call_init(int32_t n_vocab, bool render_special) +{ + std::vector prevented_tokens(n_vocab, 0); + for(llama_token token = 0; token < n_vocab; ++token) + { + const std::string piece = FileFormatTokenizeID(token, file_format, true); + if(piece.find(']') != std::string::npos) + { + prevented_tokens[token] = 1; + } + } + return llama_sampler_init(&batch_tool_call_i, new BatchToolCallSampler { + std::move(prevented_tokens), + render_special, + 0, + 0, + }); +} + static llama_sampler * batch_build_sampler(const BatchGenerateRequest & req) { llama_sampler_chain_params params = llama_sampler_chain_default_params(); @@ -4687,6 +4788,12 @@ static llama_sampler * batch_build_sampler(const BatchGenerateRequest & req) const llama_vocab * sampler_vocab = llama_model_get_vocab(llama_get_model(llama_ctx_v4)); const int32_t sampler_n_vocab = llama_vocab_n_tokens(sampler_vocab); + if(req.tool_call_fix) + { + // This sampler must not be primed with prompt tokens: its three-token window starts at generation. + llama_sampler_chain_add(chain, batch_tool_call_init(sampler_n_vocab, req.render_special)); + } + if(req.logit_biases.size()>0) { llama_sampler_chain_add(chain, llama_sampler_init_logit_bias(sampler_n_vocab, req.logit_biases.size(), req.logit_biases.data())); @@ -5122,6 +5229,7 @@ int gpttype_batch_generate_submit(const generation_inputs inputs) req->dry_base = inputs.dry_base; req->dry_allowed_length = inputs.dry_allowed_length; req->dry_penalty_last_n = inputs.dry_penalty_last_n; + req->tool_call_fix = inputs.tool_call_fix; req->allow_eos_token = inputs.allow_eos_token; req->bypass_eos_token = inputs.bypass_eos_token; req->render_special = inputs.render_special; diff --git a/koboldcpp.py b/koboldcpp.py index abd04317c..42b657f2c 100644 --- a/koboldcpp.py +++ b/koboldcpp.py @@ -2305,7 +2305,7 @@ def generate(genparams, stream_flag=False): banned_strings = coerce_ban_list(genparams.get('banned_strings', [])) # SillyTavern uses that name banned_tokens = coerce_ban_list(genparams.get('banned_tokens', banned_strings)) bypass_eos_token = genparams.get('bypass_eos', False) - tool_call_fix = genparams.get('using_openai_tools', False) + tool_call_fix = genparams.get('_generic_tool_json_array_fix', False) custom_token_bans = genparams.get('custom_token_bans', '') for tok in custom_token_bans.split(','): @@ -2543,7 +2543,7 @@ def continuous_batching_python_eligible(genparams, api_format): if dry_multiplier > 0 and 0 < dry_base < 1: utfprint("Batching disabled due to samplers set 1",2) return False - if tryparsefloat(genparams.get("smoothing_factor", 0), 0) or genparams.get("using_openai_tools", False): + if tryparsefloat(genparams.get("smoothing_factor", 0), 0): utfprint("Batching disabled due to samplers set 2",2) return False if tryparsefloat(genparams.get("top_a", 0), 0) or tryparsefloat(genparams.get("tfs", 1), 1) != 1 or tryparsefloat(genparams.get("dynatemp_range", 0), 0) < 0: @@ -4892,6 +4892,7 @@ ws ::= | " " | "\n" [ \t]{0,20} jinja_output = None request_tools = genparams.get('tools', []) genparams['_jinja_tool_formatting_used'] = False + genparams['_generic_tool_json_array_fix'] = False if use_jinja and cached_chat_template: copied_jinja_kwargs = dict(cached_jinja_kwargs or {}) # Merge user-provided chat_template_kwargs into our defaults @@ -5014,6 +5015,7 @@ ws ::= | " " | "\n" [ \t]{0,20} # Set temperature lower automatically if function calling, cannot exceed 0.5 genparams["temperature"] = (1.0 if genparams.get("temperature", 0.5) > 1.0 else genparams.get("temperature", 0.5)) genparams["using_openai_tools"] = True + genparams['_generic_tool_json_array_fix'] = True # Set grammar to llamacpp example grammar to force json response (see https://github.com/ggerganov/llama.cpp/blob/master/grammars/json_arr.gbnf) genparams["grammar"] = jsongrammar try: