1 /* Copyright 2020 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 #import <Foundation/Foundation.h> 16 17 NS_ASSUME_NONNULL_BEGIN 18 /** 19 * Struct to represent the logit and offset of the answer related to context. 20 */ 21 struct TFLPos { 22 int start; 23 int end; 24 float logit; 25 }; 26 27 /** 28 * Class for the Answer to BertQuestionAnswerer. 29 */ 30 @interface TFLQAAnswer : NSObject 31 @property(nonatomic) struct TFLPos pos; 32 @property(nonatomic) NSString* text; 33 @end 34 35 /** 36 * BertQA task API, performs tokenization for models (BERT, Albert, etc.) in 37 * preprocess and returns most possible answers. 38 * 39 * In particular, the branch of BERT models use WordPiece tokenizer, and the 40 * branch of Albert models use SentencePiece tokenizer, respectively. 41 */ 42 @interface TFLBertQuestionAnswerer : NSObject 43 44 /** 45 * Creates a BertQuestionAnswerer instance with an albert model or mobilebert 46 * model. The API expects a Bert based TFLite model with metadata containing 47 * the following information: 48 * input_process_units: for Wordpiece/Sentencepiece Tokenizer 49 * 3 input tensors with names "ids", "mask" and "segment_ids" 50 * 2 output tensors with names "end_logits" and "start_logits" 51 * Sample models: 52 * https://tfhub.dev/tensorflow/lite-model/albert_lite_base/squadv1/1 53 * https://tfhub.dev/tensorflow/lite-model/mobilebert/1/default/1 54 * @param modelPath The file path to the tflite model. 55 * @return A BertQuestionAnswerer instance. 56 */ 57 + (instancetype)questionAnswererWithModelPath:(NSString *)modelPath 58 NS_SWIFT_NAME(questionAnswerer(modelPath:)); 59 60 /** 61 * Answers question based on the context. Could be empty if no answer was found 62 * from the given context. 63 * 64 * @param context Context the question bases on. 65 * @param question Question to ask. 66 * 67 * @return A list of answers to the question, reversely sorted by the 68 * probability of each answer. 69 */ 70 - (NSArray<TFLQAAnswer*>*)answerWithContext:(NSString*)context 71 question:(NSString*)question 72 NS_SWIFT_NAME(answer(context:question:)); 73 @end 74 NS_ASSUME_NONNULL_END 75