cws
/
catbase
mirror of https://github.com/velour/catbase.git
1
0
Fork 0
Code Issues Releases Wiki Activity

update gpt2 interface

This commit is contained in:
Chris Sexton 2020-12-07 17:45:01 -05:00
parent 844e350b16
commit eefd5974d2
1 changed files with 33 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

44
plugins/gpt2/gpt2.go
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,24 +54,46 @@ func (p *GPT2Plugin) help(c bot.Connector, kind bot.Kind, message msg.Message, a
const separator = "<|endoftext|>" const separator = "<|endoftext|>"
func (p *GPT2Plugin) getGPTText(input string) (string, error) { func (p *GPT2Plugin) getGPTText(prefix string) (string, error) {
serviceURI := p.c.Get("gpt.service", "") serviceURI := p.c.Get("gpt2.service", "")
if serviceURI == "" { if serviceURI == "" {
return "", fmt.Errorf("cannot contact GPT2 service") return "", fmt.Errorf("cannot contact GPT2 service")
} }
values := url.Values{}
values.Add("text", input) args := struct {
resp, err := http.PostForm(serviceURI, values) Prefix string `json:"prefix"`
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)
} }
resp.Body.Close() defer resp.Body.Close()
txt := p.cleanup(input + string(body)) dec := json.NewDecoder(resp.Body)
return txt, nil output := struct {
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 {