main: must check pipe status on very top of program

This commit is contained in:
brian khuu 2024-05-21 14:58:34 +10:00
parent 50048f5b45
commit c1e8a6d1c0

View file

@ -120,6 +120,12 @@ static void llama_log_callback_logTee(ggml_log_level level, const char * text, v
} }
int main(int argc, char ** argv) { int main(int argc, char ** argv) {
#ifndef _MSC_VER
// Check if we have an external attachment to a file descriptor for out of band control tokens (e.g. bash `3>/dev/null` )
// Placed here to avoid file descriptor being polluted by gpt_params_parse() opening files
const bool control_token_file_descriptor_is_attached = fcntl(CONTROL_TOKEN_FILENO, F_GETFL) != -1;
#endif
gpt_params params; gpt_params params;
g_params = &params; g_params = &params;
@ -128,6 +134,16 @@ int main(int argc, char ** argv) {
} }
llama_sampling_params & sparams = params.sparams; llama_sampling_params & sparams = params.sparams;
const bool control_token_allowed_on_standard_stream = !params.conversation && sparams.grammar.empty();
#ifndef _MSC_VER
// Merge normal token stream and control token streams together only if not in conversation or grammar mode
if (control_token_allowed_on_standard_stream && !control_token_file_descriptor_is_attached) {
// Duplicate stdout file descriptor to control token file descriptor to merge the two streams
dup2(STDOUT_FILENO, CONTROL_TOKEN_FILENO);
}
#endif
#ifndef LOG_DISABLE_LOGS #ifndef LOG_DISABLE_LOGS
log_set_target(log_filename_generator("main", "log")); log_set_target(log_filename_generator("main", "log"));
LOG_TEE("Log start\n"); LOG_TEE("Log start\n");
@ -530,17 +546,6 @@ int main(int argc, char ** argv) {
exit(1); exit(1);
} }
const bool control_token_allowed_on_standard_stream = !params.conversation && sparams.grammar.empty();
#ifndef _MSC_VER
const bool control_token_descriptor_is_attached = fcntl(CONTROL_TOKEN_FILENO, F_GETFL) != -1;
if (control_token_allowed_on_standard_stream && !control_token_descriptor_is_attached) {
// Control Token File Descriptor has nothing attached to it so make control token file descriptor be an alias of stdout
// This is not done however if we are in conversation mode or grammar mode as that is typically discarded
dup2(STDOUT_FILENO, CONTROL_TOKEN_FILENO);
}
#endif
while ((n_remain != 0 && !is_antiprompt) || params.interactive) { while ((n_remain != 0 && !is_antiprompt) || params.interactive) {
// predict // predict
if (!embd.empty()) { if (!embd.empty()) {
@ -758,20 +763,18 @@ int main(int argc, char ** argv) {
// Console/Stream Output // Console/Stream Output
if (!llama_token_is_control_token(llama_get_model(ctx), id)) { if (!llama_token_is_control_token(llama_get_model(ctx), id)) {
// Stream Output Token To Standard Output // Stream Output Token To Standard Output
fflush(stdout);
fprintf(stdout, "%s", token_str.c_str()); fprintf(stdout, "%s", token_str.c_str());
} else if (!params.ctrl_token_no_out) { } else if (!params.ctrl_token_no_out) {
#ifndef _MSC_VER #ifndef _MSC_VER
if (control_token_descriptor_is_attached) { if (control_token_file_descriptor_is_attached) {
// Stream Control Token To Special Token Output. Useful for debugging control token behaviour // Stream Control Token To Special Token Output. Useful for debugging control token behaviour
ssize_t result = write(CONTROL_TOKEN_FILENO, token_str.c_str(), token_str.length()); fflush(stdout); // Ensure control token is always appended to stdout stream
(void) result; (void)! write(CONTROL_TOKEN_FILENO, token_str.c_str(), token_str.length());
} else } else
#endif #endif
if (control_token_allowed_on_standard_stream) if (control_token_allowed_on_standard_stream)
{ {
// Stream Control Token To Standard Output Stream // Stream Control Token To Standard Output Stream
fflush(stdout);
fprintf(stdout, "%s", token_str.c_str()); fprintf(stdout, "%s", token_str.c_str());
} }
} }