You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
75 lines
1.5 KiB
75 lines
1.5 KiB
package client
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
)
|
|
|
|
type Prediction struct {
|
|
Prediction string `json:"prediction"`
|
|
}
|
|
|
|
type Client struct {
|
|
baseURL string
|
|
client *http.Client
|
|
endpoint string
|
|
}
|
|
|
|
func NewClient(baseURL string) *Client {
|
|
return &Client{
|
|
baseURL: baseURL,
|
|
client: &http.Client{},
|
|
endpoint: "/predict",
|
|
}
|
|
}
|
|
|
|
type InputData struct {
|
|
Text string `json:"text"`
|
|
TopP float64 `json:"topP,omitempty"`
|
|
TopK int `json:"topK,omitempty"`
|
|
Temperature float64 `json:"temperature,omitempty"`
|
|
Tokens int `json:"tokens,omitempty"`
|
|
}
|
|
|
|
func (c *Client) Predict(text string, opts ...InputOption) (string, error) {
|
|
input := NewInputData(opts...)
|
|
input.Text = text
|
|
|
|
// encode input data to JSON format
|
|
inputBytes, err := json.Marshal(input)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
// create HTTP request
|
|
url := c.baseURL + c.endpoint
|
|
req, err := http.NewRequest("POST", url, bytes.NewBuffer(inputBytes))
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
// set request headers
|
|
req.Header.Set("Content-Type", "application/json")
|
|
|
|
// send request and get response
|
|
resp, err := c.client.Do(req)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
return "", fmt.Errorf("request failed with status %d", resp.StatusCode)
|
|
}
|
|
|
|
// decode response body to Prediction struct
|
|
var prediction Prediction
|
|
err = json.NewDecoder(resp.Body).Decode(&prediction)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
return prediction.Prediction, nil
|
|
}
|
|
|