app.py

import gradio as gr
import json


import gradio as gr

# !python -c "import torch; assert torch.cuda.get_device_capability()[0] >= 8, 'Hardware not supported for Flash Attention'"
import json
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, GemmaTokenizer,  StoppingCriteria, StoppingCriteriaList, GenerationConfig
# from google.colab import userdata
import os

model_id = "somosnlp/Sam_Diagnostic"
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)
max_seq_length=2048

# if torch.cuda.get_device_capability()[0] >= 8:
#     # print("Flash Attention")
#     attn_implementation="flash_attention_2"
# else:
#     attn_implementation=None
attn_implementation=None

tokenizer = AutoTokenizer.from_pretrained(model_id,
                                          max_length = max_seq_length)
model = AutoModelForCausalLM.from_pretrained(model_id,
                                            #  quantization_config=bnb_config,
                                             device_map = {"":0},
                                              attn_implementation = attn_implementation, # A100 o H100
                                             ).eval()



class ListOfTokensStoppingCriteria(StoppingCriteria):
    """
    Clase para definir un criterio de parada basado en una lista de tokens específicos.
    """
    def __init__(self, tokenizer, stop_tokens):
        self.tokenizer = tokenizer
        # Codifica cada token de parada y guarda sus IDs en una lista
        self.stop_token_ids_list = [tokenizer.encode(stop_token, add_special_tokens=False) for stop_token in stop_tokens]

    def __call__(self, input_ids, scores, **kwargs):
        # Verifica si los últimos tokens generados coinciden con alguno de los conjuntos de tokens de parada
        for stop_token_ids in self.stop_token_ids_list:
            len_stop_tokens = len(stop_token_ids)
            if len(input_ids[0]) >= len_stop_tokens:
                if input_ids[0, -len_stop_tokens:].tolist() == stop_token_ids:
                    return True
        return False

# Uso del criterio de parada personalizado
stop_tokens = ["<end_of_turn>"]  # Lista de tokens de parada

# Inicializa tu criterio de parada con el tokenizer y la lista de tokens de parada
stopping_criteria = ListOfTokensStoppingCriteria(tokenizer, stop_tokens)

# Añade tu criterio de parada a una StoppingCriteriaList
stopping_criteria_list = StoppingCriteriaList([stopping_criteria])

def generate_text(prompt, idioma_entrada, idioma_salida, max_length=2100):
  prompt=prompt.replace(". ", ".\n").strip()

  input_text = f'''<bos><start_of_turn>system
You are a helpful AI assistant.
Responde en formato json.
Eres un agente experto en medicina.
Lista de codigos linguisticos disponibles: ["{idioma_entrada}", "{idioma_salida}"]<end_of_turn>
<start_of_turn>user
{prompt}<end_of_turn>
<start_of_turn>model
'''

  inputs = tokenizer.encode(input_text,
                            return_tensors="pt",
                            add_special_tokens=False).to("cuda:0")
  max_new_tokens=max_length
  generation_config = GenerationConfig(
                max_new_tokens=max_new_tokens,
                temperature=0.35, #55
                #top_p=0.9,
                top_k=50, # 45
                repetition_penalty=1.,  #1.1
                do_sample=True,
            )
  outputs = model.generate(generation_config=generation_config,
                          input_ids=inputs,
                          stopping_criteria=stopping_criteria_list,)
  return tokenizer.decode(outputs[0], skip_special_tokens=False) #True

def mostrar_respuesta(pregunta, idioma_entrada, idioma_salida):
    try:
      lista_codigo_lin = {
            "español": "es",
            "ingles": "en",
        }
      # Utiliza los parámetros de idioma para obtener los códigos de idioma correspondientes.
      codigo_lin_entrada = lista_codigo_lin[idioma_entrada.lower()]
      codigo_lin_salida = lista_codigo_lin[idioma_salida.lower()]
      
      res= generate_text(pregunta, codigo_lin_entrada, codigo_lin_salida, max_length=1500)
      inicio_json = res.find('{')
      fin_json = res.rfind('}') + 1
      json_str = res[inicio_json:fin_json]
      json_obj = json.loads(json_str)
      return json_obj["description"], json_obj["medical_specialty"], json_obj["principal_diagnostic"]
    except:
      json_obj={}
      json_obj['description']='Esperando al usuario.'
      json_obj['medical_specialty']='Esperando al usuario.'
      json_obj['principal_diagnostic']='Esperando al usuario.'
      return json_obj["description"], json_obj["medical_specialty"], json_obj["principal_diagnostic"]

# Ejemplos de preguntas
ejemplos = [
    ["CHIEF COMPLAINT:,  Left wrist pain.,HISTORY OF PRESENT PROBLEM"],
    ["INDICATIONS:  ,Chest pain.,STRESS TECHNIQUE:,"],
    ["MOTIVO DE CONSULTA: Una niña de 2 meses con una historia de congestión y fiebre de 1 semana, con fiebre durante 2 días. HISTORIA DE LA ENFERMEDAD ACTUAL: La paciente es una niña de 2 meses previamente sana, que ha tenido tos y congestión durante la última semana. La madre también ha reportado dificultad para respirar, que describe como respiración rápida asociada con retracciones. La madre indica que la tos a veces es paroxística y se acompaña de vómitos post-tusivos. La paciente ha tenido pausas respiratorias cortas después de los episodios de tos. La temperatura de la paciente ha oscilado entre 102 y 104 grados. Ha tenido una disminución en la ingesta oral y menos pañales mojados. El hermano también está enfermo con síntomas de infección respiratoria, y la paciente no ha tenido diarrea. La madre informa que la niña ha comenzado a regurgitar después de las comidas. No hacía esto anteriormente. MEDICAMENTOS: Ninguno. EXPOSICIÓN AL HUMO: Ninguna. INMUNIZACIONES: Ninguna. DIETA: Similac 4 onzas cada 2 a 3 horas. ALERGIAS: No se conocen alergias a medicamentos. ANTECEDENTES MÉDICOS: La paciente nació a término. Peso al nacer fue de 6 libras 1 onza. Complicaciones postnatales: Ictericia neonatal. La paciente permaneció en el hospital durante 3 días. Las ecografías prenatales fueron reportadas como normales. HOSPITALIZACIONES ANTERIORES: Ninguna. ANTECEDENTES FAMILIARES/SOCIALES: Antecedentes familiares positivos de asma y diabetes. También hay antecedentes familiares positivos de enfermedad renal en el lado paterno de la familia. DESARROLLO: Normal. La paciente pasa la prueba de audición neonatal. REVISIÓN DE SISTEMAS: GENERALES: La paciente ha tenido fiebre, no ha tenido escalofríos. PIEL: Sin erupciones cutáneas. CABEZA/CUELLO: Congestión leve durante 1 semana. Tos, a veces paroxística, sin cianosis. La paciente se pone roja en la cara durante los episodios de tos, vómitos post-tusivos. CARDIOVASCULAR: Sin cianosis. GI: Vómitos post-tusivos, disminución de la ingesta oral. GU: Disminución de la producción de orina. ORTOPÉDICO: Sin problemas actuales. NEUROLÓGICO: Sin cambios en el estado mental. ENDOCRINO: No hay antecedentes de pérdida de peso. DESARROLLO: No hay pérdida de hitos del desarrollo. EXAMEN FÍSICO: SIGNOS VITALES: Peso de 4.8 kg, temperatura de 100.4, frecuencia cardíaca de 140, frecuencia respiratoria de 30 y saturación del 100%. GENERALES: Bebé bien aparente sin signos de malestar agudo. CABEZA/CUELLO: Fontanela anterior abierta y plana. Pupilas iguales y reactivas a la luz con reflejo rojo. Fosas nasales permeables. Mucosa oral húmeda. Faringe posterior despejada. Paladar duro intacto. Encías normales. CARDIACO: Ritmo cardíaco regular sin soplos. PULMONES: Algunos estertores leves. Sin retracciones. Sin estridor. Sin sibilancias en el examen. Taquipnea leve. EXTREMIDADES: Calientes, buena perfusión. Sin clics de cadera. NEUROLÓGICO: La paciente está alerta. Tono normal en todo el cuerpo. Reflejos tendinosos profundos 2+/4. Sin clonus. PIEL: Normal. DATOS DE LABORATORIO: El hemograma muestra un recuento de glóbulos blancos de 12.4, hemoglobina 10.1, recuento de plaquetas 611,000; 38 segmentados, 3 bastones, 42 linfocitos y 10 monocitos. Electrolitos dentro de límites normales. Proteína C reactiva 0.3. La radiografía de tórax no muestra enfermedad aguda, excepto una pequeña densidad en el área retrocardíaca en la vista posterior. El análisis de orina muestra 10 a 25 bacterias. EVALUACIÓN/PLAN: Se trata de una niña de 2 meses que presenta fiebre, tos paroxística y episodios de dificultad respiratoria. La paciente está estable en la sala de emergencias. La admitiremos en la planta de pediatría. Enviaremos un PCR de tos ferina. También seguiremos los resultados del cultivo de orina y que el análisis de orina muestra 10 a 25 bacterias. La paciente será monitoreada en busca de signos de sepsis, apnea, infección del tracto urinario y neumonía. Esperaremos un informe radiológico de la radiografía de tórax para determinar si la densidad vista en la película lateral es una variante normal o representa patología."],
]

idiomas = ["español", "ingles"]

description= """En el ámbito de la medicina, la capacidad de sintetizar información compleja y hacer derivaciones especializadas rápidas y "precisas" es crucial para el tratamiento eficaz de los pacientes. 
Aquí es donde entra en juego la tecnología avanzada, en particular los modelos de inteligencia artificial como Gemma-2B-IT. Este sistema está diseñado para transformar la manera en que los médicos acceden y utilizan la información clínica, 
mejorando significativamente la eficiencia y la "precisión" en el cuidado del paciente.

## Modelo: 
https://huggingface.co/somosnlp/Sam_Diagnostic

## Dataset formato chatml: 
https://huggingface.co/datasets/somosnlp/medical_en_es_formato_chatML_Gemma

## Datos curados: 
https://huggingface.co/datasets/somosnlp/medical_bilingual_en_es

## Implementacion rapida:
https://colab.research.google.com/drive/1gzTk-mdKefjZwzeiaHfkQ5Qnq3G10WCA?usp=sharing"""

iface = gr.Interface(
    fn=mostrar_respuesta,
    inputs=[
        gr.Textbox(label="Escribe tu transcripcion medica.", placeholder="Introduce tu consulta médica aquí..."),
        gr.Dropdown(label="Idioma de Entrada",choices=idiomas, value='español'),
        gr.Dropdown(label="Idioma de Salida", choices=idiomas, value='español'),
    ],
    outputs=[
        gr.Textbox(label="Description", lines=2),
        gr.Textbox(label="Medical specialty", lines=1),
        gr.Textbox(label="Principal diagnostic", lines=1)
    ],
    title="Consultas medicas",
    #description="Introduce tu diagnostico.",
    examples=ejemplos,
    concurrency_limit=20,
    description=description,
)

iface.queue(max_size=14).launch() # share=True,debug=True