import streamlit as st
import torch
import tempfile
import os
from transformers import WhisperProcessor, WhisperForConditionalGeneration
from datasets import load_dataset

# Configuration de l'interface Streamlit
st.title("🔊 Transcription Audio avec Whisper Fine-tuné")
st.write("Upload un fichier audio et laisse ton modèle fine-tuné faire le travail !")

# 🔹 Charger le modèle fine-tuné et le processeur
def load_model():
    model_name = "SimpleFrog/whisper_finetuned"  # Remplace par ton nom de repo sur Hugging Face
    processor = WhisperProcessor.from_pretrained(model_name)
    model = WhisperForConditionalGeneration.from_pretrained(model_name)
    model.eval()  # Mode évaluation
    return processor, model

processor, model = load_model()

# 🔹 Upload d'un fichier audio
uploaded_file = st.file_uploader("Upload un fichier audio", type=["mp3", "wav", "m4a"])

if uploaded_file is not None:
    # Sauvegarder temporairement l'audio
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_audio:
        temp_audio.write(uploaded_file.read())
        temp_audio_path = temp_audio.name

    # Charger et traiter l'audio
    st.write("📄 **Transcription en cours...**")
    
    audio_input = processor(temp_audio_path, return_tensors="pt", sampling_rate=16000)
    input_features = audio_input.input_features
    
    # Générer la transcription
    with torch.no_grad():
        predicted_ids = model.generate(input_features)
    
    # Décoder la sortie
    transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
    
    # Afficher la transcription
    st.subheader("📝 Transcription :")
    st.text_area("", transcription, height=200)

    # Supprimer le fichier temporaire après l'affichage
    os.remove(temp_audio_path)

st.write("🔹 Modèle fine-tuné utilisé :", "SimpleFrog/whisper_finetuned")