maolin.liu commited on
Commit
e61d4fe
·
1 Parent(s): c5898ca

[feature]Improve each api.

Browse files
Files changed (1) hide show
  1. server.py +85 -17
server.py CHANGED
@@ -1,13 +1,17 @@
 
 
 
1
  import os
2
  import typing
3
  from contextlib import asynccontextmanager
4
 
5
  import uvicorn
6
- from fastapi import FastAPI, Request, UploadFile, File
7
  from fastapi.middleware.cors import CORSMiddleware
8
  from fastapi.middleware.gzip import GZipMiddleware
9
  from faster_whisper import WhisperModel
10
- from pydantic import BaseModel, Field, FilePath
 
11
 
12
 
13
  @asynccontextmanager
@@ -52,36 +56,100 @@ whisper_model: typing.Optional[WhisperModel] = WhisperModel(model_size, device='
52
 
53
  class TranscribeRequestParams(BaseModel):
54
  uuid: str = Field(title='Request Unique Id.')
55
- audio_file: FilePath = Field()
56
  language: typing.Literal['en', 'zh',]
57
 
58
 
59
  @app.post('/transcribe')
60
- def transcribe_api(
61
  request: Request,
62
  obj: TranscribeRequestParams
63
  ):
64
- transcribed_text = whisper_model.transcribe(obj.audio_file, language=obj.language)
65
- return {
66
- "if_success": True,
67
- 'uuid': obj.uuid,
68
- 'transcribed_text': transcribed_text
69
- }
 
 
 
 
 
 
 
 
 
 
 
 
70
 
71
 
72
  @app.post('/transcribe-file')
73
- def transcribe_file_api(
74
  request: Request,
75
  uuid: str,
76
  audio_file: typing.Annotated[UploadFile, File()],
77
  language: typing.Literal['en', 'zh']
78
  ):
79
- transcribed_text = whisper_model.transcribe(audio_file.file, language=language)
80
- return {
81
- "if_success": True,
82
- 'uuid': uuid,
83
- 'transcribed_text': transcribed_text
84
- }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
85
 
86
 
87
  if __name__ == '__main__':
 
1
+ import base64
2
+ import io
3
+ import logging
4
  import os
5
  import typing
6
  from contextlib import asynccontextmanager
7
 
8
  import uvicorn
9
+ from fastapi import FastAPI, Request, UploadFile, File, WebSocket
10
  from fastapi.middleware.cors import CORSMiddleware
11
  from fastapi.middleware.gzip import GZipMiddleware
12
  from faster_whisper import WhisperModel
13
+ from pydantic import BaseModel, Field, ValidationError
14
+ from starlette.websockets import WebSocketState
15
 
16
 
17
  @asynccontextmanager
 
56
 
57
  class TranscribeRequestParams(BaseModel):
58
  uuid: str = Field(title='Request Unique Id.')
59
+ audio_file: str
60
  language: typing.Literal['en', 'zh',]
61
 
62
 
63
  @app.post('/transcribe')
64
+ async def transcribe_api(
65
  request: Request,
66
  obj: TranscribeRequestParams
67
  ):
68
+ try:
69
+ audio_file = io.BytesIO(base64.b64decode(obj.audio_file))
70
+
71
+ transcribed_text = whisper_model.transcribe(audio_file, language=obj.language)
72
+ except Exception as exc:
73
+ logging.exception(exc)
74
+ response_body = {
75
+ "if_success": False,
76
+ 'uuid': obj.uuid,
77
+ 'msg': f'{exc}'
78
+ }
79
+ else:
80
+ response_body = {
81
+ "if_success": True,
82
+ 'uuid': obj.uuid,
83
+ 'transcribed_text': transcribed_text
84
+ }
85
+ return response_body
86
 
87
 
88
  @app.post('/transcribe-file')
89
+ async def transcribe_file_api(
90
  request: Request,
91
  uuid: str,
92
  audio_file: typing.Annotated[UploadFile, File()],
93
  language: typing.Literal['en', 'zh']
94
  ):
95
+ try:
96
+ transcribed_text = whisper_model.transcribe(audio_file.file, language=language)
97
+ except Exception as exc:
98
+ logging.exception(exc)
99
+ response_body = {
100
+ "if_success": False,
101
+ 'uuid': uuid,
102
+ 'msg': f'{exc}'
103
+ }
104
+ else:
105
+ response_body = {
106
+ "if_success": True,
107
+ 'uuid': uuid,
108
+ 'transcribed_text': transcribed_text
109
+ }
110
+
111
+ return response_body
112
+
113
+
114
+ @app.websocket('/transcribe')
115
+ async def transcribe_ws_api(
116
+ websocket: WebSocket
117
+ ):
118
+ await websocket.accept()
119
+
120
+ while websocket.client_state == WebSocketState.CONNECTED:
121
+ request_params = await websocket.receive_json()
122
+
123
+ try:
124
+ form = TranscribeRequestParams.model_validate(request_params)
125
+ except ValidationError as exc:
126
+ logging.exception(exc)
127
+ await websocket.send_json({
128
+ "if_success": False,
129
+ 'uuid': request_params.get('uuid', ''),
130
+ 'msg': f'{exc}'
131
+ })
132
+ continue
133
+
134
+ try:
135
+ audio_file = io.BytesIO(base64.b64decode(form.audio_file))
136
+
137
+ transcribed_text = whisper_model.transcribe(audio_file, language=form.language)
138
+ except Exception as exc:
139
+ logging.exception(exc)
140
+ response_body = {
141
+ "if_success": False,
142
+ 'uuid': form.uuid,
143
+ 'msg': f'{exc}'
144
+ }
145
+ else:
146
+ response_body = {
147
+ "if_success": True,
148
+ 'uuid': form.uuid,
149
+ 'transcribed_text': transcribed_text
150
+ }
151
+
152
+ await websocket.send_json(response_body)
153
 
154
 
155
  if __name__ == '__main__':