update

Dockerfile

@@ -1,6 +1,6 @@
 FROM python:3.10.8
 
-CMD python download_private_model.py
+# CMD python download_private_model.py
 
 WORKDIR /app
 

src/FineTune/model.py

@@ -0,0 +1,306 @@
+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 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 "facebook/opt-" in model_fullname:
+        print("Using non-fast tokenizer for OPT")
+        optional_tok_kwargs['fast'] = False
+    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 ComputeScore(nn.Module):
+    def __init__(self, scoring_model_name, reference_model_name, dataset='xsum', device='cuda', cache_dir='./models'):
+        super().__init__()
+        self.device = device
+        self.reference_model_name = get_model_fullname(reference_model_name)
+        self.scoring_model_name = get_model_fullname(scoring_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('Moving model to GPU...', 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(scoring_model_name, dataset, cache_dir)
+        scoring_model = load_model(scoring_model_name, device, cache_dir)
+        if scoring_model_name in ['gemma-1b']:
+            self.peft_config = LoraConfig(
+                task_type=TaskType.CAUSAL_LM,
+                inference_mode=False,
+                r=8,
+                lora_alpha=32,
+                lora_dropout=0.1,
+                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(reference_model_name, dataset, cache_dir)
+        reference_model = load_model(reference_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 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 + 基础模型)
+        scoring_dir = os.path.join(save_directory, "scoring_model")
+        self.scoring_model.save_pretrained(scoring_dir, safe_serialization=True)
+        self.scoring_tokenizer.save_pretrained(scoring_dir)
+
+        # 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
+        )
+        model.scoring_tokenizer = AutoTokenizer.from_pretrained(scoring_dir)
+
+        # 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 get_SPO_input(self, tokenized=None, text=[""], labels=[""], training_module=False):
+        if training_module:
+            logits_score = self.scoring_model(tokenized.input_ids, attention_mask=tokenized.attention_mask).logits[:,:-1,:]
+            if self.reference_model_name != self.scoring_model_name:
+                tokenized = self.reference_tokenizer(text, return_tensors="pt", padding=True, return_token_type_ids=False, add_special_tokens=True, return_attention_mask=True).to(self.device)
+                assert torch.all(tokenized.input_ids[:, 1:] == labels), "Tokenizer is mismatch."
+            logits_ref = self.reference_model(tokenized.input_ids).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]
+                if self.reference_model_name != self.scoring_model_name:
+                    tokenized = self.reference_tokenizer(text, return_tensors="pt", padding=True, return_token_type_ids=False ,add_special_tokens=True, return_attention_mask=True).to(self.device)
+                    assert torch.all(tokenized.input_ids[:, 1:] == labels), "Tokenizer is mismatch."
+                logits_ref = self.reference_model(tokenized.input_ids).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.get_SPO_input(tokenized, original_text, 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.get_SPO_input(tokenized, sampled_text, labels,training_module=training_module)
+        
+        output = dict(crit=[train_original_crit.detach(), train_original_crit, train_sampled_crit.detach(), train_sampled_crit])
+        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}")
+
+    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.no_grad():
+            crit, _, _ = self.get_SPO_input(tokenized, text, 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)
+        
+        # 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)
+        p_value = (count + 1) / (total + 1)
+        
+        return crit, 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