import gradio as gr
from transformers import pipeline, AutoTokenizer
import torch

# Verificando se a GPU está disponível
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Carregando o modelo Whisper para transcrição de áudio
transcriber = pipeline(
    task="automatic-speech-recognition",
    model="openai/whisper-medium",
    device=device,
    chunk_length_s=30,  # Definindo chunk_length_s para 30 segundos
    stride_length_s=5,
    generate_kwargs={"language": "Portuguese", "task": "transcribe"}
)

# Carregando o tokenizer lento para o classificador
tokenizer = AutoTokenizer.from_pretrained(
    "joeddav/xlm-roberta-large-xnli",
    use_fast=False
)

# Carregando o pipeline de classificação zero-shot com o tokenizer lento
classifier = pipeline(
    "zero-shot-classification",
    model="joeddav/xlm-roberta-large-xnli",
    tokenizer=tokenizer,
    device=device
)

def transcribe_and_analyze(audio_file):
    progress = gr.Progress(track_tqdm=True)
    progress(0, desc="Iniciando transcrição...")

    # Transcrevendo o áudio com return_timestamps=True
    transcription_result = transcriber(audio_file, return_timestamps=True)
    transcription = transcription_result["text"]
    progress(50, desc="Transcrição concluída. Analisando emoções...")

    # Lista de emoções para a classificação
    emotions = ["alegria", "tristeza", "raiva", "nojo", "medo", "ansiedade", "vergonha", "tédio", "inveja"]

    # Realizando a classificação zero-shot
    classification = classifier(transcription, emotions, multi_label=True)

    # Formatando os resultados
    results = []
    for label, score in zip(classification["labels"], classification["scores"]):
        results.append(f"{label.capitalize()}: {score:.2f}")

    # Ordenando os resultados por score decrescente
    results.sort(key=lambda x: float(x.split(": ")[1]), reverse=True)

    # Unindo os resultados em uma string
    emotion_output = "\n".join(results)

    progress(100, desc="Processamento concluído.")

    return transcription, emotion_output

# Criando a interface Gradio com barra de progresso
interface = gr.Interface(
    fn=transcribe_and_analyze,
    inputs=gr.Audio(type="filepath", label="Faça upload do seu áudio"),
    outputs=[
        gr.Textbox(label="Transcrição do Áudio"),
        gr.Textbox(label="Emoções Detectadas")
    ],
    title="Voxsense 🗣️❣️",
    description="Envie um arquivo de áudio para transcrição e análise de emoções.",
    theme="default"
)

if __name__ == "__main__":
    interface.queue().launch()