From c01230640008deb5f505eb8d6b2b5438171ae756 Mon Sep 17 00:00:00 2001
From: CaIon <1808837298@qq.com>
Date: Tue, 12 Mar 2024 18:18:47 +0800
Subject: [PATCH] feat: add parameter top_k

---
 dto/text_request.go                  | 1 +
 relay/channel/ollama/relay-ollama.go | 1 +
 2 files changed, 2 insertions(+)

diff --git a/dto/text_request.go b/dto/text_request.go
index 627255a..936660e 100644
--- a/dto/text_request.go
+++ b/dto/text_request.go
@@ -14,6 +14,7 @@ type GeneralOpenAIRequest struct {
 	MaxTokens        uint            `json:"max_tokens,omitempty"`
 	Temperature      float64         `json:"temperature,omitempty"`
 	TopP             float64         `json:"top_p,omitempty"`
+	TopK             int             `json:"top_k,omitempty"`
 	Stop             any             `json:"stop,omitempty"`
 	N                int             `json:"n,omitempty"`
 	Input            any             `json:"input,omitempty"`
diff --git a/relay/channel/ollama/relay-ollama.go b/relay/channel/ollama/relay-ollama.go
index 4ddb848..0662ddd 100644
--- a/relay/channel/ollama/relay-ollama.go
+++ b/relay/channel/ollama/relay-ollama.go
@@ -25,6 +25,7 @@ func requestOpenAI2Ollama(request dto.GeneralOpenAIRequest) *OllamaRequest {
 			Temperature: request.Temperature,
 			Seed:        request.Seed,
 			Topp:        request.TopP,
+			TopK:        request.TopK,
 			Stop:        Stop,
 		},
 	}