first commit
This commit is contained in:
@@ -0,0 +1,86 @@
|
||||
package punctuation
|
||||
|
||||
import (
|
||||
"context"
|
||||
"regexp"
|
||||
"strings"
|
||||
"unicode"
|
||||
"unicode/utf8"
|
||||
)
|
||||
|
||||
type Heuristic struct{}
|
||||
|
||||
func (Heuristic) Active() bool {
|
||||
return true
|
||||
}
|
||||
|
||||
func (Heuristic) Restore(ctx context.Context, text, language string) (string, error) {
|
||||
_ = ctx
|
||||
text = strings.TrimSpace(text)
|
||||
if text == "" {
|
||||
return text, nil
|
||||
}
|
||||
text = normalizeSpaces(text)
|
||||
text = capitalizeFirst(text)
|
||||
lang := strings.ToLower(strings.TrimSpace(language))
|
||||
if lang == "ru" || lang == "rus" || lang == "russian" || lang == "auto" {
|
||||
text = heuristicRU(text)
|
||||
} else {
|
||||
text = heuristicEN(text)
|
||||
}
|
||||
return ensureTerminalPunct(text), nil
|
||||
}
|
||||
|
||||
func normalizeSpaces(s string) string {
|
||||
return strings.Join(strings.Fields(s), " ")
|
||||
}
|
||||
|
||||
func capitalizeFirst(s string) string {
|
||||
r, size := utf8.DecodeRuneInString(s)
|
||||
if r == utf8.RuneError {
|
||||
return s
|
||||
}
|
||||
return string(unicode.ToUpper(r)) + s[size:]
|
||||
}
|
||||
|
||||
var (
|
||||
reQuestionRU = regexp.MustCompile(`(?i)(^|.*\s)(как|что|где|когда|почему|зачем|кто|чей|какой|какая|какое|какие|сколько|зачем|откуда|куда|ли)(\s+[^.?!]+)$`)
|
||||
reQuestionEN = regexp.MustCompile(`(?i)^(who|what|when|where|why|how|which|whose|whom|is|are|am|was|were|do|does|did|can|could|would|will|shall|should)\b`)
|
||||
)
|
||||
|
||||
func heuristicRU(s string) string {
|
||||
if reQuestionRU.MatchString(s) && !strings.HasSuffix(s, "?") {
|
||||
return s + "?"
|
||||
}
|
||||
if !hasTerminalPunct(s) && len(strings.Fields(s)) <= 24 {
|
||||
return s + "."
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func heuristicEN(s string) string {
|
||||
lower := strings.ToLower(s)
|
||||
if reQuestionEN.MatchString(lower) && !strings.HasSuffix(s, "?") {
|
||||
return s + "?"
|
||||
}
|
||||
if !hasTerminalPunct(s) && len(strings.Fields(s)) <= 24 {
|
||||
return s + "."
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
func hasTerminalPunct(s string) bool {
|
||||
s = strings.TrimSpace(s)
|
||||
if s == "" {
|
||||
return false
|
||||
}
|
||||
r, _ := utf8.DecodeLastRuneInString(s)
|
||||
return r == '.' || r == '?' || r == '!' || r == '…'
|
||||
}
|
||||
|
||||
func ensureTerminalPunct(s string) string {
|
||||
if hasTerminalPunct(s) {
|
||||
return s
|
||||
}
|
||||
return s + "."
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
package punctuation
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestHeuristicRU_question(t *testing.T) {
|
||||
h := Heuristic{}
|
||||
out, err := h.Restore(context.Background(), "как дела", "ru")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !stringsHasSuffix(out, "?") {
|
||||
t.Fatalf("expected question mark, got %q", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHeuristicEN_period(t *testing.T) {
|
||||
h := Heuristic{}
|
||||
out, err := h.Restore(context.Background(), "hello world", "en")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !stringsHasSuffix(out, ".") {
|
||||
t.Fatalf("expected period, got %q", out)
|
||||
}
|
||||
}
|
||||
|
||||
func stringsHasSuffix(s, suffix string) bool {
|
||||
return len(s) >= len(suffix) && s[len(s)-len(suffix):] == suffix
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
//go:build xlm
|
||||
|
||||
package spwrap
|
||||
|
||||
/*
|
||||
#cgo CXXFLAGS: -std=c++17
|
||||
#cgo LDFLAGS: -lsentencepiece
|
||||
#include <stdlib.h>
|
||||
#include "sp_wrap.h"
|
||||
*/
|
||||
import "C"
|
||||
import (
|
||||
"fmt"
|
||||
"unsafe"
|
||||
)
|
||||
|
||||
type Processor struct {
|
||||
p *C.SPProcessor
|
||||
}
|
||||
|
||||
func Load(path string) (*Processor, error) {
|
||||
cpath := C.CString(path)
|
||||
defer C.free(unsafe.Pointer(cpath))
|
||||
var errMsg *C.char
|
||||
p := C.sp_load(cpath, &errMsg)
|
||||
if p == nil {
|
||||
if errMsg != nil {
|
||||
defer C.free(unsafe.Pointer(errMsg))
|
||||
return nil, fmt.Errorf("sentencepiece: %s", C.GoString(errMsg))
|
||||
}
|
||||
return nil, fmt.Errorf("sentencepiece: failed to load %s", path)
|
||||
}
|
||||
return &Processor{p: p}, nil
|
||||
}
|
||||
|
||||
func (proc *Processor) Close() {
|
||||
if proc.p != nil {
|
||||
C.sp_free(proc.p)
|
||||
proc.p = nil
|
||||
}
|
||||
}
|
||||
|
||||
func (proc *Processor) BOSID() int {
|
||||
return int(C.sp_bos_id(proc.p))
|
||||
}
|
||||
|
||||
func (proc *Processor) EOSID() int {
|
||||
return int(C.sp_eos_id(proc.p))
|
||||
}
|
||||
|
||||
func (proc *Processor) PadID() int {
|
||||
return int(C.sp_pad_id(proc.p))
|
||||
}
|
||||
|
||||
func (proc *Processor) EncodeAsIDs(text string) ([]int, error) {
|
||||
ctext := C.CString(text)
|
||||
defer C.free(unsafe.Pointer(ctext))
|
||||
var ids *C.int
|
||||
var n C.int
|
||||
var errMsg *C.char
|
||||
if C.sp_encode(proc.p, ctext, &ids, &n, &errMsg) == 0 {
|
||||
if errMsg != nil {
|
||||
defer C.free(unsafe.Pointer(errMsg))
|
||||
return nil, fmt.Errorf("sentencepiece encode: %s", C.GoString(errMsg))
|
||||
}
|
||||
return nil, fmt.Errorf("sentencepiece encode failed")
|
||||
}
|
||||
if ids == nil || n == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
defer C.free(unsafe.Pointer(ids))
|
||||
out := make([]int, int(n))
|
||||
slice := unsafe.Slice(ids, int(n))
|
||||
for i := range out {
|
||||
out[i] = int(slice[i])
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (proc *Processor) IDToPiece(id int) (string, error) {
|
||||
var errMsg *C.char
|
||||
piece := C.sp_id_to_piece(proc.p, C.int(id), &errMsg)
|
||||
if piece == nil {
|
||||
if errMsg != nil {
|
||||
defer C.free(unsafe.Pointer(errMsg))
|
||||
return "", fmt.Errorf("sentencepiece id to piece: %s", C.GoString(errMsg))
|
||||
}
|
||||
return "", fmt.Errorf("sentencepiece id to piece failed")
|
||||
}
|
||||
defer C.free(unsafe.Pointer(piece))
|
||||
return C.GoString(piece), nil
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
#include "sp_wrap.h"
|
||||
|
||||
#include <sentencepiece_processor.h>
|
||||
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
struct SPProcessor {
|
||||
sentencepiece::SentencePieceProcessor proc;
|
||||
};
|
||||
|
||||
static char *copy_err(const std::string &msg) {
|
||||
char *out = static_cast<char *>(std::malloc(msg.size() + 1));
|
||||
if (out != nullptr) {
|
||||
std::memcpy(out, msg.c_str(), msg.size() + 1);
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
SPProcessor *sp_load(const char *path, char **err) {
|
||||
if (err != nullptr) {
|
||||
*err = nullptr;
|
||||
}
|
||||
auto *p = new SPProcessor();
|
||||
const auto status = p->proc.Load(path);
|
||||
if (!status.ok()) {
|
||||
if (err != nullptr) {
|
||||
*err = copy_err(status.ToString());
|
||||
}
|
||||
delete p;
|
||||
return nullptr;
|
||||
}
|
||||
return p;
|
||||
}
|
||||
|
||||
void sp_free(SPProcessor *p) { delete p; }
|
||||
|
||||
int sp_bos_id(const SPProcessor *p) { return p->proc.bos_id(); }
|
||||
|
||||
int sp_eos_id(const SPProcessor *p) { return p->proc.eos_id(); }
|
||||
|
||||
int sp_pad_id(const SPProcessor *p) { return p->proc.pad_id(); }
|
||||
|
||||
int sp_encode(const SPProcessor *p, const char *text, int **out_ids, int *out_len, char **err) {
|
||||
if (err != nullptr) {
|
||||
*err = nullptr;
|
||||
}
|
||||
if (out_ids != nullptr) {
|
||||
*out_ids = nullptr;
|
||||
}
|
||||
if (out_len != nullptr) {
|
||||
*out_len = 0;
|
||||
}
|
||||
std::vector<int> ids;
|
||||
const auto status = p->proc.Encode(text, &ids);
|
||||
if (!status.ok()) {
|
||||
if (err != nullptr) {
|
||||
*err = copy_err(status.ToString());
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
if (ids.empty()) {
|
||||
return 1;
|
||||
}
|
||||
int *buf = static_cast<int *>(std::malloc(sizeof(int) * ids.size()));
|
||||
if (buf == nullptr) {
|
||||
if (err != nullptr) {
|
||||
*err = copy_err("malloc failed");
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
for (size_t i = 0; i < ids.size(); i++) {
|
||||
buf[i] = ids[i];
|
||||
}
|
||||
*out_ids = buf;
|
||||
*out_len = static_cast<int>(ids.size());
|
||||
return 1;
|
||||
}
|
||||
|
||||
char *sp_id_to_piece(const SPProcessor *p, int id, char **err) {
|
||||
if (err != nullptr) {
|
||||
*err = nullptr;
|
||||
}
|
||||
const std::string piece = p->proc.IdToPiece(id);
|
||||
return copy_err(piece);
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
#pragma once
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
typedef struct SPProcessor SPProcessor;
|
||||
|
||||
SPProcessor *sp_load(const char *path, char **err);
|
||||
void sp_free(SPProcessor *p);
|
||||
|
||||
int sp_bos_id(const SPProcessor *p);
|
||||
int sp_eos_id(const SPProcessor *p);
|
||||
int sp_pad_id(const SPProcessor *p);
|
||||
|
||||
int sp_encode(const SPProcessor *p, const char *text, int **out_ids, int *out_len, char **err);
|
||||
char *sp_id_to_piece(const SPProcessor *p, int id, char **err);
|
||||
|
||||
#ifdef __cplusplus
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,177 @@
|
||||
//go:build xlm
|
||||
|
||||
package punctuation
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"runtime"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
ort "github.com/yalue/onnxruntime_go"
|
||||
)
|
||||
|
||||
var (
|
||||
ortOnce sync.Once
|
||||
ortErr error
|
||||
)
|
||||
|
||||
func ensureORT() error {
|
||||
ortOnce.Do(func() {
|
||||
if p := resolveONNXRuntimeLib(); p != "" {
|
||||
ort.SetSharedLibraryPath(p)
|
||||
}
|
||||
ortErr = ort.InitializeEnvironment()
|
||||
})
|
||||
return ortErr
|
||||
}
|
||||
|
||||
func resolveONNXRuntimeLib() string {
|
||||
if p := strings.TrimSpace(os.Getenv("ONNXRUNTIME_SHARED_LIBRARY_PATH")); p != "" {
|
||||
return p
|
||||
}
|
||||
for _, p := range onnxRuntimeCandidates() {
|
||||
if st, err := os.Stat(p); err == nil && !st.IsDir() {
|
||||
return p
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func onnxRuntimeCandidates() []string {
|
||||
arch := sherpaLibArch()
|
||||
ver := sherpaLinuxModuleVersion()
|
||||
var out []string
|
||||
for _, root := range goModCacheRoots() {
|
||||
if ver != "" {
|
||||
out = append(out, filepath.Join(root,
|
||||
"github.com/k2-fsa/sherpa-onnx-go-linux@"+ver,
|
||||
"lib", arch, "libonnxruntime.so"))
|
||||
}
|
||||
}
|
||||
if exe, err := os.Executable(); err == nil {
|
||||
exeDir := filepath.Dir(exe)
|
||||
out = append(out,
|
||||
filepath.Join(exeDir, "libonnxruntime.so"),
|
||||
filepath.Join(exeDir, "lib", "libonnxruntime.so"),
|
||||
filepath.Join(exeDir, "..", "lib", "libonnxruntime.so"),
|
||||
)
|
||||
if modRoot := findModuleRoot(exeDir); modRoot != "" {
|
||||
out = append(out, filepath.Join(modRoot, "lib", "libonnxruntime.so"))
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func goModCacheRoots() []string {
|
||||
var roots []string
|
||||
seen := map[string]struct{}{}
|
||||
add := func(p string) {
|
||||
p = strings.TrimSpace(p)
|
||||
if p == "" {
|
||||
return
|
||||
}
|
||||
if _, ok := seen[p]; ok {
|
||||
return
|
||||
}
|
||||
seen[p] = struct{}{}
|
||||
roots = append(roots, p)
|
||||
}
|
||||
add(os.Getenv("GOMODCACHE"))
|
||||
if gopath := os.Getenv("GOPATH"); gopath != "" {
|
||||
for _, gp := range filepath.SplitList(gopath) {
|
||||
add(filepath.Join(gp, "pkg", "mod"))
|
||||
}
|
||||
}
|
||||
if home, err := os.UserHomeDir(); err == nil {
|
||||
add(filepath.Join(home, "go", "pkg", "mod"))
|
||||
}
|
||||
return roots
|
||||
}
|
||||
|
||||
func findModuleRoot(start string) string {
|
||||
dir := start
|
||||
for i := 0; i < 8; i++ {
|
||||
if _, err := os.Stat(filepath.Join(dir, "go.mod")); err == nil {
|
||||
return dir
|
||||
}
|
||||
parent := filepath.Dir(dir)
|
||||
if parent == dir {
|
||||
break
|
||||
}
|
||||
dir = parent
|
||||
}
|
||||
if cwd, err := os.Getwd(); err == nil {
|
||||
return findModuleRootFrom(cwd)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func findModuleRootFrom(dir string) string {
|
||||
for i := 0; i < 8; i++ {
|
||||
if _, err := os.Stat(filepath.Join(dir, "go.mod")); err == nil {
|
||||
return dir
|
||||
}
|
||||
parent := filepath.Dir(dir)
|
||||
if parent == dir {
|
||||
break
|
||||
}
|
||||
dir = parent
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func sherpaLinuxModuleVersion() string {
|
||||
for _, dir := range []string{findModuleRootFrom(mustCwd()), ""} {
|
||||
if dir == "" {
|
||||
continue
|
||||
}
|
||||
if v := readSherpaVersion(filepath.Join(dir, "go.mod")); v != "" {
|
||||
return v
|
||||
}
|
||||
}
|
||||
if exe, err := os.Executable(); err == nil {
|
||||
if root := findModuleRoot(filepath.Dir(exe)); root != "" {
|
||||
return readSherpaVersion(filepath.Join(root, "go.mod"))
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func readSherpaVersion(goModPath string) string {
|
||||
data, err := os.ReadFile(goModPath)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
for _, line := range strings.Split(string(data), "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
if !strings.Contains(line, "github.com/k2-fsa/sherpa-onnx-go-linux") {
|
||||
continue
|
||||
}
|
||||
fields := strings.Fields(line)
|
||||
if len(fields) >= 2 {
|
||||
return fields[1] // e.g. v1.13.2 — must match pkg/mod path @v1.13.2
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func sherpaLibArch() string {
|
||||
switch runtime.GOARCH {
|
||||
case "arm64":
|
||||
return "aarch64-unknown-linux-gnu"
|
||||
case "arm":
|
||||
return "arm-unknown-linux-gnueabihf"
|
||||
default:
|
||||
return "x86_64-unknown-linux-gnu"
|
||||
}
|
||||
}
|
||||
|
||||
func mustCwd() string {
|
||||
cwd, err := os.Getwd()
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return cwd
|
||||
}
|
||||
@@ -0,0 +1,154 @@
|
||||
package punctuation
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"go-whisper-api/config"
|
||||
)
|
||||
|
||||
type Restorer interface {
|
||||
Active() bool
|
||||
Restore(ctx context.Context, text, language string) (string, error)
|
||||
}
|
||||
|
||||
type Closer interface {
|
||||
Close()
|
||||
}
|
||||
|
||||
func New(cfg config.Punctuation) (Restorer, error) {
|
||||
cfg = cfg.WithDefaults()
|
||||
if !cfg.Active() {
|
||||
return Nop{}, nil
|
||||
}
|
||||
engine := strings.ToLower(strings.TrimSpace(cfg.Engine))
|
||||
switch engine {
|
||||
case "heuristic":
|
||||
return Heuristic{}, nil
|
||||
case "sherpa", "sherpa-offline", "offline":
|
||||
cfg.Engine = "sherpa"
|
||||
return newSherpaRestorer(cfg)
|
||||
case "sherpa-online", "online":
|
||||
cfg.Engine = "sherpa-online"
|
||||
return newSherpaRestorer(cfg)
|
||||
case "xlm", "xlm-roberta", "roberta":
|
||||
return newXLM(cfg)
|
||||
case "http":
|
||||
if cfg.HTTPURL == "" {
|
||||
return nil, fmt.Errorf("punctuation.http_url is required when engine=http")
|
||||
}
|
||||
return HTTP{cfg: cfg}, nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported punctuation engine %q (use: off, heuristic, xlm, sherpa, sherpa-online, http)", engine)
|
||||
}
|
||||
}
|
||||
|
||||
type Nop struct{}
|
||||
|
||||
func (Nop) Active() bool {
|
||||
return false
|
||||
}
|
||||
|
||||
func (Nop) Restore(ctx context.Context, text, language string) (string, error) {
|
||||
return text, nil
|
||||
}
|
||||
|
||||
type HTTP struct {
|
||||
cfg config.Punctuation
|
||||
client *http.Client
|
||||
}
|
||||
|
||||
func (h HTTP) Active() bool { return true }
|
||||
|
||||
func (h HTTP) Restore(ctx context.Context, text, language string) (string, error) {
|
||||
body, err := json.Marshal(map[string]string{
|
||||
"text": text,
|
||||
"language": language,
|
||||
})
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, h.cfg.HTTPURL, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
|
||||
client := h.client
|
||||
if client == nil {
|
||||
client = &http.Client{Timeout: h.cfg.Timeout()}
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
raw, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if resp.StatusCode >= 300 {
|
||||
return "", fmt.Errorf("punctuation http %s: %s", resp.Status, strings.TrimSpace(string(raw)))
|
||||
}
|
||||
var out struct {
|
||||
Text string `json:"text"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &out); err != nil {
|
||||
return strings.TrimSpace(string(raw)), nil
|
||||
}
|
||||
if out.Text == "" {
|
||||
return text, nil
|
||||
}
|
||||
return out.Text, nil
|
||||
}
|
||||
|
||||
func Apply(ctx context.Context, r Restorer, enabled bool, text, language string) (string, error) {
|
||||
if !enabled || r == nil || !r.Active() {
|
||||
return text, nil
|
||||
}
|
||||
text = strings.TrimSpace(text)
|
||||
if text == "" {
|
||||
return text, nil
|
||||
}
|
||||
return r.Restore(ctx, text, language)
|
||||
}
|
||||
|
||||
func Close(r Restorer) {
|
||||
if c, ok := r.(Closer); ok {
|
||||
c.Close()
|
||||
}
|
||||
}
|
||||
|
||||
func AutoSelect(cfg config.Punctuation) (Restorer, error) {
|
||||
cfg = cfg.WithDefaults()
|
||||
if !cfg.Active() {
|
||||
return Nop{}, nil
|
||||
}
|
||||
engine := strings.ToLower(cfg.Engine)
|
||||
if engine == "heuristic" {
|
||||
return Heuristic{}, nil
|
||||
}
|
||||
if engine == "http" {
|
||||
return New(cfg)
|
||||
}
|
||||
if engine == "xlm" || engine == "xlm-roberta" || engine == "roberta" {
|
||||
return newXLM(cfg)
|
||||
}
|
||||
if _, err := os.Stat(cfg.ModelPath()); err == nil {
|
||||
cfg.Engine = engine
|
||||
if engine == "sherpa" || engine == "sherpa-offline" || engine == "offline" || engine == "" {
|
||||
cfg.Engine = "sherpa"
|
||||
}
|
||||
return newSherpaRestorer(cfg)
|
||||
}
|
||||
if engine == "sherpa" || engine == "sherpa-online" || engine == "online" {
|
||||
return nil, fmt.Errorf("punctuation model not found at %s (run: make download-punctuation-model)", cfg.ModelPath())
|
||||
}
|
||||
return Heuristic{}, nil
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
package punctuation
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"go-whisper-api/config"
|
||||
)
|
||||
|
||||
func TestNop(t *testing.T) {
|
||||
r, err := New(config.Punctuation{Engine: "off"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
out, err := Apply(context.Background(), r, true, "hello world", "en")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if out != "hello world" {
|
||||
t.Fatalf("got %q", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestApply_disabled(t *testing.T) {
|
||||
r := Heuristic{}
|
||||
out, err := Apply(context.Background(), r, false, "hello", "en")
|
||||
if err != nil || out != "hello" {
|
||||
t.Fatalf("got %q err=%v", out, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNew_heuristic(t *testing.T) {
|
||||
r, err := New(config.Punctuation{Enabled: true, Engine: "heuristic"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !r.Active() {
|
||||
t.Fatal("expected active")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,107 @@
|
||||
//go:build sherpa
|
||||
|
||||
package punctuation
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"go-whisper-api/config"
|
||||
|
||||
sherpa "github.com/k2-fsa/sherpa-onnx-go/sherpa_onnx"
|
||||
)
|
||||
|
||||
type Sherpa struct {
|
||||
offline *sherpa.OfflinePunctuation
|
||||
online *sherpa.OnlinePunctuation
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
func newSherpaRestorer(cfg config.Punctuation) (Restorer, error) {
|
||||
return newSherpa(cfg)
|
||||
}
|
||||
|
||||
func newSherpa(cfg config.Punctuation) (*Sherpa, error) {
|
||||
cfg = cfg.WithDefaults()
|
||||
engine := strings.ToLower(cfg.Engine)
|
||||
s := &Sherpa{}
|
||||
switch engine {
|
||||
case "sherpa-online", "online":
|
||||
modelPath := cfg.ModelPath()
|
||||
vocabPath := cfg.BpeVocabPath()
|
||||
if _, err := os.Stat(modelPath); err != nil {
|
||||
return nil, fmt.Errorf("sherpa online punctuation model %q: %w", modelPath, err)
|
||||
}
|
||||
if _, err := os.Stat(vocabPath); err != nil {
|
||||
return nil, fmt.Errorf("sherpa bpe vocab %q: %w", vocabPath, err)
|
||||
}
|
||||
conf := sherpa.OnlinePunctuationConfig{}
|
||||
conf.Model.CnnBilstm = modelPath
|
||||
conf.Model.BpeVocab = vocabPath
|
||||
conf.Model.NumThreads = cfg.NumThreads
|
||||
conf.Model.Provider = "cpu"
|
||||
s.online = sherpa.NewOnlinePunctuation(&conf)
|
||||
if s.online == nil {
|
||||
return nil, fmt.Errorf("failed to create sherpa online punctuation")
|
||||
}
|
||||
default:
|
||||
modelPath := cfg.ModelPath()
|
||||
if _, err := os.Stat(modelPath); err != nil {
|
||||
return nil, fmt.Errorf("sherpa offline punctuation model %q: %w (run: make download-punctuation-model)", modelPath, err)
|
||||
}
|
||||
conf := sherpa.OfflinePunctuationConfig{}
|
||||
conf.Model.CtTransformer = modelPath
|
||||
conf.Model.NumThreads = cfg.NumThreads
|
||||
conf.Model.Provider = "cpu"
|
||||
s.offline = sherpa.NewOfflinePunctuation(&conf)
|
||||
if s.offline == nil {
|
||||
return nil, fmt.Errorf("failed to create sherpa offline punctuation")
|
||||
}
|
||||
}
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func (s *Sherpa) Active() bool {
|
||||
return true
|
||||
}
|
||||
|
||||
func (s *Sherpa) Restore(ctx context.Context, text, language string) (string, error) {
|
||||
_ = ctx
|
||||
_ = language
|
||||
text = strings.TrimSpace(text)
|
||||
if text == "" {
|
||||
return text, nil
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
var out string
|
||||
switch {
|
||||
case s.offline != nil:
|
||||
out = s.offline.AddPunct(text)
|
||||
case s.online != nil:
|
||||
out = s.online.AddPunct(text)
|
||||
default:
|
||||
return text, fmt.Errorf("sherpa punctuation not initialized")
|
||||
}
|
||||
out = strings.TrimSpace(out)
|
||||
if out == "" {
|
||||
return text, nil
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (s *Sherpa) Close() {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.offline != nil {
|
||||
sherpa.DeleteOfflinePunc(s.offline)
|
||||
s.offline = nil
|
||||
}
|
||||
if s.online != nil {
|
||||
sherpa.DeleteOnlinePunctuation(s.online)
|
||||
s.online = nil
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,13 @@
|
||||
//go:build !sherpa
|
||||
|
||||
package punctuation
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"go-whisper-api/config"
|
||||
)
|
||||
|
||||
func newSherpaRestorer(cfg config.Punctuation) (Restorer, error) {
|
||||
return nil, fmt.Errorf("punctuation engine %q requires build tag sherpa (use: make build)", cfg.Engine)
|
||||
}
|
||||
@@ -0,0 +1,319 @@
|
||||
//go:build xlm
|
||||
|
||||
package punctuation
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
"unicode"
|
||||
|
||||
"go-whisper-api/config"
|
||||
"go-whisper-api/punctuation/internal/spwrap"
|
||||
|
||||
ort "github.com/yalue/onnxruntime_go"
|
||||
)
|
||||
|
||||
type XLM struct {
|
||||
cfg config.Punctuation
|
||||
modelCfg xlmModelConfig
|
||||
sp *spwrap.Processor
|
||||
session *ort.DynamicAdvancedSession
|
||||
inputName string
|
||||
outputNames []string
|
||||
joinSBD bool
|
||||
}
|
||||
|
||||
func newXLM(cfg config.Punctuation) (*XLM, error) {
|
||||
if err := ensureORT(); err != nil {
|
||||
return nil, fmt.Errorf("onnxruntime: %w (set ONNXRUNTIME_SHARED_LIBRARY_PATH or install sherpa-onnx libs)", err)
|
||||
}
|
||||
onnxPath := cfg.ModelPath()
|
||||
if _, err := os.Stat(onnxPath); err != nil {
|
||||
return nil, fmt.Errorf("xlm onnx model not found at %s (run: make download-xlm-punctuation-model)", onnxPath)
|
||||
}
|
||||
spPath := cfg.SPModelPath()
|
||||
if _, err := os.Stat(spPath); err != nil {
|
||||
return nil, fmt.Errorf("xlm sp.model not found at %s (run: make download-xlm-punctuation-model)", spPath)
|
||||
}
|
||||
cfgPath := cfg.XLMConfigPath()
|
||||
modelCfg, err := loadXLMConfig(cfgPath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
sp, err := spwrap.Load(spPath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
inputs, outputs, err := ort.GetInputOutputInfo(onnxPath)
|
||||
if err != nil {
|
||||
sp.Close()
|
||||
return nil, fmt.Errorf("xlm model io: %w", err)
|
||||
}
|
||||
if len(inputs) == 0 || len(outputs) < 4 {
|
||||
sp.Close()
|
||||
return nil, fmt.Errorf("xlm model: unexpected inputs/outputs")
|
||||
}
|
||||
inNames := make([]string, len(inputs))
|
||||
for i, in := range inputs {
|
||||
inNames[i] = in.Name
|
||||
}
|
||||
outNames := make([]string, len(outputs))
|
||||
for i, out := range outputs {
|
||||
outNames[i] = out.Name
|
||||
}
|
||||
session, err := ort.NewDynamicAdvancedSession(onnxPath, inNames, outNames, nil)
|
||||
if err != nil {
|
||||
sp.Close()
|
||||
return nil, fmt.Errorf("xlm onnx session: %w", err)
|
||||
}
|
||||
return &XLM{
|
||||
cfg: cfg,
|
||||
modelCfg: modelCfg,
|
||||
sp: sp,
|
||||
session: session,
|
||||
inputName: inNames[0],
|
||||
outputNames: outNames,
|
||||
joinSBD: cfg.XLMJoinSentences(),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (x *XLM) Active() bool {
|
||||
return true
|
||||
}
|
||||
|
||||
func (x *XLM) Close() {
|
||||
if x.session != nil {
|
||||
_ = x.session.Destroy()
|
||||
x.session = nil
|
||||
}
|
||||
if x.sp != nil {
|
||||
x.sp.Close()
|
||||
x.sp = nil
|
||||
}
|
||||
}
|
||||
|
||||
func (x *XLM) Restore(ctx context.Context, text, language string) (string, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return "", err
|
||||
}
|
||||
text = strings.TrimSpace(normalizeXLMSpaces(text))
|
||||
if text == "" {
|
||||
return text, nil
|
||||
}
|
||||
|
||||
ids, err := x.sp.EncodeAsIDs(text)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
full := make([]int, 0, len(ids)+2)
|
||||
full = append(full, x.sp.BOSID())
|
||||
full = append(full, ids...)
|
||||
full = append(full, x.sp.EOSID())
|
||||
maxLen := x.modelCfg.MaxLength
|
||||
if maxLen <= 2 {
|
||||
maxLen = 256
|
||||
}
|
||||
if len(full) <= maxLen {
|
||||
return x.inferIDs(full)
|
||||
}
|
||||
var parts []string
|
||||
content := full[1 : len(full)-1]
|
||||
step := maxLen - 2
|
||||
for start := 0; start < len(content); {
|
||||
end := start + step
|
||||
if end > len(content) {
|
||||
end = len(content)
|
||||
}
|
||||
chunk := make([]int, 0, end-start+2)
|
||||
chunk = append(chunk, x.sp.BOSID())
|
||||
chunk = append(chunk, content[start:end]...)
|
||||
chunk = append(chunk, x.sp.EOSID())
|
||||
out, err := x.inferIDs(chunk)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if out != "" {
|
||||
parts = append(parts, out)
|
||||
}
|
||||
if end >= len(content) {
|
||||
break
|
||||
}
|
||||
start = end
|
||||
}
|
||||
return strings.TrimSpace(strings.Join(parts, " ")), nil
|
||||
}
|
||||
|
||||
func (x *XLM) inferIDs(inputIDs []int) (string, error) {
|
||||
data := make([]int64, len(inputIDs))
|
||||
for i, id := range inputIDs {
|
||||
data[i] = int64(id)
|
||||
}
|
||||
inputTensor, err := ort.NewTensor(ort.NewShape(1, int64(len(inputIDs))), data)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer inputTensor.Destroy()
|
||||
outputs := make([]ort.Value, len(x.outputNames))
|
||||
if err := x.session.Run([]ort.Value{inputTensor}, outputs); err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer destroyValues(outputs)
|
||||
pre, err := int64Row(outputs[0], len(inputIDs))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
post, err := int64Row(outputs[1], len(inputIDs))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
cap, err := boolMatrix(outputs[2], len(inputIDs))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
sbd, err := boolRow(outputs[3], len(inputIDs))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return decodeXLMSegment(x.sp, x.modelCfg, inputIDs, pre, post, cap, sbd, x.joinSBD)
|
||||
}
|
||||
|
||||
func destroyValues(vals []ort.Value) {
|
||||
for _, v := range vals {
|
||||
if v != nil {
|
||||
_ = v.Destroy()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func int64Row(v ort.Value, wantLen int) ([]int64, error) {
|
||||
switch t := v.(type) {
|
||||
case *ort.Tensor[int64]:
|
||||
d := t.GetData()
|
||||
if len(d) == wantLen {
|
||||
return d, nil
|
||||
}
|
||||
if len(d) > wantLen {
|
||||
return d[len(d)-wantLen:], nil
|
||||
}
|
||||
return nil, fmt.Errorf("int64 output short: %d < %d", len(d), wantLen)
|
||||
case *ort.Tensor[int32]:
|
||||
d := t.GetData()
|
||||
if len(d) > wantLen {
|
||||
d = d[len(d)-wantLen:]
|
||||
}
|
||||
out := make([]int64, wantLen)
|
||||
for i := 0; i < wantLen && i < len(d); i++ {
|
||||
out[i] = int64(d[i])
|
||||
}
|
||||
return out, nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unexpected int output type %T", v)
|
||||
}
|
||||
}
|
||||
|
||||
func boolRow(v ort.Value, wantLen int) ([]bool, error) {
|
||||
switch t := v.(type) {
|
||||
case *ort.Tensor[bool]:
|
||||
d := t.GetData()
|
||||
if len(d) == wantLen {
|
||||
return d, nil
|
||||
}
|
||||
if len(d) > wantLen {
|
||||
return d[len(d)-wantLen:], nil
|
||||
}
|
||||
return nil, fmt.Errorf("bool output short")
|
||||
case *ort.Tensor[float32]:
|
||||
d := t.GetData()
|
||||
out := make([]bool, wantLen)
|
||||
for i := 0; i < wantLen && i < len(d); i++ {
|
||||
out[i] = d[i] > 0.5
|
||||
}
|
||||
return out, nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unexpected bool output type %T", v)
|
||||
}
|
||||
}
|
||||
|
||||
func boolMatrix(v ort.Value, seqLen int) ([][]bool, error) {
|
||||
switch t := v.(type) {
|
||||
case *ort.Tensor[bool]:
|
||||
shape := t.GetShape()
|
||||
d := t.GetData()
|
||||
if len(shape) == 3 {
|
||||
_, sl, width := shape[0], shape[1], shape[2]
|
||||
out := make([][]bool, sl)
|
||||
for i := 0; i < int(sl); i++ {
|
||||
row := make([]bool, width)
|
||||
base := int(i) * int(width)
|
||||
copy(row, d[base:base+int(width)])
|
||||
out[i] = row
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
width := len(d) / seqLen
|
||||
if width < 1 {
|
||||
width = 1
|
||||
}
|
||||
out := make([][]bool, seqLen)
|
||||
for i := 0; i < seqLen; i++ {
|
||||
row := make([]bool, width)
|
||||
base := i * width
|
||||
if base+width <= len(d) {
|
||||
copy(row, d[base:base+width])
|
||||
}
|
||||
out[i] = row
|
||||
}
|
||||
return out, nil
|
||||
case *ort.Tensor[float32]:
|
||||
shape := t.GetShape()
|
||||
d := t.GetData()
|
||||
if len(shape) == 3 {
|
||||
_, sl, width := shape[0], shape[1], shape[2]
|
||||
out := make([][]bool, sl)
|
||||
for i := 0; i < int(sl); i++ {
|
||||
row := make([]bool, width)
|
||||
base := int(i) * int(width)
|
||||
for j := 0; j < int(width); j++ {
|
||||
row[j] = d[base+j] > 0.5
|
||||
}
|
||||
out[i] = row
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
width := len(d) / seqLen
|
||||
if width < 1 {
|
||||
width = 1
|
||||
}
|
||||
out := make([][]bool, seqLen)
|
||||
for i := 0; i < seqLen; i++ {
|
||||
row := make([]bool, width)
|
||||
base := i * width
|
||||
for j := 0; j < width && base+j < len(d); j++ {
|
||||
row[j] = d[base+j] > 0.5
|
||||
}
|
||||
out[i] = row
|
||||
}
|
||||
return out, nil
|
||||
default:
|
||||
return nil, fmt.Errorf("unexpected cap output type %T", v)
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeXLMSpaces(s string) string {
|
||||
var b strings.Builder
|
||||
prevSpace := false
|
||||
for _, r := range s {
|
||||
if unicode.IsSpace(r) {
|
||||
if !prevSpace {
|
||||
b.WriteRune(' ')
|
||||
prevSpace = true
|
||||
}
|
||||
continue
|
||||
}
|
||||
prevSpace = false
|
||||
b.WriteRune(r)
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
@@ -0,0 +1,43 @@
|
||||
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")
|
||||
}
|
||||
@@ -0,0 +1,97 @@
|
||||
//go:build xlm
|
||||
|
||||
package punctuation
|
||||
|
||||
import (
|
||||
"strings"
|
||||
|
||||
"go-whisper-api/punctuation/internal/spwrap"
|
||||
)
|
||||
|
||||
func decodeXLMSegment(
|
||||
sp *spwrap.Processor,
|
||||
cfg xlmModelConfig,
|
||||
inputIDs []int,
|
||||
prePred, postPred []int64,
|
||||
capPred [][]bool,
|
||||
sbdPred []bool,
|
||||
joinSentences bool,
|
||||
) (string, error) {
|
||||
var outputTexts []string
|
||||
current := make([]string, 0, len(inputIDs)*4)
|
||||
for tokenIdx := 1; tokenIdx < len(inputIDs)-1; tokenIdx++ {
|
||||
piece, err := sp.IDToPiece(inputIDs[tokenIdx])
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if strings.HasPrefix(piece, "▁") && len(current) > 0 {
|
||||
current = append(current, " ")
|
||||
}
|
||||
preLabel := labelAt(cfg.PreLabels, prePred, tokenIdx)
|
||||
postLabel := labelAt(cfg.PostLabels, postPred, tokenIdx)
|
||||
if preLabel != cfg.NullToken {
|
||||
current = append(current, preLabel)
|
||||
}
|
||||
charStart := 0
|
||||
if strings.HasPrefix(piece, "▁") {
|
||||
charStart = 1
|
||||
}
|
||||
runes := []rune(piece)
|
||||
for tokenCharIdx := charStart; tokenCharIdx < len(runes); tokenCharIdx++ {
|
||||
ch := string(runes[tokenCharIdx])
|
||||
if capAt(capPred, tokenIdx, tokenCharIdx) {
|
||||
ch = strings.ToUpper(ch)
|
||||
}
|
||||
current = append(current, ch)
|
||||
if postLabel == cfg.Acronym {
|
||||
current = append(current, ".")
|
||||
}
|
||||
}
|
||||
if postLabel != cfg.NullToken && postLabel != cfg.Acronym {
|
||||
current = append(current, postLabel)
|
||||
}
|
||||
if sbdAt(sbdPred, tokenIdx) {
|
||||
outputTexts = append(outputTexts, strings.Join(current, ""))
|
||||
current = current[:0]
|
||||
}
|
||||
}
|
||||
if len(current) > 0 {
|
||||
outputTexts = append(outputTexts, strings.Join(current, ""))
|
||||
}
|
||||
if len(outputTexts) == 0 {
|
||||
return "", nil
|
||||
}
|
||||
if joinSentences {
|
||||
return strings.Join(outputTexts, " "), nil
|
||||
}
|
||||
return outputTexts[0], nil
|
||||
}
|
||||
|
||||
func labelAt(labels []string, preds []int64, idx int) string {
|
||||
if idx < 0 || idx >= len(preds) {
|
||||
return labels[0]
|
||||
}
|
||||
pi := int(preds[idx])
|
||||
if pi < 0 || pi >= len(labels) {
|
||||
return labels[0]
|
||||
}
|
||||
return labels[pi]
|
||||
}
|
||||
|
||||
func capAt(capPred [][]bool, tokenIdx, charIdx int) bool {
|
||||
if tokenIdx < 0 || tokenIdx >= len(capPred) {
|
||||
return false
|
||||
}
|
||||
row := capPred[tokenIdx]
|
||||
if charIdx < 0 || charIdx >= len(row) {
|
||||
return false
|
||||
}
|
||||
return row[charIdx]
|
||||
}
|
||||
|
||||
func sbdAt(sbd []bool, tokenIdx int) bool {
|
||||
if tokenIdx < 0 || tokenIdx >= len(sbd) {
|
||||
return false
|
||||
}
|
||||
return sbd[tokenIdx]
|
||||
}
|
||||
@@ -0,0 +1,13 @@
|
||||
//go:build !xlm
|
||||
|
||||
package punctuation
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"go-whisper-api/config"
|
||||
)
|
||||
|
||||
func newXLM(cfg config.Punctuation) (Restorer, error) {
|
||||
return nil, fmt.Errorf("punctuation engine %q requires build tag xlm (run: make build-xlm)", cfg.Engine)
|
||||
}
|
||||
Reference in New Issue
Block a user