From 650a22aef11cbe1e581e5e8be3dd642092211407 Mon Sep 17 00:00:00 2001
From: mudler <mudler@c3os.io>
Date: Wed, 29 Mar 2023 18:53:24 +0200
Subject: [PATCH] Add compatibility to gpt4all models

---
 go.mod  |  2 +-
 go.sum  |  2 ++
 main.go | 14 +++++++++++++-
 3 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/go.mod b/go.mod
index 2a9b38a..ff09edf 100644
--- a/go.mod
+++ b/go.mod
@@ -6,7 +6,7 @@ require (
 	github.com/charmbracelet/bubbles v0.15.0
 	github.com/charmbracelet/bubbletea v0.23.2
 	github.com/charmbracelet/lipgloss v0.7.1
-	github.com/go-skynet/llama v0.0.0-20230325223742-a3563a2690ba
+	github.com/go-skynet/llama v0.0.0-20230329165201-84efc8db3647
 	github.com/gofiber/fiber/v2 v2.42.0
 	github.com/urfave/cli/v2 v2.25.0
 )
diff --git a/go.sum b/go.sum
index d4da439..28be5ec 100644
--- a/go.sum
+++ b/go.sum
@@ -23,6 +23,8 @@ github.com/go-skynet/llama v0.0.0-20230321172246-7be5326e18cc h1:NcmO8mA7iRZIX0Q
 github.com/go-skynet/llama v0.0.0-20230321172246-7be5326e18cc/go.mod h1:ZtYsAIud4cvP9VTTI9uhdgR1uCwaO/gGKnZZ95h9i7w=
 github.com/go-skynet/llama v0.0.0-20230325223742-a3563a2690ba h1:u6OhAqlWFHsTjfWKePdK2kP4/mTyXX5vsmKwrK5QX6o=
 github.com/go-skynet/llama v0.0.0-20230325223742-a3563a2690ba/go.mod h1:ZtYsAIud4cvP9VTTI9uhdgR1uCwaO/gGKnZZ95h9i7w=
+github.com/go-skynet/llama v0.0.0-20230329165201-84efc8db3647 h1:W6qHHD/Bv6wRXSzdv38gWMAXgw3fklHyEblfw88uEUU=
+github.com/go-skynet/llama v0.0.0-20230329165201-84efc8db3647/go.mod h1:ZtYsAIud4cvP9VTTI9uhdgR1uCwaO/gGKnZZ95h9i7w=
 github.com/gofiber/fiber/v2 v2.42.0 h1:Fnp7ybWvS+sjNQsFvkhf4G8OhXswvB6Vee8hM/LyS+8=
 github.com/gofiber/fiber/v2 v2.42.0/go.mod h1:3+SGNjqMh5VQH5Vz2Wdi43zTIV16ktlFd3x3R6O1Zlc=
 github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
diff --git a/main.go b/main.go
index 8ee79a2..fa6b4dc 100644
--- a/main.go
+++ b/main.go
@@ -36,7 +36,9 @@ func llamaFromOptions(ctx *cli.Context) (*llama.LLama, error) {
 	if ctx.Bool("alpaca") {
 		opts = append(opts, llama.EnableAlpaca)
 	}
-
+	if ctx.Bool("gpt4all") {
+		opts = append(opts, llama.EnableGPT4All)
+	}
 	return llama.New(ctx.String("model"), opts...)
 }
 
@@ -95,6 +97,11 @@ var modelFlags = []cli.Flag{
 		EnvVars: []string{"ALPACA"},
 		Value:   true,
 	},
+	&cli.BoolFlag{
+		Name:    "gpt4all",
+		EnvVars: []string{"GPT4ALL"},
+		Value:   false,
+	},
 }
 
 func main() {
@@ -168,6 +175,11 @@ echo "An Alpaca (Vicugna pacos) is a domesticated species of South American came
 						EnvVars: []string{"ALPACA"},
 						Value:   true,
 					},
+					&cli.BoolFlag{
+						Name:    "gpt4all",
+						EnvVars: []string{"GPT4ALL"},
+						Value:   false,
+					},
 					&cli.IntFlag{
 						Name:    "context-size",
 						EnvVars: []string{"CONTEXT_SIZE"},