app.py

import os
import json

import numpy as np
import ffmpeg
import whisper
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
from sklearn.tree import DecisionTreeRegressor
import torch
import youtube_dl
import pandas as pd
import streamlit as st
import altair as alt

DATA_DIR = "./data"
if not os.path.exists(DATA_DIR):
    os.makedirs(DATA_DIR)

YDL_OPTS = {
    "download_archive": os.path.join(DATA_DIR, "archive.txt"),
    "format": "bestaudio/best",
    "outtmpl": os.path.join(DATA_DIR, "%(title)s.%(ext)s"),
    "postprocessors": [
        {
            "key": "FFmpegExtractAudio",
            "preferredcodec": "mp3",
            "preferredquality": "192",
        }
    ],
}

llm = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
device = "cuda" if torch.cuda.is_available() else "cpu"


def download(url, ydl_opts):
    with youtube_dl.YoutubeDL(ydl_opts) as ydl:
        result = ydl.extract_info("{}".format(url))
        fname = ydl.prepare_filename(result)
        return fname


def transcribe(audio_path, transcript_path):
    if os.path.exists(transcript_path):
        with open(transcript_path, "r") as f:
            result = json.load(f)
    else:
        whisper_model = whisper.load_model("base")
        result = whisper_model.transcribe(audio_path)
        with open(transcript_path, "w") as f:
            json.dump(result, f)
    return result["segments"]


def compute_seg_durations(segments):
    return [s["end"] - s["start"] for s in segments]


def compute_info_densities(
    segments, seg_durations, llm, tokenizer, device, ctxt_len=512
):
    seg_encodings = [tokenizer(seg["text"], return_tensors="pt") for seg in segments]
    input_ids = [enc.input_ids.to(device) for enc in seg_encodings]
    seg_lens = [x.shape[1] for x in input_ids]
    cat_input_ids = torch.cat(input_ids, axis=1)
    end = 0
    seg_nlls = []
    n = cat_input_ids.shape[1]
    for i, seg_len in enumerate(seg_lens):
        end = min(n, end + seg_len)
        start = max(0, end - ctxt_len)
        ctxt_ids = cat_input_ids[:, start:end]
        target_ids = ctxt_ids.clone()
        target_ids[:, :-seg_len] = -100
        avg_nll = llm(ctxt_ids, labels=target_ids).loss.detach().numpy()
        nll = avg_nll * seg_len
        seg_nlls.append(nll)
    seg_nlls = np.array(seg_nlls)
    info_densities = seg_nlls / seg_durations
    return info_densities


def smooth_info_densities(info_densities, seg_durations, max_leaf_nodes, min_sec_leaf):
    min_samples_leaf = int(np.ceil(min_sec_leaf / np.mean(seg_durations)))
    tree = DecisionTreeRegressor(
        max_leaf_nodes=max_leaf_nodes, min_samples_leaf=min_samples_leaf
    )
    X = np.arange(0, len(info_densities), 1)[:, np.newaxis]
    tree.fit(X, info_densities)
    smoothed_info_densities = tree.predict(X)
    return smoothed_info_densities


def squash_segs(segments, info_densities):
    start = segments[0]["start"]
    end = None
    seg_times = []
    seg_densities = [info_densities[0]]
    for i in range(1, len(segments)):
        curr_density = info_densities[i]
        if curr_density != info_densities[i - 1]:
            seg = segments[i]
            seg_start = seg["start"]
            seg_times.append((start, seg_start))
            seg_densities.append(curr_density)
            start = seg_start
    seg_times.append((start, segments[-1]["end"]))
    return seg_times, seg_densities


def compute_speedups(info_densities):
    avg_density = np.mean(info_densities)
    speedups = avg_density / info_densities
    return speedups


def compute_actual_speedup(durations, speedups, total_duration):
    spedup_durations = durations / speedups
    spedup_total_duration = spedup_durations.sum()
    actual_speedup_factor = total_duration / spedup_total_duration
    return spedup_total_duration, actual_speedup_factor


def postprocess_speedups(
    speedups, factor, min_speedup, max_speedup, durations, total_duration, thresh=0.01
):
    assert min_speedup <= factor and factor <= max_speedup
    tuned_factor = np.array([factor / 10, factor * 10])
    actual_speedup_factor = None
    while (
        actual_speedup_factor is None
        or abs(actual_speedup_factor - factor) / factor > thresh
    ):
        mid = tuned_factor.mean()
        tuned_speedups = speedups * mid
        tuned_speedups = np.round(tuned_speedups, decimals=2)
        tuned_speedups = np.clip(tuned_speedups, min_speedup, max_speedup)
        _, actual_speedup_factor = compute_actual_speedup(
            durations, tuned_speedups, total_duration
        )
        tuned_factor[0 if actual_speedup_factor < factor else 1] = mid
    return tuned_speedups


def cat_clips(seg_times, speedups, audio_path, output_path):
    if os.path.exists(output_path):
        os.remove(output_path)
    in_file = ffmpeg.input(audio_path)
    segs = []
    for (start, end), speedup in zip(seg_times, speedups):
        seg = in_file.filter("atrim", start=start, end=end).filter("atempo", speedup)
        segs.append(seg)
    cat = ffmpeg.concat(*segs, v=0, a=1)
    cat.output(output_path).run()


def format_duration(duration):
    s = duration % 60
    m = duration // 60
    h = m // 60
    return "%02d:%02d:%02d" % (h, m, s)


def strike(url, speedup_factor, min_speedup, max_speedup, max_num_segments):

    min_speedup = max(0.5, min_speedup)  # ffmpeg limit

    name = download(url, YDL_OPTS)
    assert name.endswith(".m4a")
    name = name.split(".m4a")[0].split("/")[-1]

    audio_path = os.path.join(DATA_DIR, "%s.mp3" % name)
    transcript_path = os.path.join(DATA_DIR, "%s.json" % name)
    output_path = os.path.join(DATA_DIR, "%s_smooth.mp3" % name)

    segments = transcribe(audio_path, transcript_path)

    seg_durations = compute_seg_durations(segments)

    info_densities = compute_info_densities(
        segments, seg_durations, llm, tokenizer, device
    )

    total_duration = segments[-1]["end"] - segments[0]["start"]
    min_sec_leaf = total_duration / max_num_segments
    smoothed_info_densities = smooth_info_densities(
        info_densities, seg_durations, max_num_segments, min_sec_leaf
    )

    squashed_times, squashed_densities = squash_segs(segments, smoothed_info_densities)
    squashed_durations = np.array([end - start for start, end in squashed_times])

    speedups = compute_speedups(squashed_densities)
    speedups = postprocess_speedups(
        speedups,
        speedup_factor,
        min_speedup,
        max_speedup,
        squashed_durations,
        total_duration,
    )

    cat_clips(squashed_times, speedups, audio_path, output_path)

    spedup_total_duration, actual_speedup_factor = compute_actual_speedup(
        squashed_durations, speedups, total_duration
    )
    st.write("original duration: %s" % format_duration(total_duration))
    st.write("new duration: %s" % format_duration(spedup_total_duration))
    st.write("speedup: %0.2f" % actual_speedup_factor)

    times = np.array([(seg["start"] + seg["end"]) / 2 for seg in segments])
    times /= 60
    annotations = [seg["text"] for seg in segments]
    data = [times, info_densities / np.log(2), annotations]
    cols = ["time (minutes)", "bits per second", "transcript"]
    df = pd.DataFrame(list(zip(*data)), columns=cols)

    lines = (
        alt.Chart(df, title="information rate")
        .mark_line(color="gray", opacity=0.5)
        .encode(
            x=cols[0],
            y=cols[1],
        )
    )
    dots = (
        alt.Chart(df)
        .mark_circle(size=50, opacity=1)
        .encode(x=cols[0], y=cols[1], tooltip=["transcript"])
    )
    st.altair_chart((lines + dots).interactive(), use_container_width=True)

    times = sum([list(x) for x in squashed_times], [])
    times = np.array(times)
    times /= 60
    data = [times, np.repeat(speedups, 2)]
    cols = ["time (minutes)", "speedup"]
    df = pd.DataFrame(list(zip(*data)), columns=cols)
    st.line_chart(df, x=cols[0], y=cols[1])

    return output_path


with st.form("my_form"):
    url = st.text_input(
        "youtube url", value="https://www.youtube.com/watch?v=_3MBQm7GFIM"
    )
    speedup_factor = st.slider("speedup", min_value=1.0, max_value=10.0, value=1.5)
    min_speedup = 1
    max_speedup = st.slider("maximum speedup", min_value=1.0, max_value=10.0, value=2.0)
    speedup_factor = min(speedup_factor, max_speedup)
    max_num_segments = st.slider(
        "variance in speedup over time", min_value=2, max_value=100, value=20
    )
    submitted = st.form_submit_button("submit")
    if submitted:
        output_path = strike(
            url, speedup_factor, min_speedup, max_speedup, max_num_segments
        )
        st.audio(output_path)