Nota: es un causal LM entrenado para responder un solo dígito (0/1). Usa el chat template de Gemma y (opcionalmente) restringe la generación a {0,1}.

Datos

Dataset original del proyecto (train/val/test).
Muestra externa cochrane_sample (90% a train, 10% a val).

Métricas

split	accuracy	precision	recall	f1	roc_auc
val	1.0000	1.0000	1.0000	1.0000	1.0000

split	accuracy	precision	recall	f1	roc_auc
test	0.9767	0.9643	0.9818	0.9730	0.9983

Uso

python

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers.generation.logits_process import LogitsProcessor, LogitsProcessorList

repo_id = "Cristian11212/gemma3-4b-plaintech-original-20250921-025225"
tokenizer = AutoTokenizer.from_pretrained(repo_id, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(repo_id, torch_dtype=torch.float16 if torch.cuda.is_available() else None)
model.eval().to("cuda" if torch.cuda.is_available() else "cpu")

SYS = "You are a binary classifier. Reply with a single digit: 0 (plain) or 1 (technical). No extra text."
def build_prompt(text):
    msgs = [{"role":"system","content":SYS}, {"role":"user","content":text}]
    return tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)

zero_id = tokenizer.encode("0", add_special_tokens=False)[0]
one_id  = tokenizer.encode("1", add_special_tokens=False)[0]
class OnlyAllowTokens(LogitsProcessor):
    def __init__(self, allowed, device): self.allowed = torch.tensor(allowed, device=device)
    def __call__(self, input_ids, scores):
        mask = torch.full_like(scores, float("-inf")); mask[:, self.allowed] = 0.0; return scores + mask
lp = LogitsProcessorList([OnlyAllowTokens([zero_id, one_id], model.device)])

text = "urinary incontinence is the inability to willingly control bladder voiding..."
enc = tokenizer([build_prompt(text)], return_tensors="pt").to(model.device)
out = model.generate(**enc, max_new_tokens=1, do_sample=False, output_scores=True, return_dict_in_generate=True, logits_processor=lp)
logits = out.scores[0].float(); probs = torch.softmax(logits, dim=-1)
p1 = probs[0, one_id].item(); pred = 1 if p1 >= 0.5 else 0
print("pred:", pred, "prob_technical:", p1)