JARVIS/entrena/afina.py

#!/usr/bin/env python3
"""Fine-tuning de JARVIS con QLoRA sobre el dataset de COFRE.

    ./afina.py --revisa    comprueba que se puede entrenar aqui, sin entrenar
    ./afina.py             entrena (largo; mejor en el servidor con la 3060)

## Que hace y que NO hace el fine-tuning

Le enseña al modelo el ESTILO y las respuestas de tus apuntes con ejemplos. NO
sustituye al RAG: el RAG le da los datos frescos en el momento —y se actualiza
soltando un fichero—, mientras que esto los hornea en los pesos. Se complementan:
el fine-tuning hace que "suene" a JARVIS y conozca tus flujos sin buscar; el RAG
cubre lo que cambia y lo que es muy especifico.

Por eso el orden fue: primero RAG (hecho, funciona hoy), y esto despues.

## QLoRA en cristiano

Entrenar un 4B entero pide mas memoria de la que hay. QLoRA hace dos trucos:
  - carga el modelo base en 4 bits (cabe en poca VRAM, y se queda CONGELADO)
  - entrena solo unas matrices pequeñas nuevas (los "adaptadores LoRA"), que son
    una fraccion del tamaño

Asi se entrena un 4B en pocos GB. El resultado es un adaptador de ~50 MB que se
le pone encima al modelo base.

## Donde entrenar

  T1200 (4 GB, esta maquina)  al limite. Turing (capacidad 7.5): sin bf16 ni
                              flash-attention. Posible con lote 1 y secuencia
                              corta, lento y fragil. --revisa lo dice.
  3060 (12 GB, el servidor)   holgado. Ampere: bf16 y flash-attention. Es el
                              sitio natural, por eso el dataset se genera aparte
                              y viaja como fichero.

## Dependencias

    pip install unsloth        (trae peft, trl, bitsandbytes, accelerate)

unsloth porque es lo que mejor exprime una GPU pequeña: mismo resultado con la
mitad de VRAM y el doble de velocidad que peft+trl a pelo.
"""
import argparse
import json
import os
import sys

AQUI = os.path.dirname(os.path.abspath(__file__))
DATASET = os.path.join(AQUI, "dataset-cofre.jsonl")
SALIDA = os.path.join(AQUI, "jarvis-lora")

# El modelo base, en formato HuggingFace. ollama corre GGUF; para entrenar hace
# falta el modelo de verdad. Se prueban candidatos en orden hasta que uno
# cargue: la version pre-cuantizada de unsloth va primero porque baja menos y
# ocupa menos VRAM, que en 4 GB es lo que decide si entra.
#
# El de ollama es qwen3.5:4b (arquitectura qwen35, 4,7B, embedding 2560). Si
# ninguno carga, la fase 0 lo dice y no se pierden horas.
CANDIDATOS = [
    os.environ.get("JARVIS_BASE_HF", ""),      # lo que ponga el usuario, si pone
    "unsloth/Qwen3.5-4B-Instruct-bnb-4bit",
    "Qwen/Qwen3.5-4B-Instruct",
    "unsloth/Qwen3.5-4B-bnb-4bit",
    "Qwen/Qwen3.5-4B",
]

# Apretado para 4 GB (ver el plan). Con el 4,7B en 4 bits ocupando ~2,8 GB,
# cada megabyte de mas en secuencia o rango es el que provoca el OOM.
EPOCAS = 2
LOTE = 1
ACUMULA = 16         # lote efectivo 16 sin gastar VRAM
LR = 2e-4
MAX_SEQ = 512        # los pares Q&A son cortos; 512 no pierde casi nada
RANGO = 8            # adaptadores pequeños


def _base_que_carga():
    """El primer candidato que exista y cargue. None si ninguno."""
    from unsloth import FastLanguageModel
    for repo in CANDIDATOS:
        if not repo:
            continue
        try:
            print(f"  probando {repo} ...", flush=True)
            modelo, tok = FastLanguageModel.from_pretrained(
                model_name=repo, max_seq_length=MAX_SEQ, dtype=None,
                load_in_4bit=True)
            print(f"  ✓ carga: {repo}")
            return repo, modelo, tok
        except Exception as e:
            print(f"    no: {str(e).splitlines()[0][:80]}")
    return None, None, None


def revisa() -> int:
    """El ensayo de la fase 0: carga el 4,7B y hace UN paso de entrenamiento.

    Es el momento de la verdad. Si ese paso pasa sin quedarse sin memoria, el
    entrenamiento entero es cuestion de tiempo. Si peta, se sabe en diez minutos
    y no en la sexta hora.
    """
    print("\n  ── se puede entrenar aqui? ──\n")

    n = 0
    if os.path.exists(DATASET):
        with open(DATASET) as f:
            n = sum(1 for _ in f)
    print(f"  dataset: {n} ejemplos"
          + ("  (suficiente para el test; poco para entrenar)" if n < 500 else ""))

    try:
        import torch
    except ImportError:
        print("  torch: NO — el venv de entrenamiento no esta montado")
        return 1
    if not torch.cuda.is_available():
        print("  cuda: NO")
        return 1
    p = torch.cuda.get_device_properties(0)
    libre = torch.cuda.mem_get_info()[0] / 1e9
    print(f"  gpu: {p.name}  {p.total_memory/1e9:.1f} GB, {libre:.1f} libres, "
          f"capacidad {p.major}.{p.minor}"
          + ("  (Turing: fp16, sin flash-attn)" if p.major < 8 else ""))
    if libre < 3.3:
        print("  AVISO: hay menos de 3,3 GB libres. Cierra JARVIS y ollama antes.")

    for m in ("unsloth", "peft", "trl", "bitsandbytes"):
        try:
            __import__(m)
        except ImportError:
            print(f"  {m}: NO  ->  el venv de entrenamiento no acabo de instalar")
            return 1

    # El momento de la verdad: cargar + un paso real.
    print("\n  ── cargando el modelo y haciendo un paso de prueba ──")
    from unsloth import FastLanguageModel
    repo, modelo, tok = _base_que_carga()
    if repo is None:
        print("\n  NINGUN candidato cargo. O no existe el repo HF del qwen3.5:4b,")
        print("  o esta arquitectura (qwen35) no la soporta este transformers.")
        print("  Pon el repo correcto en JARVIS_BASE_HF y reintenta.")
        return 1

    tras_modelo = torch.cuda.memory_allocated() / 1e9
    print(f"  VRAM tras cargar el modelo: {tras_modelo:.2f} GB")

    modelo = FastLanguageModel.get_peft_model(
        modelo, r=RANGO, lora_alpha=RANGO, lora_dropout=0,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                        "gate_proj", "up_proj", "down_proj"],
        use_gradient_checkpointing="unsloth", random_state=42)

    # un lote de una frase, forward + backward, que es donde revienta si revienta
    try:
        import torch as T
        texto = tok.apply_chat_template(
            [{"role": "user", "content": "que es secretsdump"},
             {"role": "assistant", "content": "Saca los hashes NTLM, señor."}],
            tokenize=False)
        ids = tok(texto, return_tensors="pt", max_length=MAX_SEQ,
                  truncation=True).input_ids.to("cuda")
        salida = modelo(ids, labels=ids)
        salida.loss.backward()
        pico = T.cuda.max_memory_allocated() / 1e9
        print(f"  ✓ un paso de entrenamiento paso. Pico de VRAM: {pico:.2f} GB")
        print(f"\n  CABE. base a usar: {repo}")
        print("  Pon ese repo en JARVIS_BASE_HF si no es el primero, y ./afina.py\n")
        return 0
    except T.cuda.OutOfMemoryError:
        print("\n  OOM: el 4,7B no entra en esta grafica ni apretando.")
        print("  Plan B: bajar MAX_SEQ a 256, o el 3B, o el servidor (3060).")
        return 2


def entrena() -> int:
    if not os.path.exists(DATASET):
        print("  no hay dataset: ./dataset.py primero")
        return 1

    import torch
    from unsloth import FastLanguageModel
    from datasets import load_dataset
    from trl import SFTTrainer, SFTConfig

    ampere = torch.cuda.get_device_properties(0).major >= 8

    repo, modelo, tokenizer = _base_que_carga()
    if repo is None:
        print("  ningun base cargo; corre --revisa para el diagnostico")
        return 1

    modelo = FastLanguageModel.get_peft_model(
        modelo, r=RANGO, lora_alpha=RANGO, lora_dropout=0,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                        "gate_proj", "up_proj", "down_proj"],
        use_gradient_checkpointing="unsloth", random_state=42)

    datos = load_dataset("json", data_files=DATASET, split="train")

    def formatea(ej):
        return {"text": tokenizer.apply_chat_template(
            ej["messages"], tokenize=False, add_generation_prompt=False)}
    datos = datos.map(formatea)

    entrenador = SFTTrainer(
        model=modelo, tokenizer=tokenizer, train_dataset=datos,
        dataset_text_field="text", max_seq_length=MAX_SEQ,
        args=SFTConfig(
            per_device_train_batch_size=LOTE,
            gradient_accumulation_steps=ACUMULA,
            num_train_epochs=EPOCAS, learning_rate=LR,
            fp16=not ampere, bf16=ampere,
            logging_steps=10, optim="adamw_8bit",
            warmup_ratio=0.05, lr_scheduler_type="cosine",
            output_dir=os.path.join(AQUI, "checkpoints"), seed=42))

    print("  entrenando...")
    entrenador.train()

    modelo.save_pretrained(SALIDA)
    tokenizer.save_pretrained(SALIDA)
    print(f"\n  adaptador guardado: {SALIDA}")
    print("  para usarlo con ollama hay que exportarlo a GGUF y crear un")
    print("  Modelfile con ADAPTER. Ver el README de entrena/.")
    return 0


def main() -> int:
    p = argparse.ArgumentParser()
    p.add_argument("--revisa", action="store_true")
    a = p.parse_args()
    return revisa() if a.revisa else entrena()


if __name__ == "__main__":
    sys.exit(main())