atridad/himbot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update Himbot's response format and add code command

Browse Source
This commit is contained in:
Atridad Lahiji
2024-02-05 00:17:25 -07:00
parent c4cd45becb
commit 08f1b82a65
4 changed files with 127 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
lib/replicate.go
View File

@ -13,6 +13,8 @@ import (
"github.com/replicate/replicate-go"
)
var PromptPrefix = "Ready for a dose of sarcasm and wit? Himbot, your Discord assistant, is up for the challenge. Hit it with the prompt:"
func ReplicateTextGeneration(prompt string) (string, error) {
client, clientError := replicate.NewClient(replicate.WithTokenFromEnv())
if clientError != nil {
@ -20,7 +22,7 @@ func ReplicateTextGeneration(prompt string) (string, error) {
}
input := replicate.PredictionInput{
"prompt": "Respond to the following prompt as the helpful but sarcastic and witty discord assistant called Himbot: " + prompt,
"prompt": PromptPrefix + prompt,
"max_new_tokens": 4096,
"prompt_template": "<s>[INST] {prompt} [/INST]",
}
@ -59,6 +61,51 @@ func ReplicateTextGeneration(prompt string) (string, error) {
return result, nil
}
func ReplicateCodeGeneration(prompt string) (string, error) {
client, clientError := replicate.NewClient(replicate.WithTokenFromEnv())
if clientError != nil {
return "", clientError
}
input := replicate.PredictionInput{
"prompt": PromptPrefix + prompt,
"max_new_tokens": 4096,
}
webhook := replicate.Webhook{
URL: "https://example.com/webhook",
Events: []replicate.WebhookEventType{"start", "completed"},
}
prediction, predictionError := client.Run(context.Background(), "meta/codellama-70b-instruct:a279116fe47a0f65701a8817188601e2fe8f4b9e04a518789655ea7b995851bf", input, &webhook)
if predictionError != nil {
return "", predictionError
}
if prediction == nil {
return "", errors.New("there was an error generating a response based on this prompt... please reach out to @himbothyswaggins to fix this issue")
}
test, ok := prediction.([]interface{})
if !ok {
return "", errors.New("there was an error generating a response based on this prompt... please reach out to @himbothyswaggins to fix this issue")
}
strs := make([]string, len(test))
for i, v := range test {
str, ok := v.(string)
if !ok {
return "", errors.New("element is not a string")
}
strs[i] = str
}
result := strings.Join(strs, "")
return result, nil
}
func ReplicateImageGeneration(prompt string, filename string) (*bytes.Buffer, error) {
client, clientError := replicate.NewClient(replicate.WithTokenFromEnv())
if clientError != nil {