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 avançado para transcrição de áudio
transcriber = pipeline(
    task="automatic-speech-recognition",
    model="openai/whisper-large-v2",  # Modelo mais avançado
    device=device,
    chunk_length_s=30,
    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  # Desativando o tokenizer rápido
)

# 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):
    with gr.Progress(track_tqdm=True) as progress:
        progress(0, desc="Iniciando transcrição...")

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

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

        # Realizando a classificação zero-shot na transcrição
        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 de até 1 hora para transcrição e análise de emoções.",
    theme="default"
)

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