token preproc

mrsteyk · Feb 18, 2023 · 4e1d03e · 4e1d03e
1 parent bfa05f6
commit 4e1d03e
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 2 deletions.
diff --git a/python/Cargo.lock b/python/Cargo.lock
diff --git a/python/Cargo.toml b/python/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "rwkv_rs"
-version = "0.2.1"
+version = "0.2.2"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

diff --git a/python/src/lib.rs b/python/src/lib.rs
@@ -80,6 +80,12 @@ impl Rwkv {
         Ok(self.inner.rwkv().forward_raw(tokens.last().unwrap(), &mut state.inner))
     }
 
+    pub fn forward_token_preproc(&self, token: usize, state: &mut State) -> PyResult<()> {
+        let x = self.inner.rwkv().emb.get(token).ok_or(anyhow::anyhow!(RwkvError::InvalidToken(token)))?;
+        self.inner.rwkv().forward_raw_preproc(x, &mut state.inner);
+        Ok(())
+    }
+
     pub fn forward_token(&self, token: usize, state: &mut State) -> PyResult<Vec<f32>> {
         let x = self.inner.rwkv().emb.get(token).ok_or(anyhow::anyhow!(RwkvError::InvalidToken(token)))?;
         Ok(self.inner.rwkv().forward_raw(x, &mut state.inner))