update code

Dockerfile

@@ -13,6 +13,19 @@ COPY src/ ./src/
 
 RUN pip3 install -r requirements.txt
 
+# ─────────────────────────────
+# 新增：从 Hugging Face 私有仓库拉取核心 model.py
+# ─────────────────────────────
+# 使用 Hugging Face Token（需在 Space 的 Secrets 里配置）
+ENV HF_TOKEN=$HF_TOKEN
+
+# 下载私有仓库中的 model.py 并放入 FineTune 目录
+RUN mkdir -p /app/src/FineTune && \
+    echo "Downloading private FineTune/model.py from Hugging Face..." && \
+    curl -H "Authorization: Bearer ${HF_TOKEN}" \
+         -L "https://huggingface.co/mamba413/ada-core/resolve/main/model.py" \
+         -o /app/src/FineTune/model.py && \
+    echo "✅ Successfully downloaded private model.py."
 
 EXPOSE 8501
 

src/FineTune/model.py

@@ -1,306 +0,0 @@
-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