import gradio as gr
import torch
from transformers import pipeline
from timestamp import format_timestamp
MODEL_NAME = "openai/whisper-medium"
BATCH_SIZE = 8
device = 0 if torch.cuda.is_available() else "cpu"
pipe = pipeline(
task="automatic-speech-recognition",
model=MODEL_NAME,
chunk_length_s=30,
device=device,
)
def transcribe(file, task, return_timestamps):
outputs = pipe(file, batch_size=BATCH_SIZE, generate_kwargs={"task": task}, return_timestamps=True)
text = outputs["text"]
timestamps = outputs["chunks"]
if return_timestamps==True:
timestamps = [f"[{format_timestamp(chunk['timestamp'][0])} -> {format_timestamp(chunk['timestamp'][1])}] {chunk['text']}" for chunk in timestamps]
else:
timestamps = [f"{chunk['text']}" for chunk in timestamps]
text = "
".join(str(feature) for feature in timestamps)
text = f"