mirror of https://github.com/velour/catbase.git
Compare commits
2 Commits
844e350b16
...
3b2b0364c1
Author | SHA1 | Date |
---|---|---|
Chris Sexton | 3b2b0364c1 | |
Chris Sexton | eefd5974d2 |
|
@ -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 {
|
||||||
|
|
Loading…
Reference in New Issue