From 631cdcef175ad7cb38ab31210cef8c8a66b9e9b2 Mon Sep 17 00:00:00 2001
From: Etienne Perot <eperot@google.com>
Date: Thu, 15 Feb 2024 13:07:20 -0800
Subject: [PATCH] `textgen_test`: Use `instruct` code generation model.

`codellama:7b` is the bare completion model, and works better in code
completion contexts. Here, the test uses it as an instruction-following
model ("Do this."), so the `instruct` variant is the more appropriate one.

PiperOrigin-RevId: 607430293
---
 images/gpu/ollama/Dockerfile | 2 +-
 test/gpu/textgen_test.go     | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/images/gpu/ollama/Dockerfile b/images/gpu/ollama/Dockerfile
index 7b6a22fbe..da49e90b3 100644
--- a/images/gpu/ollama/Dockerfile
+++ b/images/gpu/ollama/Dockerfile
@@ -10,7 +10,7 @@ ENV PATH=$PATH:/usr/local/nvidia/bin:/bin/nvidia/bin
 RUN bash -c '                                  \
     ( ollama serve ) & serverpid="$!";         \
     sleep 5;                                   \
-    ollama pull codellama:7b &&                \
+    ollama pull codellama:7b-instruct &&       \
     ollama pull llama2-chinese:7b-chat &&      \
     kill "$serverpid" &&                       \
     wait "$serverpid"                          \
diff --git a/test/gpu/textgen_test.go b/test/gpu/textgen_test.go
index f25483639..22324ea7a 100644
--- a/test/gpu/textgen_test.go
+++ b/test/gpu/textgen_test.go
@@ -128,7 +128,7 @@ func TestLLM(t *testing.T) {
 		)
 		promptCtx, promptCancel := context.WithTimeout(ctx, 3*time.Minute)
 		prompt := ollama.Prompt{
-			Model: ollama.ZeroTemperatureModel("codellama:7b"),
+			Model: ollama.ZeroTemperatureModel("codellama:7b-instruct"),
 			Query: fmt.Sprintf(`
 				Generate a Python function that takes a string and verifies that it
 				is a valid Chinese translation of the English phrase "Hello World".