update yaml catalog schema

Signed-off-by: Dhiraj Bokde <[email protected]>
kubeflow · Jan 25, 2025 · 04fbb63 · 04fbb63
1 parent 17f9d88
commit 04fbb63
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 18 deletions.
diff --git a/internal/catalog/testdata/test-yaml-catalog.yaml b/internal/catalog/testdata/test-yaml-catalog.yaml
@@ -1,11 +1,23 @@
+source: Red Hat
 models:
-  - name: ibm-granite/granite-3.1-8b-base
+  - name: rhelai1/granite-8b-code-base
     provider: IBM
     description: A decoder-only code model designed for code generative tasks
+    longDescription: |-
+      Granite-8B-Code-Base is a decoder-only code model designed for code
+      generative tasks (e.g., code generation, code explanation, code fixing,
+      etc.). It is trained from scratch with a two-phase training strategy. In
+      phase 1, our model is trained on 4 trillion tokens sourced from 116
+      programming languages, ensuring a comprehensive understanding of
+      programming languages and syntax. In phase 2, our model is trained on 500
+      billion tokens with a carefully designed mixture of high-quality data from
+      code and natural language domains to improve the models’ ability to reason
+      and follow instructions.
     readmeLink: https://huggingface.co/ibm-granite/granite-3.1-8b-base/raw/main/README.md
     language: ["ar", "cs", "de", "en", "es", "fr", "it", "ja", "ko", "nl", "pt", "zh"]
     license: apache-2.0
     licenseLink: https://www.apache.org/licenses/LICENSE-2.0.txt
+    maturity: Generally Available
     libraryName: transformers
     tags:
       - language
@@ -14,22 +26,33 @@ models:
       - text-generation
     createTimeSinceEpoch: 1733514949000
     lastUpdateTimeSinceEpoch: 1734637721000
-  - name: ibm-granite/granite-3.1-8b-instruct
+    artifacts:
+      - protocol: oci
+        uri: oci://registry.redhat.io/rhelai1/granite-8b-code-base:1.3-1732870892
+  - name: rhelai1/granite-8b-code-instruct
     provider: IBM
     description: A fine-tuned model based on Granite 8B Code Base
+    longDescription: |-
+      Granite-8B-Code-Instruct is a 8B parameter model fine tuned from
+      Granite-8B-Code-Base on a combination of permissively licensed instruction
+      data to enhance instruction following capabilities including logical
+      reasoning and problem-solving skills.
     readmeLink: https://huggingface.co/ibm-granite/granite-3.1-8b-instruct/raw/main/README.md
     language: ["ar", "cs", "de", "en", "es", "fr", "it", "ja", "ko", "nl", "pt", "zh"]
     license: apache-2.0
     licenseLink: https://www.apache.org/licenses/LICENSE-2.0.txt
+    maturity: Generally Available
     libraryName: transformers
     baseModel:
-      - catalog: huggingface.io
-        repository: ibm-granite
-        model: granite-3.1-8b-base
+      - repository: rhelai1
+        model: granite-8b-code-base
     tags:
       - language
       - granite-3.1
     tasks:
       - text-generation
     createTimeSinceEpoch: 1733514949000
-    lastUpdateTimeSinceEpoch: 1734637721000
+    lastUpdateTimeSinceEpoch: 1734637721000
+    artifacts:
+      - protocol: oci
+        uri: oci://registry.redhat.io/rhelai1/granite-8b-code-instruct:1.3-1732870892
diff --git a/internal/catalog/yaml_catalog.go b/internal/catalog/yaml_catalog.go
@@ -14,23 +14,32 @@ type YamlBaseModel struct {
 	Model      string `yaml:"model"`
 }
 
+type YamlArtifacts struct {
+	Protocol string `yaml:"protocol"`
+	URI      string `yaml:"uri"`
+}
+
 type YamlModel struct {
-	Name                     string   `yaml:"name"`
-	Provider                 string   `yaml:"provider"`
-	Description              string   `yaml:"description"`
-	ReadmeLink               string   `yaml:"readmeLink"`
-	Language                 []string `yaml:"language"`
-	License                  string   `yaml:"license"`
-	LicenseLink              string   `yaml:"licenseLink"`
-	LibraryName              string   `yaml:"libraryName"`
-	Tags                     []string `yaml:"tags"`
-	Tasks                    []string `yaml:"tasks"`
-	CreateTimeSinceEpoch     int64    `yaml:"createTimeSinceEpoch"`
-	LastUpdateTimeSinceEpoch int64    `yaml:"lastUpdateTimeSinceEpoch"`
+	Name                     string          `yaml:"name"`
+	Provider                 string          `yaml:"provider"`
+	Description              string          `yaml:"description"`
+	LongDescription          string          `yaml:"longDescription"`
+	ReadmeLink               string          `yaml:"readmeLink"`
+	Language                 []string        `yaml:"language"`
+	License                  string          `yaml:"license"`
+	LicenseLink              string          `yaml:"licenseLink"`
+	Maturity                 string          `yaml:"maturity"`
+	LibraryName              string          `yaml:"libraryName"`
+	Tags                     []string        `yaml:"tags"`
+	Tasks                    []string        `yaml:"tasks"`
+	CreateTimeSinceEpoch     int64           `yaml:"createTimeSinceEpoch"`
+	LastUpdateTimeSinceEpoch int64           `yaml:"lastUpdateTimeSinceEpoch"`
 	BaseModel                []YamlBaseModel `yaml:"baseModel,omitempty"`
+	Artifacts                []YamlArtifacts `yaml:"artifacts"`
 }
 
 type YamlCatalog struct {
+	Source string      `yaml:"source"`
 	Models []YamlModel `yaml:"models"`
 }
 
@@ -86,6 +95,12 @@ func NewYamlCatalog(source *CatalogSource) (ModelCatalogApi, error) {
 	if err = yaml.UnmarshalStrict(bytes, &contents); err != nil {
 		return nil, fmt.Errorf("failed to parse %s file: %v", yamlCatalogPath, err)
 	}
+
+	// override catalog name from Yaml Catalog File if set
+	if len(source.Name) > 0 {
+		source.Name = contents.Source
+	}
+
 	return &yamlCatalogImpl{source: source, contents: &contents}, nil
 }