from huggingface_hub import InferenceClient
import gradio as gr

# Inicialização do cliente de inferência
client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")

# Função para formatar o prompt de entrada com o histórico de diálogo
def formatar_prompt(mensagem, historico):
    prompt = "<s>" + "".join(f"[INST] {prompt_usuario} [/INST] {resposta_bot}" for prompt_usuario, resposta_bot in historico) + 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):
    kwargs_geracao = dict(
        temperature=float(temperatura),
        max_new_tokens=int(max_tokens_novos),
        top_p=float(top_p),
        repetition_penalty=float(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)
    for resposta in stream:
        yield resposta.token.text

# 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=[
    ["Paciente de 55 anos, história familiar de doença cardíaca. Colesterol total: 260 mg/dL, HDL: 40 mg/dL, LDL: 180 mg/dL, triglicerídeos: 200 mg/dL. Interpretar resultados e sugerir intervenções.", None, None, None, None, None],
    ["Ecocardiograma em paciente de 70 anos, dispneia aos esforços. Fração de ejeção: 40%, espessamento do septo interventricular, dilatação atrial esquerda. Propor tratamento para insuficiência cardíaca.", None, None, None, None, None],
    ["Teste de esforço em esteira, paciente de 45 anos, dor torácica atípica. 10 METs atingidos, depressão do segmento ST de 2 mm no pós-exercício. Recomendar próximos passos diagnósticos e terapêuticos.", None, None, None, None, None],
    ["Holter de 24 horas, paciente de 60 anos, síncope recorrente. Episódios de pausas sinusais de até 5 segundos, arritmia sinusal respiratória. Formular plano de monitoramento e tratamento para bradicardia sintomática.", None, None, None, None, None],
    ["Cateterismo cardíaco em paciente de 65 anos, angina instável. Angiografia revela lesões significativas em duas artérias coronárias principais, FFR: 0,68. Recomendar opções de revascularização e manejo coronariano.", 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="Medchat | Automação Inteligente de Saúde",
    examples=exemplos,
    concurrency_limit=20,
).launch(show_api=False)