Node - Fix EncodeTask

This commit is contained in:
Anthony MOI
2020-01-09 14:04:20 -05:00
parent bda03ffe8c
commit de35bb4f45
2 changed files with 3 additions and 4 deletions

View File

@ -39,7 +39,7 @@ impl Task for EncodeTask {
fn perform(&self) -> Result<Self::Output, Self::Error> {
match self {
EncodeTask::Single(worker, input) => {
let mut input = unsafe { std::ptr::read(input) };
let mut input = unsafe { std::ptr::replace(&input as *const _ as *mut _, None) };
let tokenizer: &Tokenizer = unsafe { &*worker.ptr };
tokenizer
.encode(input.take().ok_or("No provided input")?)
@ -47,7 +47,7 @@ impl Task for EncodeTask {
.map(|encoding| EncodeOutput::Single(encoding))
}
EncodeTask::Batch(worker, input) => {
let mut input = unsafe { std::ptr::read(input) };
let mut input = unsafe { std::ptr::replace(&input as *const _ as *mut _, None) };
let tokenizer: &Tokenizer = unsafe { &*worker.ptr };
tokenizer
.encode_batch(input.take().ok_or("No provided input")?)

View File

@ -1,6 +1,5 @@
extern crate tokenizers as tk;
use crate::encoding::*;
use crate::models::*;
use crate::tasks::tokenizer::{EncodeTask, WorkingTokenizer};
use neon::prelude::*;
@ -16,7 +15,7 @@ pub struct Tokenizer {
impl Tokenizer {
pub fn prepare_for_task(&self) -> Option<WorkingTokenizer> {
if std::sync::Arc::strong_count(&self.running_task) > 0 {
if std::sync::Arc::strong_count(&self.running_task) > 1 {
None
} else {
unsafe {