feat(pipelines): implement parts of speech model

.gitignore

@@ -1,2 +1,3 @@
 /target
 Cargo.lock
+.vscode

README.md

@@ -21,6 +21,7 @@ methods on those classes.
 - `NERModel`
 - `QAModel`
 - `SentimentModel`
+- `POSModel`
 
 To test out these pipelines, you can try and run the `dev.ts` file. However,
 this will automatically install the necessary models so I advise you comment out

dev.ts

@@ -4,6 +4,12 @@ const SECTION = (name: string) => console.log(`\n${name}\n`);
 
 const manager = new ModelManager();
 
+SECTION("POS TAGGING MODEL");
+
+const posModel = await manager.createPOSModel();
+
+console.log(await posModel.predict(["What are the parts in this?"]));
+
 const convoModel = await manager.createConversationModel();
 
 const convoManager = await convoModel.createConversationManager();

mod.ts

@@ -15,3 +15,6 @@ export type {
   TranslateInit,
   TranslationModelInit,
 } from "./models/translation/mod.ts";
+
+export { POSModel } from "./models/pos.ts";
+export type { POSEntity } from "./models/pos.ts";

model_manager.ts

@@ -6,6 +6,7 @@ import { SentimentModel } from "./models/sentiment.ts";
 import { TranslationModel } from "./models/translation/mod.ts";
 import { ConversationModel } from "./models/conversation.ts";
 import type { TranslationModelInit } from "./models/translation/mod.ts";
+import { POSModel } from "./models/pos.ts";
 import { encode } from "./utils/encode.ts";
 import { decode } from "./utils/decode.ts";
 
@@ -42,12 +43,6 @@ const symbolDefinitions = {
     result: "isize",
     nonblocking: true,
   },
-  error_len: { parameters: [], result: "usize", nonblocking: true },
-  fill_result: {
-    parameters: ["buffer", "usize"],
-    result: "void",
-    nonblocking: true,
-  },
   create_conversation_model: {
     parameters: [],
     result: "isize",
@@ -68,6 +63,22 @@ const symbolDefinitions = {
     result: "isize",
     nonblocking: true,
   },
+  create_pos_model: {
+    parameters: [],
+    result: "isize",
+    nonblocking: true,
+  },
+  pos_predict: {
+    parameters: ["usize", "buffer", "usize"],
+    result: "isize",
+    nonblocking: true,
+  },
+  error_len: { parameters: [], result: "usize", nonblocking: true },
+  fill_result: {
+    parameters: ["buffer", "usize"],
+    result: "void",
+    nonblocking: true,
+  },
   fill_error: { parameters: ["buffer"], result: "void", nonblocking: true },
   delete_model: { parameters: ["usize"], result: "isize", nonblocking: true },
 } as const;
@@ -192,6 +203,15 @@ export class ModelManager {
     return model;
   }
 
+  async createPOSModel(): Promise<POSModel> {
+    const rid = await this.bindings.create_pos_model().then(
+      this.assertCode,
+    );
+    const model = new POSModel(this, rid);
+    this.#models.push(model);
+    return model;
+  }
+
   close() {
     this.#close();
     this.#isClosed = true;

models/pos.ts

@@ -0,0 +1,28 @@
+import type { ModelManager } from "../model_manager.ts";
+import { Model } from "../model.ts";
+import { encode } from "../utils/encode.ts";
+import { decode } from "../utils/decode.ts";
+
+export interface POSEntity {
+  word: string;
+  score: number;
+  label: string;
+}
+
+export class POSModel extends Model {
+  constructor(manager: ModelManager, rid: number) {
+    super(manager, rid);
+  }
+
+  async predict(inputs: string[]): Promise<POSEntity[]> {
+    const { bindings, assertCode } = this.manager;
+    const bytes = encode(JSON.stringify(inputs));
+    const len = await bindings
+      .pos_predict(this.rid, bytes, bytes.length)
+      .then(assertCode);
+    const buf = new Uint8Array(len);
+    await bindings.fill_result(buf, len);
+    const json = decode(buf);
+    return JSON.parse(json);
+  }
+}

models/translation/mod.ts

@@ -26,15 +26,12 @@ export class TranslationModel extends Model {
     this.#init = init;
   }
 
-  async translate(init: TranslateInit) {
-    const { sourceLanguage, targetLanguage } = init;
+  async translate(init: TranslateInit): Promise<string[]> {
     const { bindings, assertCode } = this.manager;
     const bytes = encode(JSON.stringify(init));
-    const len = await bindings.translation_translate(
-      this.rid,
-      bytes,
-      bytes.length,
-    ).then(assertCode);
+    const len = await bindings
+      .translation_translate(this.rid, bytes, bytes.length)
+      .then(assertCode);
     const buf = new Uint8Array(len);
     await bindings.fill_result(buf, len);
     const json = decode(buf);

src/allocators.rs

@@ -1,12 +1,17 @@
 use crate::create_allocator;
-use rust_bert::pipelines::{conversation, ner, question_answering, sentiment, translation};
+use rust_bert::pipelines::{
+    conversation, ner, pos_tagging, question_answering, sentiment, translation,
+    zero_shot_classification,
+};
 
 pub enum Model {
     TranslationModel(translation::TranslationModel),
     QuestionAnsweringModel(question_answering::QuestionAnsweringModel),
     NERModel(ner::NERModel),
     SentimentModel(sentiment::SentimentModel),
     ConversationModel(conversation::ConversationModel),
+    POSModel(pos_tagging::POSModel),
+    ZeroShotClassificationModel(zero_shot_classification::ZeroShotClassificationModel),
 }
 
 pub enum ModelResource {

src/lib.rs

@@ -1,14 +1,18 @@
 mod allocators;
 mod conversation;
 mod ner;
+mod pos_tagging;
 mod qa;
 mod sentiment;
 mod translation;
+mod zero_shot_classification;
 
+pub use zero_shot_classification::*;
 pub use allocators::*;
 pub use conversation::*;
 pub use ner::*;
 use once_cell::sync::Lazy;
+pub use pos_tagging::*;
 pub use qa::*;
 pub use sentiment::*;
 use std::sync::Mutex;
@@ -27,7 +31,7 @@ pub fn set_result(v: Vec<u8>) -> usize {
 /// Inspired by deno_sqlite3's `exec` helper by @littledivvy
 pub fn exec<F>(f: F) -> isize
 where
-    F: FnOnce() -> Result<isize, anyhow::Error>,
+   F: FnOnce() -> Result<isize, anyhow::Error>,
 {
     match f() {
         Ok(a) => a,

src/pos_tagging.rs

@@ -0,0 +1,32 @@
+use crate::{exec, models, set_result, Model};
+use anyhow::Context;
+use rust_bert::pipelines::pos_tagging::POSModel;
+
+#[no_mangle]
+pub extern "C" fn create_pos_model() -> isize {
+    exec(|| {
+        let model = POSModel::new(Default::default())
+            .context("Failed to load Parts of Speech Tagging model.")?;
+
+        models::allocate(Model::POSModel(model)).map(|a| a as isize)
+    })
+}
+
+#[no_mangle]
+pub extern "C" fn pos_predict(rid: usize, buf: *const u8, buf_len: usize) -> isize {
+    exec(|| {
+        let inputs: Vec<String> =
+            serde_json::from_slice(unsafe { std::slice::from_raw_parts(buf, buf_len) })?;
+
+        models::with_access(rid, |model| {
+            let model = match model {
+                Model::POSModel(m) => m,
+                _ => return Err(anyhow::anyhow!("Expected POS Model at rid '{}'.", rid)),
+            };
+
+            let outputs = model.predict(&inputs);
+            let outputs = serde_json::to_vec(&outputs).context("Failed to serialize POS tags.")?;
+            Ok(set_result(outputs) as isize)
+        })
+    })
+}

src/zero_shot_classification.rs

@@ -0,0 +1,28 @@
+use crate::{exec, models, Model};
+use anyhow::Context;
+use rust_bert::pipelines::zero_shot_classification::ZeroShotClassificationModel;
+use serde::Deserialize;
+
+#[no_mangle]
+pub extern "C" fn create_zero_shot_model() -> isize {
+    exec(|| {
+        let model = ZeroShotClassificationModel::new(Default::default())
+            .context("Failed to create zero shot classification model.")?;
+
+        models::allocate(Model::ZeroShotClassificationModel(model)).map(|a| a as isize)
+    })
+}
+
+#[derive(Deserialize)]
+pub struct ZeroShotInput {}
+
+#[no_mangle]
+pub extern "C" fn zero_shot_predict(rid: usize, buf: *const u8, buf_len: usize) -> isize {
+    exec(|| {
+        let inputs: Vec<String> =
+            serde_json::from_slice(unsafe { std::slice::from_raw_parts(buf, buf_len) })
+                .context("Failed to parse zero shot input.")?;
+
+        Ok(0)
+    })
+}