feat(uni): init for embedding

unit-mesh · Nov 14, 2023 · ab4355e · ab4355e
1 parent 50b2ce1
commit ab4355e
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 4 deletions.
diff --git a/inference_core/src/embedding/semantic.rs b/inference_core/src/embedding/semantic.rs
@@ -22,8 +22,19 @@ impl Drop for Semantic {
     }
 }
 
+pub fn init_semantic(model: Vec<u8>, tokenizer_data: Vec<u8>) -> Result<Arc<Semantic>, SemanticError> {
+    let result = Semantic::init_semantic(model, tokenizer_data)?;
+    Ok(Arc::new(result))
+}
+
 impl Semantic {
     pub async fn initialize(model: Vec<u8>, tokenizer_data: Vec<u8>) -> Result<Pin<Box<Semantic>>, SemanticError> {
+        let semantic = Self::init_semantic(model, tokenizer_data)?;
+
+        Ok(Box::pin(semantic))
+    }
+
+    pub fn init_semantic(model: Vec<u8>, tokenizer_data: Vec<u8>) -> Result<Semantic, SemanticError> {
         let environment = Arc::new(
             Environment::builder()
                 .with_name("Encode")
@@ -53,8 +64,7 @@ impl Semantic {
                 .unwrap()
                 .into(),
         };
-
-        Ok(Box::pin(semantic))
+        Ok(semantic)
     }
 
     pub fn embed(&self, sequence: &str) -> Result<Embedding, SemanticError> {

diff --git a/inference_core/src/inference.udl b/inference_core/src/inference.udl
@@ -1,4 +1,7 @@
-namespace inference { };
+namespace inference {
+    [Throws=SemanticError]
+    Semantic init_semantic(sequence<u8> model, sequence<u8> tokenizer_data);
+};
 
 [Custom]
 typedef sequence<f32> Embedding;

diff --git a/inference_core/src/lib.rs b/inference_core/src/lib.rs
@@ -9,7 +9,7 @@ pub use document::Metadata;
 pub use embedding::Embedding;
 pub use embedding::Semantic;
 pub use embedding::semantic::SemanticError;
-// pub use embedding::semantic::init_semantic;
+pub use embedding::semantic::init_semantic;
 
 pub use similarity::EmbeddingMatch;
 pub use similarity::CosineSimilarity;