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
import gradio as gr

# Função para analisar os casos clínicos
def analisar_caso(idade, sintomas, resultados_exames):
    # Aqui você pode adicionar a lógica para analisar os casos clínicos com base nas informações fornecidas
    return "Em breve, entraremos em contato com a análise do caso."

# Exemplos de prompts com títulos descritivos
exemplos=[
    ["Paciente com histórico familiar de doença cardíaca", "Análise do perfil lipídico e sugestões de intervenções.", None, None, None, None],
    ["Paciente idoso com dispneia aos esforços", "Exame ecocardiográfico e plano de tratamento para insuficiência cardíaca.", None, None, None, None],
    ["Paciente com dor torácica atípica durante teste de esforço", "Interpretação dos resultados e próximos passos diagnósticos.", None, None, None, None],
    ["Paciente com síncope recorrente", "Análise do Holter de 24 horas e plano de tratamento.", None, None, None, None],
    ["Paciente com angina instável submetido a cateterismo cardíaco", "Interpretação dos achados e opções de manejo.", None, None, None, None]
]

# Interface do Chat
gr.Interface(
    analisar_caso,
    ["text", "text", "text"],  # Tipos de entrada: texto para idade, sintomas e resultados dos exames
    "text",  # Tipo de saída: texto para a análise do caso
    title="Medchat | Automação Inteligente de Saúde",
    examples=exemplos
).launch()