Commit @47d26173bf1f002536c51d26ca56ba53e47c1dac - yjyoon/whisper_server

Fedir Zadniprovskyi 2024-07-20

feat: handle srt and vtt response formats

@47d26173bf1f002536c51d26ca56ba53e47c1dac

9f134f9

47d2617

faster_whisper_server/config.py

--- faster_whisper_server/config.py

+++ faster_whisper_server/config.py


     TEXT = "text"
     JSON = "json"
     VERBOSE_JSON = "verbose_json"
-    # NOTE: While inspecting outputs of these formats with `curl`, I noticed there's one or two "\n" inserted at the end of the response. # noqa: E501
-
-    # VTT = "vtt" # TODO
-    # 1
-    # 00:00:00,000 --> 00:00:09,220
-    # In his video on Large Language Models or LLMs, OpenAI co-founder and YouTuber Andrej Karpathy
-    #
-    # 2
-    # 00:00:09,220 --> 00:00:12,280
-    # likened LLMs to operating systems.
-    #
-    # 3
-    # 00:00:12,280 --> 00:00:13,280
-    # Karpathy said,
-    #
-    # SRT = "srt" # TODO
-    # WEBVTT
-    #
-    # 00:00:00.000 --> 00:00:09.220
-    # In his video on Large Language Models or LLMs, OpenAI co-founder and YouTuber Andrej Karpathy
-    #
-    # 00:00:09.220 --> 00:00:12.280
-    # likened LLMs to operating systems.
-    #
-    # 00:00:12.280 --> 00:00:13.280
-    # Karpathy said,
-    #
-    # 00:00:13.280 --> 00:00:19.799
-    # I see a lot of equivalence between this new LLM OS and operating systems of today.
+    SRT = "srt"
+    VTT = "vtt"
 
 
 class Device(enum.StrEnum):

9f134f9

47d2617

faster_whisper_server/core.py

--- faster_whisper_server/core.py

+++ faster_whisper_server/core.py

...	...	@@ -172,6 +172,62 @@
172	172	return "".join(segment.text for segment in segments).strip()
173	173
174	174
	175	+def srt_format_timestamp(ts: float) -> str:
	176	+ hours = ts // 3600
	177	+ minutes = (ts % 3600) // 60
	178	+ seconds = ts % 60
	179	+ milliseconds = (ts * 1000) % 1000
	180	+ return f"{int(hours):02d}:{int(minutes):02d}:{int(seconds):02d},{int(milliseconds):03d}"
	181	+
	182	+
	183	+def test_srt_format_timestamp() -> None:
	184	+ assert srt_format_timestamp(0.0) == "00:00:00,000"
	185	+ assert srt_format_timestamp(1.0) == "00:00:01,000"
	186	+ assert srt_format_timestamp(1.234) == "00:00:01,234"
	187	+ assert srt_format_timestamp(60.0) == "00:01:00,000"
	188	+ assert srt_format_timestamp(61.0) == "00:01:01,000"
	189	+ assert srt_format_timestamp(61.234) == "00:01:01,234"
	190	+ assert srt_format_timestamp(3600.0) == "01:00:00,000"
	191	+ assert srt_format_timestamp(3601.0) == "01:00:01,000"
	192	+ assert srt_format_timestamp(3601.234) == "01:00:01,234"
	193	+ assert srt_format_timestamp(23423.4234) == "06:30:23,423"
	194	+
	195	+
	196	+def vtt_format_timestamp(ts: float) -> str:
	197	+ hours = ts // 3600
	198	+ minutes = (ts % 3600) // 60
	199	+ seconds = ts % 60
	200	+ milliseconds = (ts * 1000) % 1000
	201	+ return f"{int(hours):02d}:{int(minutes):02d}:{int(seconds):02d}.{int(milliseconds):03d}"
	202	+
	203	+
	204	+def test_vtt_format_timestamp() -> None:
	205	+ assert vtt_format_timestamp(0.0) == "00:00:00.000"
	206	+ assert vtt_format_timestamp(1.0) == "00:00:01.000"
	207	+ assert vtt_format_timestamp(1.234) == "00:00:01.234"
	208	+ assert vtt_format_timestamp(60.0) == "00:01:00.000"
	209	+ assert vtt_format_timestamp(61.0) == "00:01:01.000"
	210	+ assert vtt_format_timestamp(61.234) == "00:01:01.234"
	211	+ assert vtt_format_timestamp(3600.0) == "01:00:00.000"
	212	+ assert vtt_format_timestamp(3601.0) == "01:00:01.000"
	213	+ assert vtt_format_timestamp(3601.234) == "01:00:01.234"
	214	+ assert vtt_format_timestamp(23423.4234) == "06:30:23.423"
	215	+
	216	+
	217	+def segments_to_vtt(segment: Segment, i: int) -> str:
	218	+ start = segment.start if i > 0 else 0.0
	219	+ result = f"{vtt_format_timestamp(start)} --> {vtt_format_timestamp(segment.end)}\n{segment.text}\n\n"
	220	+
	221	+ if i == 0:
	222	+ return f"WEBVTT\n\n{result}"
	223	+ else:
	224	+ return result
	225	+
	226	+
	227	+def segments_to_srt(segment: Segment, i: int) -> str:
	228	+ return f"{i + 1}\n{srt_format_timestamp(segment.start)} --> {srt_format_timestamp(segment.end)}\n{segment.text}\n\n"
	229	+
	230	+
175	231	def canonicalize_word(text: str) -> str:
176	232	text = text.lower()
177	233	# Remove non-alphabetic characters using regular expression

9f134f9

47d2617

faster_whisper_server/main.py

--- faster_whisper_server/main.py

+++ faster_whisper_server/main.py


     Task,
     config,
 )
-from faster_whisper_server.core import Segment, segments_to_text
+from faster_whisper_server.core import Segment, segments_to_srt, segments_to_text, segments_to_vtt
 from faster_whisper_server.logger import logger
 from faster_whisper_server.server_models import (
     ModelListResponse,

     segments: Iterable[Segment],
     transcription_info: TranscriptionInfo,
     response_format: ResponseFormat,
-) -> str | TranscriptionJsonResponse | TranscriptionVerboseJsonResponse:
+) -> Response:
     segments = list(segments)
     if response_format == ResponseFormat.TEXT:  # noqa: RET503
-        return segments_to_text(segments)
+        return Response(segments_to_text(segments), media_type="text/plain")
     elif response_format == ResponseFormat.JSON:
-        return TranscriptionJsonResponse.from_segments(segments)
+        return Response(
+            TranscriptionJsonResponse.from_segments(segments).model_dump_json(),
+            media_type="application/json",
+        )
     elif response_format == ResponseFormat.VERBOSE_JSON:
-        return TranscriptionVerboseJsonResponse.from_segments(segments, transcription_info)
+        return Response(
+            TranscriptionVerboseJsonResponse.from_segments(segments, transcription_info).model_dump_json(),
+            media_type="application/json",
+        )
+    elif response_format == ResponseFormat.VTT:
+        return Response(
+            "".join(segments_to_vtt(segment, i) for i, segment in enumerate(segments)), media_type="text/vtt"
+        )
+    elif response_format == ResponseFormat.SRT:
+        return Response(
+            "".join(segments_to_srt(segment, i) for i, segment in enumerate(segments)), media_type="text/plain"
+        )
 
 
 def format_as_sse(data: str) -> str:

     response_format: ResponseFormat,
 ) -> StreamingResponse:
     def segment_responses() -> Generator[str, None, None]:
-        for segment in segments:
+        for i, segment in enumerate(segments):
             if response_format == ResponseFormat.TEXT:
                 data = segment.text
             elif response_format == ResponseFormat.JSON:
                 data = TranscriptionJsonResponse.from_segments([segment]).model_dump_json()
             elif response_format == ResponseFormat.VERBOSE_JSON:
                 data = TranscriptionVerboseJsonResponse.from_segment(segment, transcription_info).model_dump_json()
+            elif response_format == ResponseFormat.VTT:
+                data = segments_to_vtt(segment, i)
+            elif response_format == ResponseFormat.SRT:
+                data = segments_to_srt(segment, i)
             yield format_as_sse(data)
 
     return StreamingResponse(segment_responses(), media_type="text/event-stream")

     response_format: Annotated[ResponseFormat, Form()] = config.default_response_format,
     temperature: Annotated[float, Form()] = 0.0,
     stream: Annotated[bool, Form()] = False,
-) -> str | TranscriptionJsonResponse | TranscriptionVerboseJsonResponse | StreamingResponse:
+) -> Response | StreamingResponse:
     whisper = load_model(model)
     segments, transcription_info = whisper.transcribe(
         file.file,

     ] = ["segment"],
     stream: Annotated[bool, Form()] = False,
     hotwords: Annotated[str | None, Form()] = None,
-) -> str | TranscriptionJsonResponse | TranscriptionVerboseJsonResponse | StreamingResponse:
+) -> Response | StreamingResponse:
     whisper = load_model(model)
     segments, transcription_info = whisper.transcribe(
         file.file,

9f134f9

47d2617

pyproject.toml

--- pyproject.toml

+++ pyproject.toml


 ]
 
 [project.optional-dependencies]
-dev = ["ruff==0.5.3", "pytest", "basedpyright==1.13.0", "pytest-xdist"]
+dev = ["ruff==0.5.3", "pytest", "webvtt-py", "srt", "basedpyright==1.13.0", "pytest-xdist"]
 
 other = ["youtube-dl @ git+https://github.com/ytdl-org/youtube-dl.git@37cea84f775129ad715b9bcd617251c831fcc980", "aider-chat==0.39.0"]
 

9f134f9

47d2617

requirements-all.txt

--- requirements-all.txt

+++ requirements-all.txt


     # via aider-chat
 semantic-version==2.10.0
     # via gradio
-setuptools==71.0.3
+setuptools==71.0.4
     # via ctranslate2
 shellingham==1.5.4
     # via typer

     # via
     #   aider-chat
     #   beautifulsoup4
+srt==3.5.3
+    # via faster-whisper-server (pyproject.toml)
 starlette==0.37.2
     # via fastapi
 streamlit==1.35.0
     # via aider-chat
-sympy==1.13.0
+sympy==1.13.1
     # via onnxruntime
 tenacity==8.3.0
     # via

     # via
     #   gradio-client
     #   uvicorn
+webvtt-py==0.5.1
+    # via faster-whisper-server (pyproject.toml)
 yarl==1.9.4
     # via
     #   aider-chat

9f134f9

47d2617

requirements-dev.txt

--- requirements-dev.txt

+++ requirements-dev.txt


     #   pandas
 onnxruntime==1.18.1
     # via faster-whisper
-openai==1.35.15
+openai==1.36.0
     # via faster-whisper-server (pyproject.toml)
 orjson==3.10.6
     # via gradio

     #   gradio
 semantic-version==2.10.0
     # via gradio
-setuptools==71.0.3
+setuptools==71.0.4
     # via ctranslate2
 shellingham==1.5.4
     # via typer

     #   openai
 soundfile==0.12.1
     # via faster-whisper-server (pyproject.toml)
+srt==3.5.3
+    # via faster-whisper-server (pyproject.toml)
 starlette==0.37.2
     # via fastapi
-sympy==1.13.0
+sympy==1.13.1
     # via onnxruntime
 tokenizers==0.19.1
     # via faster-whisper

     # via
     #   gradio-client
     #   uvicorn
+webvtt-py==0.5.1
+    # via faster-whisper-server (pyproject.toml)

9f134f9

47d2617

requirements.txt

--- requirements.txt

+++ requirements.txt


     #   pandas
 onnxruntime==1.18.1
     # via faster-whisper
-openai==1.35.15
+openai==1.36.0
     # via faster-whisper-server (pyproject.toml)
 orjson==3.10.6
     # via gradio

     # via gradio
 semantic-version==2.10.0
     # via gradio
-setuptools==71.0.3
+setuptools==71.0.4
     # via ctranslate2
 shellingham==1.5.4
     # via typer

     # via faster-whisper-server (pyproject.toml)
 starlette==0.37.2
     # via fastapi
-sympy==1.13.0
+sympy==1.13.1
     # via onnxruntime
 tokenizers==0.19.1
     # via faster-whisper

9f134f9

47d2617

tests/conftest.py

--- tests/conftest.py

+++ tests/conftest.py


 from collections.abc import Generator
 import logging
+import os
 
 from fastapi.testclient import TestClient
 from openai import OpenAI
 import pytest
 
+os.environ["WHISPER__MODEL"] = "Systran/faster-whisper-tiny.en"
 from faster_whisper_server.main import app
 
 disable_loggers = ["multipart.multipart", "faster_whisper"]

9f134f9

47d2617

tests/sse_test.py

--- tests/sse_test.py

+++ tests/sse_test.py

...	...	@@ -4,6 +4,9 @@
4	4	from fastapi.testclient import TestClient
5	5	from httpx_sse import connect_sse
6	6	import pytest
	7	+import srt
	8	+import webvtt
	9	+import webvtt.vtt
7	10
8	11	from faster_whisper_server.server_models import (
9	12	TranscriptionJsonResponse,
...	...	@@ -61,3 +64,38 @@
61	64	with connect_sse(client, "POST", endpoint, **kwargs) as event_source:
62	65	for event in event_source.iter_sse():
63	66	TranscriptionVerboseJsonResponse(**json.loads(event.data))
	67	+
	68	+
	69	+def test_transcription_vtt(client: TestClient) -> None:
	70	+ with open("audio.wav", "rb") as f:
	71	+ data = f.read()
	72	+ kwargs = {
	73	+ "files": {"file": ("audio.wav", data, "audio/wav")},
	74	+ "data": {"response_format": "vtt", "stream": False},
	75	+ }
	76	+ response = client.post("/v1/audio/transcriptions", **kwargs)
	77	+ assert response.status_code == 200
	78	+ assert response.headers["content-type"] == "text/vtt; charset=utf-8"
	79	+ text = response.text
	80	+ webvtt.from_string(text)
	81	+ text = text.replace("WEBVTT", "YO")
	82	+ with pytest.raises(webvtt.vtt.MalformedFileError):
	83	+ webvtt.from_string(text)
	84	+
	85	+
	86	+def test_transcription_srt(client: TestClient) -> None:
	87	+ with open("audio.wav", "rb") as f:
	88	+ data = f.read()
	89	+ kwargs = {
	90	+ "files": {"file": ("audio.wav", data, "audio/wav")},
	91	+ "data": {"response_format": "srt", "stream": False},
	92	+ }
	93	+ response = client.post("/v1/audio/transcriptions", **kwargs)
	94	+ assert response.status_code == 200
	95	+ assert "text/plain" in response.headers["content-type"]
	96	+
	97	+ text = response.text
	98	+ list(srt.parse(text))
	99	+ text = text.replace("1", "YO")
	100	+ with pytest.raises(srt.SRTParseError):
	101	+ list(srt.parse(text))

Add a comment

Open 0
Closed 0

List

...	...	@@ -496,7 +496,7 @@
496	496	# via aider-chat
497	497	semantic-version==2.10.0
498	498	# via gradio
499		-setuptools==71.0.3
	499	+setuptools==71.0.4
500	500	# via ctranslate2
501	501	shellingham==1.5.4
502	502	# via typer
...	...	@@ -524,11 +524,13 @@
524	524	# via
525	525	# aider-chat
526	526	# beautifulsoup4
	527	+srt==3.5.3
	528	+ # via faster-whisper-server (pyproject.toml)
527	529	starlette==0.37.2
528	530	# via fastapi
529	531	streamlit==1.35.0
530	532	# via aider-chat
531		-sympy==1.13.0
	533	+sympy==1.13.1
532	534	# via onnxruntime
533	535	tenacity==8.3.0
534	536	# via
...	...	@@ -623,6 +625,8 @@
623	625	# via
624	626	# gradio-client
625	627	# uvicorn
	628	+webvtt-py==0.5.1
	629	+ # via faster-whisper-server (pyproject.toml)
626	630	yarl==1.9.4
627	631	# via
628	632	# aider-chat

...	...	@@ -146,7 +146,7 @@
146	146	# pandas
147	147	onnxruntime==1.18.1
148	148	# via faster-whisper
149		-openai==1.35.15
	149	+openai==1.36.0
150	150	# via faster-whisper-server (pyproject.toml)
151	151	orjson==3.10.6
152	152	# via gradio
...	...	@@ -235,7 +235,7 @@
235	235	# gradio
236	236	semantic-version==2.10.0
237	237	# via gradio
238		-setuptools==71.0.3
	238	+setuptools==71.0.4
239	239	# via ctranslate2
240	240	shellingham==1.5.4
241	241	# via typer
...	...	@@ -248,9 +248,11 @@
248	248	# openai
249	249	soundfile==0.12.1
250	250	# via faster-whisper-server (pyproject.toml)
	251	+srt==3.5.3
	252	+ # via faster-whisper-server (pyproject.toml)
251	253	starlette==0.37.2
252	254	# via fastapi
253		-sympy==1.13.0
	255	+sympy==1.13.1
254	256	# via onnxruntime
255	257	tokenizers==0.19.1
256	258	# via faster-whisper
...	...	@@ -295,3 +297,5 @@
295	297	# via
296	298	# gradio-client
297	299	# uvicorn
	300	+webvtt-py==0.5.1
	301	+ # via faster-whisper-server (pyproject.toml)

...	...	@@ -15,35 +15,8 @@
15	15	TEXT = "text"
16	16	JSON = "json"
17	17	VERBOSE_JSON = "verbose_json"
18		- # NOTE: While inspecting outputs of these formats with `curl`, I noticed there's one or two "\n" inserted at the end of the response. # noqa: E501
19		-
20		- # VTT = "vtt" # TODO
21		- # 1
22		- # 00:00:00,000 --> 00:00:09,220
23		- # In his video on Large Language Models or LLMs, OpenAI co-founder and YouTuber Andrej Karpathy
24		- #
25		- # 2
26		- # 00:00:09,220 --> 00:00:12,280
27		- # likened LLMs to operating systems.
28		- #
29		- # 3
30		- # 00:00:12,280 --> 00:00:13,280
31		- # Karpathy said,
32		- #
33		- # SRT = "srt" # TODO
34		- # WEBVTT
35		- #
36		- # 00:00:00.000 --> 00:00:09.220
37		- # In his video on Large Language Models or LLMs, OpenAI co-founder and YouTuber Andrej Karpathy
38		- #
39		- # 00:00:09.220 --> 00:00:12.280
40		- # likened LLMs to operating systems.
41		- #
42		- # 00:00:12.280 --> 00:00:13.280
43		- # Karpathy said,
44		- #
45		- # 00:00:13.280 --> 00:00:19.799
46		- # I see a lot of equivalence between this new LLM OS and operating systems of today.
	18	+ SRT = "srt"
	19	+ VTT = "vtt"
47	20
48	21
49	22	class Device(enum.StrEnum):

...	...	@@ -33,7 +33,7 @@
33	33	Task,
34	34	config,
35	35	)
36		-from faster_whisper_server.core import Segment, segments_to_text
	36	+from faster_whisper_server.core import Segment, segments_to_srt, segments_to_text, segments_to_vtt
37	37	from faster_whisper_server.logger import logger
38	38	from faster_whisper_server.server_models import (
39	39	ModelListResponse,
...	...	@@ -154,14 +154,28 @@
154	154	segments: Iterable[Segment],
155	155	transcription_info: TranscriptionInfo,
156	156	response_format: ResponseFormat,
157		-) -> str \| TranscriptionJsonResponse \| TranscriptionVerboseJsonResponse:
	157	+) -> Response:
158	158	segments = list(segments)
159	159	if response_format == ResponseFormat.TEXT: # noqa: RET503
160		- return segments_to_text(segments)
	160	+ return Response(segments_to_text(segments), media_type="text/plain")
161	161	elif response_format == ResponseFormat.JSON:
162		- return TranscriptionJsonResponse.from_segments(segments)
	162	+ return Response(
	163	+ TranscriptionJsonResponse.from_segments(segments).model_dump_json(),
	164	+ media_type="application/json",
	165	+ )
163	166	elif response_format == ResponseFormat.VERBOSE_JSON:
164		- return TranscriptionVerboseJsonResponse.from_segments(segments, transcription_info)
	167	+ return Response(
	168	+ TranscriptionVerboseJsonResponse.from_segments(segments, transcription_info).model_dump_json(),
	169	+ media_type="application/json",
	170	+ )
	171	+ elif response_format == ResponseFormat.VTT:
	172	+ return Response(
	173	+ "".join(segments_to_vtt(segment, i) for i, segment in enumerate(segments)), media_type="text/vtt"
	174	+ )
	175	+ elif response_format == ResponseFormat.SRT:
	176	+ return Response(
	177	+ "".join(segments_to_srt(segment, i) for i, segment in enumerate(segments)), media_type="text/plain"
	178	+ )
165	179
166	180
167	181	def format_as_sse(data: str) -> str:
...	...	@@ -174,13 +188,17 @@
174	188	response_format: ResponseFormat,
175	189	) -> StreamingResponse:
176	190	def segment_responses() -> Generator[str, None, None]:
177		- for segment in segments:
	191	+ for i, segment in enumerate(segments):
178	192	if response_format == ResponseFormat.TEXT:
179	193	data = segment.text
180	194	elif response_format == ResponseFormat.JSON:
181	195	data = TranscriptionJsonResponse.from_segments([segment]).model_dump_json()
182	196	elif response_format == ResponseFormat.VERBOSE_JSON:
183	197	data = TranscriptionVerboseJsonResponse.from_segment(segment, transcription_info).model_dump_json()
	198	+ elif response_format == ResponseFormat.VTT:
	199	+ data = segments_to_vtt(segment, i)
	200	+ elif response_format == ResponseFormat.SRT:
	201	+ data = segments_to_srt(segment, i)
184	202	yield format_as_sse(data)
185	203
186	204	return StreamingResponse(segment_responses(), media_type="text/event-stream")
...	...	@@ -211,7 +229,7 @@
211	229	response_format: Annotated[ResponseFormat, Form()] = config.default_response_format,
212	230	temperature: Annotated[float, Form()] = 0.0,
213	231	stream: Annotated[bool, Form()] = False,
214		-) -> str \| TranscriptionJsonResponse \| TranscriptionVerboseJsonResponse \| StreamingResponse:
	232	+) -> Response \| StreamingResponse:
215	233	whisper = load_model(model)
216	234	segments, transcription_info = whisper.transcribe(
217	235	file.file,
...	...	@@ -247,7 +265,7 @@
247	265	] = ["segment"],
248	266	stream: Annotated[bool, Form()] = False,
249	267	hotwords: Annotated[str \| None, Form()] = None,
250		-) -> str \| TranscriptionJsonResponse \| TranscriptionVerboseJsonResponse \| StreamingResponse:
	268	+) -> Response \| StreamingResponse:
251	269	whisper = load_model(model)
252	270	segments, transcription_info = whisper.transcribe(
253	271	file.file,

...	...	@@ -18,7 +18,7 @@
18	18	]
19	19
20	20	[project.optional-dependencies]
21		-dev = ["ruff==0.5.3", "pytest", "basedpyright==1.13.0", "pytest-xdist"]
	21	+dev = ["ruff==0.5.3", "pytest", "webvtt-py", "srt", "basedpyright==1.13.0", "pytest-xdist"]
22	22
23	23	other = ["youtube-dl @ git+https://github.com/ytdl-org/youtube-dl.git@37cea84f775129ad715b9bcd617251c831fcc980", "aider-chat==0.39.0"]
24	24

...	...	@@ -138,7 +138,7 @@
138	138	# pandas
139	139	onnxruntime==1.18.1
140	140	# via faster-whisper
141		-openai==1.35.15
	141	+openai==1.36.0
142	142	# via faster-whisper-server (pyproject.toml)
143	143	orjson==3.10.6
144	144	# via gradio
...	...	@@ -216,7 +216,7 @@
216	216	# via gradio
217	217	semantic-version==2.10.0
218	218	# via gradio
219		-setuptools==71.0.3
	219	+setuptools==71.0.4
220	220	# via ctranslate2
221	221	shellingham==1.5.4
222	222	# via typer
...	...	@@ -231,7 +231,7 @@
231	231	# via faster-whisper-server (pyproject.toml)
232	232	starlette==0.37.2
233	233	# via fastapi
234		-sympy==1.13.0
	234	+sympy==1.13.1
235	235	# via onnxruntime
236	236	tokenizers==0.19.1
237	237	# via faster-whisper

...	...	@@ -1,10 +1,12 @@
1	1	from collections.abc import Generator
2	2	import logging
	3	+import os
3	4
4	5	from fastapi.testclient import TestClient
5	6	from openai import OpenAI
6	7	import pytest
7	8
	9	+os.environ["WHISPER__MODEL"] = "Systran/faster-whisper-tiny.en"
8	10	from faster_whisper_server.main import app
9	11
10	12	disable_loggers = ["multipart.multipart", "faster_whisper"]

Delete comment