Repository: pytorch/audio

Issue:
Optimization Suggestion for 'torchaudio.functional.merge_tokens' Function
### 🚀 The feature

Optimization of the merge_tokens function in 'torchaudio.functional.merge_tokens' to leverage PyTorch's tensor operations for improved efficiency and reduced runtime.

Proposed Implementation:
```python
def merge_tokens_optimized(tokens: torch.Tensor, scores: torch.Tensor, blank: int = 0) -> List[TokenSpan]:
    """Removes repeated tokens and blank tokens from the given CTC token sequence

    Args:
        tokens (torch.Tensor): Alignment tokens (unbatched)
        scores (torch.Tensor): Alignment scores (unbatched)
        blank (int, optional): Blank token. Defaults to 0.

    Returns:
        List[TokenSpan]: list of TokenSpan
    """

    # Compute the difference between consecutive tokens. prepend and append a -1 to make sure the first and last
    # tokens are not removed
    diff = torch.diff(
        tokens, prepend=torch.tensor([-1], device=tokens.device), append=torch.tensor([-1], device=tokens.device)
    )
    # Compute the change points and mask out the points where the new value is blank
    changes_wo_blank = torch.nonzero((diff != 0)).squeeze().tolist()

    tokens = tokens.tolist()
    # Create a TokenSpan for each change point.
    spans = [
        TokenSpan(token=token, start=start, end=end, score=scores[start:end].mean().item())
        for start, end in zip(changes_wo_blank[:-1], changes_wo_blank[1:])
        if (token := tokens[start]) != blank
    ]
    return spans
```

### Motivation, pitch

While working on some CTC token sequence operations, I noticed that the current implementation of the merge_tokens function could benefit from tensor operations, potentially offering significant runtime improvements. Specifically, the current loop-based approach might be streamlined by taking advantage of PyTorch's capabilities. My tests indicate that my proposed optimization runs approximately 7 times faster over 1000 runs while producing equivalent outputs.

### Alternatives

Proposed tensor-based optimization: This approach leverages PyTorch's tensor operations, eliminating explicit loops and potentially considerably reducing execution time.

### Additional context

To validate the accuracy of the optimized function, a method was crafted to verify the outputs of both the original and proposed functions:

```python
import torchaudio.functional as F
from torchaudio.functional import TokenSpan

def token_span_outputs_equal(span_list_1: List[TokenSpan], span_list_2: List[TokenSpan]) -> bool:
    """Compares two lists of TokenSpan objects for equality."""
    
    if len(span_list_1) != len(span_list_2):
        return False

    for span_1, span_2 in zip(span_list_1, span_list_2):
        if (
            span_1.token != span_2.token or 
            span_1.start != span_2.start or 
            span_1.end != span_2.end or 
            abs(span_1.score - span_2.score) > 1e-6  # Allowing a small tolerance for floating-point comparisons
        ):
            return False

    return True
    
   
output_original = F.merge_tokens(aligned_tokens, alignment_scores)
output_optimized = merge_tokens_optimized(aligned_tokens, alignment_scores)

assert token_span_outputs_equal(output_original, output_optimized), "Outputs are not equivalent!"
```

Benchmark code:
```python
import timeit

pytorch_time = timeit.timeit(lambda: F.merge_tokens(aligned_tokens, alignment_scores), number=1000)
new_method_time = timeit.timeit(lambda: merge_tokens_optimized(aligned_tokens, alignment_scores), number=1000)

print(f"PyTorch method: {pytorch_time}")
print(f"New method: {new_method_time}")
```

Output:
```
Pytorch method: 7.622203521430492
New method: 1.0753349959850311
```

Integrating this optimization could bring about significant improvements in both performance and maintainability, benefiting a broad spectrum of users.


Discussion:
Hi @eyalcohen308 

Thanks for the suggestion. x7 improvement sounds great. Would you like to make a PR? If not I will try to find some time adopting this later.

There are unit tests for this, which you can run as `(cd test && pytest torchaudio_unittest/functional/functional_cpu_test.py -k test_merge_repeated_tokens)`.

https://github.com/pytorch/audio/blob/0ff6d2663777a5c3df51e3ca87f327a1e597f422/test/torchaudio_unittest/functional/functional_impl.py#L1232-L1283
@mthrok I will open a PR, thanks

Resolving patch:
diff --git a/torchaudio/functional/_alignment.py b/torchaudio/functional/_alignment.py
--- a/torchaudio/functional/_alignment.py
+++ b/torchaudio/functional/_alignment.py
@@ -115,17 +115,14 @@ def merge_tokens(tokens: Tensor, scores: Tensor, blank: int = 0) -> List[TokenSp
     if len(tokens) != len(scores):
         raise ValueError("`tokens` and `scores` must be the same length.")
 
-    t_prev = blank
-    i = start = -1
-    spans = []
-    for t, token in enumerate(tokens):
-        if token != t_prev:
-            if t_prev != blank:
-                spans.append(TokenSpan(t_prev.item(), start, t, scores[start:t].mean().item()))
-            if token != blank:
-                i += 1
-                start = t
-            t_prev = token
-    if t_prev != blank:
-        spans.append(TokenSpan(t_prev.item(), start, len(tokens), scores[start:].mean().item()))
+    diff = torch.diff(
+        tokens, prepend=torch.tensor([-1], device=tokens.device), append=torch.tensor([-1], device=tokens.device)
+    )
+    changes_wo_blank = torch.nonzero((diff != 0)).squeeze().tolist()
+    tokens = tokens.tolist()
+    spans = [
+        TokenSpan(token=token, start=start, end=end, score=scores[start:end].mean().item())
+        for start, end in zip(changes_wo_blank[:-1], changes_wo_blank[1:])
+        if (token := tokens[start]) != blank
+    ]
     return spans


Test patch: