| 209 | } |
| 210 | |
| 211 | int64_t get_example_targets_batch( |
| 212 | struct llama_context * lctx, |
| 213 | struct ggml_tensor * tokens_input, |
| 214 | struct ggml_tensor * target_probs, |
| 215 | int64_t example_id, |
| 216 | const size_t * samples_offs, |
| 217 | const size_t * samples_begin, |
| 218 | const size_t * samples_size, |
| 219 | size_t samples_count, |
| 220 | const llama_token * train_data, |
| 221 | size_t n_train_data, |
| 222 | bool separate_with_eos, |
| 223 | bool separate_with_bos, |
| 224 | bool fill_with_next_samples, |
| 225 | bool sample_random_offsets |
| 226 | ) { |
| 227 | GGML_ASSERT(samples_count > 0); |
| 228 | GGML_ASSERT(tokens_input->n_dims == 2); |
| 229 | GGML_ASSERT(target_probs->n_dims == 3); |
| 230 | int64_t n_vocab = target_probs->ne[0]; |
| 231 | int64_t n_tokens = tokens_input->ne[0]; |
| 232 | int64_t n_batch = tokens_input->ne[1]; |
| 233 | GGML_ASSERT(n_vocab == target_probs->ne[0]); |
| 234 | GGML_ASSERT(n_tokens == target_probs->ne[1]); |
| 235 | GGML_ASSERT(n_batch == target_probs->ne[2]); |
| 236 | |
| 237 | int64_t used_samples = 0; |
| 238 | |
| 239 | ggml_set_f32(target_probs, 0.0f); |
| 240 | llama_token bos = llama_token_bos(llama_get_model(lctx)); |
| 241 | llama_token eos = llama_token_eos(llama_get_model(lctx)); |
| 242 | // printf("%s: example_id=%d n_batch=%d n_train_samples=%zu\n", __func__, example_id, n_batch, n_train_samples); |
| 243 | for (int k=0; k<n_batch; ++k) { |
| 244 | // printf("%s: batch %d\n", __func__, k); |
| 245 | size_t sample_idx = (example_id + used_samples) % samples_count; |
| 246 | size_t sample_offs = sample_random_offsets ? samples_offs[sample_idx] : 0; |
| 247 | size_t sample_begin = samples_begin[sample_idx]; |
| 248 | size_t sample_size = samples_size[sample_idx]; |
| 249 | ++used_samples; |
| 250 | |
| 251 | // printf("%s: sample_idx=%zu sample=%zu\n", __func__, sample_idx, sample); |
| 252 | GGML_ASSERT(sample_begin+sample_size-1 < n_train_data); |
| 253 | |
| 254 | ggml_set_i32_nd(tokens_input, 0, k, 0, 0, bos); |
| 255 | bool sample_separation_eos = !separate_with_eos; |
| 256 | bool sample_separation_bos = !separate_with_bos; |
| 257 | for (int64_t i=0; i<n_tokens; ++i) { |
| 258 | llama_token token = eos; |
| 259 | if (sample_offs >= sample_size && fill_with_next_samples) { |
| 260 | if (!sample_separation_eos) { |
| 261 | // insert eos token to separate samples |
| 262 | sample_separation_eos = true; |
| 263 | } else if (!sample_separation_bos) { |
| 264 | // insert bos token to separate samples |
| 265 | sample_separation_bos = true; |
| 266 | token = bos; |
| 267 | } else { |
| 268 | // sample separation is done, continue with next sample |
no test coverage detected