app.py

import gradio as gr
from transformers import WhisperProcessor, WhisperForConditionalGeneration
from datasets import load_dataset
import torch

# 加载 Whisper 模型和 processor
model_name = "openai/whisper-small"
processor = WhisperProcessor.from_pretrained(model_name)
model = WhisperForConditionalGeneration.from_pretrained(model_name)

# 加载数据集 bigcode/the-stack
# dataset = load_dataset("CoIR-Retrieval/CodeSearchNet-php-queries-corpus", data_dir="data", split="train[:80%]")
queries_dataset = load_dataset("CoIR-Retrieval/CodeSearchNet-php-queries-corpus", data_dir="data", split="queries")
corpus_dataset = load_dataset("CoIR-Retrieval/CodeSearchNet-php-queries-corpus", data_dir="data", split="corpus")


def transcribe(audio):
    # 处理音频进行转录
    audio_input = processor(audio, return_tensors="pt").input_values
    with torch.no_grad():
        logits = model(audio_input).logits
    predicted_ids = torch.argmax(logits, dim=-1)
    transcription = processor.batch_decode(predicted_ids)
    
    # 返回转录结果
    return transcription[0]

# Gradio 界面
iface = gr.Interface(
    fn=transcribe,
    inputs=gr.Audio( type="filepath"),
    outputs="text",
    title="Whisper Transcription for Developers",
    description="使用 Whisper 和 bigcode 数据集转录开发者相关术语。"
)

# 启动 Gradio 应用
iface.launch()