xref: /aosp_15_r20/external/executorch/extension/benchmark/apple/Benchmark/Tests/LLaMA/LLaMATests.mm (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a) 
1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9#import "ResourceTestCase.h"
10
11#import <executorch/examples/models/llama/runner/runner.h>
12
13using namespace ::executorch::extension;
14using namespace ::executorch::runtime;
15
16@interface TokensPerSecondMetric : NSObject<XCTMetric>
17
18@property(nonatomic, assign) NSUInteger tokenCount;
19
20@end
21
22@implementation TokensPerSecondMetric
23
24- (id)copyWithZone:(NSZone *)zone {
25  TokensPerSecondMetric *copy = [[[self class] allocWithZone:zone] init];
26  copy.tokenCount = self.tokenCount;
27  return copy;
28}
29
30- (NSArray<XCTPerformanceMeasurement *> *)
31    reportMeasurementsFromStartTime:
32        (XCTPerformanceMeasurementTimestamp *)startTime
33                          toEndTime:
34                              (XCTPerformanceMeasurementTimestamp *)endTime
35                              error:(NSError **)error {
36  double elapsedTime =
37      (endTime.absoluteTimeNanoSeconds - startTime.absoluteTimeNanoSeconds) /
38      (double)NSEC_PER_SEC;
39  return @[ [[XCTPerformanceMeasurement alloc]
40      initWithIdentifier:NSStringFromClass([self class])
41             displayName:@"Tokens Per Second"
42             doubleValue:(self.tokenCount / elapsedTime)
43              unitSymbol:@"t/s"] ];
44}
45
46@end
47
48@interface LLaMATests : ResourceTestCase
49@end
50
51@implementation LLaMATests
52
53+ (NSArray<NSString *> *)directories {
54  return @[
55    @"Resources",
56    @"aatp/data", // AWS Farm devices look for resources here.
57  ];
58}
59
60+ (NSDictionary<NSString *, BOOL (^)(NSString *)> *)predicates {
61  return @{
62    @"model" : ^BOOL(NSString *filename){
63      return [filename hasSuffix:@".pte"] && [filename containsString:@"llama"];
64    },
65    @"tokenizer" : ^BOOL(NSString *filename) {
66      return [filename isEqual:@"tokenizer.bin"];
67    },
68  };
69}
70
71+ (NSDictionary<NSString *, void (^)(XCTestCase *)> *)dynamicTestsForResources:
72    (NSDictionary<NSString *, NSString *> *)resources {
73  NSString *modelPath = resources[@"model"];
74  NSString *tokenizerPath = resources[@"tokenizer"];
75  return @{
76    @"generate" : ^(XCTestCase *testCase){
77      auto __block runner = std::make_unique<example::Runner>(
78          modelPath.UTF8String, tokenizerPath.UTF8String);
79      const auto status = runner->load();
80      if (status != Error::Ok) {
81        XCTFail("Load failed with error %i", status);
82        return;
83      }
84      TokensPerSecondMetric *tokensPerSecondMetric = [TokensPerSecondMetric new];
85      [testCase measureWithMetrics:@[ tokensPerSecondMetric, [XCTMemoryMetric new] ]
86                            block:^{
87                              tokensPerSecondMetric.tokenCount = 0;
88                              const auto status = runner->generate(
89                                  "Once upon a time",
90                                  128,
91                                  [=](const std::string &token) {
92                                    tokensPerSecondMetric.tokenCount++;
93                                  },
94                                  nullptr,
95                                  false);
96                              XCTAssertEqual(status, Error::Ok);
97                            }];
98    },
99  };
100}
101
102@end
103