Skip to content

Commit 6551b6a

Browse files
yhwangyhwang
authored
feat(LMES): Support custom template and prompt (trustyai-explainability#404)
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]>
1 parent bf9d8fd commit 6551b6a

File tree

9 files changed

+748
-119
lines changed

9 files changed

+748
-119
lines changed

‎Dockerfile.lmes-job

+1-1
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ RUN curl -L https://github.com/opendatahub-io/lm-evaluation-harness/archive/refs
2020

2121
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
2222

23-
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
23+
ENV PYTHONPATH=/opt/app-root/src/.local/lib/python3.11/site-packages:/opt/app-root/src:/opt/app-root/src/server
2424
ENV HF_HOME=/opt/app-root/src/hf_home
2525
ENV UNITXT_CATALOGS=/opt/app-root/src/my_catalogs
2626

‎api/lmes/v1alpha1/lmevaljob_types.go

+77-2
Original file line numberDiff line numberDiff line change
@@ -76,14 +76,69 @@ type Card struct {
7676
Custom string `json:"custom,omitempty"`
7777
}
7878

79+
type Template struct {
80+
// Unitxt template ID
81+
// +optional
82+
Name string `json:"name,omitempty"`
83+
// The name of the custom template in the custom field. Its value is a JSON string
84+
// for a custom Unitxt template. Use the documentation here: https://www.unitxt.ai/en/latest/docs/adding_template.html
85+
// to compose a custom template, store it as a JSON file by calling the
86+
// add_to_catalog API: https://www.unitxt.ai/en/latest/docs/saving_and_loading_from_catalog.html#adding-assets-to-the-catalog,
87+
// and use the JSON content as the value here.
88+
// +optional
89+
Ref string `json:"ref,omitempty"`
90+
}
91+
92+
type SystemPrompt struct {
93+
// Unitxt System Prompt id
94+
Name string `json:"name,omitempty"`
95+
// The name of the custom systemPrompt in the custom field. Its value is a custom system prompt string
96+
Ref string `json:"ref,omitempty"`
97+
}
98+
99+
type CustomArtifact struct {
100+
// Name of the custom artifact
101+
Name string `json:"name"`
102+
// Value of the custom artifact. It could be a JSON string or plain text
103+
// depending on the artifact type
104+
Value string `json:"value"`
105+
}
106+
107+
func (c *CustomArtifact) String() string {
108+
return fmt.Sprintf("%s|%s", c.Name, c.Value)
109+
}
110+
111+
type CustomArtifacts struct {
112+
Templates []CustomArtifact `json:"templates,omitempty"`
113+
SystemPrompts []CustomArtifact `json:"systemPrompts,omitempty"`
114+
}
115+
116+
func (c *CustomArtifacts) GetTemplates() []CustomArtifact {
117+
if c == nil {
118+
return nil
119+
}
120+
return c.Templates
121+
}
122+
123+
func (c *CustomArtifacts) GetSystemPrompts() []CustomArtifact {
124+
if c == nil {
125+
return nil
126+
}
127+
return c.SystemPrompts
128+
}
129+
79130
// Use a task recipe to form a custom task. It maps to the Unitxt Recipe
80131
// Find details of the Unitxt Recipe here:
81132
// https://www.unitxt.ai/en/latest/unitxt.standard.html#unitxt.standard.StandardRecipe
82133
type TaskRecipe struct {
83134
// The Unitxt dataset card
84135
Card Card `json:"card"`
85136
// The Unitxt template
86-
Template string `json:"template"`
137+
// +optional
138+
Template *Template `json:"template,omitempty"`
139+
// The Unitxt System Prompt
140+
// +optional
141+
SystemPrompt *SystemPrompt `json:"systemPrompt,omitempty"`
87142
// The Unitxt Task
88143
// +optional
89144
Task *string `json:"task,omitempty"`
@@ -109,11 +164,31 @@ type TaskList struct {
109164
TaskNames []string `json:"taskNames,omitempty"`
110165
// Task Recipes specifically for Unitxt
111166
TaskRecipes []TaskRecipe `json:"taskRecipes,omitempty"`
167+
// Custom Unitxt artifacts that can be used in a TaskRecipe
168+
CustomArtifacts *CustomArtifacts `json:"custom,omitempty"`
112169
}
113170

171+
// Use the tp_idx and sp_idx to point to the corresponding custom template
172+
// and custom system_prompt
114173
func (t *TaskRecipe) String() string {
115174
var b strings.Builder
116-
b.WriteString(fmt.Sprintf("card=%s,template=%s", t.Card.Name, t.Template))
175+
b.WriteString(fmt.Sprintf("card=%s", t.Card.Name))
176+
if t.Template != nil {
177+
if t.Template.Name != "" {
178+
b.WriteString(fmt.Sprintf(",template=%s", t.Template.Name))
179+
} else {
180+
// refer to a custom template. add "templates." prefix
181+
b.WriteString(fmt.Sprintf(",template=templates.%s", t.Template.Ref))
182+
}
183+
}
184+
if t.SystemPrompt != nil {
185+
if t.SystemPrompt.Name != "" {
186+
b.WriteString(fmt.Sprintf(",system_prompt=%s", t.SystemPrompt.Name))
187+
} else {
188+
// refer to custom system prompt. add "system_prompts." prefix
189+
b.WriteString(fmt.Sprintf(",system_prompt=system_prompts.%s", t.SystemPrompt.Ref))
190+
}
191+
}
117192
if t.Task != nil {
118193
b.WriteString(fmt.Sprintf(",task=%s", *t.Task))
119194
}

‎api/lmes/v1alpha1/zz_generated.deepcopy.go

+85
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎cmd/lmes_driver/main.go

+25-19
Original file line numberDiff line numberDiff line change
@@ -50,21 +50,25 @@ func (t *strArrayArg) String() string {
5050
}
5151

5252
var (
53-
taskRecipes strArrayArg
54-
customCards strArrayArg
55-
copy = flag.String("copy", "", "copy this binary to specified destination path")
56-
getStatus = flag.Bool("get-status", false, "Get current status")
57-
shutdown = flag.Bool("shutdown", false, "Shutdown the driver")
58-
outputPath = flag.String("output-path", OutputPath, "output path")
59-
detectDevice = flag.Bool("detect-device", false, "detect available device(s), CUDA or CPU")
60-
commPort = flag.Int("listen-port", driver.DefaultPort, "driver serves APIs on the port")
61-
downloadAssetsS3 = flag.Bool("download-assets-s3", false, "Download assets from S3")
62-
driverLog = ctrl.Log.WithName("driver")
53+
taskRecipes strArrayArg
54+
customCards strArrayArg
55+
customTemplates strArrayArg
56+
customSystemPrompts strArrayArg
57+
copy = flag.String("copy", "", "copy this binary to specified destination path")
58+
getStatus = flag.Bool("get-status", false, "Get current status")
59+
shutdown = flag.Bool("shutdown", false, "Shutdown the driver")
60+
outputPath = flag.String("output-path", OutputPath, "output path")
61+
detectDevice = flag.Bool("detect-device", false, "detect available device(s), CUDA or CPU")
62+
commPort = flag.Int("listen-port", driver.DefaultPort, "driver serves APIs on the port")
63+
downloadAssetsS3 = flag.Bool("download-assets-s3", false, "Download assets from S3")
64+
driverLog = ctrl.Log.WithName("driver")
6365
)
6466

6567
func init() {
6668
flag.Var(&taskRecipes, "task-recipe", "task recipe")
6769
flag.Var(&customCards, "custom-card", "A JSON string represents a custom card")
70+
flag.Var(&customTemplates, "custom-template", "A JSON string represents a custom template")
71+
flag.Var(&customSystemPrompts, "custom-prompt", "A string represents a custom system_prompt")
6872
}
6973

7074
func main() {
@@ -107,15 +111,17 @@ func main() {
107111
}
108112

109113
driverOpt := driver.DriverOption{
110-
Context: ctx,
111-
OutputPath: *outputPath,
112-
DetectDevice: *detectDevice,
113-
Logger: driverLog,
114-
TaskRecipes: taskRecipes,
115-
CustomCards: customCards,
116-
Args: args,
117-
CommPort: *commPort,
118-
DownloadAssetsS3: *downloadAssetsS3,
114+
Context: ctx,
115+
OutputPath: *outputPath,
116+
DetectDevice: *detectDevice,
117+
Logger: driverLog,
118+
TaskRecipes: taskRecipes,
119+
CustomCards: customCards,
120+
CustomTemplates: customTemplates,
121+
CustomSystemPrompt: customSystemPrompts,
122+
Args: args,
123+
CommPort: *commPort,
124+
DownloadAssetsS3: *downloadAssetsS3,
119125
}
120126

121127
driver, err := driver.NewDriver(&driverOpt)

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

+61-2
Original file line numberDiff line numberDiff line change
@@ -4721,6 +4721,42 @@ spec:
47214721
taskList:
47224722
description: Evaluation task list
47234723
properties:
4724+
custom:
4725+
description: Custom Unitxt artifacts that can be used in a TaskRecipe
4726+
properties:
4727+
systemPrompts:
4728+
items:
4729+
properties:
4730+
name:
4731+
description: Name of the custom artifact
4732+
type: string
4733+
value:
4734+
description: |-
4735+
Value of the custom artifact. It could be a JSON string or plain text
4736+
depending on the artifact type
4737+
type: string
4738+
required:
4739+
- name
4740+
- value
4741+
type: object
4742+
type: array
4743+
templates:
4744+
items:
4745+
properties:
4746+
name:
4747+
description: Name of the custom artifact
4748+
type: string
4749+
value:
4750+
description: |-
4751+
Value of the custom artifact. It could be a JSON string or plain text
4752+
depending on the artifact type
4753+
type: string
4754+
required:
4755+
- name
4756+
- value
4757+
type: object
4758+
type: array
4759+
type: object
47244760
taskNames:
47254761
description: TaskNames from lm-eval's task list
47264762
items:
@@ -4764,15 +4800,38 @@ spec:
47644800
numDemos:
47654801
description: Number of fewshot
47664802
type: integer
4803+
systemPrompt:
4804+
description: The Unitxt System Prompt
4805+
properties:
4806+
name:
4807+
description: Unitxt System Prompt id
4808+
type: string
4809+
ref:
4810+
description: The name of the custom systemPrompt in
4811+
the custom field. Its value is a custom system prompt
4812+
string
4813+
type: string
4814+
type: object
47674815
task:
47684816
description: The Unitxt Task
47694817
type: string
47704818
template:
47714819
description: The Unitxt template
4772-
type: string
4820+
properties:
4821+
name:
4822+
description: Unitxt template ID
4823+
type: string
4824+
ref:
4825+
description: |-
4826+
The name of the custom template in the custom field. Its value is a JSON string
4827+
for a custom Unitxt template. Use the documentation here: https://www.unitxt.ai/en/latest/docs/adding_template.html
4828+
to compose a custom template, store it as a JSON file by calling the
4829+
add_to_catalog API: https://www.unitxt.ai/en/latest/docs/saving_and_loading_from_catalog.html#adding-assets-to-the-catalog,
4830+
and use the JSON content as the value here.
4831+
type: string
4832+
type: object
47734833
required:
47744834
- card
4775-
- template
47764835
type: object
47774836
type: array
47784837
type: object

0 commit comments

Comments
 (0)