add base return_timestamps support

This commit is contained in:
Serge Zaigraeff
2025-03-03 15:59:38 +03:00
parent 8de54b7c8f
commit 450750c47b
7 changed files with 121 additions and 28 deletions
+1 -1
View File
@@ -6,7 +6,7 @@ import waitress
# Импорт классов и функций из других модулей
from .transcriber import WhisperTranscriber
from .routes import Routes
from .logger import logger
from .utils import logger
class WhisperServiceAPI:
"""Класс для API сервиса распознавания речи."""
+1 -1
View File
@@ -12,7 +12,7 @@ import uuid
from typing import Dict, Tuple
# Импорт классов и функций из других модулей
from .logger import logger
from .utils import logger
class AudioProcessor:
"""Класс для предобработки аудиофайлов перед распознаванием."""
+1 -1
View File
@@ -11,7 +11,7 @@ import requests
import abc
from typing import Dict, Tuple, Optional, BinaryIO
from .logger import logger
from .utils import logger
class AudioSource(abc.ABC):
"""Абстрактный класс для различных источников аудиофайлов.
+60 -11
View File
@@ -7,10 +7,11 @@ import os
import uuid
import tempfile
import time
import librosa
from flask import request, jsonify
from typing import Dict, Tuple
from .logger import logger
from .utils import logger
from .audio_sources import (
AudioSource,
UploadedFileSource,
@@ -35,6 +36,24 @@ class TranscriptionService:
self.config = config
self.max_file_size_mb = self.config.get("max_file_size", 100) # Default 100MB
def get_audio_duration(self, file_path: str) -> float:
"""
Определяет длительность аудиофайла в секундах.
Args:
file_path: Путь к аудиофайлу.
Returns:
Длительность в секундах.
"""
try:
y, sr = librosa.load(file_path, sr=None)
duration = librosa.get_duration(y=y, sr=sr)
return duration
except Exception as e:
logger.error(f"Ошибка при определении длительности файла: {e}")
return 0.0
def transcribe_from_source(self, source: AudioSource, params: Dict = None) -> Tuple[Dict, int]:
"""
Транскрибирует аудиофайл из указанного источника.
@@ -62,34 +81,64 @@ class TranscriptionService:
temperature = float(params.get('temperature', 0.0))
prompt = params.get('prompt', '')
# Проверяем, запрошены ли временные метки
return_timestamps = params.get('return_timestamps', self.config.get('return_timestamps', False))
# Преобразуем строковое значение в булево, если необходимо
if isinstance(return_timestamps, str):
return_timestamps = return_timestamps.lower() in ('true', 't', 'yes', 'y', '1')
# Временно изменяем настройку return_timestamps в транскрайбере
original_return_timestamps = self.transcriber.return_timestamps
self.transcriber.return_timestamps = return_timestamps
# Сохраняем файл во временный файл
temp_dir = tempfile.mkdtemp()
temp_file_path = os.path.join(temp_dir, str(uuid.uuid4()) + "_" + os.path.basename(filename))
file.save(temp_file_path)
# Определяем длительность аудиофайла
duration = self.get_audio_duration(temp_file_path)
# Для файлов из внешних источников (URL, base64), закрываем их и выполняем очистку
if hasattr(source, 'cleanup'):
file.file.close() # Закрываем файловый объект
source.cleanup() # Очищаем временные файлы источника
try:
start_time = time.time()
text = self.transcriber.process_file(temp_file_path)
result = self.transcriber.process_file(temp_file_path)
processing_time = time.time() - start_time
# Форматируем ответ в стиле OpenAI
return jsonify({
"text": text,
"processing_time": processing_time,
"response_size_bytes": len(text.encode('utf-8')),
"model": os.path.basename(self.config["model_path"])
}), 200
# Формируем ответ в зависимости от return_timestamps
if return_timestamps:
response = {
"segments": result.get("segments", []),
"text": result.get("text", ""),
"processing_time": processing_time,
"response_size_bytes": len(str(result).encode('utf-8')),
"duration_seconds": duration,
"model": os.path.basename(self.config["model_path"])
}
else:
# Если не запрашивались временные метки, result - это строка
response = {
"text": result,
"processing_time": processing_time,
"response_size_bytes": len(str(result).encode('utf-8')),
"duration_seconds": duration,
"model": os.path.basename(self.config["model_path"])
}
return jsonify(response), 200
except Exception as e:
logger.error(f"Ошибка при транскрибации: {e}")
return jsonify({"error": str(e)}), 500
finally:
# Восстанавливаем оригинальное значение return_timestamps
self.transcriber.return_timestamps = original_return_timestamps
# Очистка временных файлов
if os.path.exists(temp_file_path):
os.remove(temp_file_path)
@@ -143,7 +192,7 @@ class Routes:
file_path = data["file_path"]
source = LocalFileSource(file_path, self.config.get("max_file_size", 100))
return self.transcription_service.transcribe_from_source(source)
return self.transcription_service.transcribe_from_source(source, data)
@self.app.route('/v1/models', methods=['GET'])
def list_models():
+57 -13
View File
@@ -7,7 +7,7 @@ OpenAI для транскрибации аудиофайлов в текст.
"""
import time
from typing import Dict, Tuple
from typing import Dict, Tuple, Union
import librosa
import numpy as np
@@ -19,7 +19,7 @@ from transformers import (
)
from .audio_processor import AudioProcessor
from .logger import logger
from .utils import logger
class WhisperTranscriber:
"""Класс для распознавания речи с помощью модели Whisper."""
@@ -135,7 +135,7 @@ class WhisperTranscriber:
logger.error(f"Ошибка при загрузке аудио {file_path}: {e}")
raise
def transcribe(self, audio_path: str) -> str:
def transcribe(self, audio_path: str) -> Union[str, Dict]:
"""
Транскрибация аудиофайла.
@@ -143,7 +143,9 @@ class WhisperTranscriber:
audio_path: Путь к обработанному аудиофайлу.
Returns:
Распознанный текст.
В зависимости от параметра return_timestamps:
- Если return_timestamps=False: строка с распознанным текстом
- Если return_timestamps=True: словарь с ключами "segments" (список словарей с ключами start_time_ms, end_time_ms, text) и "text" (полный текст)
"""
logger.info(f"Начало транскрибации файла: {audio_path}")
@@ -157,12 +159,52 @@ class WhisperTranscriber:
return_timestamps=self.return_timestamps
)
transcribed_text = result.get("text", "")
logger.info(f"Транскрибация завершена: получено {len(transcribed_text)} символов текста")
# Если временные метки не запрошены, возвращаем только текст
if not self.return_timestamps:
transcribed_text = result.get("text", "")
logger.info(f"Транскрибация завершена: получено {len(transcribed_text)} символов текста")
return transcribed_text
# Если временные метки запрошены, обрабатываем и форматируем результат
segments = []
full_text = result.get("text", "")
if "chunks" in result:
# Для новых версий модели Whisper
for chunk in result["chunks"]:
start_time = chunk.get("timestamp", [0, 0])[0]
end_time = chunk.get("timestamp", [0, 0])[1]
text = chunk.get("text", "").strip()
segments.append({
"start_time_ms": int(start_time * 1000),
"end_time_ms": int(end_time * 1000),
"text": text
})
elif hasattr(result, "get") and "segments" in result:
# Для старых версий модели Whisper
for segment in result["segments"]:
start_time = segment.get("start", 0)
end_time = segment.get("end", 0)
text = segment.get("text", "").strip()
segments.append({
"start_time_ms": int(start_time * 1000),
"end_time_ms": int(end_time * 1000),
"text": text
})
else:
logger.warning("Временные метки запрошены, но не найдены в результате транскрибации")
logger.info(f"Транскрибация с временными метками завершена: получено {len(segments)} сегментов")
# Возвращаем словарь с сегментами и полным текстом
return {
"segments": segments,
"text": full_text
}
return transcribed_text
def process_file(self, input_path: str) -> str:
def process_file(self, input_path: str) -> Union[str, Dict]:
"""
Полный процесс обработки и транскрибации аудиофайла.
@@ -170,7 +212,9 @@ class WhisperTranscriber:
input_path: Путь к исходному аудиофайлу.
Returns:
Распознанный текст.
В зависимости от параметра return_timestamps:
- Если return_timestamps=False: строка с распознанным текстом
- Если return_timestamps=True: словарь с ключами "segments" и "text"
"""
start_time = time.time()
logger.info(f"Начало обработки файла: {input_path}")
@@ -182,13 +226,13 @@ class WhisperTranscriber:
processed_path, temp_files = self.audio_processor.process_audio(input_path)
# Транскрибация
text = self.transcribe(processed_path)
result = self.transcribe(processed_path)
elapsed_time = time.time() - start_time
logger.info(f"Обработка и транскрибация завершены за {elapsed_time:.2f} секунд")
return text
return result
finally:
# Очистка временных файлов
self.audio_processor.cleanup_temp_files(temp_files)
self.audio_processor.cleanup_temp_files(temp_files)
+1 -1
View File
@@ -9,4 +9,4 @@ logging.basicConfig(
logging.FileHandler("logs/transcribe.log")
]
)
logger = logging.getLogger("whisper-service")
logger = logging.getLogger("whisper-service")