-
Notifications
You must be signed in to change notification settings - Fork 7.9k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
word-level timestamps in transcribe()
#869
Changes from 1 commit
8f9357f
46ea501
cfd2b81
742d2f4
fb12414
80331c0
1d2ed66
b61e8f4
54f2901
8ce6277
812f446
cd5191f
f64d8bc
89133bd
d4f9399
040aa04
8e2756b
6c431c4
ff6cbfd
5fa4356
48537aa
8eb29c3
6ed4c11
31cd418
145f325
2b079c4
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
- Loading branch information
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,33 @@ | ||
try: | ||
import triton | ||
import triton.language as tl | ||
except ImportError: | ||
raise RuntimeError("triton import failed; try `pip install triton`") | ||
|
||
|
||
@triton.jit | ||
def dtw_kernel(cost, trace, x, x_stride, cost_stride, trace_stride, N, M, BLOCK_SIZE: tl.constexpr): | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. triton jit requires gcc/clang, python-dev, and cuda-dev at runtime. Please consider some lighter-weight alternatives. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. +1. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Thanks for the heads up; I made it fall back to pytorch and numba implementations if triton fails with RuntimeError or subprocess.CalledProcessError. I haven't tested this on a non-dev environment, so please feel free to ping me if the fallback does not work for any reason. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. +1. Thank you. |
||
offsets = tl.arange(0, BLOCK_SIZE) | ||
mask = offsets < M | ||
|
||
for k in range(1, N + M + 1): # k = i + j | ||
tl.debug_barrier() | ||
|
||
p0 = cost + (k - 1) * cost_stride | ||
p1 = cost + k * cost_stride | ||
p2 = cost + k * cost_stride + 1 | ||
|
||
c0 = tl.load(p0 + offsets, mask=mask) | ||
c1 = tl.load(p1 + offsets, mask=mask) | ||
c2 = tl.load(p2 + offsets, mask=mask) | ||
|
||
x_row = tl.load(x + (k - 1) * x_stride + offsets, mask=mask, other=0) | ||
cost_row = x_row + tl.minimum(tl.minimum(c0, c1), c2) | ||
|
||
cost_ptr = cost + (k + 1) * cost_stride + 1 | ||
tl.store(cost_ptr + offsets, cost_row, mask=mask) | ||
|
||
trace_ptr = trace + (k + 1) * trace_stride + 1 | ||
tl.store(trace_ptr + offsets, 0, mask=mask & (c0 <= c1) & (c0 <= c2)) | ||
tl.store(trace_ptr + offsets, 1, mask=mask & (c1 <= c0) & (c1 <= c2)) | ||
tl.store(trace_ptr + offsets, 2, mask=mask & (c2 <= c0) & (c2 <= c1)) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why only considering (at most) the last 6 layers?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This was because the attention weights in the later layers were more indicative of the time alignment. I've updated this part, and now it uses a mask to select which layers and heads to find the alignment.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Ok interesting. I have to check this masking trick.
I don't understand why later layers are more indicative. Is it an intuition that I am missing, or some empirical results you got from experiments?