go-whisper-api/punctuation/xlm_config.go

package punctuation

import (
    "fmt"
    "os"
    "path/filepath"

    "gopkg.in/yaml.v3"
)

type xlmModelConfig struct {
    Languages  []string `yaml:"languages"`
    MaxLength  int      `yaml:"max_length"`
    PreLabels  []string `yaml:"pre_labels"`
    PostLabels []string `yaml:"post_labels"`
    NullToken  string   `yaml:"null_token"`
    Acronym    string   `yaml:"acronym_token"`
}

func loadXLMConfig(path string) (xlmModelConfig, error) {
    data, err := os.ReadFile(path)
    if err != nil {
        return xlmModelConfig{}, err
    }
    var cfg xlmModelConfig
    if err := yaml.Unmarshal(data, &cfg); err != nil {
        return xlmModelConfig{}, fmt.Errorf("parse xlm config %s: %w", path, err)
    }
    if cfg.MaxLength <= 0 {
        cfg.MaxLength = 256
    }
    if cfg.NullToken == "" {
        cfg.NullToken = "<NULL>"
    }
    if cfg.Acronym == "" {
        cfg.Acronym = "<ACRONYM>"
    }
    return cfg, nil
}

func defaultXLMConfigPath(modelDir string) string {
    return filepath.Join(modelDir, "config.yaml")
}