feat(LMES): Support custom template and prompt

Expand the LMEvalJob CRD to support custom templates and system prompts.
This is mainly for custom unitxt task recipes. Now, users can use the
`template` and `systemPrompt` fields under the `taskRecipes` to specify
the custom template and system prompt.

Signed-off-by: Yihong Wang <[email protected]>

Author: yhwang

Dockerfile.lmes-job

@@ -20,9 +20,9 @@ RUN curl -L https://github.com/opendatahub-io/lm-evaluation-harness/archive/refs
 
 RUN python -c 'from lm_eval.tasks.unitxt import task; import os.path; print("class: !function " + task.__file__.replace("task.py", "task.Unitxt"))' > ./my_tasks/unitxt
 
-ENV PYTHONPATH=/opt/app-root/src/.local/lib/python3.11/site-packages:/opt/app-root/src/lm-evaluation-harness:/opt/app-root/src:/opt/app-root/src/server
+ENV PYTHONPATH=/opt/app-root/src/.local/lib/python3.11/site-packages:/opt/app-root/src:/opt/app-root/src/server
 ENV HF_HOME=/opt/app-root/src/hf_home
-ENV UNITXT_ARTIFACTORIES=/opt/app-root/src/my_catalogs
+ENV UNITXT_CATALOGS=/opt/app-root/src/my_catalogs
 
 USER 65532:65532
 CMD ["/opt/app-root/bin/python"]

api/lmes/v1alpha1/lmevaljob_types.go

@@ -76,14 +76,38 @@ type Card struct {
 	Custom string `json:"custom,omitempty"`
 }
 
+type Template struct {
+	// Unitxt template ID
+	// +optional
+	Name string `json:"name,omitempty"`
+	// A JSON string for a custom unitxt template.
+	// Use the documentation here: https://www.unitxt.ai/en/latest/docs/adding_template.html
+	// to compose a custom template, store it as a JSON file by calling the
+	// add_to_catalog API: https://www.unitxt.ai/en/latest/docs/saving_and_loading_from_catalog.html#adding-assets-to-the-catalog,
+	// and use the JSON content as the value here.
+	// +optional
+	Custom string `json:"custom,omitempty"`
+}
+
+type SystemPrompt struct {
+	// Unitxt System Prompt id
+	Name string `json:"name,omitempty"`
+	// A custom system prompt string
+	Custom string `json:"custom,omitempty"`
+}
+
 // Use a task recipe to form a custom task. It maps to the Unitxt Recipe
 // Find details of the Unitxt Recipe here:
 // https://www.unitxt.ai/en/latest/unitxt.standard.html#unitxt.standard.StandardRecipe
 type TaskRecipe struct {
 	// The Unitxt dataset card
 	Card Card `json:"card"`
 	// The Unitxt template
-	Template string `json:"template"`
+	// +optional
+	Template *Template `json:"template,omitempty"`
+	// The Unitxt System Prompt
+	// +optional
+	SystemPrompt *SystemPrompt `json:"systemPrompt,omitempty"`
 	// The Unitxt Task
 	// +optional
 	Task *string `json:"task,omitempty"`
@@ -111,9 +135,17 @@ type TaskList struct {
 	TaskRecipes []TaskRecipe `json:"taskRecipes,omitempty"`
 }
 
+// Use the tp_idx and sp_idx to point to the corresponding custom template
+// and custom system_prompt
 func (t *TaskRecipe) String() string {
 	var b strings.Builder
-	b.WriteString(fmt.Sprintf("card=%s,template=%s", t.Card.Name, t.Template))
+	b.WriteString(fmt.Sprintf("card=%s", t.Card.Name))
+	if t.Template != nil && t.Template.Name != "" {
+		b.WriteString(fmt.Sprintf(",template=%s", t.Template.Name))
+	}
+	if t.SystemPrompt != nil && t.SystemPrompt.Name != "" {
+		b.WriteString(fmt.Sprintf(",system_prompt=%s", t.SystemPrompt.Name))
+	}
 	if t.Task != nil {
 		b.WriteString(fmt.Sprintf(",task=%s", *t.Task))
 	}

api/lmes/v1alpha1/zz_generated.deepcopy.go

@@ -50,19 +50,23 @@ func (t *strArrayArg) String() string {
 }
 
 var (
-	taskRecipes  strArrayArg
-	customCards  strArrayArg
-	copy         = flag.String("copy", "", "copy this binary to specified destination path")
-	getStatus    = flag.Bool("get-status", false, "Get current status")
-	shutdown     = flag.Bool("shutdown", false, "Shutdown the driver")
-	outputPath   = flag.String("output-path", OutputPath, "output path")
-	detectDevice = flag.Bool("detect-device", false, "detect available device(s), CUDA or CPU")
-	driverLog    = ctrl.Log.WithName("driver")
+	taskRecipes         strArrayArg
+	customCards         strArrayArg
+	customTemplates     strArrayArg
+	customSystemPrompts strArrayArg
+	copy                = flag.String("copy", "", "copy this binary to specified destination path")
+	getStatus           = flag.Bool("get-status", false, "Get current status")
+	shutdown            = flag.Bool("shutdown", false, "Shutdown the driver")
+	outputPath          = flag.String("output-path", OutputPath, "output path")
+	detectDevice        = flag.Bool("detect-device", false, "detect available device(s), CUDA or CPU")
+	driverLog           = ctrl.Log.WithName("driver")
 )
 
 func init() {
 	flag.Var(&taskRecipes, "task-recipe", "task recipe")
 	flag.Var(&customCards, "custom-card", "A JSON string represents a custom card")
+	flag.Var(&customTemplates, "custom-template", "A JSON string represents a custom template")
+	flag.Var(&customSystemPrompts, "custom-prompt", "A string represents a custom system_prompt")
 }
 
 func main() {
@@ -105,13 +109,15 @@ func main() {
 	}
 
 	driverOpt := driver.DriverOption{
-		Context:      ctx,
-		OutputPath:   *outputPath,
-		DetectDevice: *detectDevice,
-		Logger:       driverLog,
-		TaskRecipes:  taskRecipes,
-		CustomCards:  customCards,
-		Args:         args,
+		Context:            ctx,
+		OutputPath:         *outputPath,
+		DetectDevice:       *detectDevice,
+		Logger:             driverLog,
+		TaskRecipes:        taskRecipes,
+		CustomCards:        customCards,
+		CustomTemplates:    customTemplates,
+		CustomSystemPrompt: customSystemPrompts,
+		Args:               args,
 	}
 
 	driver, err := driver.NewDriver(&driverOpt)

cmd/lmes_driver/main.go

@@ -4632,15 +4632,36 @@ spec:
                         numDemos:
                           description: Number of fewshot
                           type: integer
+                        systemPrompt:
+                          description: The Unitxt System Prompt
+                          properties:
+                            custom:
+                              description: A custom system prompt string
+                              type: string
+                            name:
+                              description: Unitxt System Prompt id
+                              type: string
+                          type: object
                         task:
                           description: The Unitxt Task
                           type: string
                         template:
                           description: The Unitxt template
-                          type: string
+                          properties:
+                            custom:
+                              description: |-
+                                A JSON string for a custom unitxt template.
+                                Use the documentation here: https://www.unitxt.ai/en/latest/docs/adding_template.html
+                                to compose a custom template, store it as a JSON file by calling the
+                                add_to_catalog API: https://www.unitxt.ai/en/latest/docs/saving_and_loading_from_catalog.html#adding-assets-to-the-catalog,
+                                and use the JSON content as the value here.
+                              type: string
+                            name:
+                              description: Unitxt template ID
+                              type: string
+                          type: object
                       required:
                       - card
-                      - template
                       type: object
                     type: array
                 type: object

config/crd/bases/trustyai.opendatahub.io_lmevaljobs.yaml

@@ -20,6 +20,7 @@ import (
 	"bufio"
 	"context"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"io/fs"
@@ -43,26 +44,30 @@ var (
 
 const (
 	// put the domain socket under /tmp. may move to emptydir to share across containers
-	socketPath             = "/tmp/ta-lmes-driver.sock"
-	DefaultTaskRecipesPath = "/opt/app-root/src/my_tasks"
-	DefaultCatalogPath     = "/opt/app-root/src/my_catalogs"
-	TaskRecipePrefix       = "tr"
-	CustomCardPrefix       = "custom"
-	ShutdownURI            = "/Shutdown"
-	GetStatusURI           = "/GetStatus"
+	socketPath               = "/tmp/ta-lmes-driver.sock"
+	DefaultTaskRecipesPath   = "/opt/app-root/src/my_tasks"
+	DefaultCatalogPath       = "/opt/app-root/src/my_catalogs"
+	TaskRecipePrefix         = "tr"
+	CustomCardPrefix         = "custom"
+	CustomTemplatePrefix     = "tp"
+	CustomSystemPromptPrefix = "sp"
+	ShutdownURI              = "/Shutdown"
+	GetStatusURI             = "/GetStatus"
 )
 
 type DriverOption struct {
-	Context         context.Context
-	OutputPath      string
-	DetectDevice    bool
-	TaskRecipesPath string
-	TaskRecipes     []string
-	CatalogPath     string
-	CustomCards     []string
-	Logger          logr.Logger
-	Args            []string
-	SocketPath      string
+	Context            context.Context
+	OutputPath         string
+	DetectDevice       bool
+	TaskRecipesPath    string
+	TaskRecipes        []string
+	CatalogPath        string
+	CustomCards        []string
+	CustomTemplates    []string
+	CustomSystemPrompt []string
+	Logger             logr.Logger
+	Args               []string
+	SocketPath         string
 }
 
 type Driver interface {
@@ -313,9 +318,19 @@ func (d *driverImpl) exec() error {
 		return fmt.Errorf("failed to create task recipes: %v", err)
 	}
 
+	if err := d.prepDir4CustomArtifacts(); err != nil {
+		return fmt.Errorf("failed to create the directories for custom artifacts: %v", err)
+	}
+
 	if err := d.createCustomCards(); err != nil {
 		return fmt.Errorf("failed to create custom cards: %v", err)
 	}
+	if err := d.createCustomTemplates(); err != nil {
+		return fmt.Errorf("failed to create custom templates: %v", err)
+	}
+	if err := d.createCustomSystemPrompts(); err != nil {
+		return fmt.Errorf("failed to create custom system_prompts: %v", err)
+	}
 
 	// Detect available devices if needed
 	if err := d.detectDevice(); err != nil {
@@ -478,6 +493,15 @@ func (d *driverImpl) createTaskRecipes() error {
 	return nil
 }
 
+func (d *driverImpl) prepDir4CustomArtifacts() error {
+	subDirs := []string{"cards", "templates", "system_prompts"}
+	var errs []error
+	for _, dir := range subDirs {
+		errs = append(errs, mkdirIfNotExist(filepath.Join(d.Option.CatalogPath, dir)))
+	}
+	return errors.Join(errs...)
+}
+
 func (d *driverImpl) createCustomCards() error {
 	for i, customCard := range d.Option.CustomCards {
 		err := os.WriteFile(
@@ -491,3 +515,42 @@ func (d *driverImpl) createCustomCards() error {
 	}
 	return nil
 }
+
+func (d *driverImpl) createCustomTemplates() error {
+	for i, customTemplate := range d.Option.CustomTemplates {
+		err := os.WriteFile(
+			filepath.Join(d.Option.CatalogPath, "templates", fmt.Sprintf("%s_%d.json", CustomTemplatePrefix, i)),
+			[]byte(customTemplate),
+			0666,
+		)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (d *driverImpl) createCustomSystemPrompts() error {
+	for i, systemPrompt := range d.Option.CustomSystemPrompt {
+		err := os.WriteFile(
+			filepath.Join(d.Option.CatalogPath, "system_prompts", fmt.Sprintf("%s_%d.json", CustomSystemPromptPrefix, i)),
+			[]byte(fmt.Sprintf(`{ "__type__": "textual_system_prompt", "text": "%s" }`, systemPrompt)),
+			0666,
+		)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func mkdirIfNotExist(path string) error {
+	fi, err := os.Stat(path)
+	if err == nil && !fi.IsDir() {
+		return fmt.Errorf("%s is a file. can not create a directory", path)
+	}
+	if os.IsNotExist(err) {
+		return os.MkdirAll(path, 0770)
+	}
+	return nil
+}