Skip to content

Commit cb66b4b

Browse files
authored
Feature/generate completion info (#259)
* Make GenerateCompletionInfo properties public * Added public init to GenerateCompletionInfo
1 parent 046cf15 commit cb66b4b

File tree

1 file changed

+17
-5
lines changed

1 file changed

+17
-5
lines changed

Libraries/MLXLMCommon/Evaluate.swift

Lines changed: 17 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -780,18 +780,18 @@ public func generate(
780780
/// Represents metadata and statistics related to token generation.
781781
///
782782
/// Provides information about the number of tokens processed during both the prompt and generation phases, as well as the time taken for each phase.
783-
public struct GenerateCompletionInfo {
783+
public struct GenerateCompletionInfo: Sendable {
784784
/// The number of tokens included in the input prompt.
785-
let promptTokenCount: Int
785+
public let promptTokenCount: Int
786786

787787
/// The number of tokens generated by the language model.
788-
let generationTokenCount: Int
788+
public let generationTokenCount: Int
789789

790790
/// The time interval (in seconds) taken to process the input prompt.
791-
let promptTime: TimeInterval
791+
public let promptTime: TimeInterval
792792

793793
/// The time interval (in seconds) taken to generate the output tokens.
794-
let generationTime: TimeInterval
794+
public let generationTime: TimeInterval
795795

796796
/// The number of tokens processed per second during the prompt phase.
797797
public var promptTokensPerSecond: Double {
@@ -802,6 +802,18 @@ public struct GenerateCompletionInfo {
802802
public var tokensPerSecond: Double {
803803
Double(generationTokenCount) / generationTime
804804
}
805+
806+
public init(
807+
promptTokenCount: Int,
808+
generationTokenCount: Int,
809+
promptTime: TimeInterval,
810+
generationTime: TimeInterval
811+
) {
812+
self.promptTokenCount = promptTokenCount
813+
self.generationTokenCount = generationTokenCount
814+
self.promptTime = promptTime
815+
self.generationTime = generationTime
816+
}
805817
}
806818

807819
/// Represents the different stages or outputs of the token generation process.

0 commit comments

Comments
 (0)