Skip to content

Commit 6e27406

Browse files
committed
embedding: assign n_ubatch value, print error on n_batch overflow
1 parent ad3a050 commit 6e27406

File tree

1 file changed

+5
-1
lines changed

1 file changed

+5
-1
lines changed

examples/embedding/embedding.cpp

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,8 @@ int main(int argc, char ** argv) {
6161
}
6262

6363
params.embedding = true;
64+
// For BERT models, batch size must be equal to ubatch size
65+
params.n_ubatch = params.n_batch;
6466

6567
print_build_info();
6668

@@ -114,7 +116,9 @@ int main(int argc, char ** argv) {
114116
for (const auto & prompt : prompts) {
115117
auto inp = ::llama_tokenize(ctx, prompt, true, false);
116118
if (inp.size() > n_batch) {
117-
inp.resize(n_batch);
119+
fprintf(stderr, "%s: error: number of tokens in input line (%lld) exceeds batch size (%lld), increase batch size and re-run\n",
120+
__func__, (long long int) inp.size(), (long long int) n_batch);
121+
return 1;
118122
}
119123
inputs.push_back(inp);
120124
}

0 commit comments

Comments
 (0)