import { LayersModel , Tensor3D } from '@tensorflow/tfjs-node' ; import { join } from 'path' ; import { Dataset } from '../../dataset' ; import { generateEmbeddingsMatrix , WordEmbeddingStorage , } from '../../embeddings' ; import { IndexedItems , LabelEncoder , OneHotEncoder } from '../../encode' ; import { ModelEvaluation , ModelManager } from '../../model' ; import { Preprocessor } from '../../preprocess' ; import { ConfigStorage , ModelStorage , NlpTaskConfig } from '../../storage' ; import { unique } from '../../utils' ; import { createBiLstmModel , NER_TEMPLATE , NerModelParameters } from './models' ; import { Entity , NEUTRAL_ENTITY , PredictionProcessor , Processor , } from './process' ; export interface NerConfig extends NlpTaskConfig { entities : string [ ] ; } export class NerConfigStorage extends ConfigStorage < NerConfig > { } export class BotonicNer { private readonly config : Readonly < NerConfig > ; private readonly processor : Processor ; private readonly predictionProcessor : PredictionProcessor ; private modelManager : ModelManager ; constructor ( config , preprocessor ) { this . config = { ... config , entities : unique ( [ NEUTRAL_ENTITY ] . concat ( config . entities ) ) , } ; this . processor = new Processor ( preprocessor , new LabelEncoder ( new IndexedItems ( this . config . vocabulary ) ) , new OneHotEncoder ( new IndexedItems ( this . config . entities ) ) ) ; this . predictionProcessor = new PredictionProcessor ( this . config . entities ) ; } static async load ( path , preprocessor ) { const config = new NerConfigStorage ( ) . load ( path ) ; const ner = new BotonicNer ( config , preprocessor ) ; const model = await new ModelStorage ( ) . load ( path ) ; ner . modelManager = new ModelManager ( model ) ; return ner ; } async createModel ( template , storage , params ? ) { const embeddingsMatrix = await generateEmbeddingsMatrix ( storage , this . config . vocabulary ) ; switch ( template ) { case NER_TEMPLATE . BILSTM : return createBiLstmModel ( this . config . maxLength , this . config . entities , embeddingsMatrix , params ) ; default : throw new Error ( `"${ String ( template ) }" is an invalid model template.` ) ; } } setModel ( model ) { this . modelManager = new ModelManager ( model ) ; } async train ( dataset , epochs , batchSize ) { const { x , y } = this . processor . process ( dataset . samples ) ; await this . modelManager . train ( x , y , { epochs , batchSize } ) ; } async evaluate ( dataset ) { const { x , y } = this . processor . process ( dataset . samples ) ; return await this . modelManager . evaluate ( x , y ) ; } recognizeEntities ( text ) { const { sequence , input } = this . processor . generateInput ( text ) ; const prediction = this . modelManager . predict ( input ) as Tensor3D ; return this . predictionProcessor . process ( sequence , prediction ) ; } async saveModel ( path ) { path = join ( path , this . config . locale ) ; new NerConfigStorage ( ) . save ( this . config , path ) ; await new ModelStorage ( ) . save ( this . modelManager . model , path ) ; } }