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-small",
    device=device
)

# 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):
    """
    Recebe um arquivo de áudio, transcreve e analisa as emoções presentes.
    """
    # Transcrevendo o áudio
    transcription = transcriber(audio_file)["text"]
    
    # Lista de emoções para a classificação
    emotions = ["medo", "rancor", "ódio", "raiva", "felicidade", "amor", "tristeza", "ansiedade", "inveja", "sarcasmo", "vergonha"]
    
    # 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}: {score:.2f}")
    
    # Unindo os resultados em uma string
    emotion_output = "\n".join(results)
    
    return transcription, emotion_output

# Criando a interface Gradio
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()