File size: 2,001 Bytes
72c0672 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | extern crate tokenizers as tk;
use crate::models::Model;
use napi::bindgen_prelude::*;
use std::sync::{Arc, RwLock};
use tokenizers::models::bpe::{BpeBuilder, BPE};
use tokenizers::models::wordlevel::{WordLevel, WordLevelBuilder};
use tokenizers::models::wordpiece::{WordPiece, WordPieceBuilder};
pub struct BPEFromFilesTask {
pub(crate) builder: Option<BpeBuilder>,
}
impl Task for BPEFromFilesTask {
type Output = BPE;
type JsValue = Model;
fn compute(&mut self) -> Result<Self::Output> {
self
.builder
.take()
.ok_or(Error::from_reason("Empty builder".to_string()))?
.build()
.map_err(|e| Error::from_reason(format!("{}", e)))
}
fn resolve(&mut self, _env: Env, output: Self::Output) -> Result<Self::JsValue> {
Ok(Model {
model: Some(Arc::new(RwLock::new(output.into()))),
})
}
}
pub struct WordPieceFromFilesTask {
pub(crate) builder: Option<WordPieceBuilder>,
}
impl Task for WordPieceFromFilesTask {
type Output = WordPiece;
type JsValue = Model;
fn compute(&mut self) -> Result<Self::Output> {
self
.builder
.take()
.ok_or(Error::from_reason("Empty builder".to_string()))?
.build()
.map_err(|e| Error::from_reason(format!("{}", e)))
}
fn resolve(&mut self, _env: Env, output: Self::Output) -> Result<Self::JsValue> {
Ok(Model {
model: Some(Arc::new(RwLock::new(output.into()))),
})
}
}
pub struct WordLevelFromFilesTask {
pub(crate) builder: Option<WordLevelBuilder>,
}
impl Task for WordLevelFromFilesTask {
type Output = WordLevel;
type JsValue = Model;
fn compute(&mut self) -> Result<Self::Output> {
self
.builder
.take()
.ok_or(Error::from_reason("Empty builder".to_string()))?
.build()
.map_err(|e| Error::from_reason(format!("{}", e)))
}
fn resolve(&mut self, _env: Env, output: Self::Output) -> Result<Self::JsValue> {
Ok(Model {
model: Some(Arc::new(RwLock::new(output.into()))),
})
}
}
|