Skip to content

Commit 3626fe1

Browse files
Peiyang-SongPeiyang-Song
committed
Chores: format
1 parent b36077f commit 3626fe1

1 file changed

Lines changed: 29 additions & 23 deletions

File tree

‎cpp/ct2.cpp‎

Lines changed: 29 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -207,7 +207,8 @@ extern "C" lean_obj_res ct2_generate(
207207
return reinterpret_cast<lean_obj_res>(output);
208208
}
209209

210-
extern "C" uint8_t init_ct2_encoder(b_lean_obj_arg model_path, b_lean_obj_arg _device) {
210+
extern "C" uint8_t init_ct2_encoder(b_lean_obj_arg model_path,
211+
b_lean_obj_arg _device) {
211212
const char *dir = lean_string_cstr(model_path);
212213
if (!exists(dir)) {
213214
return false;
@@ -216,7 +217,7 @@ extern "C" uint8_t init_ct2_encoder(b_lean_obj_arg model_path, b_lean_obj_arg _d
216217
delete p_encoder;
217218
}
218219

219-
ctranslate2::Device device =
220+
ctranslate2::Device device =
220221
ctranslate2::str_to_device(lean_string_cstr(_device));
221222

222223
p_encoder = new ctranslate2::Encoder(dir, device);
@@ -257,7 +258,8 @@ extern "C" lean_obj_res ct2_encode(b_lean_obj_arg _input_tokens) {
257258
return arr;
258259
}
259260

260-
extern "C" uint8_t init_premise_embeddings(b_lean_obj_arg embeddings_path, b_lean_obj_arg _device) {
261+
extern "C" uint8_t init_premise_embeddings(b_lean_obj_arg embeddings_path,
262+
b_lean_obj_arg _device) {
261263
const char *emb_path = lean_string_cstr(embeddings_path);
262264
if (!exists(emb_path)) {
263265
return false;
@@ -266,9 +268,10 @@ extern "C" uint8_t init_premise_embeddings(b_lean_obj_arg embeddings_path, b_lea
266268
delete premise_embeddings;
267269
}
268270

269-
ctranslate2::Device device =
271+
ctranslate2::Device device =
270272
ctranslate2::str_to_device(lean_string_cstr(_device));
271-
device = ctranslate2::Device::CPU; // [TODO]: We should remove this line when everything can work well on CUDA.
273+
device = ctranslate2::Device::CPU; // [TODO]: We should remove this line when
274+
// everything can work well on CUDA.
272275

273276
auto d = npy::read_npy<double>(emb_path);
274277
std::vector<double> data = d.data;
@@ -285,8 +288,7 @@ extern "C" uint8_t init_premise_embeddings(b_lean_obj_arg embeddings_path, b_lea
285288
std::transform(shape.begin(), shape.end(), shape_i64.begin(),
286289
[](unsigned long ul) { return static_cast<int64_t>(ul); });
287290

288-
premise_embeddings =
289-
new ctranslate2::StorageView(shape_i64, data_f, device);
291+
premise_embeddings = new ctranslate2::StorageView(shape_i64, data_f, device);
290292
return true;
291293
}
292294

@@ -321,7 +323,7 @@ extern "C" uint8_t is_premise_dictionary_initialized(lean_object *) {
321323
return is_premise_dictionary_initialized_aux();
322324
}
323325

324-
extern "C" lean_obj_res ct2_retrieve(b_lean_obj_arg _encoded_state) {
326+
extern "C" lean_obj_res ct2_retrieve(b_lean_obj_arg _encoded_state) {
325327
const lean_array_object *p_arr = lean_to_array(_encoded_state);
326328

327329
assert(static_cast<int64_t>(p_arr->m_size) == premise_embeddings->dim(1));
@@ -336,8 +338,9 @@ extern "C" lean_obj_res ct2_retrieve(b_lean_obj_arg _encoded_state) {
336338
std::vector<int64_t> state_embedding_shape{
337339
static_cast<int64_t>(p_arr->m_size), 1};
338340

339-
ctranslate2::StorageView *state_embedding = new ctranslate2::StorageView(
340-
state_embedding_shape, state_embedding_data, premise_embeddings->device());
341+
ctranslate2::StorageView *state_embedding =
342+
new ctranslate2::StorageView(state_embedding_shape, state_embedding_data,
343+
premise_embeddings->device());
341344

342345
int k = 10;
343346
ctranslate2::ops::MatMul matmul(false, false, 1.0);
@@ -346,31 +349,34 @@ extern "C" lean_obj_res ct2_retrieve(b_lean_obj_arg _encoded_state) {
346349
std::vector<int64_t> probs_shape{premise_embeddings->dim(0), 1};
347350

348351
ctranslate2::StorageView *probs =
349-
new ctranslate2::StorageView(probs_shape, ctranslate2::DataType::FLOAT32, premise_embeddings->device());\
350-
352+
new ctranslate2::StorageView(probs_shape, ctranslate2::DataType::FLOAT32,
353+
premise_embeddings->device());
354+
351355
matmul(*premise_embeddings, *state_embedding, *probs);
352356
probs->resize({premise_embeddings->dim(0)});
353-
354-
ctranslate2::StorageView *topk_values =
355-
new ctranslate2::StorageView({k}, ctranslate2::DataType::FLOAT32, premise_embeddings->device());
356-
ctranslate2::StorageView *topk_indices =
357-
new ctranslate2::StorageView({k}, ctranslate2::DataType::INT32, premise_embeddings->device());
357+
358+
ctranslate2::StorageView *topk_values = new ctranslate2::StorageView(
359+
{k}, ctranslate2::DataType::FLOAT32, premise_embeddings->device());
360+
ctranslate2::StorageView *topk_indices = new ctranslate2::StorageView(
361+
{k}, ctranslate2::DataType::INT32, premise_embeddings->device());
358362
topk(*probs, *topk_values, *topk_indices);
359-
363+
360364
lean_array_object *output =
361365
reinterpret_cast<lean_array_object *>(lean_alloc_array(k, k));
362366
int *topk_indices_ptr = topk_indices->data<int>();
363367
float *topk_values_ptr = topk_values->data<float>();
364-
368+
365369
for (int i = 0; i < k; i++) {
366-
assert (topk_indices_ptr[i] >= 0 && topk_indices_ptr[i] < premise_embeddings->dim(0));
367-
std::string this_premise =
370+
assert(topk_indices_ptr[i] >= 0 &&
371+
topk_indices_ptr[i] < premise_embeddings->dim(0));
372+
// [NOTE]: This is where the server crash occurs on CUDA.
373+
std::string this_premise =
368374
(*premise_dictionary)[std::to_string(*(topk_indices_ptr + i))];
369375

370376
output->m_data[i] =
371377
lean_mk_pair(lean_mk_string(this_premise.c_str()),
372-
lean_box_float(static_cast<double>(topk_values_ptr[i])));
378+
lean_box_float(static_cast<double>(topk_values_ptr[i])));
373379
}
374-
380+
375381
return reinterpret_cast<lean_obj_res>(output);
376382
}

0 commit comments

Comments
 (0)