import gradio as gr
import pandas as pd
import io

# Load the model
model = gr.load("models/luigisaetta/whisper-atcosim3")

# Function to process the model output and convert to CSV
def process_output(audio):
    # Get the transcription from the model
    transcription = model(audio)
    
    # Assuming the transcription includes timestamps in some way
    # If not, you would need to generate timestamps (this depends on your model's output format)
    # Here, we assume the transcription is a list of dictionaries with 'start', 'end', and 'text'
    
    # Convert the transcription to a DataFrame
    df = pd.DataFrame(transcription)
    
    # Create a CSV file in memory
    csv_buffer = io.StringIO()
    df.to_csv(csv_buffer, index=False)
    csv_buffer.seek(0)
    
    return csv_buffer.getvalue()

# Create the Gradio interface
iface = gr.Interface(
    fn=process_output, 
    inputs=gr.Audio(type="file"), 
    outputs=gr.File(type="csv"), 
    title="Audio Transcription to CSV",
    description="Upload an audio file and get the transcription in a CSV format with timestamps."
)

# Launch the app
iface.launch()