jeans/ollama
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

WIP

Browse Source
This commit is contained in:
ParthSareen 2024-12-18 10:49:22 -08:00
parent 11acb85ff3
commit e3dd90102d
1 changed files with 4 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
server/routes_tokenize_test.go
View File

@ -23,9 +23,7 @@ func (ml *mockModelLoader) LoadModel(name string, params llama.ModelParams) (*lo
return ml.LoadModelFn(name, params)
}
return &loadedModel{
model: mockModel{},
}, nil
return nil, nil
}
type mockModel struct {
@ -34,24 +32,18 @@ type mockModel struct {
TokenToPieceFn func(token int) string
}
func (m *mockModel) Tokenize(text string, addBos bool, addEos bool) ([]int, error) {
func (mockModel) Tokenize(text string, addBos bool, addEos bool) ([]int, error) {
return []int{1, 2, 3}, nil
}
func (m *mockModel) TokenToPiece(token int) string {
func (mockModel) TokenToPiece(token int) string {
return fmt.Sprint(token)
}
func TestTokenizeHandler(t *testing.T) {
gin.SetMode(gin.TestMode)
mockLoader := mockModelLoader{
LoadModelFn: func(name string, params llama.ModelParams) (*loadedModel, error) {
return &loadedModel{
model: mockModel{},
}, nil
},
}
mockModel := mockModel{}
s := Server{
sched: &Scheduler{