zhangshuai/mlx-swift-examples
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore(llm-tool): add the top_p option in the llm-tool (#41)

Browse Source 
* chore: add top p option in llm-tool
* chore: wire up the top p with async generate
This commit is contained in:
Anchen
2024-04-04 01:54:54 +11:00
committed by GitHub
parent b3eb428c60
commit 2d0fdfe3a9
4 changed files with 13 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
Tools/llm-tool/LLMTool.swift
View File

@@ -29,6 +29,9 @@ struct LLMArguments: ParsableArguments {
@Option(name: .shortAndLong, help: "The sampling temperature")
var temperature: Float = 0.6
@Option(name: .shortAndLong, help: "The top p sampling")
var topP: Float = 0.9
@Option(name: .long, help: "The PRNG seed")
var seed: UInt64 = 0
@@ -127,7 +130,7 @@ struct SyncGenerator: AsyncParsableCommand {
var printed = 0
for token in TokenIterator(
prompt: MLXArray(promptTokens), model: model, temp: args.temperature)
prompt: MLXArray(promptTokens), model: model, temp: args.temperature, topP: args.topP)
{
if tokens.isEmpty {
eval(token)
@@ -205,7 +208,7 @@ struct AsyncGenerator: AsyncParsableCommand {
var printed = 0
let (task, channel) = generate(
prompt: MLXArray(promptTokens), model: model, temp: args.temperature)
prompt: MLXArray(promptTokens), model: model, temp: args.temperature, topP: args.topP)
for await token in channel {
if tokens.isEmpty {