279 lines
8.0 KiB
Go
279 lines
8.0 KiB
Go
package stage
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"sort"
|
|
"strings"
|
|
|
|
"golang.org/x/sync/errgroup"
|
|
|
|
"gitea.maximumdirect.net/eric/narratio/internal/adapters/whisperx"
|
|
"gitea.maximumdirect.net/eric/narratio/internal/artifacts"
|
|
"gitea.maximumdirect.net/eric/narratio/internal/manifest"
|
|
)
|
|
|
|
type transcribeStage struct{}
|
|
|
|
func (transcribeStage) Name() string { return "transcribe" }
|
|
|
|
func (transcribeStage) Run(ctx context.Context, env *Env, m *manifest.Manifest) (*StageResult, error) {
|
|
if env == nil || env.Config == nil {
|
|
return nil, fmt.Errorf("transcribe: stage environment config is required")
|
|
}
|
|
if env.ArtifactStore == nil {
|
|
return nil, fmt.Errorf("transcribe: artifact store is required")
|
|
}
|
|
if env.ManifestStore == nil {
|
|
return nil, fmt.Errorf("transcribe: manifest store is required")
|
|
}
|
|
if env.Config.Pipeline == nil || env.Config.Session == nil {
|
|
return nil, fmt.Errorf("transcribe: resolved config must include pipeline and session")
|
|
}
|
|
if m == nil {
|
|
return nil, fmt.Errorf("transcribe: session manifest is required")
|
|
}
|
|
|
|
var sessionID string
|
|
if m != nil {
|
|
sessionID = strings.TrimSpace(m.SessionID)
|
|
}
|
|
if sessionID == "" {
|
|
sessionID = strings.TrimSpace(env.Config.Session.SessionID)
|
|
}
|
|
if sessionID == "" {
|
|
return nil, fmt.Errorf("transcribe: session id is required")
|
|
}
|
|
|
|
paths := sessionPathsForEnv(env, sessionID)
|
|
reconciliation, err := reconcileTranscribeSpeakers(ctx, env, m)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("transcribe: reconcile speakers: %w", err)
|
|
}
|
|
state := newTranscribeExecutionState(env, m, paths, reconciliation.current)
|
|
if err := state.checkpointStartingCollection(ctx); err != nil {
|
|
return nil, fmt.Errorf("transcribe: checkpoint reconciled speaker collection before execution: %w", err)
|
|
}
|
|
|
|
jobs := append([]transcribeReconciliationJob(nil), reconciliation.work...)
|
|
if len(jobs) > 0 && env.WhisperX == nil {
|
|
return nil, fmt.Errorf("transcribe: whisperx adapter is required")
|
|
}
|
|
if len(jobs) > 0 {
|
|
runLayout, err := resolveRunStageLayout(env, m, paths, sessionID, "transcribe")
|
|
if err != nil {
|
|
return nil, fmt.Errorf("transcribe: resolve run-stage layout: %w", err)
|
|
}
|
|
for index := range jobs {
|
|
jobs[index].runOutputPath, err = runLocalPathForCanonical(runLayout, paths, jobs[index].canonicalOutputPath)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("transcribe: resolve run-local output path for %q: %w", jobs[index].speakerID, err)
|
|
}
|
|
}
|
|
}
|
|
|
|
concurrency := 1
|
|
if env.Config.Pipeline.WhisperX.Concurrency != nil && *env.Config.Pipeline.WhisperX.Concurrency > 0 {
|
|
concurrency = *env.Config.Pipeline.WhisperX.Concurrency
|
|
}
|
|
if len(jobs) > 0 && concurrency > len(jobs) {
|
|
concurrency = len(jobs)
|
|
}
|
|
|
|
group, groupCtx := errgroup.WithContext(ctx)
|
|
group.SetLimit(concurrency)
|
|
for _, job := range jobs {
|
|
job := job
|
|
group.Go(func() error {
|
|
if err := groupCtx.Err(); err != nil {
|
|
return err
|
|
}
|
|
if err := state.checkpointFailure(); err != nil {
|
|
return err
|
|
}
|
|
res, runErr := env.WhisperX.Transcribe(groupCtx, whisperx.TranscribeRequest{
|
|
SpeakerID: job.speakerID,
|
|
AudioPath: job.audioPath,
|
|
OutputRawTranscriptPath: job.runOutputPath,
|
|
})
|
|
if runErr != nil {
|
|
return fmt.Errorf("speaker %q (%s): %w", job.speakerID, filepath.Base(job.audioPath), runErr)
|
|
}
|
|
|
|
meta := map[string]any{
|
|
"attempts": res.Attempts,
|
|
"http_status": res.HTTPStatus,
|
|
"duration_ms": res.Duration.Milliseconds(),
|
|
"output_path": job.runOutputPath,
|
|
}
|
|
|
|
if err := state.acceptSpeaker(ctx, job, res.OutputRawTranscriptPath, meta); err != nil {
|
|
return fmt.Errorf("speaker %q: %w", job.speakerID, err)
|
|
}
|
|
return nil
|
|
})
|
|
}
|
|
groupErr := group.Wait()
|
|
if err := state.checkpointFailure(); err != nil {
|
|
return nil, err
|
|
}
|
|
completedCount := state.completedCount()
|
|
if err := ctx.Err(); err != nil {
|
|
return nil, fmt.Errorf("transcribe: canceled after completing %d of %d audio jobs: %w", completedCount, len(jobs), err)
|
|
}
|
|
if groupErr != nil {
|
|
return nil, groupErr
|
|
}
|
|
if completedCount != len(jobs) {
|
|
return nil, fmt.Errorf("transcribe: completed %d of %d planned audio jobs", completedCount, len(jobs))
|
|
}
|
|
outputs, outputPaths, runOutputPaths, orderedPerFile, err := state.successDetails()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("transcribe: assemble completed outputs: %w", err)
|
|
}
|
|
|
|
retries := 0
|
|
if env.Config.Pipeline.WhisperX.Retries != nil {
|
|
retries = *env.Config.Pipeline.WhisperX.Retries
|
|
}
|
|
|
|
return &StageResult{
|
|
Outputs: outputs,
|
|
Metadata: map[string]any{
|
|
"stage": "transcribe",
|
|
"audio_files_count": len(reconciliation.jobs),
|
|
"language": env.Config.Pipeline.WhisperX.Language,
|
|
"concurrency": concurrency,
|
|
"retries": retries,
|
|
"retry_delay": env.Config.Pipeline.WhisperX.RetryDelay,
|
|
"timeout": env.Config.Pipeline.WhisperX.Timeout,
|
|
"run_output_paths": runOutputPaths,
|
|
"output_paths": outputPaths,
|
|
"per_file": orderedPerFile,
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
type transcribeJob struct {
|
|
speakerID string
|
|
audioPath string
|
|
canonicalOutputPath string
|
|
runOutputPath string
|
|
}
|
|
|
|
func buildTranscribeJobs(m *manifest.Manifest, paths artifacts.SessionPaths) ([]transcribeJob, error) {
|
|
audioFiles, err := discoverPreparedAudio(m, paths)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
jobs := make([]transcribeJob, 0, len(audioFiles))
|
|
seenSpeaker := map[string]string{}
|
|
for _, audioPath := range audioFiles {
|
|
base := strings.TrimSuffix(filepath.Base(audioPath), filepath.Ext(audioPath))
|
|
if base == "" {
|
|
return nil, fmt.Errorf("could not derive speaker id from %q", audioPath)
|
|
}
|
|
if previous, exists := seenSpeaker[base]; exists {
|
|
return nil, fmt.Errorf("duplicate speaker/audio basename %q from %q and %q", base, previous, audioPath)
|
|
}
|
|
seenSpeaker[base] = audioPath
|
|
jobs = append(jobs, transcribeJob{
|
|
speakerID: base,
|
|
audioPath: audioPath,
|
|
canonicalOutputPath: filepath.Join(paths.TranscriptsRawDir, base+".json"),
|
|
})
|
|
}
|
|
return jobs, nil
|
|
}
|
|
|
|
func discoverPreparedAudio(m *manifest.Manifest, paths artifacts.SessionPaths) ([]string, error) {
|
|
fromManifest := make([]string, 0)
|
|
if m != nil {
|
|
for _, in := range m.Inputs {
|
|
if in.Kind != "audio" {
|
|
continue
|
|
}
|
|
p := strings.TrimSpace(in.Path)
|
|
if p == "" {
|
|
continue
|
|
}
|
|
fromManifest = append(fromManifest, artifacts.ResolveSessionLocalPathForRead(paths, p))
|
|
}
|
|
}
|
|
if len(fromManifest) > 0 {
|
|
sort.Strings(fromManifest)
|
|
if err := validateAudioFiles(fromManifest); err != nil {
|
|
return nil, err
|
|
}
|
|
return fromManifest, nil
|
|
}
|
|
|
|
entries, err := os.ReadDir(paths.AudioDir)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("read audio directory %q: %w", paths.AudioDir, err)
|
|
}
|
|
|
|
out := make([]string, 0)
|
|
for _, entry := range entries {
|
|
if entry.IsDir() {
|
|
continue
|
|
}
|
|
p := filepath.Join(paths.AudioDir, entry.Name())
|
|
if !isFlac(p) {
|
|
continue
|
|
}
|
|
out = append(out, p)
|
|
}
|
|
sort.Strings(out)
|
|
if err := validateAudioFiles(out); err != nil {
|
|
return nil, err
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func validateAudioFiles(paths []string) error {
|
|
if len(paths) == 0 {
|
|
return nil
|
|
}
|
|
seen := map[string]struct{}{}
|
|
for _, p := range paths {
|
|
clean := filepath.Clean(strings.TrimSpace(p))
|
|
if clean == "" {
|
|
return fmt.Errorf("audio path is required")
|
|
}
|
|
if !isFlac(clean) {
|
|
return fmt.Errorf("audio file %q must have .flac extension", clean)
|
|
}
|
|
info, err := os.Stat(clean)
|
|
if err != nil {
|
|
return fmt.Errorf("audio file %q not found: %w", clean, err)
|
|
}
|
|
if !info.Mode().IsRegular() {
|
|
return fmt.Errorf("audio file %q is not a regular file", clean)
|
|
}
|
|
if _, exists := seen[clean]; exists {
|
|
return fmt.Errorf("duplicate audio file path %q", clean)
|
|
}
|
|
seen[clean] = struct{}{}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func validateTranscriptJSONFile(path string) error {
|
|
if err := requireFile(path, "raw transcript"); err != nil {
|
|
return err
|
|
}
|
|
data, err := readExternalResult(path, "whisperx transcript result")
|
|
if err != nil {
|
|
return err
|
|
}
|
|
var payload any
|
|
if err := json.Unmarshal(data, &payload); err != nil {
|
|
return fmt.Errorf("decode json: %w", err)
|
|
}
|
|
return nil
|
|
}
|