record-asr/internal/transcribe/pipeline.go

// Package transcribe ties the silero VAD and the SenseVoice recognizer together
// into a single "audio file -> transcript" pipeline.
//
// The VAD is stateful, so segmentation of one file cannot interleave with
// another; the manager guards VAD access with a mutex. Recognition of the
// resulting speech segments is stateless and could be parallelized.
package transcribe

import (
	"fmt"
	"strings"

	sherpa "github.com/k2-fsa/sherpa-onnx-go/sherpa_onnx"

	"git.misaka.ren/m1saka/record-asr/internal/audio"
	"git.misaka.ren/m1saka/record-asr/internal/models"
)

// Segment is one VAD-detected speech region with its transcript.
type Segment struct {
	Text  string  `json:"text"`
	Start float64 `json:"start"` // seconds from file start
	End   float64 `json:"end"`   // seconds from file start
	Lang  string  `json:"lang"`
}

// Result is the full transcript for one audio file.
type Result struct {
	Text     string    `json:"text"`      // all segments joined by newline
	Lang     string    `json:"lang"`      // dominant detected language
	Emotion  string    `json:"emotion"`   // dominant detected emotion
	Duration float64   `json:"duration"`  // audio duration in seconds
	Segments []Segment `json:"segments"`
}

// New returns a Pipeline bound to the given model manager.
func New(m *models.Manager) *Pipeline {
	return &Pipeline{m: m}
}

// Pipeline wraps a model manager to provide file transcription.
type Pipeline struct {
	m *models.Manager
}

// Transcribe decodes the audio file at path, runs VAD to isolate speech, and
// recognizes each segment with SenseVoice.
func (p *Pipeline) Transcribe(path string) (*Result, error) {
	samples, sampleRate, err := audio.Decode(path)
	if err != nil {
		return nil, err
	}
	if sampleRate != audio.TargetSampleRate {
		return nil, fmt.Errorf("internal: expected 16kHz samples, got %d", sampleRate)
	}
	duration := float64(len(samples)) / float64(audio.TargetSampleRate)

	segments, err := p.segment(samples)
	if err != nil {
		return nil, err
	}

	res := &Result{
		Duration: duration,
		Segments: make([]Segment, 0, len(segments)),
	}
	var texts []string
	for _, seg := range segments {
		// Recognize, splitting extremely long segments (should not happen
		// given the VAD MaxSpeechDuration cap, but guards against pathological
		// inputs that exceed SenseVoice's ~30s comfort window).
		parts := splitForRecognize(seg.Samples, 28*audio.TargetSampleRate)
		var segTexts []string
		var segLang, segEmotion string
		segStart := float64(seg.Start) / float64(audio.TargetSampleRate)
		for _, part := range parts {
			txt, lang, emo, perr := p.recognize(part)
			if perr != nil {
				return nil, fmt.Errorf("recognize segment at %.2fs: %w", segStart, perr)
			}
			if txt != "" {
				segTexts = append(segTexts, txt)
			}
			if segLang == "" {
				segLang, segEmotion = lang, emo
			}
		}
		segText := strings.Join(segTexts, " ")
		segEnd := segStart + float64(len(seg.Samples))/float64(audio.TargetSampleRate)
		if segText == "" {
			segText = " "
		}
		res.Segments = append(res.Segments, Segment{
			Text:  segText,
			Start: segStart,
			End:   segEnd,
			Lang:  segLang,
		})
		texts = append(texts, segText)
		if res.Lang == "" {
			res.Lang, res.Emotion = segLang, segEmotion
		}
	}
	res.Text = strings.Join(texts, "\n")
	return res, nil
}

// segment runs the silero VAD over the full sample buffer and returns the
// completed speech segments. The VAD is stateful and shared, so this is the
// only place that touches it, under the manager's lock, and it resets the
// detector before use so no state leaks between files.
func (p *Pipeline) segment(samples []float32) ([]sherpa.SpeechSegment, error) {
	if p.m.VAD == nil {
		return nil, fmt.Errorf("VAD not loaded")
	}
	p.m.VadMu().Lock()
	defer p.m.VadMu().Unlock()

	vad := p.m.VAD
	vad.Reset()

	// Feed in window-sized chunks; drain completed segments after each feed.
	w := 512 // silero window at 16kHz
	if len(samples) > 0 && w > len(samples) {
		w = len(samples)
	}
	var out []sherpa.SpeechSegment
	drain := func() {
		for !vad.IsEmpty() {
			seg := vad.Front()
			if seg == nil {
				break
			}
			// Copy because Front()'s backing memory is freed on Pop().
			cp := sherpa.SpeechSegment{Start: seg.Start, Samples: append([]float32(nil), seg.Samples...)}
			out = append(out, cp)
			vad.Pop()
		}
	}
	for i := 0; i < len(samples); i += w {
		end := i + w
		if end > len(samples) {
			end = len(samples)
		}
		vad.AcceptWaveform(samples[i:end])
		drain()
	}
	vad.Flush()
	drain()
	return out, nil
}

// recognize runs SenseVoice over a single contiguous speech chunk.
func (p *Pipeline) recognize(samples []float32) (text, lang, emotion string, err error) {
	if len(samples) == 0 {
		return "", "", "", nil
	}
	if p.m.Recognizer == nil {
		return "", "", "", fmt.Errorf("recognizer not loaded")
	}
	stream := sherpa.NewOfflineStream(p.m.Recognizer)
	if stream == nil {
		return "", "", "", fmt.Errorf("create offline stream failed")
	}
	defer sherpa.DeleteOfflineStream(stream)
	stream.AcceptWaveform(audio.TargetSampleRate, samples)
	p.m.Recognizer.Decode(stream)
	r := stream.GetResult()
	if r == nil {
		return "", "", "", nil
	}
	return r.Text, r.Lang, r.Emotion, nil
}

// splitForRecognize breaks a sample slice into chunks no longer than maxLen to
// stay within the recognizer's comfortable input window.
func splitForRecognize(samples []float32, maxLen int) [][]float32 {
	if maxLen <= 0 || len(samples) <= maxLen {
		return [][]float32{samples}
	}
	var parts [][]float32
	for i := 0; i < len(samples); i += maxLen {
		end := i + maxLen
		if end > len(samples) {
			end = len(samples)
		}
		parts = append(parts, samples[i:end])
	}
	return parts
}