Better handling of feed prompt callbacks

Dognam · May 11, 2023 · 82f1c68 · 82f1c68
1 parent 5502f08
commit 82f1c68
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 19 deletions.
diff --git a/crates/llm-base/src/inference_session.rs b/crates/llm-base/src/inference_session.rs
@@ -187,7 +187,7 @@ impl InferenceSession {
             parameters,
             request.prompt,
             output_request,
-            TokenUtf8Buffer::adapt_callback(&mut callback),
+            feed_prompt_callback(&mut callback),
         )?;
         stats.feed_prompt_duration = start_at.elapsed().unwrap();
         stats.prompt_tokens = self.n_past;
@@ -634,6 +634,18 @@ pub enum InferenceFeedback {
     Halt,
 }
 
+/// Adapt an [InferenceResponse] callback so that it can be used in a call to
+/// [InferenceSession::feed_prompt].
+pub fn feed_prompt_callback<'a, E: std::error::Error + 'static>(
+    mut callback: impl FnMut(InferenceResponse) -> Result<InferenceFeedback, E> + 'a,
+) -> impl FnMut(&[u8]) -> Result<InferenceFeedback, E> + 'a {
+    let mut buffer = TokenUtf8Buffer::new();
+    move |token| match buffer.push(token) {
+        Some(tokens) => callback(InferenceResponse::PromptToken(tokens)),
+        None => Ok(InferenceFeedback::Continue),
+    }
+}
+
 fn scratch_buffers() -> [ggml::Buffer; 2] {
     [
         ggml::Buffer::new(SCRATCH_SIZE),

diff --git a/crates/llm-base/src/lib.rs b/crates/llm-base/src/lib.rs
@@ -21,7 +21,7 @@ pub use ggml;
 pub use ggml::Type as ElementType;
 
 pub use inference_session::{
-    InferenceFeedback, InferenceRequest, InferenceResponse, InferenceSession,
+    feed_prompt_callback, InferenceFeedback, InferenceRequest, InferenceResponse, InferenceSession,
     InferenceSessionConfig, InferenceSnapshot, InferenceStats, ModelKVMemoryType, SnapshotError,
 };
 pub use loader::{

diff --git a/crates/llm-base/src/util.rs b/crates/llm-base/src/util.rs
@@ -20,8 +20,6 @@ macro_rules! mulf {
 use memmap2::{Mmap, MmapAsRawDesc, MmapOptions};
 use thiserror::Error;
 
-use crate::inference_session::{InferenceFeedback, InferenceResponse};
-
 /// Used to buffer incoming tokens until they produce a valid string of UTF-8 text.
 ///
 /// Tokens are *not* valid UTF-8 by themselves. However, the LLM will produce valid UTF-8
@@ -61,18 +59,6 @@ impl TokenUtf8Buffer {
             }
         }
     }
-
-    /// Adapt an [InferenceResponse] callback so that it can be used in a `&[u8]`
-    /// context.
-    pub fn adapt_callback<'a, E: std::error::Error + 'static>(
-        mut callback: impl FnMut(InferenceResponse) -> Result<InferenceFeedback, E> + 'a,
-    ) -> impl FnMut(&[u8]) -> Result<crate::inference_session::InferenceFeedback, E> + 'a {
-        let mut buffer = Self::new();
-        move |token| match buffer.push(token) {
-            Some(tokens) => callback(InferenceResponse::PromptToken(tokens)),
-            None => Ok(InferenceFeedback::Continue),
-        }
-    }
 }
 
 #[derive(Error, Debug)]

diff --git a/crates/llm/examples/vicuna-chat.rs b/crates/llm/examples/vicuna-chat.rs
@@ -1,6 +1,6 @@
 use llm_base::{
-    InferenceFeedback, InferenceRequest, InferenceResponse, InferenceStats, LoadProgress,
-    TokenUtf8Buffer,
+    feed_prompt_callback, InferenceFeedback, InferenceRequest, InferenceResponse, InferenceStats,
+    LoadProgress,
 };
 use rustyline::error::ReadlineError;
 use spinoff::{spinners::Dots2, Spinner};
@@ -50,7 +50,7 @@ fn main() {
             &Default::default(),
             format!("{persona}\n{history}").as_str(),
             &mut Default::default(),
-            TokenUtf8Buffer::adapt_callback(prompt_callback),
+            feed_prompt_callback(prompt_callback),
         )
         .expect("Failed to ingest initial prompt.");