from huggingface_hub import InferenceClient
import gradio as gr

client = InferenceClient(
    "mistralai/Mixtral-8x7B-Instruct-v0.1"
)

# Função para formatar a prompt de entrada com o histórico de diálogo
def formatar_prompt(mensagem, historico):
    prompt = "<s>"
    for prompt_usuario, resposta_bot in historico:
        prompt += f"[INST] {prompt_usuario} [/INST]"
        prompt += f" {resposta_bot}</s> "
    prompt += f"[INST] {mensagem} [/INST]"
    return prompt

# Função para gerar resposta do modelo
def gerar(
    prompt, historico, prompt_sistema, temperatura=0.9, max_tokens_novos=256, top_p=0.95, penalidade_repeticao=1.0,
):
    temperatura = float(temperatura)
    if temperatura < 1e-2:
        temperatura = 1e-2
    top_p = float(top_p)

    kwargs_geracao = dict(
        temperature=temperatura,
        max_new_tokens=max_tokens_novos,
        top_p=top_p,
        repetition_penalty=penalidade_repeticao,
        do_sample=True,
        seed=42,
    )

    prompt_formatada = formatar_prompt(f"{prompt_sistema}, {prompt}", historico)
    stream = client.text_generation(prompt_formatada, **kwargs_geracao, stream=True, details=True, return_full_text=False)
    output = ""

    for resposta in stream:
        output += resposta.token.text
        yield output
    return output

# Inputs adicionais para o modelo
inputs_adicionais=[
    gr.Textbox(
        label="Prompt do Sistema",
        max_lines=1,
        interactive=True,
    ),
    gr.Slider(
        label="Temperatura",
        value=0.9,
        minimum=0.0,
        maximum=1.0,
        step=0.05,
        interactive=True,
        info="Valores mais altos produzem saídas mais diversas",
    ),
    gr.Slider(
        label="Máximo de Novos Tokens",
        value=256,
        minimum=0,
        maximum=1048,
        step=64,
        interactive=True,
        info="O número máximo de novos tokens",
    ),
    gr.Slider(
        label="Top-p (amostragem de núcleo)",
        value=0.90,
        minimum=0.0,
        maximum=1,
        step=0.05,
        interactive=True,
        info="Valores mais altos amostram mais tokens de baixa probabilidade",
    ),
    gr.Slider(
        label="Penalidade de Repetição",
        value=1.2,
        minimum=1.0,
        maximum=2.0,
        step=0.05,
        interactive=True,
        info="Penalize tokens repetidos",
    )
]

# Exemplos de prompts
exemplos=[["Estou planejando férias no Japão. Você pode sugerir um itinerário de uma semana incluindo lugares imperdíveis e culinárias locais para experimentar?", None, None, None, None, None, ],
          ["Você pode escrever uma história curta sobre um detetive viajante do tempo que resolve mistérios históricos?", None, None, None, None, None,],
          ["Estou tentando aprender francês. Você pode fornecer algumas frases comuns que seriam úteis para um iniciante, juntamente com suas pronúncias?", None, None, None, None, None,],
          ["Eu tenho frango, arroz e pimentões na minha cozinha. Você pode sugerir uma receita fácil que eu possa fazer com esses ingredientes?", None, None, None, None, None,],
          ["Você pode explicar como o algoritmo QuickSort funciona e fornecer uma implementação em Python?", None, None, None, None, None,],
          ["Quais são algumas características únicas do Rust que o destacam em comparação com outras linguagens de programação de sistemas como C ++?", None, None, None, None, None,],
         ]

# Interface do Chat
gr.ChatInterface(
    fn=gerar,
    chatbot=gr.Chatbot(show_label=False, show_share_button=False, show_copy_button=True, likeable=True, layout="panel"),
    additional_inputs=inputs_adicionais,
    title="Mixtral 46.7B",
    examples=exemplos,
    concurrency_limit=20,
).launch(show_api=False)