Halt inference

coasys · Nov 3, 2023 · 9f817b9 · 9f817b9
1 parent 2b5edc9
commit 9f817b9
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/cli/src/eve.rs b/cli/src/eve.rs
@@ -253,15 +253,23 @@ pub async fn run(command: EveCommands) -> Result<()> {
                 &mut Default::default(),
                 // output callback
                 |t| {
+                    let mut cont = true;
                     match t {
                         InferenceResponse::PromptToken(t) | InferenceResponse::InferredToken(t) | llm::InferenceResponse::SnapshotToken(t) => {
+                            if t == "Eve:" || t == "User:" {
+                                cont = false;
+                            }
                             print_token(t);
                         }
                         _ => {}
                     }
                     std::io::stdout().flush().unwrap();
 
-                    Ok(llm::InferenceFeedback::Continue)
+                    if cont {
+                        Ok(llm::InferenceFeedback::Continue)
+                    } else {
+                        Ok(llm::InferenceFeedback::Halt)
+                    }
                 }
             );