Compare commits

..

No commits in common. "3b2b0364c10720a2ee13ba67708d8bf10f6d7247" and "844e350b16c21360ecc59c385ef5b1798d423791" have entirely different histories.

1 changed files with 11 additions and 33 deletions

View File

@ -1,10 +1,10 @@
package gpt2 package gpt2
import ( import (
"bytes"
"encoding/json"
"fmt" "fmt"
"io/ioutil"
"net/http" "net/http"
"net/url"
"strings" "strings"
"github.com/velour/catbase/bot" "github.com/velour/catbase/bot"
@ -54,46 +54,24 @@ func (p *GPT2Plugin) help(c bot.Connector, kind bot.Kind, message msg.Message, a
const separator = "<|endoftext|>" const separator = "<|endoftext|>"
func (p *GPT2Plugin) getGPTText(prefix string) (string, error) { func (p *GPT2Plugin) getGPTText(input string) (string, error) {
serviceURI := p.c.Get("gpt2.service", "") serviceURI := p.c.Get("gpt.service", "")
if serviceURI == "" { if serviceURI == "" {
return "", fmt.Errorf("cannot contact GPT2 service") return "", fmt.Errorf("cannot contact GPT2 service")
} }
values := url.Values{}
args := struct { values.Add("text", input)
Prefix string `json:"prefix"` resp, err := http.PostForm(serviceURI, values)
Length int `json:"length"`
Temperature float64 `json:"temperature"`
TopP float64 `json:"top_p"`
TopK float64 `json:"top_k"`
}{
Prefix: prefix,
Length: p.c.GetInt("gpt2.length", 50),
Temperature: p.c.GetFloat64("gpt2.temperature", 0.7),
TopK: p.c.GetFloat64("gpt2.topk", 0),
TopP: p.c.GetFloat64("gpt2.topp", 0),
}
values, _ := json.Marshal(args)
resp, err := http.Post(serviceURI, "application/json", bytes.NewBuffer(values))
if err != nil { if err != nil {
return "", fmt.Errorf("error retrieving GPT2 response: %s", err) return "", fmt.Errorf("error retrieving GPT2 response: %s", err)
} }
body, err := ioutil.ReadAll(resp.Body)
if err != nil { if err != nil {
return "", fmt.Errorf("error reading GPT2 response: %s", err) return "", fmt.Errorf("error reading GPT2 response: %s", err)
} }
defer resp.Body.Close() resp.Body.Close()
dec := json.NewDecoder(resp.Body) txt := p.cleanup(input + string(body))
output := struct { return txt, nil
Text string `json:"text"`
}{}
err = dec.Decode(&output)
if err != nil {
return "", err
}
return p.cleanup(output.Text), nil
} }
func (p *GPT2Plugin) cleanup(txt string) string { func (p *GPT2Plugin) cleanup(txt string) string {