Include speaker name in transcription text

mattermost · Nov 1, 2023 · 4018fd4 · 4018fd4
1 parent cb2b277
commit 4018fd4
Show file tree

Hide file tree

Showing 3 changed files with 100 additions and 11 deletions.
diff --git a/cmd/transcriber/call/tracks.go b/cmd/transcriber/call/tracks.go
@@ -210,7 +210,9 @@ func (t *Transcriber) handleClose() error {
 	}
 	defer f.Close()
 
-	if err := tr.WebVTT(f); err != nil {
+	if err := tr.WebVTT(f, transcribe.WebVTTOptions{
+		OmitSpeaker: false,
+	}); err != nil {
 		return fmt.Errorf("failed to write WebVTT file: %w", err)
 	}
 

diff --git a/cmd/transcriber/transcribe/output.go b/cmd/transcriber/transcribe/output.go
@@ -44,7 +44,11 @@ func (t Transcription) interleave() []namedSegment {
 	return nss
 }
 
-func (t Transcription) WebVTT(w io.Writer) error {
+type WebVTTOptions struct {
+	OmitSpeaker bool
+}
+
+func (t Transcription) WebVTT(w io.Writer, opts WebVTTOptions) error {
 	_, err := fmt.Fprintf(w, "WEBVTT\n")
 	if err != nil {
 		return fmt.Errorf("failed to write: %w", err)
@@ -54,7 +58,11 @@ func (t Transcription) WebVTT(w io.Writer) error {
 		if err != nil {
 			return fmt.Errorf("failed to write: %w", err)
 		}
-		_, err = fmt.Fprintf(w, "<v %s>%s\n", s.Speaker, s.Text)
+		tmpl := "<v %[1]s>(%[1]s) %[2]s\n"
+		if opts.OmitSpeaker {
+			tmpl = "%[2]s\n"
+		}
+		_, err = fmt.Fprintf(w, tmpl, s.Speaker, s.Text)
 		if err != nil {
 			return fmt.Errorf("failed to write: %w", err)
 		}

diff --git a/cmd/transcriber/transcribe/output_test.go b/cmd/transcriber/transcribe/output_test.go
@@ -208,7 +208,7 @@ func TestWebVTT(t *testing.T) {
 	t.Run("empty", func(t *testing.T) {
 		var tr Transcription
 		var b strings.Builder
-		err := tr.WebVTT(&b)
+		err := tr.WebVTT(&b, WebVTTOptions{})
 		require.NoError(t, err)
 		require.Equal(t, "WEBVTT\n", b.String())
 	})
@@ -266,24 +266,103 @@ func TestWebVTT(t *testing.T) {
 		expected := `WEBVTT
 
 00:00:00.000 --> 00:00:01.000
-<v SpeakerA>A1
+<v SpeakerA>(SpeakerA) A1
 
 00:00:02.000 --> 00:00:03.000
-<v SpeakerA>A2
+<v SpeakerA>(SpeakerA) A2
 
 00:00:03.000 --> 00:00:04.000
-<v SpeakerB>B1
+<v SpeakerB>(SpeakerB) B1
 
 00:00:04.000 --> 00:00:05.000
-<v SpeakerA>A3
+<v SpeakerA>(SpeakerA) A3
 
 00:00:05.000 --> 00:00:06.000
-<v SpeakerA>A4
+<v SpeakerA>(SpeakerA) A4
 
 00:00:06.000 --> 00:00:07.000
-<v SpeakerB>B2
+<v SpeakerB>(SpeakerB) B2
 `
-		err := tr.WebVTT(&b)
+		err := tr.WebVTT(&b, WebVTTOptions{
+			OmitSpeaker: false,
+		})
+		require.NoError(t, err)
+		require.Equal(t, expected, b.String())
+	})
+
+	t.Run("omit speaker", func(t *testing.T) {
+		tr := Transcription{
+			TrackTranscription{
+				Speaker: "SpeakerA",
+				Segments: []Segment{
+					{
+						StartTS: 0,
+						EndTS:   1000,
+						Text:    "A1",
+					},
+					{
+						StartTS: 2000,
+						EndTS:   3000,
+						Text:    "A2",
+					},
+				},
+			},
+			TrackTranscription{
+				Speaker: "SpeakerA",
+				Segments: []Segment{
+					{
+						StartTS: 4000,
+						EndTS:   5000,
+						Text:    "A3",
+					},
+					{
+						StartTS: 5000,
+						EndTS:   6000,
+						Text:    "A4",
+					},
+				},
+			},
+			TrackTranscription{
+				Speaker: "SpeakerB",
+				Segments: []Segment{
+					{
+						StartTS: 3000,
+						EndTS:   4000,
+						Text:    "B1",
+					},
+					{
+						StartTS: 6000,
+						EndTS:   7000,
+						Text:    "B2",
+					},
+				},
+			},
+		}
+
+		var b strings.Builder
+		expected := `WEBVTT
+
+00:00:00.000 --> 00:00:01.000
+A1
+
+00:00:02.000 --> 00:00:03.000
+A2
+
+00:00:03.000 --> 00:00:04.000
+B1
+
+00:00:04.000 --> 00:00:05.000
+A3
+
+00:00:05.000 --> 00:00:06.000
+A4
+
+00:00:06.000 --> 00:00:07.000
+B2
+`
+		err := tr.WebVTT(&b, WebVTTOptions{
+			OmitSpeaker: true,
+		})
 		require.NoError(t, err)
 		require.Equal(t, expected, b.String())
 	})