Simplify inference API.

Cargo.toml

@@ -9,20 +9,20 @@ keywords = ["deep-learning", "language", "model", "rwkv"]
 license = "MIT OR Apache-2.0"
 name = "web-rwkv"
 repository = "https://github.com/cryscan/web-rwkv"
-version = "0.8.9"
+version = "0.8.10"
 
 [dependencies]
 ahash = "0.8"
 anyhow = "1.0"
 bytemuck = { version = "1.13", features = ["extern_crate_alloc"] }
-derive-getters = "0.3"
+derive-getters = "0.4"
 document-features = "0.2.8"
 flume = { version = "0.11.0" }
 futures = "0.3"
 gpp = "0.6.2"
 half = { version = "2.2", features = ["bytemuck", "serde"] }
 instant = { version = "0.1", features = ["inaccurate", "wasm-bindgen"] }
-itertools = "0.12"
+itertools = "0.13"
 log = "0.4"
 regex = "1.10"
 rustc-hash = "1.1.0"

examples/rt-batch.rs

@@ -31,7 +31,7 @@ use web_rwkv::{
         loader::{Loader, Lora},
         model::{Build, ContextAutoLimits, ModelBuilder, ModelInfo, ModelVersion, Quant},
         softmax::softmax,
-        v4, v5, v6, JobRuntime, Submission,
+        v4, v5, v6, JobRuntime,
     },
     tokenizer::Tokenizer,
 };
@@ -294,11 +294,7 @@ async fn main() -> Result<()> {
         }
 
         let input = inference.clone();
-        let (sender, receiver) = tokio::sync::oneshot::channel();
-        let submission = Submission { input, sender };
-
-        let _ = runtime.send(submission).await;
-        let (input, output) = receiver.await.unwrap();
+        let (input, output) = runtime.infer(input).await;
         inference = input;
 
         let output = output.iter().map(|batch| batch.0.clone()).collect_vec();

examples/rt-chat.rs

@@ -23,7 +23,7 @@ use web_rwkv::{
             State,
         },
         softmax::softmax_one,
-        v4, v5, v6, JobRuntime, Submission,
+        v4, v5, v6, JobRuntime,
     },
     tensor::{TensorCpu, TensorInit, TensorShape},
     tokenizer::Tokenizer,
@@ -293,11 +293,7 @@ async fn main() -> Result<()> {
 
     loop {
         let input = inference.clone();
-        let (sender, receiver) = tokio::sync::oneshot::channel();
-        let submission = Submission { input, sender };
-
-        let _ = runtime.send(submission).await;
-        let (input, output) = receiver.await.unwrap();
+        let (input, output) = runtime.infer(input).await;
         inference = input;
 
         if output[0].size() > 0 {
@@ -353,11 +349,7 @@ async fn main() -> Result<()> {
 
         loop {
             let input = inference.clone();
-            let (sender, receiver) = tokio::sync::oneshot::channel();
-            let submission = Submission { input, sender };
-
-            let _ = runtime.send(submission).await;
-            let (input, output) = receiver.await.unwrap();
+            let (input, output) = runtime.infer(input).await;
             inference = input;
 
             let output = output[0].0.clone();

examples/rt-gen.rs

@@ -23,7 +23,7 @@ use web_rwkv::{
         loader::{Loader, Lora},
         model::{Build, ContextAutoLimits, ModelBuilder, ModelInfo, ModelVersion, Quant},
         softmax::softmax_one,
-        v4, v5, v6, JobRuntime, Submission,
+        v4, v5, v6, JobRuntime,
     },
     tokenizer::Tokenizer,
 };
@@ -205,14 +205,8 @@ async fn main() -> Result<()> {
 
     let num_token = 500;
     for _ in 0..num_token {
-        let (sender, receiver) = tokio::sync::oneshot::channel();
         let input = prompt.clone();
-        let submission = Submission { input, sender };
-
-        let _ = runtime.send(submission).await;
-        let Ok((input, output)) = receiver.await else {
-            break;
-        };
+        let (input, output) = runtime.infer(input).await;
         prompt = input;
 
         let output = output[0].0.clone();

src/runtime/mod.rs

@@ -1,7 +1,6 @@
 use std::future::Future;
 
 use anyhow::Result;
-use web_rwkv_derive::Deref;
 
 pub mod infer;
 pub mod loader;
@@ -41,9 +40,9 @@ pub trait JobBuilder<J: Job>: Send + Clone + 'static {
 }
 
 #[derive(Debug)]
-pub struct Submission<I, O> {
-    pub input: I,
-    pub sender: tokio::sync::oneshot::Sender<(I, O)>,
+struct Submission<I, O> {
+    input: I,
+    sender: tokio::sync::oneshot::Sender<(I, O)>,
 }
 
 pub trait JobInput: Send + 'static {
@@ -56,7 +55,7 @@ pub trait JobInput: Send + 'static {
     fn chunk(&self) -> Self::Chunk;
 }
 
-#[derive(Debug, Clone, Deref)]
+#[derive(Debug, Clone)]
 pub struct JobRuntime<I, O>(tokio::sync::mpsc::Sender<Submission<I, O>>);
 
 #[allow(clippy::type_complexity)]
@@ -127,8 +126,6 @@ where
                 }
                 let iter = iter.as_mut().expect("iter should be assigned");
 
-                // let remain = queue.len() + candidates.len().max(1) - 1;
-                // let predict = MAX_QUEUE_SIZE - MAX_QUEUE_SIZE.min(remain);
                 for info in iter.take(predict) {
                     #[cfg(feature = "trace")]
                     tracing::event!(
@@ -176,4 +173,13 @@ where
         }
         Ok(())
     }
+
+    /// Perform (partial) inference and return the remaining input and (perhaps partial) output.
+    /// The amount of input processed during one call is bound by the input chunk size.
+    pub async fn infer(&self, input: I) -> (I, O) {
+        let (sender, receiver) = tokio::sync::oneshot::channel();
+        let submission = Submission { input, sender };
+        let _ = self.0.send(submission).await;
+        receiver.await.expect("receive infer output error")
+    }
 }