Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion pkg/apis/ome/v1beta1/model.go
Original file line number Diff line number Diff line change
Expand Up @@ -286,7 +286,7 @@ const (
// ModelCapability enum
// TODO: Remove legacy capabilities
//
// +kubebuilder:validation:Enum=TEXT_GENERATION;TEXT_SUMMARIZATION;TEXT_EMBEDDINGS;TEXT_RERANK;CHAT;VISION;EMBEDDING;RERANK;TEXT_TO_TEXT;TEXT_TO_AUDIO;TEXT_TO_IMAGE;TEXT_TO_VIDEO;IMAGE_TEXT_TO_TEXT;IMAGE_TEXT_TO_AUDIO;IMAGE_TEXT_TO_IMAGE;IMAGE_TEXT_TO_VIDEO;VIDEO_TEXT_TO_AUDIO;AUDIO_TO_TEXT;AUDIO_TO_AUDIO;AUDIO_TEXT_TO_TEXT;AUDIO_TRANSLATION
// +kubebuilder:validation:Enum=TEXT_GENERATION;TEXT_SUMMARIZATION;TEXT_EMBEDDINGS;TEXT_RERANK;CHAT;VISION;EMBEDDING;RERANK;TEXT_TO_TEXT;TEXT_TO_AUDIO;TEXT_TO_IMAGE;TEXT_TO_VIDEO;IMAGE_TEXT_TO_TEXT;IMAGE_TEXT_TO_AUDIO;IMAGE_TEXT_TO_IMAGE;IMAGE_TEXT_TO_VIDEO;VIDEO_TEXT_TO_AUDIO;VIDEO_TEXT_TO_TEXT;AUDIO_TO_TEXT;AUDIO_TO_AUDIO;AUDIO_TEXT_TO_TEXT;AUDIO_TRANSLATION
type ModelCapability string

const (
Expand All @@ -310,6 +310,7 @@ const (
ModelCapabilityImageTextToImage ModelCapability = "IMAGE_TEXT_TO_IMAGE"
ModelCapabilityImageTextToVideo ModelCapability = "IMAGE_TEXT_TO_VIDEO"
ModelCapabilityVideoTextToAudio ModelCapability = "VIDEO_TEXT_TO_AUDIO"
ModelCapabilityVideoTextToText ModelCapability = "VIDEO_TEXT_TO_TEXT"
ModelCapabilityAudioToText ModelCapability = "AUDIO_TO_TEXT"
ModelCapabilityAudioToAudio ModelCapability = "AUDIO_TO_AUDIO"
ModelCapabilityAudioTextToText ModelCapability = "AUDIO_TEXT_TO_TEXT"
Expand Down
4 changes: 3 additions & 1 deletion pkg/modelconfig/capability.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ const (
CapabilityTextToAudio Capability = "TEXT_TO_AUDIO"
CapabilityImageTextToAudio Capability = "IMAGE_TEXT_TO_AUDIO"
CapabilityVideoTextToAudio Capability = "VIDEO_TEXT_TO_AUDIO"
CapabilityVideoTextToText Capability = "VIDEO_TEXT_TO_TEXT"
CapabilityAudioToText Capability = "AUDIO_TO_TEXT"
CapabilityAudioToAudio Capability = "AUDIO_TO_AUDIO"
CapabilityAudioTextToText Capability = "AUDIO_TEXT_TO_TEXT"
Expand Down Expand Up @@ -172,7 +173,8 @@ func nemotronHNanoRule(hf HuggingFaceModel) []Capability {
return []Capability{
CapabilityImageTextToText,
CapabilityTextToText,
CapabilityAudioToText,
CapabilityAudioTextToText,
CapabilityVideoTextToText,
}
}

Expand Down
6 changes: 4 additions & 2 deletions pkg/modelconfig/capability_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -198,7 +198,8 @@ func TestClassifyCapabilities_NemotronH_Nano(t *testing.T) {
want := []Capability{
CapabilityImageTextToText,
CapabilityTextToText,
CapabilityAudioToText,
CapabilityAudioTextToText,
CapabilityVideoTextToText,
}
if got := classifyCapabilities(m); !equalCaps(got, want) {
t.Errorf("got %v, want %v", got, want)
Expand Down Expand Up @@ -413,7 +414,8 @@ func TestClassifyCapabilities_RegressionSet(t *testing.T) {
want: []Capability{
CapabilityImageTextToText,
CapabilityTextToText,
CapabilityAudioToText,
CapabilityAudioTextToText,
CapabilityVideoTextToText,
},
},
{
Expand Down
1 change: 1 addition & 0 deletions pkg/modelparser/capability.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ var capabilityToOME = map[modelconfig.Capability]v1beta1.ModelCapability{
modelconfig.CapabilityTextToAudio: v1beta1.ModelCapabilityTextToAudio,
modelconfig.CapabilityImageTextToAudio: v1beta1.ModelCapabilityImageTextToAudio,
modelconfig.CapabilityVideoTextToAudio: v1beta1.ModelCapabilityVideoTextToAudio,
modelconfig.CapabilityVideoTextToText: v1beta1.ModelCapabilityVideoTextToText,
modelconfig.CapabilityAudioToText: v1beta1.ModelCapabilityAudioToText,
modelconfig.CapabilityAudioToAudio: v1beta1.ModelCapabilityAudioToAudio,
modelconfig.CapabilityAudioTextToText: v1beta1.ModelCapabilityAudioTextToText,
Expand Down
1 change: 1 addition & 0 deletions pkg/modelparser/capability_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ func TestCapabilityToOME(t *testing.T) {
{modelconfig.CapabilityTextToAudio, v1beta1.ModelCapabilityTextToAudio},
{modelconfig.CapabilityImageTextToAudio, v1beta1.ModelCapabilityImageTextToAudio},
{modelconfig.CapabilityVideoTextToAudio, v1beta1.ModelCapabilityVideoTextToAudio},
{modelconfig.CapabilityVideoTextToText, v1beta1.ModelCapabilityVideoTextToText},
{modelconfig.CapabilityAudioToText, v1beta1.ModelCapabilityAudioToText},
{modelconfig.CapabilityAudioToAudio, v1beta1.ModelCapabilityAudioToAudio},
{modelconfig.CapabilityAudioTextToText, v1beta1.ModelCapabilityAudioTextToText},
Expand Down