in SwiftChat/ContentView.swift [68:116]
func run() {
guard let languageModel = languageModel else { return }
@Sendable func showOutput(currentGeneration: String, progress: Double, completedTokensPerSecond: Double? = nil) {
Task { @MainActor in
// Temporary hack to remove start token returned by llama tokenizers
var response = currentGeneration.deletingPrefix("<s> ")
// Strip prompt
guard response.count > prompt.count else { return }
response = response[prompt.endIndex...].replacingOccurrences(of: "\\n", with: "\n")
// Format prompt + response with different colors
var styledPrompt = AttributedString(prompt)
styledPrompt.foregroundColor = .black
var styledOutput = AttributedString(response)
styledOutput.foregroundColor = .accentColor
outputText = styledPrompt + styledOutput
if let tps = completedTokensPerSecond {
status = .ready(tps)
} else {
status = .generating(progress)
}
}
}
Task.init {
status = .generating(0)
var tokensReceived = 0
let begin = Date()
do {
let output = try await languageModel.generate(config: config, prompt: prompt) { inProgressGeneration in
tokensReceived += 1
showOutput(currentGeneration: inProgressGeneration, progress: Double(tokensReceived)/Double(config.maxNewTokens))
}
let completionTime = Date().timeIntervalSince(begin)
let tokensPerSecond = Double(tokensReceived) / completionTime
showOutput(currentGeneration: output, progress: 1, completedTokensPerSecond: tokensPerSecond)
print("Took \(completionTime)")
} catch {
print("Error \(error)")
Task { @MainActor in
status = .failed("\(error)")
}
}
}
}