Various improvements in server

server/conf/conf.go

@@ -17,6 +17,7 @@ func InitConf() types.Conf {
 	viper.AddConfigPath(".")
 	viper.SetDefault("origins", []string{"localhost"})
 	viper.SetDefault("cmd_api_key", nil)
+	viper.SetDefault("models", map[string]string{})
 	viper.SetDefault("features", []string{})
 	err := viper.ReadInConfig() // Find and read the config file
 	if err != nil {             // Handle errors reading the config file
@@ -25,12 +26,14 @@ func InitConf() types.Conf {
 	or := viper.GetStringSlice("origins")
 	ak := viper.GetString("api_key")
 	cmdak := viper.GetString("cmd_api_key")
+	models := viper.GetStringMapString("models")
 	ft := viper.GetStringSlice("features")
 	return types.Conf{
 		Origins:   or,
 		ApiKey:    ak,
 		CmdApiKey: &cmdak,
 		Features:  ft,
+		Models:    models,
 	}
 }
 

server/httpserver/tasks.go

@@ -43,7 +43,11 @@ func ExecuteTaskHandler(c echo.Context) error {
 		}
 		return c.NoContent(http.StatusBadRequest)
 	}
-	ok, task, err := files.ReadTask(tp, state.ModelSize)
+	ms, ok := state.ModelsConf[taskName]
+	if !ok {
+		ms = "default"
+	}
+	ok, task, err := files.ReadTask(tp, ms)
 	if err != nil {
 		log.Println(err)
 		return c.NoContent(http.StatusInternalServerError)
@@ -66,7 +70,7 @@ func ExecuteTaskHandler(c echo.Context) error {
 	select {
 	case res, ok := <-ch:
 		if ok {
-			if state.IsVerbose {
+			if state.IsDebug {
 				fmt.Println("-------- result ----------")
 				for key, value := range res.Data {
 					fmt.Printf("%s: %v\n", key, value)

server/lm/infer.go

@@ -37,19 +37,21 @@ func InferTask(
 			MsgType: types.ErrorMsgType,
 		}
 	}
-	if state.IsVerbose {
+	if state.IsDebug {
 		//fmt.Println("Inference params:")
 		//fmt.Println(params)
 		fmt.Println("---------- prompt ----------")
 		fmt.Println(finalPrompt)
 		fmt.Println("----------------------------")
-		fmt.Println("Thinking ..")
+	}
+	if state.IsVerbose {
+		fmt.Println("Ingesting prompt ..")
 	}
 	if state.IsDebug {
 		fmt.Println("Inference params:")
-		fmt.Printf("%+v\n\n", task.InferParams)
+		fmt.Printf("%+v\n", task.InferParams)
 	}
-	fmt.Println("Stop", stop)
+	//fmt.Println("Stop", stop)
 	for _, s := range stop {
 		st, ok := task.InferParams["stop"]
 		if st == nil || !ok {

server/lm/ollama.go

@@ -49,7 +49,7 @@ func ollamaInfer(
 			thinkingElapsed = time.Since(startThinking)
 			if state.IsVerbose {
 				fmt.Println("Thinking time:", thinkingElapsed)
-				fmt.Println("Emitting ..")
+				fmt.Println("Emitting ..\n")
 			}
 			smsg := types.StreamedMessage{
 				Num:     ntokens,
@@ -99,7 +99,7 @@ func ollamaInfer(
 		// the inference was not aborted
 		emittingElapsed := time.Since(startEmitting)
 		if state.IsVerbose {
-			fmt.Println("Emitting time:", emittingElapsed)
+			fmt.Println("\n\nEmitting time:", emittingElapsed)
 		}
 		tpsRaw := float64(ntokens) / emittingElapsed.Seconds()
 		s := fmt.Sprintf("%.2f", tpsRaw)

server/main.go

@@ -29,7 +29,7 @@ func main() {
 	}
 	state.IsVerbose = !*quiet
 	config := conf.InitConf()
-	err := state.Init(config.Features)
+	err := state.Init(config.Features, config.Models)
 	if err != nil {
 		log.Fatal("Error initializing state", err)
 	}

server/state/state.go

@@ -15,11 +15,12 @@ var ContinueInferingController = true
 var IsInfering = false
 
 // models
-var ModelSize = "xsmall"
+var ModelsConf = make((map[string]string))
 
 var Tasks = map[string]string{}
 
-func Init(features []string) error {
+func Init(features []string, models map[string]string) error {
+	ModelsConf = models
 	ts, err := files.InitTasks(features)
 	if err != nil {
 		return err

server/types/types.go

@@ -5,6 +5,7 @@ type Conf struct {
 	ApiKey    string
 	CmdApiKey *string
 	Features  []string
+	Models    map[string]string
 }
 
 // ModelConf represents the configuration for a language model.