license import * as tfliteWebAPIClient from '../tflite_web_api_client' ; import { BertNLClassifier as TaskLibraryBertNLClassifier } from '../types/bert_nl_classifier' ; import { BaseTaskLibraryClient , Class } from './common' ; export interface BertNLClassifierOptions { /**
     * Max number of tokens to pass to the model.
     *
     * Default to 128.
     */ maxSeqLen ? : number ; } /**
 * Client for BertNLClassifier TFLite Task Library.
 *
 * It is a wrapper around the underlying javascript API to make it more
 * convenient to use. See comments in the corresponding type declaration file in
 * src/types for more info.
 */ export class BertNLClassifier extends BaseTaskLibraryClient { constructor ( protected instance ) { super ( instance ) ; } static async create ( model , options ? ) { const protoOptions = new tfliteWebAPIClient . tfweb . BertNLClassifierOptions ( ) ; if ( options ) { if ( options . maxSeqLen ) { protoOptions . setMaxSeqLen ( options . maxSeqLen ) ; } } const instance = await tfliteWebAPIClient . tfweb . BertNLClassifier . create ( model , protoOptions ) ; return new BertNLClassifier ( instance ) ; } classify ( input ) { return this . instance . classify ( input ) . map ( category => { return { probability : category . score , className : category . className , } ; } ) ; } }