Merge pull request #1 from httpjamesm:require-stop-string-slice

feat: stop string slices
httpjamesm · Sep 11, 2023 · d6af8c5 · d6af8c5
2 parents 58f867e + 7798aee
commit d6af8c5
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 12 deletions.
diff --git a/README.md b/README.md
@@ -17,7 +17,7 @@ resp, err := client.NewInference(InferenceConfig{
     Model:     "togethercomputer/RedPajama-INCITE-7B-Instruct",
     Prompt:    "The capital of France is",
     MaxTokens: 128,
-    Stop:      &stopString,
+    Stop:      &stopStrings,
 })
 if err != nil {
     panic(err)

diff --git a/inference.go b/inference.go
@@ -6,15 +6,15 @@ import (
 )
 
 type InferenceConfig struct {
-	Model             string   `json:"model"`
-	Prompt            string   `json:"prompt"`
-	MaxTokens         int32    `json:"max_tokens"`
-	Stop              *string  `json:"stop"`
-	Temperature       *float32 `json:"temperature"`
-	TopP              *float32 `json:"top_p"`
-	TopK              *int32   `json:"top_k"`
-	RepetitionPenalty *float32 `json:"repetition_penalty"`
-	LogProbs          *int32   `json:"logprobs"`
+	Model             string    `json:"model"`
+	Prompt            string    `json:"prompt"`
+	MaxTokens         int32     `json:"max_tokens"`
+	Stop              *[]string `json:"stop"`
+	Temperature       *float32  `json:"temperature"`
+	TopP              *float32  `json:"top_p"`
+	TopK              *int32    `json:"top_k"`
+	RepetitionPenalty *float32  `json:"repetition_penalty"`
+	LogProbs          *int32    `json:"logprobs"`
 }
 
 type inferenceRequestBody struct {

diff --git a/inference_test.go b/inference_test.go
@@ -11,13 +11,13 @@ import (
 func TestNewInference(t *testing.T) {
 	client := NewClient(os.Getenv("TOGETHERAI_API_KEY"))
 
-	stopString := "*"
+	stopStrings := []string{"*"}
 
 	respBody, err := client.NewInference(InferenceConfig{
 		Model:     "togethercomputer/RedPajama-INCITE-7B-Instruct",
 		Prompt:    "The capital of France is",
 		MaxTokens: 128,
-		Stop:      &stopString,
+		Stop:      &stopStrings,
 	})
 	assert.NoError(t, err)
 	assert.NotNil(t, respBody)