import { PADDING_TOKEN } from '@botonic/nlp/lib/preprocess/constants' ; import { NEUTRAL_ENTITY } from '@botonic/nlp/lib/tasks/ner/process/constants' ; import { Entity , Prediction } from '@botonic/nlp/lib/tasks/ner/process/types' ; import { Tensor3D } from '@tensorflow/tfjs' ; export class PredictionProcessor { constructor ( private readonly entities ) { } process ( sequence , prediction ) { const confidences = prediction . arraySync ( ) [ 0 ] ; return this . getEntities ( sequence , confidences ) ; } private getEntities ( sequence , confidences ) { const entities = sequence . map ( ( token , idx ) => this . generateEntity ( token , confidences [ idx ] ) ) ; return this . filterEntities ( entities ) ; } private generateEntity ( token , confidences ) { const confidence = Math . max ( ... confidences ) ; const label = this . entities [ confidences . indexOf ( confidence ) ] ; const predictions = this . getPredictions ( confidences ) ; return { text : token , label , confidence , predictions } ; } private getPredictions ( confidences ) { return confidences . map ( ( c , i ) => { return { label : this . entities [ i ] , confidence : c } ; } ) ; } private filterEntities ( entities ) { return entities . filter ( entity => entity . text != PADDING_TOKEN && entity . label != NEUTRAL_ENTITY ) ; } }