Pipeline NLP Clínico con HuggingFace Transformers

# novamed/inferencia.py — Pipeline NLP clínico con Transformers v5 # Instalación: uv pip install "transformers[torch]==5.12.0" accelerate==1.14.0 from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification from transformers import AutoModelForTokenClassification, Pipeline from huggingface_hub import login import os import torch # ✓ Seguridad: token desde variable de entorno, nunca hardcodeado login(token=os.environ["HF_TOKEN"]) DEVICE = 0 if torch.cuda.is_available() else -1 DTYPE = torch.bfloat16 if DEVICE == 0 else torch.float32 # ── 1. Clasificador de urgencia (BERT fine-tuneado) ────────────────────────── triage_pipe = pipeline( "text-classification", model="novamed/bert-triage-v2", tokenizer="novamed/bert-triage-v2", device=DEVICE, torch_dtype=DTYPE, top_k=None, # devuelve scores de todas las clases ) # ── 2. NER clínico (entidades biomédicas en español) ───────────────────────── ner_pipe = pipeline( "token-classification", model="PlanTL-GOB-ES/bsc-bio-ehr-es", aggregation_strategy="simple", device=DEVICE, ) # ── 3. Resumen para enfermería (mT5 fine-tuneado en español) ───────────────── summ_pipe = pipeline( "summarization", model="IIC/mt5-spanish-mlsum", device=DEVICE, torch_dtype=DTYPE, ) def analizar_nota(texto: str) -> dict: """Analiza una nota médica devolviendo triaje, entidades y resumen.""" return { "urgencia" : triage_pipe(texto, truncation=True, max_length=512), "entidades" : ner_pipe(texto), "resumen" : summ_pipe(texto, max_length=80, min_length=20)[0]["summary_text"], }

from datasets import load_dataset, ClassLabel from transformers import ( AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer, DataCollatorWithPadding, ) import evaluate import numpy as np LABELS = ["baja", "media", "alta", "critica"] MODEL = "bert-base-multilingual-cased" # base para fine-tuning HUB_ID = "novamed/bert-triage-v2" # ── Dataset ────────────────────────────────────────────────────────────────── ds = load_dataset("parquet", data_files="data/notas_etiquetadas.parquet") ds = ds["train"].train_test_split(test_size=0.2, seed=42) tokenizer = AutoTokenizer.from_pretrained(MODEL) def tokenize(batch): return tokenizer(batch["texto"], truncation=True, max_length=512) ds = ds.map(tokenize, batched=True) # ── Modelo ─────────────────────────────────────────────────────────────────── model = AutoModelForSequenceClassification.from_pretrained( MODEL, num_labels=len(LABELS), id2label={i: l for i, l in enumerate(LABELS)}, label2id={l: i for i, l in enumerate(LABELS)}, ) # ── Training Args ───────────────────────────────────────────────────────────── args = TrainingArguments( output_dir="./checkpoints/triage", num_train_epochs=5, per_device_train_batch_size=32, per_device_eval_batch_size=64, learning_rate=2e-5, bf16=True, # bfloat16 en A10G eval_strategy="epoch", save_strategy="best", load_best_model_at_end=True, metric_for_best_model="f1", push_to_hub=True, hub_model_id=HUB_ID, report_to="none", # sin W&B en primera pasada ) metric = evaluate.load("f1") def compute_metrics(eval_pred): logits, labels = eval_pred preds = np.argmax(logits, axis=-1) return metric.compute(predictions=preds, references=labels, average="macro") trainer = Trainer( model=model, args=args, train_dataset=ds["train"], eval_dataset=ds["test"], tokenizer=tokenizer, data_collator=DataCollatorWithPadding(tokenizer), compute_metrics=compute_metrics, ) trainer.train() trainer.push_to_hub() # sube el mejor checkpoint al Hub privado