Update main.py

main.py

@@ -1,7 +1,34 @@
-from fastapi import FastAPI
+import torch
+from fastapi import BackgroundTasks, FastAPI
+from transformers import pipeline
+
+MODEL_NAME = "openai/whisper-large-v2"
+
+device = 0 if torch.cuda.is_available() else "cpu"
+
+pipe = pipeline(
+    task="automatic-speech-recognition",
+    model=MODEL_NAME,
+    chunk_length_s=30,
+    device=device,
+)
+
+all_special_ids = pipe.tokenizer.all_special_ids
+transcribe_token_id = all_special_ids[-5]
+translate_token_id = all_special_ids[-6]
+
+def transcribe(file):
+    pipe.model.config.forced_decoder_ids = [[2, transcribe_token_id]]
+    text = pipe(file, return_timestamps=True)
+    return text
 
 app = FastAPI()
 
+@app.get("/transcribe")
+def transcribe(file: UploadFile = File(...), background_tasks: BackgroundTasks):
+    background_tasks.add_task(transcribe, file)
+    return {"text": "Processing file..."}
+
 @app.get("/")
 def read_root():
-    return {"Hello": "World!"}
+    return {"text": "World!"}