import gradio as gr
from huggingface_hub import InferenceClient

# Função para gerar resposta do modelo
def gerar(prompt, historico, prompt_sistema, modelo, temperatura=0.9, max_tokens_novos=256, top_p=0.95, penalidade_repeticao=1.0):
    # Inicialização do cliente de inferência
    if modelo == "Modelo 1":
        client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.2")
    elif modelo == "Modelo 2":
        client = InferenceClient("microsoft/Phi-3-mini-4k-instruct")
    elif modelo == "Modelo 3":
        client = InferenceClient("google/gemma-1.1-7b-it")
    else:
        return "Modelo não reconhecido."

        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,
    )
    # Função para formatar o prompt de entrada com o histórico de diálogo
    prompt_formatada = "<s>" + "".join(f"[INST] {prompt_usuario} [/INST] {resposta_bot}" for prompt_usuario, resposta_bot in historico) + f"[INST] {prompt_sistema} [/INST]"
    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"),
    gr.Dropdown(label="Modelo", choices=["Modelo 1", "Modelo 2", "Modelo 3"], default="Modelo 1", interactive=True)
]

# 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,
    inputs=inputs_adicionais,
    outputs=gr.Textbox(label="Resposta do Modelo"),
    title="Medchat | Automação Inteligente de Saúde",
    examples=exemplos,
    layout="vertical",
    enable_queue=True,
    allow_screenshot=True,
    allow_flagging=True,
    flagging_options=["Inapropriado", "Incorreto", "Ofensivo"],
    server_name="0.0.0.0",
    server_port=8080,
    server_name_description="Endereço IP",
    server_port_description="Porta",
    interpretation="markdown",
    keep_history=True,
).launch()