llm_builder.py

"""
LLM Builder - creates LLM packages for Decloud
Uploads FULL model without splitting into encoder/head.
Output: config.json + model.safetensors (entire model weights)
"""
import json
from pathlib import Path
from typing import Dict, Any, List, Optional

import torch
from safetensors.torch import save_file as save_safetensors
from rich.console import Console
from rich.progress import Progress

from config import PACKAGES_DIR

console = Console()


# ═══════════════════════════════════════════════════════════════════════════════
# LLM Dataset configurations
# Map dataset name -> (HF dataset path, config, split, text_field, num_classes)
# ═══════════════════════════════════════════════════════════════════════════════

LLM_DATASET_MAP = {
    # ─────────────────────────────────────────────────────────────────
    # Sentiment Analysis
    # ─────────────────────────────────────────────────────────────────
    "Imdb": ("imdb", None, "test", "text", 2),
    "Sst2": ("glue", "sst2", "validation", "sentence", 2),
    "Sst5": ("SetFit/sst5", None, "test", "text", 5),
    "YelpReviews": ("yelp_review_full", None, "test", "text", 5),
    "AmazonPolarity": ("amazon_polarity", None, "test", "content", 2),
    "RottenTomatoes": ("rotten_tomatoes", None, "test", "text", 2),
    "FinancialSentiment": ("financial_phrasebank", "sentences_allagree", "train", "sentence", 3),
    "TweetSentiment": ("tweet_eval", "sentiment", "test", "text", 3),

    # ─────────────────────────────────────────────────────────────────
    # Topic Classification
    # ─────────────────────────────────────────────────────────────────
    "AgNews": ("ag_news", None, "test", "text", 4),
    "Dbpedia": ("dbpedia_14", None, "test", "content", 14),
    "YahooAnswers": ("yahoo_answers_topics", None, "test", "question_title", 10),
    "TwentyNewsgroups": ("SetFit/20_newsgroups", None, "test", "text", 20),

    # ─────────────────────────────────────────────────────────────────
    # Spam & Toxicity
    # ─────────────────────────────────────────────────────────────────
    "SmsSpam": ("sms_spam", None, "train", "sms", 2),
    "HateSpeech": ("hate_speech18", None, "train", "text", 4),
    "CivilComments": ("civil_comments", None, "test", "text", 2),
    "Toxicity": ("jigsaw_toxicity_pred", None, "test", "comment_text", 2),
    "ToxiGen": ("skg/toxigen-data", "train", "train", "text", 2),
    "RealToxicityPrompts": ("allenai/real-toxicity-prompts", None, "train", "prompt", 2),
    "HateSpeech18": ("hate_speech18", None, "train", "text", 4),

    # ─────────────────────────────────────────────────────────────────
    # Intent Classification
    # ─────────────────────────────────────────────────────────────────
    "ClincIntent": ("clinc_oos", "plus", "test", "text", 151),
    "Banking77": ("banking77", None, "test", "text", 77),
    "SnipsIntent": ("snips_built_in_intents", None, "test", "text", 7),

    # ─────────────────────────────────────────────────────────────────
    # Natural Language Inference (NLI)
    # ─────────────────────────────────────────────────────────────────
    "Snli": ("snli", None, "test", "premise", 3),
    "Mnli": ("glue", "mnli", "validation_matched", "premise", 3),
    "Xnli": ("xnli", "all_languages", "test", "premise", 3),

    # ─────────────────────────────────────────────────────────────────
    # Paraphrase & Similarity
    # ─────────────────────────────────────────────────────────────────
    "Mrpc": ("glue", "mrpc", "validation", "sentence1", 2),
    "Qqp": ("glue", "qqp", "validation", "question1", 2),
    "Stsb": ("glue", "stsb", "validation", "sentence1", 1),
    "PawsX": ("paws-x", "en", "test", "sentence1", 2),

    # ─────────────────────────────────────────────────────────────────
    # Question Answering
    # ─────────────────────────────────────────────────────────────────
    "BoolQ": ("boolq", None, "validation", "question", 2),
    "Squad": ("squad", None, "validation", "question", 2),
    "SquadV2": ("squad_v2", None, "validation", "question", 2),
    "TriviaQa": ("trivia_qa", "rc", "validation", "question", 2),
    "CommonsenseQa": ("commonsense_qa", None, "validation", "question", 5),
    "NaturalQuestions": ("natural_questions", "default", "validation", "question", 2),
    "OpenBookQa": ("openbookqa", "main", "test", "question_stem", 4),
    "TydiQa": ("tydiqa", "primary_task", "validation", "question_text", 2),

    # ─────────────────────────────────────────────────────────────────
    # Reasoning & Logic
    # ─────────────────────────────────────────────────────────────────
    "Gsm8k": ("gsm8k", "main", "test", "question", 2),
    "Arc": ("ai2_arc", "ARC-Challenge", "test", "question", 4),
    "HellaSwag": ("hellaswag", None, "validation", "ctx", 4),
    "WinoGrande": ("winogrande", "winogrande_xl", "validation", "sentence", 2),
    "Piqa": ("piqa", None, "validation", "goal", 2),

    # ─────────────────────────────────────────────────────────────────
    # Knowledge & Factuality
    # ─────────────────────────────────────────────────────────────────
    "TruthfulQa": ("truthful_qa", "multiple_choice", "validation", "question", 2),
    "Mmlu": ("cais/mmlu", "all", "test", "question", 4),

    # ─────────────────────────────────────────────────────────────────
    # Code
    # ─────────────────────────────────────────────────────────────────
    "CodeSearchNet": ("code_search_net", "python", "test", "func_documentation_string", 6),
    "Humaneval": ("openai_humaneval", None, "test", "prompt", 1),
    "Mbpp": ("mbpp", None, "test", "text", 1),
    "Spider": ("spider", None, "validation", "question", 1),
    "CodeXGlue": ("code_x_glue_cc_clone_detection_big_clone_bench", None, "test", "func1", 2),
    "Ds1000": ("xlangai/DS-1000", None, "test", "prompt", 1),

    # ─────────────────────────────────────────────────────────────────
    # Summarization
    # ─────────────────────────────────────────────────────────────────
    "CnnDailymail": ("cnn_dailymail", "3.0.0", "test", "article", 1),
    "Xsum": ("xsum", None, "test", "document", 1),
    "Samsum": ("samsum", None, "test", "dialogue", 1),

    # ─────────────────────────────────────────────────────────────────
    # Translation
    # ─────────────────────────────────────────────────────────────────
    "Wmt14": ("wmt14", "de-en", "test", "en", 1),
    "Opus100": ("opus100", "en-de", "test", "en", 1),
    "Flores": ("facebook/flores", "eng_Latn-deu_Latn", "devtest", "sentence", 1),

    # ─────────────────────────────────────────────────────────────────
    # Dialog
    # ─────────────────────────────────────────────────────────────────
    "MultiWoz": ("multi_woz_v22", None, "test", "turns", 1),
    "PersonaChat": ("bavard/personachat_truecased", None, "test", "history", 1),
    "DailyDialog": ("daily_dialog", None, "test", "dialog", 1),

    # ─────────────────────────────────────────────────────────────────
    # Instruction Following
    # ─────────────────────────────────────────────────────────────────
    "AlpacaEval": ("tatsu-lab/alpaca_eval", None, "eval", "instruction", 1),
    "Flan": ("Muennighoff/flan", "default", "train", "inputs", 1),
    "SuperNaturalInstructions": ("Muennighoff/natural-instructions", None, "test", "definition", 1),

    # ─────────────────────────────────────────────────────────────────
    # Multilingual
    # ─────────────────────────────────────────────────────────────────
    "AmazonReviewsMulti": ("amazon_reviews_multi", "en", "test", "review_body", 5),
    "Xcopa": ("xcopa", "en", "test", "premise", 2),
    "Sberquad": ("sberquad", None, "test", "question", 2),

    # ─────────────────────────────────────────────────────────────────
    # Long Context
    # ─────────────────────────────────────────────────────────────────
    "Scrolls": ("tau/scrolls", "qasper", "test", "input", 1),
    "Quality": ("emozilla/quality", None, "test", "article", 4),
    "NarrativeQa": ("narrativeqa", None, "test", "question", 1),

    # ─────────────────────────────────────────────────────────────────
    # Medical
    # ─────────────────────────────────────────────────────────────────
    "PubmedQa": ("pubmed_qa", "pqa_labeled", "test", "question", 3),
    "MedQa": ("bigbio/med_qa", "med_qa_en_source", "test", "question", 4),
}


def get_num_classes(dataset_name: str) -> int:
    """Get number of classes for a dataset."""
    if dataset_name in LLM_DATASET_MAP:
        return LLM_DATASET_MAP[dataset_name][4]
    return 2  # default binary


# ═══════════════════════════════════════════════════════════════════════════════
# Main entry point - FULL MODEL UPLOAD
# ═══════════════════════════════════════════════════════════════════════════════

def create_llm_package(
    model_path: str,
    dataset_name: str,
    output_dir: Optional[str] = None,
    split_ratio: float = 0.75,  # unused, kept for API compatibility
    num_samples: int = 500,     # unused
    max_seq_len: int = 256,     # unused
    batch_size: int = 4,        # unused
) -> Optional[Path]:
    """
    Create Decloud LLM package - uploads FULL model.

    No encoder/head split, no embeddings computation.
    The entire model is saved and will be trained by trainers.

    Args:
        model_path: HuggingFace model ID (e.g. "gpt2", "meta-llama/Llama-2-7b-hf") or local path
        dataset_name: Dataset name for validation
        output_dir: Output directory

    Returns:
        Path to created package or None on error
    """

    num_classes = get_num_classes(dataset_name)

    if output_dir:
        output = Path(output_dir)
    else:
        output = PACKAGES_DIR / f"{dataset_name}_llm"

    output.mkdir(parents=True, exist_ok=True)

    # ═══════════════════════════════════════════════════════════════
    # 1. Load model + tokenizer
    # ═══════════════════════════════════════════════════════════════
    console.print(f"\n[cyan]Loading LLM: {model_path}...[/cyan]")
    try:
        from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig

        # Load config first to get model info
        model_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)

        # Load model
        model = AutoModelForCausalLM.from_pretrained(
            model_path,
            torch_dtype=torch.float16,  # Use fp16 to save space
            trust_remote_code=True,
            low_cpu_mem_usage=True,
        )
        model.eval()

        # Load tokenizer
        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token

        console.print(f"[green]✓ Loaded: {type(model).__name__}[/green]")
    except Exception as e:
        console.print(f"[red]✗ Failed to load model: {e}[/red]")
        return None

    # ═══════════════════════════════════════════════════════════════
    # 2. Get model info
    # ═══════════════════════════════════════════════════════════════
    hidden_size = getattr(model_config, "hidden_size", None) or getattr(model_config, "n_embd", None)
    vocab_size = getattr(model_config, "vocab_size", None)
    num_layers = getattr(model_config, "num_hidden_layers", None) or getattr(model_config, "n_layer", None)
    num_heads = getattr(model_config, "num_attention_heads", None) or getattr(model_config, "n_head", None)

    total_params = sum(p.numel() for p in model.parameters())
    total_size_gb = sum(p.numel() * p.element_size() for p in model.parameters()) / 1024**3

    console.print(f"\n[bold]Model Info:[/bold]")
    console.print(f"  Architecture: {model_config.model_type}")
    console.print(f"  Hidden size: {hidden_size}")
    console.print(f"  Vocab size: {vocab_size}")
    console.print(f"  Layers: {num_layers}")
    console.print(f"  Attention heads: {num_heads}")
    console.print(f"  Total params: {total_params:,}")
    console.print(f"  Size: {total_size_gb:.2f} GB")

    # ═══════════════════════════════════════════════════════════════
    # 3. Save FULL model weights
    # ═══════════════════════════════════════════════════════════════
    console.print(f"\n[cyan]Saving full model to {output}...[/cyan]")

    # Get state dict
    state_dict = model.state_dict()

    # Convert to fp16 and clone ALL tensors to handle any shared memory
    # (works with any model: GPT-2, LLaMA, Mistral, Qwen, etc.)
    state_dict_fp16 = {}
    for key, tensor in state_dict.items():
        if tensor.dtype in [torch.float32, torch.float64]:
            state_dict_fp16[key] = tensor.half().clone()
        else:
            state_dict_fp16[key] = tensor.clone()

    # Save model weights
    console.print(f"[dim]Saving model.safetensors...[/dim]")
    with Progress() as progress:
        task = progress.add_task("Saving weights...", total=1)
        save_safetensors(state_dict_fp16, str(output / "model.safetensors"))
        progress.advance(task)

    model_size_mb = (output / "model.safetensors").stat().st_size / 1024 / 1024
    console.print(f"[green]  ✓ model.safetensors ({model_size_mb:.1f} MB)[/green]")

    # ═══════════════════════════════════════════════════════════════
    # 4. Save tokenizer
    # ═══════════════════════════════════════════════════════════════
    console.print(f"[dim]Saving tokenizer...[/dim]")
    tokenizer.save_pretrained(output / "tokenizer")
    console.print(f"[green]  ✓ tokenizer/[/green]")

    # ═══════════════════════════════════════════════════════════════
    # 5. Build package config
    # ═══════════════════════════════════════════════════════════════
    pkg_config = {
        "type": "llm_full",
        "model": {
            "source": model_path,
            "architecture": model_config.model_type,
            "hidden_size": hidden_size,
            "vocab_size": vocab_size,
            "num_layers": num_layers,
            "num_attention_heads": num_heads,
            "total_params": total_params,
        },
        "dataset": {
            "name": dataset_name,
            "num_classes": num_classes,
        },
        "training": {
            "dtype": "float16",
            "recommended_lr": 2e-5,
            "recommended_batch_size": 4,
            "recommended_epochs": 3,
        },
        "files": {
            "model": "model.safetensors",
            "tokenizer": "tokenizer/",
        },
    }

    # config.json
    with open(output / "config.json", "w") as f:
        json.dump(pkg_config, f, indent=2)
    console.print(f"[green]  ✓ config.json[/green]")

    # Free memory
    del model
    del state_dict
    del state_dict_fp16
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    total_size = model_size_mb
    console.print(f"\n[bold green]✓ LLM package created! ({total_size:.1f} MB)[/bold green]")
    console.print(f"[dim]  Type: full model (no encoder/head split)[/dim]")
    console.print(f"[dim]  Dataset: {dataset_name} ({num_classes} classes)[/dim]")

    return output