Update apis/chat_api.py

apis/chat_api.py

@@ -16,9 +16,10 @@ import base64
 import re
 import requests
 from utils.enver import enver
+import shutil
 
 
-from fastapi import FastAPI, Response
+from fastapi import FastAPI, Response, File, UploadFile
 from fastapi.encoders import jsonable_encoder
 from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import BaseModel, Field
@@ -141,6 +142,51 @@ class ChatAPIApp:
         }
         json_compatible_item_data = jsonable_encoder(item_response)
         return JSONResponse(content=json_compatible_item_data)
+
+    class WhisperPostItem(BaseModel):
+        model: str = Field(
+            default="whisper-small",
+            description="(str) `Whisper model`",
+        )
+        lang: str = Field(
+            default="en",
+            description="(str) `transcribe to`",
+        )
+        audio_file: UploadFile = File(
+            description="Source Audio File",
+        )
+    def whisper_transcribe(self, item:WhisperPostItem):
+        MODEL_MAP = {
+        "whisper-small": "openai/whisper-small",
+        "whisper-medium": "openai/whisper-medium",
+        "whisper-large": "openai/whisper-large",   
+        "default": "openai/whisper-small",
+        }
+        if item.model in MODEL_MAP.keys():
+            target_model = item.model
+        else:
+            target_model = "default"
+
+        real_name = MODEL_MAP[target_model]
+        device = 0 if torch.cuda.is_available() else "cpu"
+        pipe = pipeline(
+           task="automatic-speech-recognition",
+           model=real_name,
+           chunk_length_s=30,
+           device=device,
+        )
+        time_start = time.time()
+        pipe.model.config.forced_decoder_ids = pipe.tokenizer.get_decoder_prompt_ids(language=item.lang, task="transcribe")
+        text = pipe(item.audio_file)["text"]
+        time_end = time.time()
+        item_response = {
+            "statue": 200,
+            "result": text,
+            "start": str(time_start),
+            "end": str(time_end)
+        }
+        json_compatible_item_data = jsonable_encoder(item_response)
+        return JSONResponse(content=json_compatible_item_data)
     class TranslateAiPostItem(BaseModel):
         model: str = Field(
             default="t5-base",
@@ -255,6 +301,10 @@ class ChatAPIApp:
                 summary="translate text with ai",
             )(self.translate_ai_completions)
 
+            self.app.post(
+                prefix + "/transcribe",
+                summary="transcribe audio to text",
+            )(self.whisper_transcribe)
             
             self.app.post(
                 prefix + "/detect",