scripts/promote_merged_model.py

#!/usr/bin/env python3
"""
Faz merge da base Qwen2.5 + LoRA e envia o modelo full para o Hub (ex.: eda-llm-qwen2.5-merged).

Requer HF_TOKEN com permissao de escrita. GPU recomendada para 1.5B em fp16.
"""

from __future__ import annotations

import argparse
import os
import sys
import tempfile
from pathlib import Path

try:
    import torch
    from huggingface_hub import HfApi, create_repo, login
    from peft import PeftModel
    from transformers import AutoModelForCausalLM, AutoTokenizer
except Exception as e:
    err = str(e)
    if "numpy.dtype" in err or "binary incompatibility" in err:
        print(
            "Incompatibilidade entre numpy e pandas/sklearn (comum apos upgrades).\n"
            "Corrija e rode de novo:\n"
            "  pip install --upgrade --force-reinstall numpy pandas scikit-learn\n"
            "Ou use um venv limpo com: pip install -r requirements.txt\n",
            file=sys.stderr,
        )
    raise


def _require_token() -> str:
    token = os.environ.get("HF_TOKEN")
    if not token:
        print("Defina HF_TOKEN ou use huggingface-cli login.", file=sys.stderr)
        sys.exit(1)
    return token


def merge_and_upload(
    base_model_id: str,
    lora_id: str,
    merged_repo_id: str,
    *,
    dtype: str,
    token: str,
) -> None:
    torch_dtype = torch.float16 if dtype == "float16" else torch.bfloat16

    print(f"Carregando base: {base_model_id}")
    base = AutoModelForCausalLM.from_pretrained(
        base_model_id,
        torch_dtype=torch_dtype,
        device_map="auto",
        trust_remote_code=True,
    )
    print(f"Carregando LoRA: {lora_id}")
    model = PeftModel.from_pretrained(base, lora_id)
    print("Merge LoRA na base...")
    merged = model.merge_and_unload()

    tokenizer = AutoTokenizer.from_pretrained(base_model_id, trust_remote_code=True)

    with tempfile.TemporaryDirectory(prefix="eda_merged_") as tmp:
        out = Path(tmp)
        print(f"Salvando merge em {out}...")
        merged.save_pretrained(out, safe_serialization=True)
        tokenizer.save_pretrained(out)

        print(f"Upload para {merged_repo_id}...")
        api = HfApi(token=token)
        create_repo(
            repo_id=merged_repo_id,
            repo_type="model",
            token=token,
            exist_ok=True,
        )
        api.upload_folder(
            folder_path=str(out),
            repo_id=merged_repo_id,
            repo_type="model",
        )

    print("Concluido.")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Merge Qwen base + LoRA e push para o repositorio merged no Hub.",
    )
    parser.add_argument(
        "--base-model",
        default=os.environ.get("MODEL_NAME", "Qwen/Qwen2.5-1.5B-Instruct"),
        help="Modelo base no Hub (default: Qwen/Qwen2.5-1.5B-Instruct ou MODEL_NAME).",
    )
    parser.add_argument(
        "--lora-repo",
        default=os.environ.get("OUTPUT_REPO", "beAnalytic/eda-llm-qwen2.5-lora"),
        help="Repo ID do LoRA treinado (default: OUTPUT_REPO ou beAnalytic/eda-llm-qwen2.5-lora).",
    )
    parser.add_argument(
        "--merged-repo",
        default="beAnalytic/eda-llm-qwen2.5-merged",
        help="Repo ID do modelo mergeado (full weights).",
    )
    parser.add_argument(
        "--dtype",
        choices=("float16", "bfloat16"),
        default="float16",
        help="Dtype do merge (default float16).",
    )
    args = parser.parse_args()

    token = _require_token()
    login(token=token, add_to_git_credential=False)

    try:
        merge_and_upload(
            args.base_model,
            args.lora_repo,
            args.merged_repo,
            dtype=args.dtype,
            token=token,
        )
    except Exception as e:
        print(f"Erro: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()