diff --git a/images/gpu/ollama/Dockerfile b/images/gpu/ollama/Dockerfile index ff7d72db3..a859f3f46 100644 --- a/images/gpu/ollama/Dockerfile +++ b/images/gpu/ollama/Dockerfile @@ -12,8 +12,9 @@ ENV OLLAMA_HOST=0.0.0.0:11434 RUN bash -c ' \ ( ollama serve ) & serverpid="$!"; \ sleep 5; \ - ollama pull codellama:7b-instruct && \ + ollama pull codellama:7b-instruct && \ ollama pull llama2-chinese:7b-chat && \ + ollama pull llava:7b-v1.6 && \ kill "$serverpid" && \ wait "$serverpid" \ ' diff --git a/images/gpu/ollama/bench/Dockerfile b/images/gpu/ollama/bench/Dockerfile index 8a97fa832..fffa251eb 100644 --- a/images/gpu/ollama/bench/Dockerfile +++ b/images/gpu/ollama/bench/Dockerfile @@ -6,7 +6,7 @@ ENV OLLAMA_ORIGINS=* ENV OLLAMA_HOST=0.0.0.0:11434 # Pre-install models useful for benchmarking. -# These are huge (total ~115 GiB), but necessary to benchmark +# These are huge (total ~120 GiB), but necessary to benchmark # models of various sizes. They are in their own image file to # keep the test-only image lighter by comparison. RUN bash -c ' \ @@ -22,6 +22,7 @@ RUN bash -c ' \ ollama pull gemma:2b-instruct && \ ollama pull gemma:7b-instruct && \ ollama pull llava:7b-v1.6 && \ + ollama pull llava:34b-v1.6 && \ kill "$serverpid" && \ wait "$serverpid" \ ' diff --git a/test/gpu/BUILD b/test/gpu/BUILD index 9caaa2845..4b302ca3e 100644 --- a/test/gpu/BUILD +++ b/test/gpu/BUILD @@ -32,6 +32,7 @@ go_test( go_test( name = "textgen_test", srcs = ["textgen_test.go"], + embedsrcs = ["gvisor.png"], tags = [ "local", "noguitar", diff --git a/test/gpu/gvisor.png b/test/gpu/gvisor.png new file mode 100644 index 000000000..5919dba68 Binary files /dev/null and b/test/gpu/gvisor.png differ diff --git a/test/gpu/ollama/ollama.go b/test/gpu/ollama/ollama.go index fd92d932c..8326105df 100644 --- a/test/gpu/ollama/ollama.go +++ b/test/gpu/ollama/ollama.go @@ -436,11 +436,22 @@ type Prompt struct { // Common leading whitespace will be removed. Query string + // images is a set of attached images. + // Use AddImage to add an image. + images [][]byte + // Context is the conversational context to follow up on, if any. // This is returned from `Response`. Context ConversationContext } +// AddImage attaches an image to the prompt. +// Returns itself for chainability. +func (p *Prompt) AddImage(data []byte) *Prompt { + p.images = append(p.images, data) + return p +} + // CleanQuery removes common whitespace from query lines, and all // leading/ending whitespace-only lines. // It is useful to be able to specify query string as indented strings @@ -529,6 +540,7 @@ func (p *Prompt) WithHotterModel() *Prompt { type PromptJSON struct { Model string `json:"model"` Prompt string `json:"prompt,omitempty"` + Images []string `json:"images"` Stream bool `json:"stream"` Context ConversationContext `json:"context"` Options map[string]any `json:"options"` @@ -541,9 +553,14 @@ func (p *Prompt) json() PromptJSON { if p.KeepModelAlive != 0 { keepAlive = p.KeepModelAlive.String() } + images := make([]string, len(p.images)) + for i, image := range p.images { + images[i] = base64.StdEncoding.EncodeToString(image) + } return PromptJSON{ Model: p.Model.Name, Prompt: p.CleanQuery(), + Images: images, Stream: true, Context: p.Context, Options: p.Model.Options, diff --git a/test/gpu/textgen_test.go b/test/gpu/textgen_test.go index 0ac7cf07a..44768ae00 100644 --- a/test/gpu/textgen_test.go +++ b/test/gpu/textgen_test.go @@ -17,6 +17,7 @@ package textgen_test import ( "context" + _ "embed" "errors" "fmt" "strings" @@ -28,6 +29,9 @@ import ( "gvisor.dev/gvisor/test/gpu/ollama" ) +//go:embed gvisor.png +var gVisorPNG []byte + // extractCode extracts code between two code block markers. func extractCode(response, codeBlockDelim string) (string, error) { if !strings.Contains(response, codeBlockDelim) { @@ -205,4 +209,26 @@ func TestLLM(t *testing.T) { } t.Logf("Translation verification succeeded with code:\n\n%s\n\n", pythonCode) }) + t.Run("ocr", func(t *testing.T) { + const textInImage = "gVisor" + promptCtx, promptCancel := context.WithTimeout(ctx, 3*time.Minute) + prompt := ollama.Prompt{ + Model: ollama.ZeroTemperatureModel("llava:7b-v1.6"), + Query: "What is the text written in this image?", + } + prompt.AddImage(gVisorPNG) + response, err := llm.PromptUntil(promptCtx, &prompt, func(prompt *ollama.Prompt, response *ollama.Response) (*ollama.Prompt, error) { + defer prompt.Model.RaiseTemperature() + text := strings.TrimSpace(response.Text()) + if !strings.Contains(strings.ToLower(text), strings.ToLower(textInImage)) { + return prompt, fmt.Errorf("text does not contain %q: %q", textInImage, text) + } + return prompt, nil + }) + promptCancel() + if err != nil { + t.Fatalf("OCR failed: %v", err) + } + t.Logf("OCR response for gVisor logo: %q", response.Text()) + }) }