503 lines
19 KiB
Python
503 lines
19 KiB
Python
"""
|
||
recognition.py — модуль распознавания аудиофайла с ответами ученика ЕГЭ (аудирование, английский язык).
|
||
|
||
Зависимости:
|
||
pip install faster-whisper
|
||
|
||
Использование:
|
||
from recognition import transcribe, extract_answers
|
||
|
||
# Полный pipeline: аудио -> транскрипт -> структурированные ответы
|
||
result = transcribe("student_answers.mp3")
|
||
answers = extract_answers(result.text)
|
||
print(answers)
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import re
|
||
import logging
|
||
from dataclasses import dataclass, field
|
||
from pathlib import Path
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Константы
|
||
# ---------------------------------------------------------------------------
|
||
|
||
# Модели faster-whisper по убыванию скорости / возрастанию качества:
|
||
# tiny, base, small, medium, large-v2, large-v3
|
||
DEFAULT_MODEL = "medium"
|
||
|
||
# Подсказка для Whisper — описывает формат ответов ученика.
|
||
# Критически важна для правильного распознавания "one/two/three" как цифр
|
||
# и "A equals 3" как ответа на задание 1.
|
||
WHISPER_PROMPT = (
|
||
"Student answers to EGE English listening exam. "
|
||
"Task one matching: speaker A answer three, speaker B answer one, "
|
||
"speaker C answer five, speaker D answer seven, speaker E answer two, speaker F answer four. "
|
||
"Tasks two through nine True False Not Stated: "
|
||
"task two true, task three false, task four not stated. "
|
||
"Tasks ten through eighteen multiple choice one two or three: "
|
||
"task ten two, task eleven one, task twelve three."
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Структуры данных
|
||
# ---------------------------------------------------------------------------
|
||
|
||
@dataclass
|
||
class TranscriptResult:
|
||
"""Результат транскрипции аудиофайла."""
|
||
text: str # Полный текст транскрипта
|
||
language: str # Определённый язык ("en")
|
||
duration_seconds: float # Длительность аудио в секундах
|
||
segments: list[dict] = field(default_factory=list) # Детальные сегменты с таймкодами
|
||
model_used: str = DEFAULT_MODEL
|
||
|
||
|
||
@dataclass
|
||
class StudentAnswers:
|
||
"""Структурированные ответы ученика, извлечённые из транскрипта."""
|
||
|
||
# Задание 1: соответствие A–F → цифра 1–7
|
||
task1: dict[str, str] = field(default_factory=dict) # {"A": "3", "B": "1", ...}
|
||
|
||
# Задания 2–9: True(1) / False(2) / Not Stated(3)
|
||
task2_9: dict[int, str] = field(default_factory=dict) # {2: "1", 3: "2", ...}
|
||
|
||
# Задания 10–18: выбор из вариантов 1/2/3
|
||
task10_18: dict[int, str] = field(default_factory=dict) # {10: "2", 11: "1", ...}
|
||
|
||
# Задания, которые не удалось распознать
|
||
unrecognized: list[str] = field(default_factory=list)
|
||
|
||
def to_dict(self) -> dict:
|
||
return {
|
||
"task1": self.task1,
|
||
"task2_9": {str(k): v for k, v in self.task2_9.items()},
|
||
"task10_18": {str(k): v for k, v in self.task10_18.items()},
|
||
"unrecognized": self.unrecognized,
|
||
}
|
||
|
||
def summary(self) -> str:
|
||
"""Читаемое представление для вывода агенту / в лог."""
|
||
lines = ["=== Распознанные ответы ученика ==="]
|
||
|
||
if self.task1:
|
||
lines.append("\nЗадание 1 (соответствие):")
|
||
for letter in "ABCDEF":
|
||
ans = self.task1.get(letter, "?")
|
||
lines.append(f" {letter} → {ans}")
|
||
|
||
if self.task2_9:
|
||
lines.append("\nЗадания 2–9 (True/False/Not Stated):")
|
||
labels = {"1": "True", "2": "False", "3": "Not Stated"}
|
||
for task_num in range(2, 10):
|
||
ans = self.task2_9.get(task_num, "?")
|
||
label = labels.get(ans, ans)
|
||
lines.append(f" Задание {task_num}: {ans} ({label})")
|
||
|
||
if self.task10_18:
|
||
lines.append("\nЗадания 10–18 (выбор):")
|
||
for task_num in range(10, 19):
|
||
ans = self.task10_18.get(task_num, "?")
|
||
lines.append(f" Задание {task_num}: {ans}")
|
||
|
||
if self.unrecognized:
|
||
lines.append(f"\nНе распознано: {', '.join(self.unrecognized)}")
|
||
|
||
return "\n".join(lines)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Транскрипция
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def transcribe(
|
||
audio_path: str | Path,
|
||
model_size: str = DEFAULT_MODEL,
|
||
device: str = "auto",
|
||
compute_type: str = "auto",
|
||
language: str = "en",
|
||
beam_size: int = 5,
|
||
) -> TranscriptResult:
|
||
"""
|
||
Транскрибирует аудиофайл с ответами ученика.
|
||
|
||
Args:
|
||
audio_path: Путь к аудиофайлу (MP3, WAV, M4A, OGG, WEBM, FLAC).
|
||
model_size: Размер модели Whisper: tiny/base/small/medium/large-v2/large-v3.
|
||
medium — хороший баланс скорость/качество для ЕГЭ.
|
||
large-v3 — максимальное качество, медленнее.
|
||
device: "auto" | "cpu" | "cuda". "auto" выберет GPU если доступен.
|
||
compute_type: "auto" | "int8" | "float16" | "float32".
|
||
"auto" подберёт оптимальный тип для устройства.
|
||
language: Язык аудио. "en" для ответов на английском.
|
||
beam_size: Ширина луча beam search. 5 — стандарт, выше = точнее но медленнее.
|
||
|
||
Returns:
|
||
TranscriptResult с текстом, языком, длительностью и сегментами.
|
||
|
||
Raises:
|
||
FileNotFoundError: Если аудиофайл не найден.
|
||
RuntimeError: Если faster-whisper не установлен.
|
||
"""
|
||
try:
|
||
from faster_whisper import WhisperModel
|
||
except ImportError:
|
||
raise RuntimeError(
|
||
"faster-whisper не установлен. Установите: pip install faster-whisper"
|
||
)
|
||
|
||
audio_path = Path(audio_path)
|
||
if not audio_path.exists():
|
||
raise FileNotFoundError(f"Аудиофайл не найден: {audio_path}")
|
||
|
||
# Автовыбор устройства и типа вычислений
|
||
resolved_device, resolved_compute = _resolve_device(device, compute_type)
|
||
|
||
logger.info(
|
||
"Загрузка модели %s на %s (%s)...",
|
||
model_size, resolved_device, resolved_compute
|
||
)
|
||
|
||
model = WhisperModel(
|
||
model_size,
|
||
device=resolved_device,
|
||
compute_type=resolved_compute,
|
||
)
|
||
|
||
logger.info("Транскрибирую: %s", audio_path.name)
|
||
|
||
segments_gen, info = model.transcribe(
|
||
str(audio_path),
|
||
language=language,
|
||
beam_size=beam_size,
|
||
initial_prompt=WHISPER_PROMPT,
|
||
word_timestamps=False,
|
||
vad_filter=True, # Фильтрация тишины — полезно для записей с паузами
|
||
vad_parameters={
|
||
"min_silence_duration_ms": 500, # Паузы >0.5с считаются тишиной
|
||
"speech_pad_ms": 200,
|
||
},
|
||
)
|
||
|
||
# Материализуем генератор сегментов
|
||
segments = []
|
||
full_text_parts = []
|
||
|
||
for seg in segments_gen:
|
||
segments.append({
|
||
"start": round(seg.start, 2),
|
||
"end": round(seg.end, 2),
|
||
"text": seg.text.strip(),
|
||
})
|
||
full_text_parts.append(seg.text.strip())
|
||
|
||
full_text = " ".join(full_text_parts)
|
||
|
||
logger.info(
|
||
"Транскрипция завершена. Длительность: %.1f сек, слов ~%d",
|
||
info.duration, len(full_text.split())
|
||
)
|
||
|
||
return TranscriptResult(
|
||
text=full_text,
|
||
language=info.language,
|
||
duration_seconds=round(info.duration, 1),
|
||
segments=segments,
|
||
model_used=model_size,
|
||
)
|
||
|
||
|
||
def _resolve_device(device: str, compute_type: str) -> tuple[str, str]:
|
||
"""Определяет оптимальное устройство и тип вычислений."""
|
||
if device != "auto" and compute_type != "auto":
|
||
return device, compute_type
|
||
|
||
# Проверяем наличие CUDA
|
||
try:
|
||
from torch import cuda
|
||
has_cuda = cuda.is_available()
|
||
except ImportError:
|
||
has_cuda = False
|
||
|
||
if device == "auto":
|
||
device = "cuda" if has_cuda else "cpu"
|
||
|
||
if compute_type == "auto":
|
||
if device == "cuda":
|
||
compute_type = "float16" # GPU: float16 быстрее и точнее чем int8
|
||
else:
|
||
compute_type = "int8" # CPU: int8 значительно быстрее float32
|
||
|
||
return device, compute_type
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Извлечение ответов из транскрипта
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def extract_answers(transcript_text: str) -> StudentAnswers:
|
||
"""
|
||
Извлекает структурированные ответы из текста транскрипта.
|
||
|
||
Обрабатывает разные форматы речи ученика:
|
||
- "A three", "speaker A answer three", "A equals 3", "A — три"
|
||
- "task two true", "number two true", "two — true", "2 true"
|
||
- "task ten two", "ten — 2", "question ten answer two"
|
||
|
||
Args:
|
||
transcript_text: Текст транскрипта от Whisper.
|
||
|
||
Returns:
|
||
StudentAnswers со структурированными ответами.
|
||
"""
|
||
answers = StudentAnswers()
|
||
text = transcript_text.lower().strip()
|
||
|
||
_extract_task1(text, answers)
|
||
_extract_task2_9(text, answers)
|
||
_extract_task10_18(text, answers)
|
||
|
||
logger.debug("Извлечено ответов: task1=%d, task2_9=%d, task10_18=%d, нераспознано=%d",
|
||
len(answers.task1), len(answers.task2_9),
|
||
len(answers.task10_18), len(answers.unrecognized))
|
||
|
||
return answers
|
||
|
||
|
||
def _extract_task1(text: str, answers: StudentAnswers) -> None:
|
||
"""
|
||
Задание 1: соответствие A–F → цифра 1–7.
|
||
Примеры: "A three", "speaker A answer 3", "A equals three", "A — 3"
|
||
"""
|
||
# Числа словами → цифры
|
||
word_to_digit = {
|
||
"one": "1", "two": "2", "three": "3", "four": "4",
|
||
"five": "5", "six": "6", "seven": "7",
|
||
# На случай русских ответов
|
||
"один": "1", "два": "2", "три": "3", "четыре": "4",
|
||
"пять": "5", "шесть": "6", "семь": "7",
|
||
}
|
||
|
||
for letter in "abcdef":
|
||
# Паттерн: буква, затем необязательный разделитель, затем цифра/слово
|
||
pattern = (
|
||
rf"(?:speaker\s+)?{letter}"
|
||
rf"(?:\s+(?:answer|equals|is|—|-|:))?"
|
||
rf"\s+"
|
||
rf"({_digit_or_word_pattern(1, 7)})"
|
||
)
|
||
match = re.search(pattern, text)
|
||
if match:
|
||
raw = match.group(1).strip()
|
||
digit = word_to_digit.get(raw, raw) if not raw.isdigit() else raw
|
||
if digit in [str(i) for i in range(1, 8)]:
|
||
answers.task1[letter.upper()] = digit
|
||
else:
|
||
answers.unrecognized.append(f"1{letter.upper()}")
|
||
else:
|
||
answers.unrecognized.append(f"1{letter.upper()}")
|
||
|
||
|
||
def _extract_task2_9(text: str, answers: StudentAnswers) -> None:
|
||
"""
|
||
Задания 2–9: True(1) / False(2) / Not Stated(3).
|
||
Примеры: "task two true", "number 3 false", "four not stated", "5 — 2"
|
||
"""
|
||
tfs_map = {
|
||
"true": "1", "1": "1",
|
||
"false": "2", "2": "2",
|
||
"not stated": "3", "not_stated": "3", "3": "3",
|
||
}
|
||
|
||
num_words = {
|
||
"two": 2, "three": 3, "four": 4, "five": 5,
|
||
"six": 6, "seven": 7, "eight": 8, "nine": 9,
|
||
"2": 2, "3": 3, "4": 4, "5": 5,
|
||
"6": 6, "7": 7, "8": 8, "9": 9,
|
||
}
|
||
|
||
for word, num in num_words.items():
|
||
pattern = (
|
||
rf"(?:task|number|question|задание)?\s*{re.escape(word)}"
|
||
rf"(?:\s+(?:is|answer|—|-|:))?"
|
||
rf"\s+"
|
||
rf"(true|false|not\s+stated|not_stated|[123])"
|
||
)
|
||
match = re.search(pattern, text)
|
||
if match:
|
||
raw = match.group(1).strip().replace(" ", "_")
|
||
digit = tfs_map.get(raw) or tfs_map.get(raw.replace("_", " "))
|
||
if digit:
|
||
answers.task2_9[num] = digit
|
||
else:
|
||
answers.unrecognized.append(str(num))
|
||
else:
|
||
answers.unrecognized.append(str(num))
|
||
|
||
|
||
def _extract_task10_18(text: str, answers: StudentAnswers) -> None:
|
||
"""
|
||
Задания 10–18: выбор из вариантов 1/2/3.
|
||
Примеры: "task ten two", "eleven — 1", "question 12 answer three"
|
||
"""
|
||
word_to_digit_choice = {
|
||
"one": "1", "two": "2", "three": "3",
|
||
"1": "1", "2": "2", "3": "3",
|
||
}
|
||
|
||
num_words = {
|
||
"ten": 10, "eleven": 11, "twelve": 12, "thirteen": 13,
|
||
"fourteen": 14, "fifteen": 15, "sixteen": 16,
|
||
"seventeen": 17, "eighteen": 18,
|
||
"10": 10, "11": 11, "12": 12, "13": 13,
|
||
"14": 14, "15": 15, "16": 16, "17": 17, "18": 18,
|
||
}
|
||
|
||
for word, num in num_words.items():
|
||
pattern = (
|
||
rf"(?:task|number|question|задание)?\s*{re.escape(word)}"
|
||
rf"(?:\s+(?:is|answer|—|-|:))?"
|
||
rf"\s+"
|
||
rf"({_digit_or_word_pattern(1, 3)})"
|
||
)
|
||
match = re.search(pattern, text)
|
||
if match:
|
||
raw = match.group(1).strip()
|
||
digit = word_to_digit_choice.get(raw)
|
||
if digit:
|
||
answers.task10_18[num] = digit
|
||
else:
|
||
answers.unrecognized.append(str(num))
|
||
else:
|
||
answers.unrecognized.append(str(num))
|
||
|
||
|
||
def _digit_or_word_pattern(min_val: int, max_val: int) -> str:
|
||
"""Строит regex-паттерн для диапазона цифр и их словесных форм."""
|
||
digits = [str(i) for i in range(min_val, max_val + 1)]
|
||
words = {
|
||
1: "one", 2: "two", 3: "three", 4: "four", 5: "five",
|
||
6: "six", 7: "seven", 8: "eight", 9: "nine", 10: "ten",
|
||
11: "eleven", 12: "twelve", 13: "thirteen", 14: "fourteen",
|
||
15: "fifteen", 16: "sixteen", 17: "seventeen", 18: "eighteen",
|
||
}
|
||
options = digits + [words[i] for i in range(min_val, max_val + 1) if i in words]
|
||
return "(" + "|".join(sorted(options, key=len, reverse=True)) + ")"
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Удобный pipeline
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def process_audio(
|
||
audio_path: str | Path,
|
||
model_size: str = DEFAULT_MODEL,
|
||
device: str = "auto",
|
||
verbose: bool = False,
|
||
) -> tuple[TranscriptResult, StudentAnswers]:
|
||
"""
|
||
Полный pipeline: аудиофайл → транскрипт → структурированные ответы.
|
||
|
||
Args:
|
||
audio_path: Путь к аудиофайлу.
|
||
model_size: Размер модели Whisper (medium по умолчанию).
|
||
device: "auto" | "cpu" | "cuda".
|
||
verbose: Если True — выводить промежуточные результаты в stdout.
|
||
|
||
Returns:
|
||
Кортеж (TranscriptResult, StudentAnswers).
|
||
|
||
Example:
|
||
result, answers = process_audio("student.mp3")
|
||
print(answers.summary())
|
||
# Передать answers.to_dict() агенту для сверки с ключами
|
||
"""
|
||
if verbose:
|
||
print(f"[1/2] Транскрибирую {Path(audio_path).name} (модель: {model_size})...")
|
||
|
||
transcript = transcribe(audio_path, model_size=model_size, device=device)
|
||
|
||
if verbose:
|
||
print(f" Длительность: {transcript.duration_seconds} сек")
|
||
print(f" Транскрипт: {transcript.text[:120]}...")
|
||
print("[2/2] Извлекаю ответы...")
|
||
|
||
answers = extract_answers(transcript.text)
|
||
|
||
if verbose:
|
||
print(answers.summary())
|
||
|
||
return transcript, answers
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# CLI
|
||
# ---------------------------------------------------------------------------
|
||
|
||
if __name__ == "__main__":
|
||
import argparse
|
||
import json
|
||
import sys
|
||
|
||
logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
|
||
|
||
parser = argparse.ArgumentParser(
|
||
description="Распознавание аудиоответов ЕГЭ (аудирование, английский язык)"
|
||
)
|
||
parser.add_argument("audio", help="Путь к аудиофайлу")
|
||
parser.add_argument(
|
||
"--model", default=DEFAULT_MODEL,
|
||
choices=["tiny", "base", "small", "medium", "large-v2", "large-v3"],
|
||
help=f"Размер модели Whisper (по умолчанию: {DEFAULT_MODEL})"
|
||
)
|
||
parser.add_argument(
|
||
"--device", default="auto",
|
||
choices=["auto", "cpu", "cuda"],
|
||
help="Устройство для инференса (по умолчанию: auto)"
|
||
)
|
||
parser.add_argument(
|
||
"--output", choices=["summary", "json", "transcript"],
|
||
default="summary",
|
||
help="Формат вывода: summary (читаемый), json (машинный), transcript (сырой текст)"
|
||
)
|
||
parser.add_argument(
|
||
"--transcript-only", action="store_true",
|
||
help="Только транскрипция без извлечения ответов"
|
||
)
|
||
|
||
args = parser.parse_args()
|
||
|
||
try:
|
||
if args.transcript_only or args.output == "transcript":
|
||
result = transcribe(args.audio, model_size=args.model, device=args.device)
|
||
print(result.text)
|
||
else:
|
||
result, answers = process_audio(
|
||
args.audio,
|
||
model_size=args.model,
|
||
device=args.device,
|
||
verbose=(args.output == "summary"),
|
||
)
|
||
|
||
if args.output == "json":
|
||
output = {
|
||
"transcript": result.text,
|
||
"language": result.language,
|
||
"duration_seconds": result.duration_seconds,
|
||
"model_used": result.model_used,
|
||
"answers": answers.to_dict(),
|
||
}
|
||
print(json.dumps(output, ensure_ascii=False, indent=2))
|
||
elif args.output == "summary":
|
||
# verbose=True уже вывел всё в process_audio
|
||
pass
|
||
|
||
except (FileNotFoundError, RuntimeError) as e:
|
||
print(f"Ошибка: {e}", file=sys.stderr)
|
||
sys.exit(1)
|