utils/text.py

from typing import List
from pathlib import Path
import pandas as pd


def get_sentence_data(filename: str, timestamp_dict: dict) -> pd.DataFrame:
    """Extracts the sentences from the output dictionary of whisper inference

    Parameters
    ----------
    filename : str
        Name of the audio analyzed
    timestamp_dict : dict
        Output dictionary from whisper inference

    Returns
    -------
    pd.DataFrame
        DataFrame containing audio filename, start, end and duration of sentences with
        its transcriptions.
    """
    sentence_df = pd.DataFrame(
        columns=["Audio file", "Sentence", "Start", "End", "Duration"]
    )
    for sentence_i in timestamp_dict["segments"]:
        sentence_i = pd.DataFrame(
            {
                "Audio file": [filename],
                "Sentence": [str(sentence_i["text"])],
                "Start": [sentence_i["start"]],
                "End": [sentence_i["end"]],
                "Duration": [sentence_i["end"] - sentence_i["start"]],
            }
        )
        sentence_df = pd.concat([sentence_df, sentence_i], ignore_index=True)
    return sentence_df


def get_word_data(filename: str, timestamp_dict: dict):
    """Extracts the words from the output dictionary of whisper inference

    Parameters
    ----------
    filename : str
        Name of the audio analyzed
    timestamp_dict : dict
        Output dictionary from whisper inference

    Returns
    -------
    pd.DataFrame
        DataFrame containing audio filename, start, end and duration of words with
        its transcriptions.
    """
    word_df = pd.DataFrame(columns=["Audio file", "Word", "Start", "End", "Duration"])
    for sentence_i in timestamp_dict["segments"]:
        for word_i in sentence_i["words"]:
            word_i_df = pd.DataFrame(
                {
                    "Audio file": [filename],
                    "Word": [str(word_i["text"])],
                    "Start": [word_i["start"]],
                    "End": [word_i["end"]],
                    "Duration": [word_i["end"] - word_i["start"]],
                }
            )
            word_df = pd.concat([word_df, word_i_df], ignore_index=True)
    return word_df


def filter_dataframe_by_audiofile(timestamps_df: pd.DataFrame, audio_file: str) -> List:
    """Generates a list from timestamps_df with the timestamps belonging to audio_file

    Parameters
    ----------
    timestamps_df : pd.DataFrame
        Dataframe containing timestamps
    audio_file : str
        Name of the audio file.

    Returns
    -------
    List
        List of tuples containing the start and end of each stamp.
        E.g: [(start_1, end_2), ..., (start_n, end_n)]
    """
    audio_df = timestamps_df[timestamps_df["Audio file"] == audio_file]
    return list(zip(audio_df["Start"], audio_df["End"]))


def get_utterances_transcriptions(timestamps_df: pd.DataFrame) -> List[str]:
    """Gives column with transcriptions

    Parameters
    ----------
    timestamps_df : pd.DataFrame
        DataFrame with transcriptions

    Returns
    -------
    List[str]
        List of the transcriptions
    """
    return timestamps_df.iloc[:, 1].tolist()


def save_transcriptions_segments(
    audio_path: Path, transcriptions_list: List[str], destination: Path
) -> None:
    """Save transcription segments to text files.

    Parameters
    ----------
    audio_path : Path
        Path to the audio file.
    transcriptions_list : List[str]
        List of transcriptions.
    destination : Path
        Destination path for the text files.
    """
    for i, transcription_i in enumerate(transcriptions_list):
        transcription_i_path = destination / f"{audio_path.stem}-{i}.txt"
        with open(str(transcription_i_path), "w") as file:
            file.write(transcription_i)


def generate_transcriptions_splits(
    audio_path: Path, timestamps_df: pd.DataFrame, destination: Path
):
    """Generate and save transcription splits based on timestamps.

    Parameters
    ----------
    audio_path : Path
        Path to the audio file.
    timestamps_df : pd.DataFrame
        DataFrame containing timestamps.
    destination : Path
        Destination path for the text files.
    """
    transcriptions_list = get_utterances_transcriptions(timestamps_df)
    save_transcriptions_segments(audio_path, transcriptions_list, destination)