|
| 1 | +package whisper_test |
| 2 | + |
| 3 | +import ( |
| 4 | + "testing" |
| 5 | + |
| 6 | + "github.com/ggerganov/whisper.cpp/bindings/go/pkg/whisper" |
| 7 | + assert "github.com/stretchr/testify/assert" |
| 8 | +) |
| 9 | + |
| 10 | +func TestNew(t *testing.T) { |
| 11 | + assert := assert.New(t) |
| 12 | + t.Run("valid model path", func(t *testing.T) { |
| 13 | + model, err := whisper.New(ModelPath) |
| 14 | + assert.NoError(err) |
| 15 | + assert.NotNil(model) |
| 16 | + defer model.Close() |
| 17 | + |
| 18 | + }) |
| 19 | + |
| 20 | + t.Run("invalid model path", func(t *testing.T) { |
| 21 | + invalidModelPath := "invalid-model-path.bin" |
| 22 | + model, err := whisper.New(invalidModelPath) |
| 23 | + assert.Error(err) |
| 24 | + assert.Nil(model) |
| 25 | + }) |
| 26 | +} |
| 27 | + |
| 28 | +func TestClose(t *testing.T) { |
| 29 | + assert := assert.New(t) |
| 30 | + |
| 31 | + model, err := whisper.New(ModelPath) |
| 32 | + assert.NoError(err) |
| 33 | + assert.NotNil(model) |
| 34 | + |
| 35 | + err = model.Close() |
| 36 | + assert.NoError(err) |
| 37 | +} |
| 38 | + |
| 39 | +func TestNewContext(t *testing.T) { |
| 40 | + assert := assert.New(t) |
| 41 | + |
| 42 | + model, err := whisper.New(ModelPath) |
| 43 | + assert.NoError(err) |
| 44 | + assert.NotNil(model) |
| 45 | + defer model.Close() |
| 46 | + |
| 47 | + context, err := model.NewContext() |
| 48 | + assert.NoError(err) |
| 49 | + assert.NotNil(context) |
| 50 | +} |
| 51 | + |
| 52 | +func TestIsMultilingual(t *testing.T) { |
| 53 | + assert := assert.New(t) |
| 54 | + |
| 55 | + model, err := whisper.New(ModelPath) |
| 56 | + assert.NoError(err) |
| 57 | + assert.NotNil(model) |
| 58 | + defer model.Close() |
| 59 | + |
| 60 | + isMultilingual := model.IsMultilingual() |
| 61 | + |
| 62 | + // This returns false since |
| 63 | + // the model 'models/ggml-small.en.bin' |
| 64 | + // that is loaded is not multilingual |
| 65 | + assert.False(isMultilingual) |
| 66 | +} |
| 67 | + |
| 68 | +func TestLanguages(t *testing.T) { |
| 69 | + assert := assert.New(t) |
| 70 | + |
| 71 | + model, err := whisper.New(ModelPath) |
| 72 | + assert.NoError(err) |
| 73 | + assert.NotNil(model) |
| 74 | + defer model.Close() |
| 75 | + |
| 76 | + expectedLanguages := []string{ |
| 77 | + "en", "zh", "de", "es", "ru", "ko", "fr", "ja", "pt", "tr", "pl", |
| 78 | + "ca", "nl", "ar", "sv", "it", "id", "hi", "fi", "vi", "he", "uk", |
| 79 | + "el", "ms", "cs", "ro", "da", "hu", "ta", "no", "th", "ur", "hr", |
| 80 | + "bg", "lt", "la", "mi", "ml", "cy", "sk", "te", "fa", "lv", "bn", |
| 81 | + "sr", "az", "sl", "kn", "et", "mk", "br", "eu", "is", "hy", "ne", |
| 82 | + "mn", "bs", "kk", "sq", "sw", "gl", "mr", "pa", "si", "km", "sn", |
| 83 | + "yo", "so", "af", "oc", "ka", "be", "tg", "sd", "gu", "am", "yi", |
| 84 | + "lo", "uz", "fo", "ht", "ps", "tk", "nn", "mt", "sa", "lb", "my", |
| 85 | + "bo", "tl", "mg", "as", "tt", "haw", "ln", "ha", "ba", "jw", "su", |
| 86 | + } |
| 87 | + |
| 88 | + actualLanguages := model.Languages() |
| 89 | + |
| 90 | + assert.Equal(expectedLanguages, actualLanguages) |
| 91 | +} |
0 commit comments