license import * as use from '@tensorflow-models/universal-sentence-encoder' ; import * as tfconv from '@tensorflow/tfjs-converter' ; import * as tf from '@tensorflow/tfjs-core' ; export { version } from './version' ; declare interface ModelInputs extends tf . NamedTensorMap { Placeholder_1 : tf . Tensor ; Placeholder : tf . Tensor ; } param threshold param toxicityLabels export async function load ( threshold , toxicityLabels ) { const model = new ToxicityClassifier ( threshold , toxicityLabels ) ; await model . load ( ) ; return model ; } export class ToxicityClassifier { private tokenizer : use . Tokenizer ; private model : tfconv . GraphModel ; private labels : string [ ] ; private threshold : number ; private toxicityLabels : string [ ] ; constructor ( threshold = 0.85 , toxicityLabels = [ ] ) { this . threshold = threshold ; this . toxicityLabels = toxicityLabels ; } async loadModel ( ) { return tfconv . loadGraphModel ( 'https://tfhub.dev/tensorflow/tfjs-model/toxicity/1/default/1' , { fromTFHub : true } ) ; } async loadTokenizer ( ) { return use . loadTokenizer ( ) ; } async load ( ) { const [ model , tokenizer ] = await Promise . all ( [ this . loadModel ( ) , this . loadTokenizer ( ) ] ) ; this . model = model ; this . tokenizer = tokenizer ; this . labels = model . outputs . map ( ( d ) => d . name . split ( '/' ) [ 0 ] ) ; if ( this . toxicityLabels . length === 0 ) { this . toxicityLabels = this . labels ; } else { tf . util . assert ( this . toxicityLabels . every ( d => this . labels . indexOf ( d ) > - 1 ) , ( ) => `toxicityLabels argument must contain only items from the ` + `model heads ${ this . labels . join ( ', ' ) }, ` + `got ${ this . toxicityLabels . join ( ', ' ) }` ) ; } } param inputs async classify ( inputs ) { if ( typeof inputs === 'string' ) { inputs = [ inputs ] ; } const encodings = inputs . map ( d => this . tokenizer . encode ( d ) ) ; const indicesArr = encodings . map ( ( arr , i ) => arr . map ( ( d , index ) => [ i , index ] ) ) ; let flattenedIndicesArr = [ ] ; for ( let i = 0 ; i < indicesArr . length ; i ++ ) { flattenedIndicesArr = flattenedIndicesArr . concat ( indicesArr [ i ] as Array < [ number , number ] > ) ; } const indices = tf . tensor2d ( flattenedIndicesArr , [ flattenedIndicesArr . length , 2 ] , 'int32' ) ; const values = tf . tensor1d ( tf . util . flatten ( encodings ) as number [ ] , 'int32' ) ; const modelInputs = { Placeholder_1 : indices , Placeholder : values } ; const labels = await this . model . executeAsync ( modelInputs ) ; indices . dispose ( ) ; values . dispose ( ) ; return ( labels as tf . Tensor2D [ ] ) . map ( ( d , i ) => ( { data : d , headIndex : i } ) ) . filter ( ( d ) => this . toxicityLabels . indexOf ( this . labels [ d . headIndex ] ) > - 1 ) . map ( ( d ) => { const prediction = d . data . dataSync ( ) as Float32Array ; const results = [ ] ; for ( let input = 0 ; input < inputs . length ; input ++ ) { const probabilities = prediction . slice ( input * 2 , input * 2 + 2 ) ; let match = null ; if ( Math . max ( probabilities [ 0 ] , probabilities [ 1 ] ) > this . threshold ) { match = probabilities [ 0 ] < probabilities [ 1 ] ; } results . push ( { probabilities , match } ) ; } return { label : this . labels [ d . headIndex ] , results } ; } ) ; } }