Skip to content

Commit c9e2b8f

Browse files
authored
Make initializer and some methods public in GenerateResult and GenerateParameters (#239)
* Make the initializer public for GenerateResult and sampler and processor methods of GenerateParameters.
1 parent 95556b1 commit c9e2b8f

File tree

2 files changed

+31
-11
lines changed

2 files changed

+31
-11
lines changed

Libraries/MLXLMCommon/Evaluate.swift

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,7 @@ public struct GenerateParameters: Sendable {
7979
self.repetitionContextSize = repetitionContextSize
8080
}
8181

82-
func sampler() -> LogitSampler {
82+
public func sampler() -> LogitSampler {
8383
if temperature == 0 {
8484
return ArgMaxSampler()
8585
} else if topP > 0 && topP < 1 {
@@ -89,7 +89,7 @@ public struct GenerateParameters: Sendable {
8989
}
9090
}
9191

92-
func processor() -> LogitProcessor? {
92+
public func processor() -> LogitProcessor? {
9393
if let repetitionPenalty, repetitionContextSize > 0 {
9494
return RepetitionContext(
9595
repetitionPenalty: repetitionPenalty, repetitionContextSize: repetitionContextSize)
@@ -379,6 +379,26 @@ public struct TokenIterator: Sequence, IteratorProtocol {
379379

380380
/// Result of a call to ``generate(input:parameters:context:didGenerate:)``.
381381
public struct GenerateResult: Sendable {
382+
383+
/// Initializes a new `GenerateResult` instance.
384+
///
385+
/// - Parameters:
386+
/// - inputText: The input text used for generation.
387+
/// - tokens: The array of tokens generated.
388+
/// - output: The generated output string.
389+
/// - promptTime: The time taken to prompt the input.
390+
/// - generateTime: The time taken to generate the output.
391+
public init(
392+
inputText: LMInput.Text, tokens: [Int], output: String, promptTime: TimeInterval,
393+
generateTime: TimeInterval
394+
) {
395+
self.inputText = inputText
396+
self.tokens = tokens
397+
self.output = output
398+
self.promptTime = promptTime
399+
self.generateTime = generateTime
400+
}
401+
382402
/// input (prompt, images, etc.)
383403
public let inputText: LMInput.Text
384404

Package.resolved

Lines changed: 9 additions & 9 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)