smoores.dev

Storyteller is now a sprawling ecosystem of software, with a full stack web application, native apps for Android and iOS, KOReader plugins, and upcoming macOS, watchOS, and tvOS apps. We're alpha and beta testing v3 releases of the above, which will bring updated UIs, a huge suite of new library management features, and more. We're really excited to show you what we've been working on!

But at the core of Storyteller is its alignment algorithm. Storyteller can take an ebook and an audiobook that you provide it and align them, finding where each word of the ebook is spoken in the audiobook. It does this automatically, without any input from you. Then it uses the EPUB specification's built-in audio synchronization system, called Media Overlays, to embed the audio and synchronization info into the EPUB. That lets you read your book immersively, with your reader app highlighting each sentence as it's narrated by the narrator, just like in the demo above (which is using real Storyteller alignments!).

Initially, that's all Storyteller was: one Python script. It took one audiobook file and one ebook file, and it output a new ebook file with audio synchronization metadata. At the time, there were a very small number of ereader apps (and zero ereader devices) that could consume these files, which used the EPUB Media Overlay specification for their functionality. I would run the script on my computer, copy the resulting EPUB to my phone, and then use BookFusion's fledging Media Overlay support to read and listen to my books.

Back then, I was completely unfamiliar with the field of forced alignment that I was unwittingly sprinting into headlong. My first alignment algorithm was a wobbly construction of clunky, nested while loops. I felt like I was stumbling through the dark, aware that there was probably light at the end, but unable to see it.

Challenges

Still, I found a few insights along the way. Existing forced aligners, even ones designed for this specific task of aligning ebooks and audiobooks, struggled with a few challenges common to books, and even my first fledgling attempt handled these (to varying degrees of success):

Chapter order

Ebooks and audiobooks may (and often do) have different chapter orders. For example, content that would be considered frontmatter in an ebook, like a dedication, may be read at the end of the audiobook instead, since audiobooks often attempt to start with the content immediately.

Chapter existence

Each format will almost certainly have chapters that the other is missing entirely. Appendices, forewards, tables of contents — these are all almost always skipped in audiobooks. And audiobooks often have small chapters that don't exist in ebooks, as well.

Skipped spans

Sometimes smaller spans of content will be skipped in the audiobook narration, or the audiobook will contain content that isn't in the ebook, like a description of an image or graphic.

Alternate word choices

Sometimes audiobook directors or narrators will intentionally choose a different word or phrase when the original is hard to speak fluently or sounds awkward when read aloud. Also, sometimes they make mistakes!

Of these, even fairly basic forced alignment systems can generally handle alternate word choices without issue. And skipped spans can be challenging, especially when the audio skips spans in the text, but the results usually aren't disastrous, just imperfect.

But the missing and reordered chapters can be dealbreakers for many forced aligners. Tools that existed before Storyteller, like the very cool syncabook, required that users identify which ebook chapters correspond to which audiobook chapter in advance. This is both very manual and rather challenging, as many audiobooks don't even have chapter metadata or proper per-chapter files.

I wanted to do better, and that meant solving this problem. I needed a search algorithm.

Prerequisite: Boundary search

Before we can even look at the actual forced alignment problem, we need to find (roughly) where a given chapter of text can be found in the audio, if it can be found at all.

As prerequisite problems go, this one is… uh… kinda rough. We haven't done any alignment yet, so we don't know anything about the verbal content of the audio. And even if we had a perfect transcription (which we don't have any way of getting¹ — doing this is the forced alignment problem we need to solve later), we can't just scan the transcript for the contents of the chapter, because even a perfect transcription will deviate from the baseline ebook text.

So we can't do the easy thing. But while we can't get a full, accurate transcription of the audio, we can get some textual representation of it. We can use the Massively Multilingual Speech² model to generate CTC emissions, and then greedily decode those emissions to produce text.

... I will now explain the prior jargon. We're gonna go pretty deep. There will be graphics.

CTC, Wav2Vec 2.0, and MMS

Connectionist Temporal Classification (CTC, and yes, it does sound like something out of Dune) has been a staple of automatic speech recognition and forced alignment for over a decade. It's essentially a loss function: the function used by machine learning models to evaluate their output and train themselves. In order to work with this loss function, a model must contain a "CTC head", a layer that outputs "CTC emissions." Emissions are an intermediate representation used by CTC — they'll be discussed in depth in a moment.

Because any model using a CTC head will produce the same shape of output (the aforementioned CTC emissions), there are also standard algorithms for further decoding emissions into text. The two we care about for our use case are "unconstrained greedy decoding" and "Viterbi forced alignment". We'll explain these in detail as we get to them.

So we have a way to turn our model's internal representation into emissions, via our CTC decoder. Wav2Vec 2.0 goes the other way — it's a pretrained encoder, responsible for turning audio data into the internal representation that the machine learning model can operate on.

The model itself, which incorporates the Wav2Vec 2.0 encoder, the CTC decoder, and is then fine-tuned on some corpus of data so that it can "learn" the weights that minimize the CTC loss function, is Massively Multilingual Speech, or MMS.

We can take some audio, feed it into MMS in chunks, and get out some CTC emissions. The emissions themselves are a two-dimensional matrix: one vector of character probabilities³ per frame of audio, where a frame is 20ms of audio.

Decoding without labels

Now that we have our emissions, we need to solve our prerequisite problem: finding where each chapter starts and stops in the audio. One nice feature of emissions is that they're regular — since each emission vector represents one 20ms frame of audio, if we can find which frame a chapter starts in, we also know what millisecond it starts in.

In order to search for text, we need something that we can compare text to. Our emissions don't really fit this bill, at the moment. But we can extract text out of our emissions, can't we? What if we just walked through our emission vectors, and, for each one, we took the character with the highest probability? We wouldn't get a good transcription in any sense, but we would get some text, and a lot of it would probably be correct.

I'm calling this algorithm "unconstrained greedy decoding." Unconstrained because we didn't provide a baseline text to try to decode to, and greedy because at each step, we take the best probability, and never reconsider previous steps. Here’s what it looks like when we run in on the entire section from the demo at the start of this post:

The resulting text looks quite a bit like our ebook's text! And we can make them look even more similar by conditioning the ebook text: removing punctuation, collapsing whitespace, and lowercasing each character. We can even convert numerals to their spelled forms, e.g. "2,000" to "two thousand". I talked more about how we can do this conditioning without losing track of where the text came from in the original XHTML in a previous post.

RANSAC’d n-grams

You've probably noticed that our greedy decoding doesn't perfectly match our query. And this is only a very small sample — most audiobooks will have several deviations from the ebook text, as we discussed earlier, and most greedy decodings will have loads of transcription errors. So we can't just scan through the document until we find our exact query text.

Instead, we need to break up our document and query into pieces small enough that many of them are likely to match between the two. These are called "n-grams." For our purposes, a gram will be equivalent to a character, and our "n" will be 10. In both our document and query, we will record every single 10-letter span, along with the position it starts at. Many of these will exist in both texts — we can use those matches to locate the query in the document!

This algorithm has a few really nice features:

1. It's incredibly robust to noisy decodings. Whether the audiobook narration has many deviations from the ebook text, or the greedy decoding just did an especially poor job of estimating the spoken content, even if only 10% of the n-grams match, that's still thousands of points we can use for finding our line.
2. It gives us lots of additional information. We'll dive into this more later, but we can use information from this algorithm, like the local rate of speech and the location of known inliers, to implement our actual forced alignment pass.
3. It's really efficient!

Forced alignment

Now we know where our chapters start and end in the audiobook emissions, thanks to our "RANSAC'd n-grams" boundary search. This lets us move on to a more straightforward forced alignment algorithm: the CTC Viterbi algorithm.

Forced alignment is usually framed as a global optimization problem⁴. We have some loss function, like "the sum of all of the chosen characters' probabilities," and we want to maximize the output of that function across our whole chapter. Generally speaking, global optimization problems like this are hard to compute efficiently. Considering all of the possible outcomes — in order to compare them and find the one with the best score — tends to be very expensive. But we have two tricks that will help us out considerably here:

Match anchors

When we computed our matches earlier, we found several n-grams that existed in both the chapter and the audio. When we used RANSAC to find the inliers, we were left with matches that we were pretty sure actually corresponded to real matches between the chapter and audio. If we add one additional constraint — that we only consider matches that are globally unique in their respective documents — we can be quite sure that they're real matches. Because we know that the chapter and audio definitely line up at those points, we only need to run our forced alignment algorithm on the frames between those anchors. This means that we can split up our chapter and audio at the anchor points and only run forced alignment on one segment at a time.

If we look for one of these unique match anchors roughly every 2,000 characters, then we only need to run our forced alignment algorithm on roughly 2,000 characters at a time, instead of the entire chapter at once, which might be tens of thousands of characters or more!

Viterbi

The other "trick" we have is Viterbi. Viterbi is a bottom-up dynamic programming algorithm, which means that it finds its solution by first finding the solutions to its sub-problems, and building up to its ultimate solution from its sub-solutions. The key insight that allows us to use dynamic programming on our CTC alignment problem is this:

For states A, B, and C, if B lies on the shortest path from A to C, then the section of that shortest path from A to C that runs from A to B must also be the shortest path from A to B. If it wasn't, then we could swap out that path with a shorter path from A to B, which would make our overall path from A to C shorter as well.

To explain a little further, let's take just the very first word in our chapter, "All." We saw above that our emissions contain multiple frames for each character in this word, and each frame vector has probabilities for each character in our vocabulary (the letter "a" through "z", plus a "blank" token for when nothing is being spoken or the model can't distinguish the current character).

We want to find a path through our frames that spells out the word "all". A path consists of states and transitions. The states are our chapter text. When we put together our possible states, we insert a blank token between each letter. This lets us represent double letters, like the "ll" in "all", as "l → blank → l" in our path. So our state sequence looks like "blank → a → blank → l → blank → l → blank".

The transitions we allow are: stay on this state; move to the next state; or skip the next blank state and move to the one after. The skip transition is only allowed between to non-equal characters, so we can skip the blank between "a" and "l", but not between the two "l"s.

Our path needs to walk through each state using the transitions we defined above, so it's allowed to skip blank states between non-equal characters. It's also allowed to skip the leading and trailing blanks.

One way to do this is to enumerate all possible paths, add up their scores, and pick the best one. We have 7 frames and 7 states, and there are 210 possible paths through our frames. For one word. And this explodes exponentially as we increase our frame and state sizes from 7 to, say, 70,000.

This is not really feasible. Instead, let’s look at the Viterbi approach. Below are our 7 frames again, this time showing the actual logarithmic probabilities of each of the characters we're aligning:

First, we populate a 7-number vector, one for each state. This will be our "current scores" vector. It starts out with negative infinity for all scores except the first two, which are our possible starting states: index 0 is our starting blank, and index 1 is our "a". For these, we use the scores directly from the first frame of our emissions: the blank had -8.56 and the "a" had -0.00.

Then we iterate through our frames. At each frame, we look at the next frames scores to determine which transition to take. For our leading blank, we can either stay on the blank (whose score is a -0.00 in the next frame), or move to the "a" (whose score is -9.16 in the next frame). The better score is the -0.00, so the better transition is to "stay". So we record a "stay" for the transition from token 0 in frame 0, and we put -8.56 + -0.00 = -8.56 in our "current scores" vector for the leading blank state.

For our "a", we can either stay on the "a" (whose score is a -9.16 in the next frame), move to the next blank (whose score is a -0.00 in the next frame), or skip the next blank and move directly to the first "l" (whose score is a -10.47 in the next frame). The best score is the move to the next blank, so we record a "next" for the transition from token 1 in frame 0, and we put -0.00 + -0.00 = -0.00 in our "current scores" vector for the "a" state.

If we continue this pattern through the frames, we’ll end up with a "current scores" vector that contains the scores for the best paths to each state, and a transitions history that records how we got to each of those states. Then we simply pick the terminal state (the second "l" or the final blank) with the highest score and follow the transitions backwards to determine the complete path to that state.

This is the CTC Viterbi algorithm. In order to determine the optimal path to a given state, we only need to know the optimal score of the path to the previous state, because the shortest path to the current state must start with the shortest path to the previous state.

And all told, that means that instead of comparing 210 paths, we only need to compute 2! Exactly one optimal path ends at each of the valid ending frames, the second "l" and the final blank. Our best path is the one with the best score (in this case, the path to the second "l").

Now all we need to do is loop through the frames and keep track of which frame each token starts and ends on. The first "l", for example, is spoken over two frames, the third and fourth frames. And since frames are all 20ms long, we can use that to determine the start and end time for each letter.

Zooming out a bit, this lets us determine the start and end time for each word and sentence in the book. If you want to see it for yourself, you can use the new --ctc flag in stalign, or the brand new CTC aligner option in the Storyteller v3 beta! If you're not using the beta yet, no worries — it'll be available to everyone soon!

Pretty neat, right?