Update model.py

src/FineTune/model.py

@@ -7,6 +7,10 @@ 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:
@@ -29,9 +33,6 @@ def get_model_fullname(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:
@@ -60,12 +61,12 @@ def get_sampling_discrepancy_analytic(logits_ref, logits_score, labels):
     
     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'):
+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(reference_model_name)
-        self.scoring_model_name = get_model_fullname(scoring_model_name)
+        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)
@@ -76,16 +77,16 @@ class ComputeScore(nn.Module):
             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)
+            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(scoring_model_name, dataset, cache_dir)
-        scoring_model = load_model(scoring_model_name, device, cache_dir)
-        if scoring_model_name in ['gemma-1b']:
+        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,
@@ -105,8 +106,8 @@ class ComputeScore(nn.Module):
         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_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():
@@ -146,7 +147,6 @@ class ComputeScore(nn.Module):
         # 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 = {}
@@ -185,7 +185,6 @@ class ComputeScore(nn.Module):
             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")
@@ -207,21 +206,15 @@ class ComputeScore(nn.Module):
         print(f"✅ Model loaded from {load_directory}")
         return model
     
-    def get_SPO_input(self, tokenized=None, text=[""], labels=[""], training_module=False):
+    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,:]
-            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,:]
+            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]
-                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,:]
+                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
 
@@ -231,13 +224,14 @@ class ComputeScore(nn.Module):
         
         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)
+        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.get_SPO_input(tokenized, sampled_text, labels,training_module=training_module)
+        train_sampled_crit, _, _ = self.compute_stats(tokenized, labels, training_module=training_module)
         
-        output = dict(crit=[train_original_crit.detach(), train_original_crit, train_sampled_crit.detach(), train_sampled_crit])
+        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):
@@ -255,7 +249,7 @@ class ComputeScore(nn.Module):
 
         # 直接覆盖 buffer，避免 delattr 带来的问题
         self._buffers[distr_name] = null_distr
-        print(f"✅ Null distribution on {domain} with shape: {self._buffers[distr_name].shape}")
+        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):
         """
@@ -273,8 +267,8 @@ class ComputeScore(nn.Module):
         ).to(self.device)
         labels = tokenized.input_ids[:, 1:] 
         
-        with torch.no_grad():
-            crit, _, _ = self.get_SPO_input(tokenized, text, labels, training_module=False)
+        with torch.inference_mode():
+            crit, _, _ = self.compute_stats(tokenized, labels, training_module=False)
         
         # 获取对应domain的null distribution
         distr_name = f"null_distr_{domain}"
@@ -283,16 +277,19 @@ class ComputeScore(nn.Module):
                 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)
-        p_value = (count + 1) / (total + 1)
-        
-        return crit, p_value
+        # 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):
         """