Fix RepititionPenalty from int to float (#7)

* Fix `RepititionPenalty` from int to float --------- Co-authored-by: Robby <h0rv@users.noreply.github.com>
IBM · Jan 28, 2024 · 1aa4aaf · 1aa4aaf
1 parent 62936df
commit 1aa4aaf
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/models/generate_option.go b/models/generate_option.go
@@ -28,7 +28,7 @@ type GenerateOptions struct {
 	TopP                *float64       `json:"top_p,omitempty"`
 	TopK                *uint          `json:"top_k,omitempty"`
 	RandomSeed          *uint          `json:"random_seed,omitempty"`
-	RepetitionPenalty   *uint          `json:"repetition_penalty,omitempty"`
+	RepetitionPenalty   *float64       `json:"repetition_penalty,omitempty"`
 	MinNewTokens        *uint          `json:"min_new_tokens,omitempty"`
 	MaxNewTokens        *uint          `json:"max_new_tokens,omitempty"`
 	StopSequences       *[]string      `json:"stop_sequences,omitempty"`
@@ -73,7 +73,7 @@ func WithRandomSeed(randomSeed uint) GenerateOption {
 	}
 }
 
-func WithRepetitionPenalty(repetitionPenalty uint) GenerateOption {
+func WithRepetitionPenalty(repetitionPenalty float64) GenerateOption {
 	return func(opts *GenerateOptions) {
 		opts.RepetitionPenalty = &repetitionPenalty
 	}