Added option tests to iOS Language Detector Tests

google-ai-edge · Dec 22, 2023 · 9f27248 · 9f27248
1 parent 7fab4db
commit 9f27248
Showing 1 changed file with 66 additions and 0 deletions.
diff --git a/mediapipe/tasks/ios/test/text/language_detector/MPPLanguageDetectorTests.mm b/mediapipe/tasks/ios/test/text/language_detector/MPPLanguageDetectorTests.mm
@@ -128,6 +128,72 @@ - (void)testClassifyWithL2CModelSucceeds {
       approximatelyEqualsExpectedLanguagePredictions:expectedZhLanguagePredictions];
 }
 
+- (void)testClassifyWithMaxResultsSucceeds {
+  MPPLanguageDetectorOptions *options =
+      [self languageDetectorOptionsWithModelFileInfo:kLanguageDetectorModelFileInfo];
+  options.maxResults = 1;
+  MPPLanguageDetector *languageDetector = [self createLanguageDetectorWithOptionsSucceeds:options];
+
+  NSString *zhText = @"分久必合合久必分";
+  NSArray<MPPLanguagePrediction *> *expectedZhLanguagePredictions = @[
+    [[MPPLanguagePrediction alloc] initWithLanguageCode:@"zh" probability:0.505424f],
+  ];
+
+  [self assertResultsOfDetectLanguageOfText:zhText
+                               usingLanguageDetector:languageDetector
+      approximatelyEqualsExpectedLanguagePredictions:expectedZhLanguagePredictions];
+}
+
+- (void)testClassifyWithScoreThresholdSucceeds {
+  MPPLanguageDetectorOptions *options =
+      [self languageDetectorOptionsWithModelFileInfo:kLanguageDetectorModelFileInfo];
+  options.scoreThreshold = 0.5f;
+  MPPLanguageDetector *languageDetector = [self createLanguageDetectorWithOptionsSucceeds:options];
+
+  NSString *zhText = @"分久必合合久必分";
+  NSArray<MPPLanguagePrediction *> *expectedZhLanguagePredictions = @[
+    [[MPPLanguagePrediction alloc] initWithLanguageCode:@"zh" probability:0.505424f],
+  ];
+
+  [self assertResultsOfDetectLanguageOfText:zhText
+                               usingLanguageDetector:languageDetector
+      approximatelyEqualsExpectedLanguagePredictions:expectedZhLanguagePredictions];
+}
+
+- (void)testClassifyWithCategoryAllowListSucceeds {
+  MPPLanguageDetectorOptions *options =
+      [self languageDetectorOptionsWithModelFileInfo:kLanguageDetectorModelFileInfo];
+  options.categoryAllowlist = @[ @"zh" ];
+
+  MPPLanguageDetector *languageDetector = [self createLanguageDetectorWithOptionsSucceeds:options];
+
+  NSString *zhText = @"分久必合合久必分";
+  NSArray<MPPLanguagePrediction *> *expectedZhLanguagePredictions = @[
+    [[MPPLanguagePrediction alloc] initWithLanguageCode:@"zh" probability:0.505424f],
+  ];
+
+  [self assertResultsOfDetectLanguageOfText:zhText
+                               usingLanguageDetector:languageDetector
+      approximatelyEqualsExpectedLanguagePredictions:expectedZhLanguagePredictions];
+}
+
+- (void)testClassifyWithCategoryDenyListSucceeds {
+  MPPLanguageDetectorOptions *options =
+      [self languageDetectorOptionsWithModelFileInfo:kLanguageDetectorModelFileInfo];
+  options.categoryDenylist = @[ @"zh" ];
+
+  MPPLanguageDetector *languageDetector = [self createLanguageDetectorWithOptionsSucceeds:options];
+
+  NSString *zhText = @"分久必合合久必分";
+  NSArray<MPPLanguagePrediction *> *expectedZhLanguagePredictions = @[
+    [[MPPLanguagePrediction alloc] initWithLanguageCode:@"ja" probability:0.481617f],
+  ];
+
+  [self assertResultsOfDetectLanguageOfText:zhText
+                               usingLanguageDetector:languageDetector
+      approximatelyEqualsExpectedLanguagePredictions:expectedZhLanguagePredictions];
+}
+
 #pragma mark Assert Segmenter Results
 - (void)assertResultsOfDetectLanguageOfText:(NSString *)text
                              usingLanguageDetector:(MPPLanguageDetector *)languageDetector