Skip to content

Commit 8d0033a

Browse files
committed
common, server : add top-a sampler
This improves compatibility with clients like AI Horde.
1 parent 5b09797 commit 8d0033a

File tree

7 files changed

+41
-5
lines changed

7 files changed

+41
-5
lines changed

common/common.cpp

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -376,6 +376,12 @@ bool gpt_params_parse_ex(int argc, char ** argv, gpt_params & params) {
376376
break;
377377
}
378378
sparams.min_p = std::stof(argv[i]);
379+
} else if (arg == "--top-a") {
380+
if (++i >= argc) {
381+
invalid_param = true;
382+
break;
383+
}
384+
sparams.top_a = std::stof(argv[i]);
379385
} else if (arg == "--temp") {
380386
if (++i >= argc) {
381387
invalid_param = true;
@@ -984,6 +990,7 @@ void gpt_print_usage(int /*argc*/, char ** argv, const gpt_params & params) {
984990
printf(" --top-k N top-k sampling (default: %d, 0 = disabled)\n", sparams.top_k);
985991
printf(" --top-p N top-p sampling (default: %.1f, 1.0 = disabled)\n", (double)sparams.top_p);
986992
printf(" --min-p N min-p sampling (default: %.1f, 0.0 = disabled)\n", (double)sparams.min_p);
993+
printf(" --top-a N top-a sampling (default: %.1f, 0.0 = disabled)\n", (double)sparams.top_a);
987994
printf(" --tfs N tail free sampling, parameter z (default: %.1f, 1.0 = disabled)\n", (double)sparams.tfs_z);
988995
printf(" --typical N locally typical sampling, parameter p (default: %.1f, 1.0 = disabled)\n", (double)sparams.typical_p);
989996
printf(" --repeat-last-n N last n tokens to consider for penalize (default: %d, 0 = disabled, -1 = ctx_size)\n", sparams.penalty_last_n);
@@ -1157,6 +1164,7 @@ std::vector<llama_sampler_type> sampler_types_from_names(const std::vector<std::
11571164
{"top_p", llama_sampler_type::TOP_P},
11581165
{"typical_p", llama_sampler_type::TYPICAL_P},
11591166
{"min_p", llama_sampler_type::MIN_P},
1167+
{"top_a", llama_sampler_type::TOP_A},
11601168
{"tfs_z", llama_sampler_type::TFS_Z},
11611169
{"temperature", llama_sampler_type::TEMPERATURE}
11621170
};
@@ -1170,6 +1178,7 @@ std::vector<llama_sampler_type> sampler_types_from_names(const std::vector<std::
11701178
{"typical-p", llama_sampler_type::TYPICAL_P},
11711179
{"typical", llama_sampler_type::TYPICAL_P},
11721180
{"min-p", llama_sampler_type::MIN_P},
1181+
{"top-a", llama_sampler_type::TOP_A},
11731182
{"tfs-z", llama_sampler_type::TFS_Z},
11741183
{"tfs", llama_sampler_type::TFS_Z},
11751184
{"temp", llama_sampler_type::TEMPERATURE}
@@ -1205,6 +1214,7 @@ std::vector<llama_sampler_type> sampler_types_from_chars(const std::string & nam
12051214
{'p', llama_sampler_type::TOP_P},
12061215
{'y', llama_sampler_type::TYPICAL_P},
12071216
{'m', llama_sampler_type::MIN_P},
1217+
{'a', llama_sampler_type::TOP_A},
12081218
{'f', llama_sampler_type::TFS_Z},
12091219
{'t', llama_sampler_type::TEMPERATURE}
12101220
};
@@ -1227,6 +1237,7 @@ std::string sampler_type_to_name_string(llama_sampler_type sampler_type) {
12271237
case llama_sampler_type::TYPICAL_P: return "typical_p";
12281238
case llama_sampler_type::TOP_P: return "top_p";
12291239
case llama_sampler_type::MIN_P: return "min_p";
1240+
case llama_sampler_type::TOP_A: return "top_a";
12301241
case llama_sampler_type::TEMPERATURE: return "temperature";
12311242
default : return "";
12321243
}
@@ -1773,6 +1784,7 @@ void dump_non_result_info_yaml(FILE * stream, const gpt_params & params, const l
17731784
fprintf(stream, "top_k: %d # default: 40\n", sparams.top_k);
17741785
fprintf(stream, "top_p: %f # default: 0.95\n", sparams.top_p);
17751786
fprintf(stream, "min_p: %f # default: 0.0\n", sparams.min_p);
1787+
fprintf(stream, "top_a: %f # default: 0.0\n", sparams.top_a);
17761788
fprintf(stream, "typical_p: %f # default: 1.0\n", sparams.typical_p);
17771789
fprintf(stream, "verbose_prompt: %s # default: false\n", params.verbose_prompt ? "true" : "false");
17781790
fprintf(stream, "display_prompt: %s # default: true\n", params.display_prompt ? "true" : "false");

common/sampling.cpp

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -91,10 +91,10 @@ std::string llama_sampling_print(const llama_sampling_params & params) {
9191

9292
snprintf(result, sizeof(result),
9393
"\trepeat_last_n = %d, repeat_penalty = %.3f, frequency_penalty = %.3f, presence_penalty = %.3f\n"
94-
"\ttop_k = %d, tfs_z = %.3f, top_p = %.3f, min_p = %.3f, typical_p = %.3f, temp = %.3f\n"
94+
"\ttop_k = %d, tfs_z = %.3f, top_p = %.3f, min_p = %.3f, top_a = %.3f, typical_p = %.3f, temp = %.3f\n"
9595
"\tmirostat = %d, mirostat_lr = %.3f, mirostat_ent = %.3f",
9696
params.penalty_last_n, params.penalty_repeat, params.penalty_freq, params.penalty_present,
97-
params.top_k, params.tfs_z, params.top_p, params.min_p, params.typical_p, params.temp,
97+
params.top_k, params.tfs_z, params.top_p, params.min_p, params.top_a, params.typical_p, params.temp,
9898
params.mirostat, params.mirostat_eta, params.mirostat_tau);
9999

100100
return std::string(result);
@@ -128,6 +128,7 @@ static void sampler_queue(
128128
const int32_t top_k = params.top_k;
129129
const float top_p = params.top_p;
130130
const float min_p = params.min_p;
131+
const float top_a = params.top_a;
131132
const float tfs_z = params.tfs_z;
132133
const float typical_p = params.typical_p;
133134
const std::vector<llama_sampler_type> & samplers_sequence = params.samplers_sequence;
@@ -139,6 +140,7 @@ static void sampler_queue(
139140
case llama_sampler_type::TYPICAL_P: llama_sample_typical (ctx_main, &cur_p, typical_p, min_keep); break;
140141
case llama_sampler_type::TOP_P : llama_sample_top_p (ctx_main, &cur_p, top_p, min_keep); break;
141142
case llama_sampler_type::MIN_P : llama_sample_min_p (ctx_main, &cur_p, min_p, min_keep); break;
143+
case llama_sampler_type::TOP_A : llama_sample_top_a (ctx_main, &cur_p, top_a, min_keep); break;
142144
case llama_sampler_type::TEMPERATURE:
143145
if (dynatemp_range > 0) {
144146
float dynatemp_min = std::max(0.0f, temp - dynatemp_range);

common/sampling.h

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ enum class llama_sampler_type : char {
1313
TOP_K = 'k',
1414
TOP_P = 'p',
1515
MIN_P = 'm',
16+
TOP_A = 'a',
1617
TFS_Z = 'f',
1718
TYPICAL_P = 'y',
1819
TEMPERATURE = 't'
@@ -26,6 +27,7 @@ typedef struct llama_sampling_params {
2627
int32_t top_k = 40; // <= 0 to use vocab size
2728
float top_p = 0.95f; // 1.0 = disabled
2829
float min_p = 0.05f; // 0.0 = disabled
30+
float top_a = 0.00f; // 0.0 = disabled
2931
float tfs_z = 1.00f; // 1.0 = disabled
3032
float typical_p = 1.00f; // 1.0 = disabled
3133
float temp = 0.80f; // <= 0.0 to sample greedily, 0.0 to not output probabilities
@@ -46,6 +48,7 @@ typedef struct llama_sampling_params {
4648
llama_sampler_type::TYPICAL_P,
4749
llama_sampler_type::TOP_P,
4850
llama_sampler_type::MIN_P,
51+
llama_sampler_type::TOP_A,
4952
llama_sampler_type::TEMPERATURE
5053
};
5154

examples/server/README.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -213,6 +213,8 @@ node index.js
213213

214214
`min_p`: The minimum probability for a token to be considered, relative to the probability of the most likely token (default: 0.05).
215215

216+
`top_a`: Limit the next token selection to a subset of tokens with a probability above a*P^2, where P is the most probable token (default: 0.0, 0.0 = disabled).
217+
216218
`n_predict`: Set the maximum number of tokens to predict when generating text. **Note:** May exceed the set limit slightly if the last token is a partial multibyte character. When 0, no tokens will be generated but the prompt is evaluated into the cache. (default: -1, -1 = infinity).
217219

218220
`n_keep`: Specify the number of tokens from the prompt to retain when the context size is exceeded and tokens need to be discarded.

examples/server/server.cpp

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -816,6 +816,7 @@ struct server_context {
816816
slot.sparams.top_k = json_value(data, "top_k", default_sparams.top_k);
817817
slot.sparams.top_p = json_value(data, "top_p", default_sparams.top_p);
818818
slot.sparams.min_p = json_value(data, "min_p", default_sparams.min_p);
819+
slot.sparams.top_a = json_value(data, "top_a", default_sparams.top_a);
819820
slot.sparams.tfs_z = json_value(data, "tfs_z", default_sparams.tfs_z);
820821
slot.sparams.typical_p = json_value(data, "typical_p", default_sparams.typical_p);
821822
slot.sparams.temp = json_value(data, "temperature", default_sparams.temp);
@@ -1194,6 +1195,7 @@ struct server_context {
11941195
{"top_k", slot.sparams.top_k},
11951196
{"top_p", slot.sparams.top_p},
11961197
{"min_p", slot.sparams.min_p},
1198+
{"top_a", slot.sparams.top_a},
11971199
{"tfs_z", slot.sparams.tfs_z},
11981200
{"typical_p", slot.sparams.typical_p},
11991201
{"repeat_last_n", slot.sparams.penalty_last_n},

llama.cpp

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10725,7 +10725,7 @@ void llama_sample_top_p(struct llama_context * ctx, llama_token_data_array * can
1072510725
}
1072610726
}
1072710727

10728-
void llama_sample_min_p(struct llama_context * ctx, llama_token_data_array * candidates, float p, size_t min_keep) {
10728+
static void llama_sample_min_p_pow(struct llama_context * ctx, llama_token_data_array * candidates, float p, float pow, size_t min_keep) {
1072910729
if (p <= 0.0f || !candidates->size) {
1073010730
return;
1073110731
}
@@ -10742,7 +10742,7 @@ void llama_sample_min_p(struct llama_context * ctx, llama_token_data_array * can
1074210742
for (size_t i = 0; i < candidates->size; ++i) {
1074310743
max_logit = std::max(max_logit, candidates->data[i].logit);
1074410744
}
10745-
const float min_logit = max_logit + logf(p); // min logit for p_i >= p * p_max
10745+
const float min_logit = max_logit + logf(p) * pow; // min logit for p_i >= p * p_max^pow
1074610746

1074710747
for (size_t i = 0; i < candidates->size; ++i) {
1074810748
if (candidates->data[i].logit >= min_logit) {
@@ -10768,7 +10768,7 @@ void llama_sample_min_p(struct llama_context * ctx, llama_token_data_array * can
1076810768
candidates->sorted = true;
1076910769
}
1077010770

10771-
const float min_logit = candidates->data[0].logit + logf(p); // min logit for p_i >= p * p_max
10771+
const float min_logit = candidates->data[0].logit + logf(p) * pow; // min logit for p_i >= p * p_max^pow
1077210772
size_t i = 1; // first token always matches
1077310773

1077410774
for (; i < candidates->size; ++i) {
@@ -10786,6 +10786,14 @@ void llama_sample_min_p(struct llama_context * ctx, llama_token_data_array * can
1078610786
}
1078710787
}
1078810788

10789+
void llama_sample_min_p(struct llama_context * ctx, llama_token_data_array * candidates, float p, size_t min_keep) {
10790+
llama_sample_min_p_pow(ctx, candidates, p, 1.f, min_keep);
10791+
}
10792+
10793+
void llama_sample_top_a(struct llama_context * ctx, llama_token_data_array * candidates, float a, size_t min_keep) {
10794+
llama_sample_min_p_pow(ctx, candidates, a, 2.f, min_keep);
10795+
}
10796+
1078910797
void llama_sample_tail_free(struct llama_context * ctx, llama_token_data_array * candidates, float z, size_t min_keep) {
1079010798
if (z >= 1.0f || candidates->size <= 2) {
1079110799
return;

llama.h

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -813,6 +813,13 @@ extern "C" {
813813
float p,
814814
size_t min_keep);
815815

816+
/// @details Top-A sampling as described in https://github.com/BlinkDL/RWKV-LM/tree/4cb363e5aa31978d801a47bc89d28e927ab6912e#the-top-a-sampling-method
817+
LLAMA_API void llama_sample_top_a(
818+
struct llama_context * ctx,
819+
llama_token_data_array * candidates,
820+
float a,
821+
size_t min_keep);
822+
816823
/// @details Tail Free Sampling described in https://www.trentonbricken.com/Tail-Free-Sampling/.
817824
LLAMA_API void llama_sample_tail_free(
818825
struct llama_context * ctx,

0 commit comments

Comments
 (0)