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 ml.LoadModelFn(name, params)
} }
return &loadedModel{ return nil, nil
model: mockModel{},
}, nil
} }
type mockModel struct { type mockModel struct {
@ -34,24 +32,18 @@ type mockModel struct {
TokenToPieceFn func(token int) string 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 return []int{1, 2, 3}, nil
} }
func (m *mockModel) TokenToPiece(token int) string { func (mockModel) TokenToPiece(token int) string {
return fmt.Sprint(token) return fmt.Sprint(token)
} }
func TestTokenizeHandler(t *testing.T) { func TestTokenizeHandler(t *testing.T) {
gin.SetMode(gin.TestMode) gin.SetMode(gin.TestMode)
mockLoader := mockModelLoader{ mockModel := mockModel{}
LoadModelFn: func(name string, params llama.ModelParams) (*loadedModel, error) {
return &loadedModel{
model: mockModel{},
}, nil
},
}
s := Server{ s := Server{
sched: &Scheduler{ sched: &Scheduler{