ggerganov commited on
Commit
d06a6c2
·
1 Parent(s): 83926f7

main : print colors + no timestamps

Browse files
Files changed (1) hide show
  1. main.cpp +16 -2
main.cpp CHANGED
@@ -154,12 +154,26 @@ void whisper_print_segment_callback(struct whisper_context * ctx, void * user_da
154
 
155
  if (params.no_timestamps) {
156
  if (params.print_colors) {
157
- // TODO
 
 
 
 
 
 
 
 
 
 
 
 
 
 
158
  } else {
159
  const char * text = whisper_full_get_segment_text(ctx, i);
160
  printf("%s", text);
161
- fflush(stdout);
162
  }
 
163
  } else {
164
  const int64_t t0 = whisper_full_get_segment_t0(ctx, i);
165
  const int64_t t1 = whisper_full_get_segment_t1(ctx, i);
 
154
 
155
  if (params.no_timestamps) {
156
  if (params.print_colors) {
157
+ for (int j = 0; j < whisper_full_n_tokens(ctx, i); ++j) {
158
+ if (params.print_special_tokens == false) {
159
+ const whisper_token id = whisper_full_get_token_id(ctx, i, j);
160
+ if (id >= whisper_token_eot(ctx)) {
161
+ continue;
162
+ }
163
+ }
164
+
165
+ const char * text = whisper_full_get_token_text(ctx, i, j);
166
+ const float p = whisper_full_get_token_p (ctx, i, j);
167
+
168
+ const int col = std::max(0, std::min((int) k_colors.size(), (int) (std::pow(p, 3)*float(k_colors.size()))));
169
+
170
+ printf("%s%s%s", k_colors[col].c_str(), text, "\033[0m");
171
+ }
172
  } else {
173
  const char * text = whisper_full_get_segment_text(ctx, i);
174
  printf("%s", text);
 
175
  }
176
+ fflush(stdout);
177
  } else {
178
  const int64_t t0 = whisper_full_get_segment_t0(ctx, i);
179
  const int64_t t1 = whisper_full_get_segment_t1(ctx, i);