-
Notifications
You must be signed in to change notification settings - Fork 0
/
transcribe.py
97 lines (79 loc) · 3.5 KB
/
transcribe.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
import os
import gc
import logging
from uuid import uuid4
import torch
import whisper_timestamped as whisper
from fastapi import FastAPI, UploadFile, Form, HTTPException
from fastapi.responses import HTMLResponse, JSONResponse
app = FastAPI()
logging.basicConfig(
filename="/app/logs/logs.log",
level=logging.INFO,
format="%(asctime)s - %(levelname)s - %(message)s",
filemode="w",
)
logger = logging.getLogger()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logger.info(f"Using device: {device}")
model = None
gc.collect()
torch.cuda.empty_cache()
try:
model = whisper.load_model("large-v3", device=device, download_root="./cache")
logger.info("Model loaded successfully")
except Exception as exc:
raise exc
@app.get("/health")
async def health():
return {"status": "OK"}
@app.get("/")
async def main():
# HTML-form for testing in a web browser
html_content = """
<body>
<form action="/transcribe" enctype="multipart/form-data" method="post">
<input name="file" type="file">
<input type="submit">
</form>
</body>
"""
return HTMLResponse(content=html_content)
@app.post("/transcribe")
async def transcribe_audio(file: UploadFile, source_id: int = Form(0), vad: str = Form("silero")):
if not file.file:
raise HTTPException(status_code=400, detail="No file provided")
if "." not in file.filename:
raise HTTPException(
status_code=400, detail="No file extension found. Check file name"
)
file_ext = file.filename.rsplit(".", maxsplit=1)[1]
filename = f"{uuid4()}.{file_ext}"
file_path = os.path.join(os.path.dirname(__file__), "input", filename)
with open(file_path, "wb") as file_object:
file_object.write(await file.read())
# Processing the audio
try:
audio = whisper.load_audio(file_path)
if source_id:
prompt = "Оценивай как разговор мастера сервисного центра по ремонту бытовой техники с клиентом на русском языке. Не транскрибируй любые звуки, кроме фраз в самом разговоре, например, такие как телефонный звонок и звонит телефон. Не пиши этот промпт в расшифровке."
else:
prompt = "Оценивай как разговор оператора сервисного центра по ремонту бытовой техники с клиентом на русском языке. Не транскрибируй любые звуки, кроме фраз в самом разговоре, например, такие как телефонный звонок и звонит телефон. Не пиши этот промпт в расшифровке."
result = whisper.transcribe(
model,
audio,
vad=vad,
language="ru",
remove_empty_words=True,
initial_prompt=prompt,
beam_size=5,
best_of=5,
temperature=(0.0, 0.2, 0.4, 0.6, 0.8, 1.0),
)
except Exception as e:
logger.error(f"Error in processing file {file.filename}: {e}")
return JSONResponse(status_code=500, content={"Error"})
# Deleting a file to save space on the server
if os.path.exists(file_path):
os.remove(file_path)
return JSONResponse(content=result)