Initial Batch Embedding

This commit is contained in:
Roy Han
2024-06-18 17:34:36 -07:00
parent 26d0bf9236
commit c22d54895a
4 changed files with 93 additions and 10 deletions

View File

@@ -608,7 +608,7 @@ type mockLlm struct {
pingResp error
waitResp error
completionResp error
embeddingResp []float64
embeddingResp [][]float64
embeddingRespErr error
tokenizeResp []int
tokenizeRespErr error
@@ -626,7 +626,8 @@ func (s *mockLlm) WaitUntilRunning(ctx context.Context) error { return s.waitRes
func (s *mockLlm) Completion(ctx context.Context, req llm.CompletionRequest, fn func(llm.CompletionResponse)) error {
return s.completionResp
}
func (s *mockLlm) Embedding(ctx context.Context, prompt string) ([]float64, error) {
func (s *mockLlm) Embedding(ctx context.Context, prompts interface{}) ([][]float64, error) {
return s.embeddingResp, s.embeddingRespErr
}
func (s *mockLlm) Tokenize(ctx context.Context, content string) ([]int, error) {