FlaskAI/pkg/model/loader.go

package model

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"text/template"

	gptj "github.com/go-skynet/go-gpt4all-j.cpp"
	llama "github.com/go-skynet/go-llama.cpp"
)

type ModelLoader struct {
	modelPath        string
	mu               sync.Mutex
	models           map[string]*llama.LLama
	gptmodels        map[string]*gptj.GPTJ
	promptsTemplates map[string]*template.Template
}

func NewModelLoader(modelPath string) *ModelLoader {
	return &ModelLoader{modelPath: modelPath, gptmodels: make(map[string]*gptj.GPTJ), models: make(map[string]*llama.LLama), promptsTemplates: make(map[string]*template.Template)}
}

func (ml *ModelLoader) ListModels() ([]string, error) {
	files, err := ioutil.ReadDir(ml.modelPath)
	if err != nil {
		return []string{}, err
	}

	models := []string{}
	for _, file := range files {
		if strings.HasSuffix(file.Name(), ".bin") {
			models = append(models, strings.TrimRight(file.Name(), ".bin"))
		}
	}

	return models, nil
}

func (ml *ModelLoader) TemplatePrefix(modelName string, in interface{}) (string, error) {
	ml.mu.Lock()
	defer ml.mu.Unlock()

	m, ok := ml.promptsTemplates[modelName]
	if !ok {
		// try to find a s.bin
		modelBin := fmt.Sprintf("%s.bin", modelName)
		m, ok = ml.promptsTemplates[modelBin]
		if !ok {
			return "", fmt.Errorf("no prompt template available")
		}
	}

	var buf bytes.Buffer

	if err := m.Execute(&buf, in); err != nil {
		return "", err
	}
	return buf.String(), nil
}

func (ml *ModelLoader) loadTemplate(modelName, modelFile string) error {
	modelTemplateFile := fmt.Sprintf("%s.tmpl", modelFile)

	// Check if the model path exists
	if _, err := os.Stat(modelTemplateFile); err != nil {
		return nil
	}

	dat, err := os.ReadFile(modelTemplateFile)
	if err != nil {
		return err
	}

	// Parse the template
	tmpl, err := template.New("prompt").Parse(string(dat))
	if err != nil {
		return err
	}
	ml.promptsTemplates[modelName] = tmpl

	return nil
}

func (ml *ModelLoader) LoadGPTJModel(modelName string) (*gptj.GPTJ, error) {
	ml.mu.Lock()
	defer ml.mu.Unlock()

	// Check if we already have a loaded model
	modelFile := filepath.Join(ml.modelPath, modelName)

	if m, ok := ml.gptmodels[modelFile]; ok {
		return m, nil
	}

	// Check if the model path exists
	if _, err := os.Stat(modelFile); os.IsNotExist(err) {
		// try to find a s.bin
		modelBin := fmt.Sprintf("%s.bin", modelFile)
		if _, err := os.Stat(modelBin); os.IsNotExist(err) {
			return nil, err
		} else {
			modelName = fmt.Sprintf("%s.bin", modelName)
			modelFile = modelBin
		}
	}

	// Load the model and keep it in memory for later use
	model, err := gptj.New(modelFile)
	if err != nil {
		return nil, err
	}

	// If there is a prompt template, load it
	if err := ml.loadTemplate(modelName, modelFile); err != nil {
		return nil, err
	}

	ml.gptmodels[modelFile] = model
	return model, err
}

func (ml *ModelLoader) LoadLLaMAModel(modelName string, opts ...llama.ModelOption) (*llama.LLama, error) {
	ml.mu.Lock()
	defer ml.mu.Unlock()

	// Check if we already have a loaded model
	modelFile := filepath.Join(ml.modelPath, modelName)
	if m, ok := ml.models[modelFile]; ok {
		return m, nil
	}
	// TODO: This needs refactoring, it's really bad to have it in here
	// Check if we have a GPTJ model loaded instead
	if _, ok := ml.gptmodels[modelFile]; ok {
		return nil, fmt.Errorf("this model is a GPTJ one")
	}

	// Check if the model path exists
	if _, err := os.Stat(modelFile); os.IsNotExist(err) {
		// try to find a s.bin
		modelBin := fmt.Sprintf("%s.bin", modelFile)
		if _, err := os.Stat(modelBin); os.IsNotExist(err) {
			return nil, err
		} else {
			modelName = fmt.Sprintf("%s.bin", modelName)
			modelFile = modelBin
		}
	}

	// Load the model and keep it in memory for later use
	model, err := llama.New(modelFile, opts...)
	if err != nil {
		return nil, err
	}

	// If there is a prompt template, load it
	if err := ml.loadTemplate(modelName, modelFile); err != nil {
		return nil, err
	}

	ml.models[modelFile] = model
	return model, err
}