whisper / app.py
artyomboyko's picture
Update app.py
779f9f6 verified
raw
history blame contribute delete
5.73 kB
# -*- coding: utf-8 -*-
import torch
import gradio as gr
import yt_dlp as youtube_dl
from transformers import pipeline
from transformers.pipelines.audio_utils import ffmpeg_read
from huggingface_hub import login
import tempfile
import os
import time
# --- Конфигурация ---
HF_TOKEN = os.getenv("HF_ACCESS_TOKEN")
MODEL_NAME = "artyomboyko/whisper-small-ru-v4"
BATCH_SIZE = 8
# FILE_LIMIT_MB = 1000
# --- Инициализация ---
try:
if HF_TOKEN:
login(token=HF_TOKEN)
print("Успешный вход в Hugging Face Hub.")
else:
print("HF_ACCESS_TOKEN не найден. Запуск в оффлайн-режиме или с публичными моделями.")
except Exception as e:
print(f"Не удалось войти в Hugging Face Hub: {e}")
# Определение устройства для вычислений (GPU или CPU)
device = 0 if torch.cuda.is_available() else "cpu"
if device == 0:
print("GPU доступен, вычисления будут производиться на GPU.")
else:
print("GPU не доступен, вычисления будут производиться на CPU.")
# Загрузка модели распознавания речи
try:
pipe = pipeline(
task="automatic-speech-recognition",
model=MODEL_NAME,
chunk_length_s=30,
device=device,
)
print(f"Модель {MODEL_NAME} успешно загружена.")
except Exception as e:
print(f"Ошибка при загрузке модели {MODEL_NAME}: {e}")
exit()
# --- Функции ---
def transcribe(filepath, task):
"""
Транскрибирует аудиофайл с помощью загруженной модели Whisper.
Args:
filepath (str): Путь к аудиофайлу.
task (str): Задача для модели ("transcribe" или "translate").
Returns:
str: Распознанный текст.
"""
if filepath is None:
# Используем gr.Warning вместо gr.Error для неблокирующего уведомления
gr.Warning("Аудиофайл не предоставлен! Пожалуйста, загрузите или запишите аудио перед отправкой запроса.")
return "Ошибка: Аудиофайл не предоставлен." # Возвращаем строку ошибки
print(f"Начало транскрибации файла: {filepath}, Задача: {task}")
try:
result = pipe(filepath, batch_size=BATCH_SIZE, generate_kwargs={"task": task}, return_timestamps=True)
text = result["text"]
print("Транскрибация завершена успешно.")
return text
except Exception as e:
print(f"Ошибка во время транскрибации: {e}")
# Используем gr.Error для вывода серьезной ошибки
raise gr.Error(f"Произошла ошибка при обработке аудио: {e}")
# --- Создание интерфейса Gradio ---
with gr.Blocks(theme=gr.themes.Default(primary_hue="blue", secondary_hue="neutral")) as demo:
gr.Markdown(
f"""
# Whisper: Транскрибация аудио
Транскрибируйте длинные аудиозаписи (с микрофона или из файла) одним нажатием кнопки!
Демо использует модель OpenAI Whisper [{MODEL_NAME}](https://huggingface.co/{MODEL_NAME})
через библиотеку 🤗 Transformers для распознавания речи.
"""
)
# Создание вкладок
with gr.Tabs():
# Вкладка 1: Транскрибация с микрофона
with gr.TabItem("Микрофон"):
with gr.Row():
mic_input = gr.Audio(sources=["microphone"], type="filepath", label="Запись с микрофона")
task_mic = gr.Radio(["transcribe", "translate"], label="Задача", value="transcribe")
output_mic = gr.Textbox(label="Результат", lines=7) # Увеличили количество строк для текста
mic_button = gr.Button("Транскрибировать с микрофона")
# Вкладка 2: Транскрибация из файла
with gr.TabItem("Аудиофайл"):
with gr.Row():
file_input = gr.Audio(sources=["upload"], type="filepath", label="Загрузить аудиофайл")
task_file = gr.Radio(["transcribe", "translate"], label="Задача", value="transcribe")
output_file = gr.Textbox(label="Результат", lines=7)
file_button = gr.Button("Транскрибировать файл")
# --- Привязка функций к кнопкам ---
# Связываем кнопку "Транскрибировать с микрофона" с функцией transcrib
mic_button.click(
fn=transcribe,
inputs=[mic_input, task_mic],
outputs=output_mic
)
# Связываем кнопку "Транскрибировать файл" с функцией transcribe
file_button.click(
fn=transcribe,
inputs=[file_input, task_file],
outputs=output_file
)
# --- Запуск приложения ---
print("Запуск интерфейса Gradio...")
demo.queue().launch(debug=True, share=False)