import gradio as gr
from transformers import pipeline
import requests
import json
import edge_tts
import asyncio
import tempfile
import os
import inflect
from huggingface_hub import InferenceClient
import re
import time
from streaming_stt_nemo import Model

number_to_word = inflect.engine()

default_lang = "fr"

engines = { default_lang: Model(default_lang) }

def transcribe(audio):
    lang = "fr"
    model = engines[lang]
    text = model.stt_file(audio)[0]
    return text

client1 = InferenceClient("dbddv01/gpt2-french-small")

system_instructions1 = "<s>[SYSTEM] Répondez comme le vrai Jarvis JARVIS, fabriqué par 'Tony Stark'. La conversation doit être très courte, claire, amicale et concise. Le texte fourni est une demande de réponse spécifique de votre part, en tant qu'assistant virtuel. Il vous demande de fournir des réponses amicales comme si vous étiez le personnage Jarvis, créé par « Tony Stark ». On s'attend à ce que j'évite les présentations et que je réponde directement à la requête, en développant tous les aspects de la demande. En tant qu'assistant doté d'une IA, votre tâche consiste à générer des réponses qui semblent avoir été créées par l'utilisateur et non par un assistant doté d'une IA.[USER]"

def model(text):
    generate_kwargs = dict(
        temperature=0.7,
        max_new_tokens=512,
        top_p=0.95,
        repetition_penalty=1,
        do_sample=True,
        seed=42,
    )
    
    formatted_prompt = system_instructions1 + text + "[JARVIS]"
    stream = client1.text_generation(
        formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
    output = ""
    for response in stream:
        if not response.token.text == "</s>":
            output += response.token.text

    return output

def number_to_words(str):
    words = str.split(' ')
    result = []
    
    for word in words:
        if( any(char.isdigit() for char in word) ):
            word = number_to_word.number_to_words(word)

        result.append(word)

    final_result = ' '.join(result).replace('point', '')
    return final_result

async def respond(audio):
    user = transcribe(audio)
    reply = model(user)
    reply2 = number_to_words(reply)
    communicate = edge_tts.Communicate(reply2)
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_file:
        tmp_path = tmp_file.name
        await communicate.save(tmp_path)
    yield tmp_path

DESCRIPTION = """ # <center><b>JARVIS⚡</b></center>
        ### <center>A personal Assistant of Tony Stark for YOU
        ### <center>Voice Chat with your personal Assistant</center>
        """

MORE = """ ## TRY Other Models
        ### Instant Video: Create Amazing Videos in 5 Second -> https://huggingface.co/spaces/KingNish/Instant-Video
        ### Instant Image: 4k images in 5 Second -> https://huggingface.co/spaces/KingNish/Instant-Image
        """

BETA = """ ### Voice Chat (BETA)"""

FAST = """## Fastest Model"""

Complex = """## Best in Complex Question"""

Detail = """## Best for Detailed Generation or Long Answers"""

base_loaded = "mistralai/Mixtral-8x7B-Instruct-v0.1"

client1 = InferenceClient(base_loaded)

system_instructions1 = "<s>[SYSTEM] Répondez comme le vrai Jarvis JARVIS, fabriqué par 'Tony Stark'. La conversation doit être très courte, claire, amicale et concise. Le texte fourni est une demande de réponse spécifique de votre part, en tant qu'assistant virtuel. Il vous demande de fournir des réponses amicales comme si vous étiez le personnage Jarvis, créé par « Tony Stark ». On s'attend à ce que j'évite les présentations et que je réponde directement à la requête, en développant tous les aspects de la demande. En tant qu'assistant doté d'une IA, votre tâche consiste à générer des réponses qui semblent avoir été créées par l'utilisateur et non par un assistant doté d'une IA.[USER]"

async def generate1(prompt):
    generate_kwargs = dict(
        temperature=0.7,
        max_new_tokens=512,
        top_p=0.95,
        repetition_penalty=1,
        do_sample=False,
    )
    formatted_prompt = system_instructions1 + prompt + "[JARVIS]"
    stream = client1.text_generation(
        formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=True)
    output = ""
    for response in stream:
        if not response.token.text == "</s>":
            output += response.token.text

    communicate = edge_tts.Communicate(output)
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp_file:
        tmp_path = tmp_file.name
        await communicate.save(tmp_path)
    yield tmp_path

with gr.Blocks(css="style.css") as demo:    
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        user_input = gr.Audio(label="Voice Chat (BETA)", type="filepath")
        output_audio = gr.Audio(label="JARVIS", type="filepath",
                        interactive=False,
                        autoplay=True,
                        elem_classes="audio")
    with gr.Row():
        translate_btn = gr.Button("Response")
        translate_btn.click(fn=respond, inputs=user_input,
                            outputs=output_audio, api_name=False)
    gr.Markdown(FAST)
    with gr.Row():
        user_input = gr.Textbox(label="Prompt", value="What is Wikipedia")
        input_text = gr.Textbox(label="Input Text", elem_id="important")
        output_audio = gr.Audio(label="JARVIS", type="filepath",
                        interactive=False,
                        autoplay=True,
                        elem_classes="audio")
    with gr.Row():
        translate_btn = gr.Button("Response")
        translate_btn.click(fn=generate1, inputs=user_input,
                            outputs=output_audio, api_name="translate")  

gr.Markdown(MORE)

if __name__ == "__main__":
    demo.queue(max_size=200).launch()