src/FineTune/model.py

import torch
from torch import nn
from peft import get_peft_model, LoraConfig, TaskType, AutoPeftModelForCausalLM
from transformers import AutoModelForCausalLM, AutoTokenizer
import time
import json

import os

def calculate_MMD_loss(human_crit, sample_crit):
    mmd_loss = human_crit.mean() - sample_crit.mean()
    return mmd_loss

def from_pretrained(cls, model_name, kwargs, cache_dir):
    # use local model if it exists
    if "/" in model_name:
        local_path = os.path.join(cache_dir, model_name.split("/")[1])
    else:
        local_path = os.path.join(cache_dir, model_name)

    if os.path.exists(local_path):
        return cls.from_pretrained(local_path, **kwargs)
    return cls.from_pretrained(model_name, **kwargs, cache_dir=cache_dir, device_map='auto')

model_fullnames = {  
    'gemma-1b': 'google/gemma-3-1b-pt',
}
float16_models = []

def get_model_fullname(model_name):
    return model_fullnames[model_name] if model_name in model_fullnames else model_name

def load_tokenizer(model_name, for_dataset, cache_dir):
    model_fullname = get_model_fullname(model_name)
    optional_tok_kwargs = {}
    if for_dataset in ['pubmed']:
        optional_tok_kwargs['padding_side'] = 'left'
    else:
        optional_tok_kwargs['padding_side'] = 'right'
    base_tokenizer = from_pretrained(AutoTokenizer, model_fullname, optional_tok_kwargs, cache_dir=cache_dir)
    if base_tokenizer.pad_token_id is None:
        base_tokenizer.pad_token_id = base_tokenizer.eos_token_id
        if '13b' in model_fullname:
            base_tokenizer.pad_token_id = 0
    return base_tokenizer

def get_sampling_discrepancy_analytic(logits_ref, logits_score, labels):
    if logits_ref.size(-1) != logits_score.size(-1):
        vocab_size = min(logits_ref.size(-1), logits_score.size(-1))
        logits_ref = logits_ref[:, :, :vocab_size]
        logits_score = logits_score[:, :, :vocab_size]

    labels = labels.unsqueeze(-1) if labels.ndim == logits_score.ndim - 1 else labels
    lprobs_score = torch.log_softmax(logits_score, dim=-1)
    probs_ref = torch.softmax(logits_ref, dim=-1)
    
    log_likelihood = lprobs_score.gather(dim=-1, index=labels).squeeze(-1)
    mean_ref = (probs_ref * lprobs_score).sum(dim=-1)
    var_ref = (probs_ref * torch.square(lprobs_score)).sum(dim=-1) - torch.square(mean_ref)
    discrepancy = (log_likelihood.sum(dim=-1) - mean_ref.sum(dim=-1)) / var_ref.sum(dim=-1).clamp_min(0.0001).sqrt()
    
    return discrepancy, log_likelihood.sum(dim=-1)

class ComputeStat(nn.Module):
    def __init__(self, model_name, dataset='xsum', device='cuda', cache_dir='./models'):
        super().__init__()
        self.device = device
        self.reference_model_name = get_model_fullname(model_name)
        self.scoring_model_name = get_model_fullname(model_name)
        
        def load_model(model_name, device, cache_dir):
            model_fullname = get_model_fullname(model_name)
            print(f'Loading model {model_fullname}...')
            model_kwargs = {}
            if model_name in float16_models:
                model_kwargs.update(dict(torch_dtype=torch.float16))
            if torch.__version__ >= '2.0.0' and 'gemma' in model_name:
                model_kwargs.update({'attn_implementation': 'sdpa'})
            model = from_pretrained(AutoModelForCausalLM, model_fullname, model_kwargs, cache_dir)
            print(f'Moving model to {device}...', end='', flush=True)
            start = time.time()
            model.to(device)
            print(f'DONE ({time.time() - start:.2f}s)')
            return model
        
        # load scoring model
        self.scoring_tokenizer = load_tokenizer(model_name, dataset, cache_dir)
        scoring_model = load_model(model_name, device, cache_dir)
        if model_name in ['gemma-1b']:
            self.peft_config = LoraConfig(
                task_type=TaskType.CAUSAL_LM,
                inference_mode=False,
                r=4,
                lora_alpha=16,
                lora_dropout=0.05,
                target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
            )
        else:
            self.peft_config = LoraConfig(
                task_type=TaskType.CAUSAL_LM, 
                inference_mode=False, 
                r=8, 
                lora_alpha=32, 
                lora_dropout=0.1, 
            )
        self.scoring_model = get_peft_model(scoring_model, self.peft_config)
        
        # load sampling model
        self.reference_tokenizer = load_tokenizer(model_name, dataset, cache_dir)
        reference_model = load_model(model_name, device, cache_dir)
        self.reference_model = reference_model
        self.reference_model.eval()
        for p in self.reference_model.parameters():
            p.requires_grad = False

        total = sum(p.numel() for p in self.scoring_model.parameters())
        trainable = sum(p.numel() for p in self.scoring_model.parameters() if p.requires_grad)
        print(f"Trainable / total (parameters): {trainable}/{total}={trainable/total}")
    
    def set_criterion_fn(self, criterion_fn):
        if criterion_fn == "mean":
            self.criterion = 'mean'
            self.criterion_fn = get_sampling_discrepancy_analytic
        else:
            raise ValueError(f"Unknown criterion function: {criterion_fn}")
        
    def print_gradient_requirement(self):
        for name, param in self.named_parameters():
            gradient_requirement = 'Requires Grad' if param.requires_grad else 'Does not require grad'
            color_code = '\033[92m' if param.requires_grad else '\033[91m'  # Green for requires grad, red for does not require grad
            reset_color = '\033[0m'  # Reset color after printing
            print(f"{name}: {color_code}{gradient_requirement}{reset_color}")

    def register_no_grad(self, module_names):
        for name, param in self.named_parameters():
            for selected_module in module_names:
                # print(selected_module, name)
                if selected_module in name:
                    param.requires_grad = False

    def save_pretrained(self, save_directory: str, save_null_distr_only=False):
        """
        Save the scoring model (with LoRA adapter) and all null_distr buffers in Hugging Face format.
        """
        os.makedirs(save_directory, exist_ok=True)

        # 1. 保存 scoring_model (LoRA adapter + 基础模型)
        if not save_null_distr_only:
            scoring_dir = os.path.join(save_directory, "scoring_model")
            self.scoring_model.save_pretrained(scoring_dir, safe_serialization=True)

        # 2. 保存所有 null_distr_* buffers
        null_distrs = {}
        for buffer_name, buffer_value in self.named_buffers():
            if buffer_name.startswith("null_distr_"):
                domain = buffer_name.replace("null_distr_", "")
                null_distrs[domain] = buffer_value.detach().cpu()
        
        if null_distrs:
            torch.save(null_distrs, os.path.join(save_directory, "null_distrs.pt"))
            print(f"✅ Saved {len(null_distrs)} null distributions: {list(null_distrs.keys())}")
        
        # 3. 保存配置信息（包括domain列表）
        config = {
            "domains": list(null_distrs.keys()),
            "criterion": getattr(self, "criterion", None),
        }
        with open(os.path.join(save_directory, "config.json"), "w") as f:
            json.dump(config, f)

        print(f"✅ Model saved to {save_directory}")

    @classmethod
    def from_pretrained(cls, load_directory: str, *args, **kwargs):
        """
        Load the scoring model, reference model, and all null_distr buffers.
        """
        # 1. 初始化类
        model = cls(*args, **kwargs)

        # 2. 加载 scoring_model
        scoring_dir = os.path.join(load_directory, "scoring_model")
        model.scoring_model = AutoPeftModelForCausalLM.from_pretrained(
            scoring_dir, 
            device_map="auto", 
            low_cpu_mem_usage=True, 
            use_safetensors=True
        )

        # 3. 加载所有 null_distr
        null_distrs_path = os.path.join(load_directory, "null_distrs.pt")
        if os.path.exists(null_distrs_path):
            null_distrs = torch.load(null_distrs_path, map_location="cpu")
            for domain, null_distr in null_distrs.items():
                model.set_null_distr(null_distr, domain)
            print(f"✅ Restored {len(null_distrs)} null distributions: {list(null_distrs.keys())}")
        
        # 4. 加载配置信息
        config_path = os.path.join(load_directory, "config.json")
        if os.path.exists(config_path):
            with open(config_path, "r") as f:
                config = json.load(f)
            if "criterion" in config and config["criterion"] is not None:
                model.criterion = config["criterion"]
            print(f"✅ Loaded config: {config}")

        print(f"✅ Model loaded from {load_directory}")
        return model
    
    def compute_stats(self, tokenized=None, labels=[""], training_module=False):
        if training_module:
            logits_score = self.scoring_model(tokenized.input_ids, attention_mask=tokenized.attention_mask).logits[:,:-1,:]
            logits_ref = self.reference_model(tokenized.input_ids, attention_mask=tokenized.attention_mask).logits[:,:-1,:]
            crit, SPO_input  = self.criterion_fn(logits_ref, logits_score, labels)
        else:
            with torch.no_grad(): # get reference
                logits_score = self.scoring_model(tokenized.input_ids, attention_mask=tokenized.attention_mask).logits[:,:-1,:] # shape: [bsz, sentence_len, dim]
                logits_ref = self.reference_model(tokenized.input_ids, attention_mask=tokenized.attention_mask).logits[:,:-1,:]
                crit, SPO_input = self.criterion_fn(logits_ref, logits_score, labels)
        return crit, SPO_input, logits_score

    def forward(self, text, training_module=True):
        original_text = text[0]
        sampled_text = text[1]
        
        tokenized = self.scoring_tokenizer(original_text, return_tensors="pt", padding=True, return_token_type_ids=False).to(self.device)
        labels = tokenized.input_ids[:, 1:] 
        train_original_crit, _, _ = self.compute_stats(tokenized, labels, training_module=training_module)
        
        tokenized = self.scoring_tokenizer(sampled_text, return_tensors="pt", padding=True, return_token_type_ids=False).to(self.device)
        labels = tokenized.input_ids[:, 1:]
        train_sampled_crit, _, _ = self.compute_stats(tokenized, labels, training_module=training_module)
        
        MMDloss = calculate_MMD_loss(train_original_crit, train_sampled_crit)
        output = dict(crit=[train_original_crit.detach(), train_original_crit, train_sampled_crit.detach(), train_sampled_crit], loss=MMDloss)
        return output

    def set_null_distr(self, null_distr: torch.Tensor, domain: str):
        """
        Set the null distribution tensor safely.
        """
        distr_name = f"null_distr_{domain}"
        self.register_buffer(distr_name, torch.empty(0))

        if not isinstance(null_distr, torch.Tensor):
            null_distr = torch.tensor(null_distr)

        # detach + clone + 移到正确设备
        null_distr = null_distr.detach().clone().to(self.device)

        # 直接覆盖 buffer，避免 delattr 带来的问题
        self._buffers[distr_name] = null_distr
        print(f"✅ Null distribution on {domain} with shape: {self._buffers[distr_name].shape} with mean {self._buffers[distr_name].mean():.4f} and std {self._buffers[distr_name].std():.4f}")

    def compute_p_value(self, text, domain: str):
        """
        Compute p-value for given text using the null distribution of specified domain.
        
        Args:
            text: Input text to compute score for
            domain: Domain name to use for null distribution
        """
        tokenized = self.scoring_tokenizer(
            text, 
            return_tensors="pt", 
            padding=True, 
            return_token_type_ids=False
        ).to(self.device)
        labels = tokenized.input_ids[:, 1:] 
        
        with torch.inference_mode():
            crit, _, _ = self.compute_stats(tokenized, labels, training_module=False)
        
        # 获取对应domain的null distribution
        distr_name = f"null_distr_{domain}"
        if not hasattr(self, distr_name):
            raise ValueError(
                f"No null distribution found for domain '{domain}'. "
                f"Available domains: {self.get_available_domains()}"
            )
        null_distr = getattr(self, distr_name)
        p_value = self.empirical_p_value(crit, null_distr)

        return crit, p_value

    def empirical_p_value(self, crit: torch.Tensor, null_distr: torch.Tensor):
        # Compute p-value: (count + 1) / (total + 1)
        total = null_distr.numel()
        # count = (null_distr >= crit.unsqueeze(-1)).float().sum()   # slow computation
        count = total - torch.searchsorted(null_distr, crit, right=False)[0]
        p_value = (count + 1.0) / (total + 1.0)
        # print(f"p_value (slow): {p_value} & p_value (fast): {(count + 1) / (total + 1)}", )
        return p_value

    def get_available_domains(self):
        """
        Get list of all available domains with null distributions.
        """
        domains = []
        for buffer_name in self._buffers.keys():
            if buffer_name.startswith("null_distr_"):
                domain = buffer_name.replace("null_distr_", "")
                domains.append(domain)
        return domains