Skip to content

Commit

Permalink
Enhance user input handling for llama-run
Browse files Browse the repository at this point in the history
The main motivation for this change is it was not handing ctrl-d
correctly. Modify `read_user_input` to handle EOF, "/bye" command,
and empty input cases. Introduce `get_user_input` function to
manage user input loop and handle different return cases.

Signed-off-by: Eric Curtin <[email protected]>
  • Loading branch information
ericcurtin committed Jan 8, 2025
1 parent 99a3755 commit 64db527
Showing 1 changed file with 35 additions and 3 deletions.
38 changes: 35 additions & 3 deletions examples/run/run.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -801,7 +801,20 @@ static int generate(LlamaData & llama_data, const std::string & prompt, std::str

static int read_user_input(std::string & user) {
std::getline(std::cin, user);
return user.empty(); // Should have data in happy path
if (std::cin.eof()) {
printf("\n");
return 1;
}

if (user == "/bye") {
return 1;
}

if (user.empty()) {
return 2;
}

return 0; // Should have data in happy path
}

// Function to generate a response based on the prompt
Expand Down Expand Up @@ -868,15 +881,34 @@ static bool is_stdout_a_terminal() {
#endif
}

// Function to tokenize the prompt
// Function to handle user input
static int get_user_input(std::string & user_input, const std::string & user) {
while (true) {
const int ret = handle_user_input(user_input, user);
if (ret == 1) {
return 1;
}

if (ret == 2) {
continue;
}

break;
}

return 0;
}

// Main chat loop function
static int chat_loop(LlamaData & llama_data, const std::string & user) {
int prev_len = 0;
llama_data.fmtted.resize(llama_n_ctx(llama_data.context.get()));
static const bool stdout_a_terminal = is_stdout_a_terminal();
while (true) {
// Get user input
std::string user_input;
while (handle_user_input(user_input, user)) {
if (get_user_input(user_input, user) == 1) {
return 0;
}

add_message("user", user.empty() ? user_input : user, llama_data);
Expand Down

0 comments on commit 64db527

Please sign in to comment.