question_and_answer.ts

/**
 * @license
 * Copyright 2020 Google LLC. All Rights Reserved.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 * =============================================================================
 */
import*astfconvfrom'@tensorflow/tfjs-converter';
import*astffrom'@tensorflow/tfjs-core';

import{BertTokenizer,CLS_INDEX,loadTokenizer,SEP_INDEX,Token}from'./bert_tokenizer';

constMODEL_URL='https://tfhub.dev/tensorflow/tfjs-model/mobilebert/1';
constINPUT_SIZE=384;
constMAX_ANSWER_LEN=32;
constMAX_QUERY_LEN=64;
constMAX_SEQ_LEN=384;
constPREDICT_ANSWER_NUM=5;
constOUTPUT_OFFSET=1;
// This is threshold value for determining if a question is irrelevant to the
// context. This value comes from the QnA model, and is generated by the
// training process based on the SQUaD 2.0 dataset.
constNO_ANSWER_THRESHOLD=4.3980759382247925;

exportinterfaceQuestionAndAnswer{
/**
 * Given the question and context, find the best answers.
 * @param question the question to find answers for.
 * @param context context where the answers are looked up from.
 * @return array of answers
 */
findAnswers(question: string,context: string): Promise<Answer[]>;
}

/**
 * MobileBert model loading is configurable using the following config
 * dictionary.
 *
 * `modelUrl`: An optional string that specifies custom url of the model. This
 * is useful for area/countries that don't have access to the model hosted on
 * GCP.
 */
exportinterfaceModelConfig{
/**
 * An optional string that specifies custom url of the model. This
 * is useful for area/countries that don't have access to the model hosted on
 * GCP.
 */
modelUrl: string;
/**
 * Wheter the url is from tfhub.
 */
fromTFHub?: boolean;
}

/**
 * Answer object returned by the model.
 * `text`: string, the text of the answer.
 * `startIndex`: number, the index of the starting character of the answer in
 * the passage.
 * `endIndex`: number, index of the last character of the answer text.
 * `score`: number, indicates the confident
 * level.
 */
exportinterfaceAnswer{
text: string;
startIndex: number;
endIndex: number;
score: number;
}

interfaceFeature{
inputIds: number[];
inputMask: number[];
segmentIds: number[];
origTokens: Token[];
tokenToOrigMap: {[key: number]: number};
}

interfaceAnswerIndex{
start: number;
end: number;
score: number;
}

classQuestionAndAnswerImplimplementsQuestionAndAnswer{
privatemodel: tfconv.GraphModel;
privatetokenizer: BertTokenizer;

constructor(privatemodelConfig: ModelConfig){
if(this.modelConfig==null){
this.modelConfig={modelUrl: MODEL_URL,fromTFHub: true};
}
if(this.modelConfig.fromTFHub==null){
this.modelConfig.fromTFHub=false;
}
}
privateprocess(
query: string,context: string,maxQueryLen: number,maxSeqLen: number,
docStride=128): Feature[]{
// always add the question mark to the end of the query.
query=query.replace(/\?/g,'');
query=query.trim();
query=query+'?';

constqueryTokens=this.tokenizer.tokenize(query);
if(queryTokens.length>maxQueryLen){
thrownewError(
`The length of question token exceeds the limit (${maxQueryLen}).`);
}

constorigTokens=this.tokenizer.processInput(context.trim());
consttokenToOrigIndex: number[]=[];
constallDocTokens: number[]=[];
for(leti=0;i<origTokens.length;i++){
consttoken=origTokens[i].text;
constsubTokens=this.tokenizer.tokenize(token);
for(letj=0;j<subTokens.length;j++){
constsubToken=subTokens[j];
tokenToOrigIndex.push(i);
allDocTokens.push(subToken);
}
}
// The -3 accounts for [CLS], [SEP] and [SEP]
constmaxContextLen=maxSeqLen-queryTokens.length-3;

// We can have documents that are longer than the maximum sequence
// length. To deal with this we do a sliding window approach, where we
// take chunks of the up to our max length with a stride of
// `doc_stride`.
constdocSpans: Array<{start: number,length: number}>=[];
letstartOffset=0;
while(startOffset<allDocTokens.length){
letlength=allDocTokens.length-startOffset;
if(length>maxContextLen){
length=maxContextLen;
}
docSpans.push({start: startOffset, length});
if(startOffset+length===allDocTokens.length){
break;
}
startOffset+=Math.min(length,docStride);
}

constfeatures=docSpans.map(docSpan=>{
consttokens=[];
constsegmentIds=[];
consttokenToOrigMap: {[index: number]: number}={};
tokens.push(CLS_INDEX);
segmentIds.push(0);
for(leti=0;i<queryTokens.length;i++){
constqueryToken=queryTokens[i];
tokens.push(queryToken);
segmentIds.push(0);
}
tokens.push(SEP_INDEX);
segmentIds.push(0);
for(leti=0;i<docSpan.length;i++){
constsplitTokenIndex=i+docSpan.start;
constdocToken=allDocTokens[splitTokenIndex];
tokens.push(docToken);
segmentIds.push(1);
tokenToOrigMap[tokens.length]=tokenToOrigIndex[splitTokenIndex];
}
tokens.push(SEP_INDEX);
segmentIds.push(1);
constinputIds=tokens;
constinputMask=inputIds.map(id=>1);
while((inputIds.length<maxSeqLen)){
inputIds.push(0);
inputMask.push(0);
segmentIds.push(0);
}
return{inputIds, inputMask, segmentIds, origTokens, tokenToOrigMap};
});
returnfeatures;
}

asyncload(){
this.model=awaittfconv.loadGraphModel(
this.modelConfig.modelUrl,{fromTFHub: this.modelConfig.fromTFHub});
// warm up the backend
constbatchSize=1;
constinputIds=tf.ones([batchSize,INPUT_SIZE],'int32');
constsegmentIds=tf.ones([1,INPUT_SIZE],'int32');
constinputMask=tf.ones([1,INPUT_SIZE],'int32');
this.model.execute({
input_ids: inputIds,
segment_ids: segmentIds,
input_mask: inputMask,
global_step: tf.scalar(1,'int32')
});

this.tokenizer=awaitloadTokenizer();
}

/**
 * Given the question and context, find the best answers.
 * @param question the question to find answers for.
 * @param context context where the answers are looked up from.
 * @return array of answers
 */
asyncfindAnswers(question: string,context: string): Promise<Answer[]>{
if(question==null||context==null){
thrownewError(
'The input to findAnswers call is null, '+
'please pass a string as input.');
}

constfeatures=
this.process(question,context,MAX_QUERY_LEN,MAX_SEQ_LEN);
constinputIdArray=features.map(f=>f.inputIds);
constsegmentIdArray=features.map(f=>f.segmentIds);
constinputMaskArray=features.map(f=>f.inputMask);
constglobalStep=tf.scalar(1,'int32');
constbatchSize=features.length;
constresult=tf.tidy(()=>{
constinputIds=
tf.tensor2d(inputIdArray,[batchSize,INPUT_SIZE],'int32');
constsegmentIds=
tf.tensor2d(segmentIdArray,[batchSize,INPUT_SIZE],'int32');
constinputMask=
tf.tensor2d(inputMaskArray,[batchSize,INPUT_SIZE],'int32');
returnthis.model.execute(
{
input_ids: inputIds,
segment_ids: segmentIds,
input_mask: inputMask,
global_step: globalStep
},
['start_logits','end_logits'])as[tf.Tensor2D,tf.Tensor2D];
});
constlogits=awaitPromise.all([result[0].array(),result[1].array()]);
// dispose all intermediate tensors
globalStep.dispose();
result[0].dispose();
result[1].dispose();

constanswers=[];
for(leti=0;i<batchSize;i++){
answers.push(this.getBestAnswers(
logits[0][i],logits[1][i],features[i].origTokens,
features[i].tokenToOrigMap,context,i));
}

returnanswers.reduce((flatten,array)=>flatten.concat(array),[])
.sort((logitA,logitB)=>logitB.score-logitA.score)
.slice(0,PREDICT_ANSWER_NUM);
}

/**
 * Find the Best N answers & logits from the logits array and input feature.
 * @param startLogits start index for the answers
 * @param endLogits end index for the answers
 * @param origTokens original tokens of the passage
 * @param tokenToOrigMap token to index mapping
 */
getBestAnswers(
startLogits: number[],endLogits: number[],origTokens: Token[],
tokenToOrigMap: {[key: string]: number},context: string,
docIndex=0): Answer[]{
// Model uses the closed interval [start, end] for indices.
conststartIndexes=this.getBestIndex(startLogits);
constendIndexes=this.getBestIndex(endLogits);
constorigResults: AnswerIndex[]=[];
startIndexes.forEach(start=>{
endIndexes.forEach(end=>{
if(tokenToOrigMap[start+OUTPUT_OFFSET]&&tokenToOrigMap[end+OUTPUT_OFFSET]&&end>=start){
constlength=end-start+1;
if(length<MAX_ANSWER_LEN){
origResults.push(
{start, end,score: startLogits[start]+endLogits[end]});
}
}
});
});

origResults.sort((a,b)=>b.score-a.score);
constanswers: Answer[]=[];
for(leti=0;i<origResults.length;i++){
if(i>=PREDICT_ANSWER_NUM||
origResults[i].score<NO_ANSWER_THRESHOLD){
break;
}

letconvertedText='';
letstartIndex=0;
letendIndex=0;
if(origResults[i].start>0){
[convertedText,startIndex,endIndex]=this.convertBack(
origTokens,tokenToOrigMap,origResults[i].start,
origResults[i].end,context);
}else{
convertedText='';
}
answers.push({
text: convertedText,
score: origResults[i].score,
 startIndex,
 endIndex
});
}
returnanswers;
}

/** Get the n-best logits from a list of all the logits. */
getBestIndex(logits: number[]): number[]{
consttmpList=[];
for(leti=0;i<MAX_SEQ_LEN;i++){
tmpList.push([i,i,logits[i]]);
}
tmpList.sort((a,b)=>b[2]-a[2]);

constindexes=[];
for(leti=0;i<PREDICT_ANSWER_NUM;i++){
indexes.push(tmpList[i][0]);
}

returnindexes;
}

/** Convert the answer back to original text form. */
convertBack(
origTokens: Token[],tokenToOrigMap: {[key: string]: number},
start: number,end: number,context: string): [string,number,number]{
// Shifted index is: index of logits + offset.
constshiftedStart=start+OUTPUT_OFFSET;
constshiftedEnd=end+OUTPUT_OFFSET;
conststartIndex=tokenToOrigMap[shiftedStart];
constendIndex=tokenToOrigMap[shiftedEnd];
conststartCharIndex=origTokens[startIndex].index;

constendCharIndex=endIndex<origTokens.length-1 ?
origTokens[endIndex+1].index-1 :
origTokens[endIndex].index+origTokens[endIndex].text.length;

return[
context.slice(startCharIndex,endCharIndex+1).trim(),startCharIndex,
endCharIndex
];
}
}

exportasyncfunctionload(modelConfig?: ModelConfig):
Promise<QuestionAndAnswer>{
constmobileBert=newQuestionAndAnswerImpl(modelConfig);
awaitmobileBert.load();
returnmobileBert;
}