From a5b98e6fda697f410191315ad4d231608b5bc54d Mon Sep 17 00:00:00 2001
From: bel <bel@bel.bel>
Date: Sat, 17 Jun 2023 10:31:35 -0600
Subject: [PATCH] gr

---
 vicuna-tools.d/main.go | 47 +++++++++++++++++++++++++++++++++---------
 1 file changed, 37 insertions(+), 10 deletions(-)

diff --git a/vicuna-tools.d/main.go b/vicuna-tools.d/main.go
index 830d863..004b23d 100644
--- a/vicuna-tools.d/main.go
+++ b/vicuna-tools.d/main.go
@@ -16,6 +16,7 @@ import (
 	"os/exec"
 	"os/signal"
 	"path"
+	"strconv"
 	"strings"
 	"sync"
 	"syscall"
@@ -30,6 +31,7 @@ var (
 			semaphore sync.Mutex
 			WD        string
 			Command   string
+			N         int
 		}
 	}
 
@@ -87,7 +89,8 @@ func config(ctx context.Context) {
 	flag.IntVar(&Config.Port, "p", 37070, "port to listen on")
 	flag.StringVar(&Config.ChatBot.SessionD, "chatbot-session-d", d, "dir to store chat bot sessions")
 	flag.StringVar(&Config.ChatBot.WD, "chatbot-working-d", "./llama.cpp", "working directory for chatbot")
-	flag.StringVar(&Config.ChatBot.Command, "chatbot-cmd", "./main -m ./models/ggml-vic7b-uncensored-q5_1.bin -n 256 --repeat_penalty 1.0", "chatbot cmd prefix")
+	flag.StringVar(&Config.ChatBot.Command, "chatbot-cmd", "./main -m ./models/ggml-vic7b-uncensored-q5_1.bin --repeat_penalty 1.0", "chatbot cmd prefix")
+	flag.IntVar(&Config.ChatBot.N, "chatbot-n", 64, "chatbot items to gen")
 	flag.BoolVar(&Config.Debug, "debug", false, "debug mode")
 	flag.Parse()
 }
@@ -291,6 +294,7 @@ func handleAPIChatBotPut(w http.ResponseWriter, r *http.Request) error {
 
 	promptF := path.Join(sessionD, "prompt.txt")
 	inputF := path.Join(sessionD, "input.txt")
+	cacheF := path.Join(sessionD, "cache.bin")
 	if err := copyFile(inputF, promptF); err != nil {
 		return err
 	}
@@ -299,13 +303,19 @@ func handleAPIChatBotPut(w http.ResponseWriter, r *http.Request) error {
 	}
 
 	commands := strings.Fields(Config.ChatBot.Command)
-	commands = append(commands, "-f", inputF)
+	commands = append(commands,
+		"-f", inputF,
+		"--prompt-cache-all",
+		"--prompt-cache", cacheF,
+		"-n", strconv.Itoa(Config.ChatBot.N),
+	)
 	command := exec.CommandContext(
 		r.Context(),
 		commands[0],
 		commands[1:]...,
 	)
 	command.Dir = Config.ChatBot.WD
+	command.Stderr = os.Stderr
 
 	stdout, err := command.StdoutPipe()
 	if err != nil {
@@ -316,14 +326,6 @@ func handleAPIChatBotPut(w http.ResponseWriter, r *http.Request) error {
 		io.Copy(buff, stdout)
 	}()
 
-	stderr, err := command.StderrPipe()
-	if err != nil {
-		return err
-	}
-	go func() {
-		io.Copy(os.Stderr, stderr)
-	}()
-
 	Config.ChatBot.semaphore.Lock()
 	defer Config.ChatBot.semaphore.Unlock()
 
@@ -333,6 +335,31 @@ func handleAPIChatBotPut(w http.ResponseWriter, r *http.Request) error {
 
 	log.Printf("OUTPUT: %s", buff.Bytes())
 	return errors.New("not impl: move input to prompt")
+
+	/*
+		go func() {
+			Config.ChatBot.semaphore.Lock()
+			defer Config.ChatBot.semaphore.Unlock()
+
+			commands := strings.Fields(Config.ChatBot.Command)
+			commands = append(commands,
+				"-f", promptF,
+				"--batch_size", "8",
+				"--prompt-cache", cacheF,
+				"--file", promptF,
+				"--n_predict", "1",
+			)
+			command := exec.CommandContext(
+				r.Context(),
+				commands[0],
+				commands[1:]...,
+			)
+			command.Dir = Config.ChatBot.WD
+			command.Run()
+		}()
+	*/
+
+	return nil
 }
 
 func copyFile(toF, fromF string) error {