license import { util } from '@tensorflow/tfjs-core' ; /**
 * The StringNGramsOp class creates ngrams from ragged string data.
 * The constructor contains all attributes related to the operation such as
 * padding widths and strings, and the compute function can be used to
 * compute the ngrams for different ragged tensor inputs.
 */ class StringNGramsOp { private separator : Uint8Array ; private nGramWidths : number [ ] ; private padWidth : number ; private leftPad : Uint8Array ; private rightPad : Uint8Array ; private preserveShort : boolean ; constructor ( separator , nGramWidths , leftPad , rightPad , padWidth , preserveShortSequences ) { this . separator = util . encodeString ( separator ) ; this . nGramWidths = nGramWidths ; this . leftPad = util . encodeString ( leftPad ) ; this . rightPad = util . encodeString ( rightPad ) ; this . padWidth = padWidth ; this . preserveShort = preserveShortSequences ; } private getPadWidth ( nGramWidth ) { return Math . min ( this . padWidth < 0 ? nGramWidth - 1 : this . padWidth , nGramWidth - 1 ) ; } private getNumNGrams ( length , nGramWidth ) { const padWidth = this . getPadWidth ( nGramWidth ) ; return Math . max ( 0 , ( ( length + 2 * padWidth ) - nGramWidth ) + 1 ) ; } private createNGrams ( data , splitIndex , output , outputStartIndex , numNGrams , nGramWidth ) { for ( let nGramIndex = 0 ; nGramIndex < numNGrams ; ++ nGramIndex ) { const padWidth = this . getPadWidth ( nGramWidth ) ; const leftPadding = Math . max ( 0 , padWidth - nGramIndex ) ; const rightPadding = Math . max ( 0 , padWidth - ( numNGrams - ( nGramIndex + 1 ) ) ) ; const numTokens = nGramWidth - ( leftPadding + rightPadding ) ; const dataStartIndex = splitIndex + ( leftPadding > 0 ? 0 : nGramIndex - padWidth ) ; let nGramSize = 0 ; nGramSize += leftPadding * this . leftPad . length ; for ( let n = 0 ; n < numTokens ; ++ n ) { nGramSize += data [ dataStartIndex + n ] . length ; } nGramSize += rightPadding * this . rightPad . length ; const numSeparators = leftPadding + rightPadding + numTokens - 1 ; nGramSize += numSeparators * this . separator . length ; output [ outputStartIndex + nGramIndex ] = new Uint8Array ( nGramSize ) ; const nGram = output [ outputStartIndex + nGramIndex ] ; let nextNGramIndex = 0 ; const appendToNGram = ( str ) => str . forEach ( ( value ) => nGram [ nextNGramIndex ++ ] = value ) ; for ( let n = 0 ; n < leftPadding ; ++ n ) { appendToNGram ( this . leftPad ) ; appendToNGram ( this . separator ) ; } for ( let n = 0 ; n < numTokens - 1 ; ++ n ) { appendToNGram ( data [ dataStartIndex + n ] ) ; appendToNGram ( this . separator ) ; } if ( numTokens > 0 ) { appendToNGram ( data [ dataStartIndex + numTokens - 1 ] ) ; for ( let n = 0 ; n < rightPadding ; ++ n ) { appendToNGram ( this . separator ) ; appendToNGram ( this . rightPad ) ; } } else { for ( let n = 0 ; n < rightPadding - 1 ; ++ n ) { appendToNGram ( this . rightPad ) ; appendToNGram ( this . separator ) ; } appendToNGram ( this . rightPad ) ; } } } public compute ( data , splits ) { const inputDataSize = data . length ; const splitsSize = splits . length ; if ( splitsSize > 0 ) { let prevSplit = splits [ 0 ] ; if ( prevSplit !== 0 ) { throw new Error ( `First split value must be 0, got ${ prevSplit }` ) ; } for ( let i = 1 ; i < splitsSize ; ++ i ) { let validSplits = splits [ i ] >= prevSplit ; validSplits = validSplits && ( splits [ i ] <= inputDataSize ) ; if ( ! validSplits ) { throw new Error ( `Invalid split value ${ splits [ i ] }, must be in [${ prevSplit }, ${ inputDataSize }]` ) ; } prevSplit = splits [ i ] ; } if ( prevSplit !== inputDataSize ) { throw new Error ( `Last split value must be data size. Expected ${ inputDataSize }, got ${ prevSplit }` ) ; } } const numBatchItems = splitsSize - 1 ; const nGramsSplits = util . getArrayFromDType ( 'int32' , splitsSize ) ; if ( inputDataSize === 0 || splitsSize === 0 ) { const empty = new Array ( inputDataSize ) ; for ( let i = 0 ; i <= numBatchItems ; ++ i ) { nGramsSplits [ i ] = 0 ; } return [ empty , nGramsSplits ] ; } nGramsSplits [ 0 ] = 0 ; for ( let i = 1 ; i <= numBatchItems ; ++ i ) { const length = splits [ i ] - splits [ i - 1 ] ; let numNGrams = 0 ; this . nGramWidths . forEach ( ( nGramWidth ) => { numNGrams += this . getNumNGrams ( length , nGramWidth ) ; } ) ; if ( this . preserveShort && length > 0 && numNGrams === 0 ) { numNGrams = 1 ; } nGramsSplits [ i ] = nGramsSplits [ i - 1 ] + numNGrams ; } const nGrams = new Array ( nGramsSplits [ numBatchItems ] ) ; for ( let i = 0 ; i < numBatchItems ; ++ i ) { const splitIndex = splits [ i ] ; let outputStartIdx = nGramsSplits [ i ] ; this . nGramWidths . forEach ( ( nGramWidth ) => { const length = splits [ i + 1 ] - splits [ i ] ; const numNGrams = this . getNumNGrams ( length , nGramWidth ) ; this . createNGrams ( data , splitIndex , nGrams , outputStartIdx , numNGrams , nGramWidth ) ; outputStartIdx += numNGrams ; } ) ; if ( this . preserveShort && outputStartIdx === nGramsSplits [ i ] ) { const dataLength = splits [ i + 1 ] - splits [ i ] ; if ( dataLength === 0 ) { continue ; } const nGramWidth = dataLength + 2 * this . padWidth ; const numNGrams = 1 ; this . createNGrams ( data , splitIndex , nGrams , outputStartIdx , numNGrams , nGramWidth ) ; } } return [ nGrams , nGramsSplits ] ; } } export function stringNGramsImpl ( data , dataSplits , separator , nGramWidths , leftPad , rightPad , padWidth , preserveShortSequences ) { return new StringNGramsOp ( separator , nGramWidths , leftPad , rightPad , padWidth , preserveShortSequences ) . compute ( data , dataSplits ) ; }