#importing all the required libraries
import gradio as gr 
import librosa
from transformers import AutoFeatureExtractor, pipeline

#Loading and fixing the audio input
def load_and_fix_data(input_file, model_sampling_rate):
    speech, sample_rate = librosa.load(input_file)
    if len(speech.shape) > 1:
        speech = speech[:, 0] + speech[:, 1]
    if sample_rate != model_sampling_rate:
        speech = librosa.resample(speech, sample_rate, model_sampling_rate)
    return speech

#Loading the feature extractor and instantiating the pipeline
model_name1 = "jonatasgrosman/wav2vec2-xls-r-1b-spanish"
feature_extractor = AutoFeatureExtractor.from_pretrained(model_name1)
sampling_rate = feature_extractor.sampling_rate
asr = pipeline("automatic-speech-recognition", model=model_name1)

#Instantiating a pipeline for harassment detection (text classification)
model_name2 = "hackathon-pln-es/Detect-Acoso-Twitter-Es"
classifier = pipeline("text-classification", model = model_name2)

#Defining a function for speech-to-text-conversion
def speech_to_text(input_file):
    speech = load_and_fix_data(input_file, sampling_rate)
    transcribed_text = asr(speech, chunk_length_s=15, stride_length_s=1)["text"]
    return transcribed_text
    
#Defining a function for Harassment detection (text classification)
def harassment_detector(transcribed_text):
    harassment_detection = classifier(transcribed_text)[0]["label"]
    return harassment_detection

#Defining a function which outputs audio transcription and the output of harassment detection module
new_line = "\n\n\n"
def asr_and_harassment_detection(input_file):
    transcribed_text = speech_to_text(input_file)
    harassment_detection = harassment_detector(transcribed_text)
    return f"Audio Transcription :{transcribed_text} {new_line} Audio content is: {harassment_detection}"


inputs=[gr.inputs.Audio(source="microphone", type="filepath", label="Record your audio")]
outputs=[gr.outputs.Textbox(label="Predicción")]
examples=[["audio2.wav"], ["sample_audio.wav"], ["test1.wav"], ["test2.wav"]]
title="Spanish-Audio-Transcription-based-Harassment-Detection"

description = """ This is a Gradio demo for Spanish audio transcription-based harassment detection. To use this, simply provide an audio input (audio recording or via microphone), which will subsequently be transcribed and classified as Harassment/non-harassment pertaining to audio (transcription) with the help of pre-trained models.


Pre-trained model used for Spanish ASR: [jonatasgrosman/wav2vec2-xls-r-1b-spanish](https://huggingface.co/jonatasgrosman/wav2vec2-xls-r-1b-spanish)


Pre-trained model used for Harassment Detection: [hackathon-pln-es/Detect-Acoso-Twitter-Es](https://huggingface.co/hackathon-pln-es/Detect-Acoso-Twitter-Es)"""

gr.Interface(
    asr_and_harassment_detection,
    inputs=inputs,
    outputs=outputs,
    examples=examples,
    title=title,
    description=description,
    layout="horizontal",
    theme="huggingface",
).launch(enable_queue=True)