ollama test: Add image processing capability test.

This asks the LLM to OCR the gVisor logo.

Response: `The image you've provided appears to be a logo or branding for a
company named "gVisor." The logo features a stylized astronaut with a space
theme, suggesting that the company may be related to space exploration,
satellite technology, or similar fields.`
PiperOrigin-RevId: 612241427

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"                          \
 '

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"                          \
 '

test/gpu/BUILD

@@ -32,6 +32,7 @@ go_test(
 go_test(
     name = "textgen_test",
     srcs = ["textgen_test.go"],
+    embedsrcs = ["gvisor.png"],
     tags = [
         "local",
         "noguitar",

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,

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())
+	})
 }