h_brainstorming.py

import json
import os
from typing import Any, Dict, List, cast

from groq import Groq
from openai import OpenAI

from utils.other.endpoints import timeit

os.environ['GROQ_API_KEY'] = ''

# filename = '../audioSamples/empty.wav'


# filename = '../audioSamples/1719787052529-temp.wav'
filename = 'data/more/18-42-24-841626.wav'


# filename = 'data/more/18-45-32-069108.wav'


def execute() -> None:
    client = OpenAI()
    with open(filename, "rb") as file:
        transcription = client.audio.transcriptions.create(
            file=(filename, file.read()),
            model="whisper-1",
            timestamp_granularities=["segment"],
            response_format="verbose_json",
            language="en",
            temperature=0.0,
        )
        segments = transcription.model_dump_json()
        print(segments)
        data: List[Dict[str, Any]] = json.loads(segments).get('segments', [])
        for segment in data:
            print(segment['start'], segment['end'], segment['text'])


@timeit
def execute_groq():
    client = Groq()
    with open(filename, "rb") as file:
        transcription = client.audio.transcriptions.create(
            file=(filename, file.read()),
            model="whisper-large-v3",
            response_format="text",
            language="en",
            temperature=0.0,
        )
        # print(transcription)
        return transcription


@timeit
def diarization(content: str):
    # client = OpenAI()
    client = Groq()
    system_prompt = '''You are a helpful assistant for correcting transcriptions of conversations.\
    Correct any spelling discrepancies in the transcribed text, add necessary punctuation such as periods, commas, \
    and capitalization, and most important differentiate contextually within the multiple speakers in the conversation.
    
    The output should be formatted as a JSON instance that conforms to the JSON schema below.
    
    As an example, for the schema {"properties": {"foo": {"title": "Foo", "description": "a list of strings", "type": "array", "items": {"type": "string"}}}, "required": ["foo"]}
    the object {"foo": ["bar", "baz"]} is a well-formatted instance of the schema. The object {"properties": {"foo": ["bar", "baz"]}} is not well-formatted.
    
    Here is the output schema:
    ```
    {"properties": {"segments": {"title": "Segments", "description": "The segments of the conversation", "default": [], "type": "array", "items": {"$ref": "#/definitions/Segment"}}}, "definitions": {"Segment": {"title": "Segment", "type": "object", "properties": {"speaker": {"title": "Speaker", "description": "The speaker id for this segment", "default": "SPEAKER_00", "type": "string"}, "text": {"title": "Text", "description": "The text of the segment", "default": "", "type": "string"}}}}}
    ```'''.replace('  ', '').strip()

    response = client.chat.completions.create(
        # model="gpt-4o",
        model="llama3-70b-8192",
        temperature=0,
        messages=[{"role": "system", "content": system_prompt}, {"role": "user", "content": content}],
    )
    return response.choices[0].message.content


import torch  # type: ignore[reportMissingImports]  # torch not installed in dev venv

# torch ships without type stubs; alias as Any to avoid cascading unknown-member warnings.
_torch: Any = cast(Any, torch)

_torch.set_num_threads(1)

model, utils = _torch.hub.load(repo_or_dir='snakers4/silero-vad', model='silero_vad')
get_speech_timestamps, _, read_audio, _, _ = utils


@timeit
def has_audio() -> bool:
    wav = read_audio(filename)
    speech_timestamps = get_speech_timestamps(wav, model, sampling_rate=8000)
    return len(speech_timestamps) > 0


def retrieve_proper_segment_points(file_path: str) -> List[Any]:
    wav = read_audio(file_path)
    speech_timestamps = get_speech_timestamps(wav, model, sampling_rate=8000)
    if not speech_timestamps:
        return [None, None]
    return [speech_timestamps[0]['start'] / 1000, speech_timestamps[-1]['end'] / 1000]


if __name__ == '__main__':
    # execute()
    files = sorted(os.listdir('../audioSamples'), key=lambda x: x)
    print('Files:', files)
    for path in os.listdir('../audioSamples'):
        filename = f'../audioSamples/{path}'
        transcription = execute_groq()
        # transcription = fal()
        print(diarization(cast(str, transcription)))
    # has_audio()
    # print(retrieve_proper_segment_points(filename))