Eric Rakestraw
b20438acf0
All checks were successful
ci/woodpecker/tag/release Pipeline was successful
308 lines
8.4 KiB
Go
308 lines
8.4 KiB
Go
package normalize
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
"strings"
|
|
)
|
|
|
|
// InputShape identifies which top-level input shape was parsed.
|
|
type InputShape string
|
|
|
|
const (
|
|
ShapeObjectWithSegments InputShape = "object_with_segments"
|
|
ShapeBareSegmentsArray InputShape = "bare_segments_array"
|
|
)
|
|
|
|
// ParsedTranscript is the validated normalize input model.
|
|
type ParsedTranscript struct {
|
|
Shape InputShape
|
|
InputSegmentCount int
|
|
Stats NormalizeStats
|
|
Segments []InputSegment
|
|
}
|
|
|
|
// InputSegment is a validated segment from normalize input.
|
|
type InputSegment struct {
|
|
InputIndex int
|
|
OriginalID *int
|
|
StartPresent bool
|
|
EndPresent bool
|
|
SpeakerPresent bool
|
|
TextPresent bool
|
|
Start float64
|
|
End float64
|
|
Speaker string
|
|
Text string
|
|
Categories []string
|
|
Source string
|
|
SourceSegmentIndex *int
|
|
SourceRef string
|
|
DerivedFrom []string
|
|
OverlapGroupID *int
|
|
}
|
|
|
|
// NormalizeStats captures deterministic repair/drop outcomes from normalize input processing.
|
|
type NormalizeStats struct {
|
|
TimingFieldsRepaired int
|
|
TimingOrderSwapped int
|
|
SpeakerFilled int
|
|
SegmentsDroppedText int
|
|
}
|
|
|
|
type inputSegmentPayload struct {
|
|
ID *int `json:"id"`
|
|
Start *float64 `json:"start"`
|
|
End *float64 `json:"end"`
|
|
Speaker *string `json:"speaker"`
|
|
Text *string `json:"text"`
|
|
Categories []string `json:"categories"`
|
|
Source string `json:"source"`
|
|
SourceSegmentIndex *int `json:"source_segment_index"`
|
|
SourceRef string `json:"source_ref"`
|
|
DerivedFrom []string `json:"derived_from"`
|
|
OverlapGroupID *int `json:"overlap_group_id"`
|
|
}
|
|
|
|
// ParseFile parses normalize input JSON from file path.
|
|
func ParseFile(path string) (ParsedTranscript, error) {
|
|
file, err := os.Open(path)
|
|
if err != nil {
|
|
return ParsedTranscript{}, err
|
|
}
|
|
defer file.Close()
|
|
|
|
return ParseReader(file)
|
|
}
|
|
|
|
// ParseReader parses normalize input JSON from a reader.
|
|
func ParseReader(reader io.Reader) (ParsedTranscript, error) {
|
|
var raw json.RawMessage
|
|
decoder := json.NewDecoder(reader)
|
|
decoder.UseNumber()
|
|
if err := decoder.Decode(&raw); err != nil {
|
|
return ParsedTranscript{}, fmt.Errorf("decode normalize input JSON: %w", err)
|
|
}
|
|
if err := ensureSingleValue(decoder); err != nil {
|
|
return ParsedTranscript{}, err
|
|
}
|
|
|
|
trimmed := bytes.TrimSpace(raw)
|
|
if len(trimmed) == 0 {
|
|
return ParsedTranscript{}, fmt.Errorf("normalize input is empty")
|
|
}
|
|
|
|
switch trimmed[0] {
|
|
case '{':
|
|
return parseObjectShape(trimmed)
|
|
case '[':
|
|
segments, stats, inputCount, err := parseSegmentsArray(trimmed)
|
|
if err != nil {
|
|
return ParsedTranscript{}, err
|
|
}
|
|
return ParsedTranscript{
|
|
Shape: ShapeBareSegmentsArray,
|
|
InputSegmentCount: inputCount,
|
|
Stats: stats,
|
|
Segments: segments,
|
|
}, nil
|
|
default:
|
|
return ParsedTranscript{}, fmt.Errorf("normalize input must be a top-level object with \"segments\" or a top-level segment array")
|
|
}
|
|
}
|
|
|
|
func ensureSingleValue(decoder *json.Decoder) error {
|
|
var extra json.RawMessage
|
|
err := decoder.Decode(&extra)
|
|
if err == io.EOF {
|
|
return nil
|
|
}
|
|
if err == nil {
|
|
return fmt.Errorf("normalize input must contain exactly one top-level JSON value")
|
|
}
|
|
return fmt.Errorf("decode normalize input JSON: %w", err)
|
|
}
|
|
|
|
func parseObjectShape(raw []byte) (ParsedTranscript, error) {
|
|
var object map[string]json.RawMessage
|
|
if err := json.Unmarshal(raw, &object); err != nil {
|
|
return ParsedTranscript{}, fmt.Errorf("decode normalize object input: %w", err)
|
|
}
|
|
|
|
segmentsRaw, exists := object["segments"]
|
|
if !exists {
|
|
return ParsedTranscript{}, fmt.Errorf("normalize object input must contain a \"segments\" field")
|
|
}
|
|
|
|
segments, stats, inputCount, err := parseSegmentsArray(segmentsRaw)
|
|
if err != nil {
|
|
return ParsedTranscript{}, err
|
|
}
|
|
|
|
return ParsedTranscript{
|
|
Shape: ShapeObjectWithSegments,
|
|
InputSegmentCount: inputCount,
|
|
Stats: stats,
|
|
Segments: segments,
|
|
}, nil
|
|
}
|
|
|
|
func parseSegmentsArray(raw []byte) ([]InputSegment, NormalizeStats, int, error) {
|
|
var segmentValues []json.RawMessage
|
|
if err := json.Unmarshal(raw, &segmentValues); err != nil {
|
|
return nil, NormalizeStats{}, 0, fmt.Errorf("normalize input \"segments\" must be an array")
|
|
}
|
|
|
|
segments := make([]InputSegment, len(segmentValues))
|
|
for index, segmentRaw := range segmentValues {
|
|
segment, err := decodeSegment(index, segmentRaw)
|
|
if err != nil {
|
|
return nil, NormalizeStats{}, 0, err
|
|
}
|
|
segments[index] = segment
|
|
}
|
|
normalized, stats, err := normalizeSegments(segments)
|
|
return normalized, stats, len(segmentValues), err
|
|
}
|
|
|
|
func decodeSegment(index int, raw []byte) (InputSegment, error) {
|
|
var payload inputSegmentPayload
|
|
if err := json.Unmarshal(raw, &payload); err != nil {
|
|
return InputSegment{}, fmt.Errorf("segment %d: invalid segment object: %w", index, err)
|
|
}
|
|
|
|
var start float64
|
|
if payload.Start != nil {
|
|
start = *payload.Start
|
|
}
|
|
var end float64
|
|
if payload.End != nil {
|
|
end = *payload.End
|
|
}
|
|
|
|
speaker := ""
|
|
if payload.Speaker != nil {
|
|
speaker = strings.TrimSpace(*payload.Speaker)
|
|
}
|
|
|
|
text := ""
|
|
if payload.Text != nil {
|
|
text = *payload.Text
|
|
}
|
|
|
|
return InputSegment{
|
|
InputIndex: index,
|
|
OriginalID: payload.ID,
|
|
StartPresent: payload.Start != nil,
|
|
EndPresent: payload.End != nil,
|
|
SpeakerPresent: payload.Speaker != nil,
|
|
TextPresent: payload.Text != nil,
|
|
Start: start,
|
|
End: end,
|
|
Speaker: speaker,
|
|
Text: text,
|
|
Categories: append([]string(nil), payload.Categories...),
|
|
Source: payload.Source,
|
|
SourceSegmentIndex: payload.SourceSegmentIndex,
|
|
SourceRef: payload.SourceRef,
|
|
DerivedFrom: append([]string(nil), payload.DerivedFrom...),
|
|
OverlapGroupID: payload.OverlapGroupID,
|
|
}, nil
|
|
}
|
|
|
|
func normalizeSegments(segments []InputSegment) ([]InputSegment, NormalizeStats, error) {
|
|
stats := NormalizeStats{}
|
|
filtered := make([]InputSegment, 0, len(segments))
|
|
for _, segment := range segments {
|
|
if strings.TrimSpace(segment.Text) == "" {
|
|
stats.SegmentsDroppedText++
|
|
continue
|
|
}
|
|
filtered = append(filtered, segment)
|
|
}
|
|
|
|
if len(filtered) == 0 {
|
|
return filtered, stats, nil
|
|
}
|
|
|
|
hasStart := make([]bool, len(filtered))
|
|
hasEnd := make([]bool, len(filtered))
|
|
for index, segment := range filtered {
|
|
hasStart[index] = segment.StartPresent
|
|
hasEnd[index] = segment.EndPresent
|
|
}
|
|
for index := range filtered {
|
|
if !hasStart[index] && !hasEnd[index] {
|
|
midpoint := inferTimestampFromNeighbors(filtered, hasStart, hasEnd, index)
|
|
filtered[index].Start = midpoint
|
|
filtered[index].End = midpoint
|
|
stats.TimingFieldsRepaired += 2
|
|
hasStart[index] = true
|
|
hasEnd[index] = true
|
|
continue
|
|
}
|
|
if hasStart[index] && !hasEnd[index] {
|
|
filtered[index].End = filtered[index].Start
|
|
stats.TimingFieldsRepaired++
|
|
hasEnd[index] = true
|
|
}
|
|
if !hasStart[index] && hasEnd[index] {
|
|
filtered[index].Start = filtered[index].End
|
|
stats.TimingFieldsRepaired++
|
|
hasStart[index] = true
|
|
}
|
|
}
|
|
|
|
for index := range filtered {
|
|
if filtered[index].Start < 0 {
|
|
return nil, NormalizeStats{}, fmt.Errorf("segment %d has start %v; start must be >= 0", filtered[index].InputIndex, filtered[index].Start)
|
|
}
|
|
if filtered[index].End < filtered[index].Start {
|
|
filtered[index].Start, filtered[index].End = filtered[index].End, filtered[index].Start
|
|
stats.TimingOrderSwapped++
|
|
}
|
|
if strings.TrimSpace(filtered[index].Speaker) == "" {
|
|
filtered[index].Speaker = "Unknown_Speaker"
|
|
stats.SpeakerFilled++
|
|
}
|
|
}
|
|
|
|
return filtered, stats, nil
|
|
}
|
|
|
|
func inferTimestampFromNeighbors(segments []InputSegment, hasStart []bool, hasEnd []bool, index int) float64 {
|
|
prevEnd, hasPrev := nearestPreviousEnd(segments, hasEnd, index)
|
|
nextStart, hasNext := nearestNextStart(segments, hasStart, index)
|
|
if hasPrev && hasNext {
|
|
return (prevEnd + nextStart) / 2
|
|
}
|
|
if hasPrev {
|
|
return prevEnd
|
|
}
|
|
if hasNext {
|
|
return nextStart
|
|
}
|
|
return 0
|
|
}
|
|
|
|
func nearestPreviousEnd(segments []InputSegment, hasEnd []bool, index int) (float64, bool) {
|
|
for i := index - 1; i >= 0; i-- {
|
|
if hasEnd[i] {
|
|
return segments[i].End, true
|
|
}
|
|
}
|
|
return 0, false
|
|
}
|
|
|
|
func nearestNextStart(segments []InputSegment, hasStart []bool, index int) (float64, bool) {
|
|
for i := index + 1; i < len(segments); i++ {
|
|
if hasStart[i] {
|
|
return segments[i].Start, true
|
|
}
|
|
}
|
|
return 0, false
|
|
}
|