Implement concat for ModelOutput.

cryscan · Feb 2, 2024 · f6bef59 · f6bef59
1 parent 6d905d8
commit f6bef59
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/Cargo.toml b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "web-rwkv"
-version = "0.6.0"
+version = "0.6.1"
 edition = "2021"
 authors = ["Zhenyuan Zhang <[email protected]>"]
 license = "MIT OR Apache-2.0"

diff --git a/src/model/mod.rs b/src/model/mod.rs
@@ -92,6 +92,19 @@ pub enum ModelOutput {
     Full(Vec<Vec<f32>>),
 }
 
+impl ModelOutput {
+    pub fn concat(self, other: Self) -> Self {
+        match (self, other) {
+            (Self::None, y) => y,
+            (x, Self::None) => x,
+            (Self::Last(x), Self::Last(y)) => Self::Full(vec![x, y]),
+            (Self::Last(x), Self::Full(y)) => Self::Full([vec![x], y].concat()),
+            (Self::Full(x), Self::Last(y)) => Self::Full([x, vec![y]].concat()),
+            (Self::Full(x), Self::Full(y)) => Self::Full([x, y].concat()),
+        }
+    }
+}
+
 #[derive(Debug, Default, Clone, Copy)]
 pub enum OutputType {
     /// Only the prediction of the last token.