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

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

# Carregando o modelo Whisper pequeno para transcrição de áudio
transcriber = pipeline(
    task="automatic-speech-recognition",
    model="openai/whisper-medium",  # Modelo menor para maior velocidade
    device=device,
    chunk_length_s=30,  # Ajuste conforme necessário
    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):
    import numpy as np

    # Dividindo o áudio em chunks
    total_duration = transcriber.feature_extractor.sampling_rate * 30  # 30 segundos
    audio_input, _ = transcriber.feature_extractor(audio_file, return_tensors="pt").values()
    audio_length = audio_input.shape[1]
    num_chunks = int(np.ceil(audio_length / total_duration))
    
    transcription = ""
    
    for i in range(num_chunks):
        start_time = i * total_duration
        end_time = min((i + 1) * total_duration, audio_length)
        chunk = audio_input[:, int(start_time):int(end_time)]
        
        # Transcrevendo o chunk
        chunk_transcription = transcriber(chunk, generate_kwargs={"language": "Portuguese", "task": "transcribe"})["text"]
        transcription += chunk_transcription + " "
        
        # Atualizando a barra de progresso
        progress = int((i + 1) / num_chunks * 100)
        yield f"Transcrevendo... {progress}% concluído", ""
    
    # Após a transcrição completa, analisar as emoções
    yield "Análise de emoções em andamento...", ""
    
    # 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 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)
    
    emotion_output = "\n".join(results)
    
    yield 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()