update code

requirements.txt

@@ -3,7 +3,7 @@ gradio==5.31.0
 pandas==2.3.1
 torch==2.8.0
 numpy==2.1.3
-transformers==4.55.2
+transformers==4.57.1
 peft==0.17.1
 tqdm
 scikit-learn

src/FineTune/.gitignore

@@ -163,3 +163,5 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+!ckpt/domain_clf

src/FineTune/ckpt/domain_clf/README.md

@@ -0,0 +1,206 @@
+---
+base_model: google/gemma-3-1b-pt
+library_name: peft
+tags:
+- base_model:adapter:google/gemma-3-1b-pt
+- lora
+- transformers
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]
+### Framework versions
+
+- PEFT 0.19.1

src/FineTune/ckpt/domain_clf/adapter_config.json

@@ -0,0 +1,48 @@
+{
+  "alora_invocation_tokens": null,
+  "alpha_pattern": {},
+  "arrow_config": null,
+  "auto_mapping": null,
+  "base_model_name_or_path": "google/gemma-3-1b-pt",
+  "bias": "none",
+  "corda_config": null,
+  "ensure_weight_tying": false,
+  "eva_config": null,
+  "exclude_modules": null,
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": true,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 32,
+  "lora_bias": false,
+  "lora_dropout": 0.05,
+  "lora_ga_config": null,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": [
+    "classifier",
+    "score"
+  ],
+  "peft_type": "LORA",
+  "peft_version": "0.19.1",
+  "qalora_group_size": 16,
+  "r": 8,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "k_proj",
+    "o_proj",
+    "q_proj",
+    "v_proj"
+  ],
+  "target_parameters": null,
+  "task_type": "SEQ_CLS",
+  "trainable_token_indices": null,
+  "use_bdlora": null,
+  "use_dora": false,
+  "use_qalora": false,
+  "use_rslora": false
+}

src/FineTune/ckpt/domain_clf/adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4c99e0b144d6e36d98b25f1b6c456e0d413b2de2e1a1be424283c44d88196bed
+size 6009592

src/FineTune/ckpt/domain_clf/label_names.json

@@ -0,0 +1 @@
+{"label_names": ["Academia", "Finance", "Government", "Knowledge", "Legislation", "Medicine", "News", "UserReview"], "base_model": "gemma-1b"}

src/FineTune/model.py

@@ -1,16 +1,34 @@
 import torch
 from torch import nn
-from peft import get_peft_model, LoraConfig, TaskType, AutoPeftModelForCausalLM
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import (
+    get_peft_model,
+    LoraConfig,
+    TaskType,
+    AutoPeftModelForCausalLM,
+    AutoPeftModelForSequenceClassification,
+)
+from transformers import (
+    AutoModelForCausalLM,
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+)
 import time
 import json
+import random
 
 import os
 
+try:
+    from transformers import AdamW
+except ImportError:
+    from torch.optim import AdamW
+
+
 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, device=None):
     # use local model if it exists
     if "/" in model_name:
@@ -32,14 +50,29 @@ def from_pretrained(cls, model_name, kwargs, cache_dir, device=None):
         remote_kwargs["device_map"] = {"": device}
     return cls.from_pretrained(model_name, **remote_kwargs)
 
-model_fullnames = {  
+
+model_fullnames = {
+    # default base for the witness functions
     'gemma-1b': 'google/gemma-3-1b-pt',
+    'gemma-4b': 'google/gemma-3-4b-pt',
+    'qwen-1.5b': 'Qwen/Qwen2.5-1.5B',
+    'falcon-1b': 'tiiuae/Falcon3-1B-Base',
+    'phi-1b': 'microsoft/phi-1',
 }
 float16_models = []
 
+# Dtype for the *trainable* scoring model under gemma-1b. fp32 keeps the LoRA
+# weight updates exact (bf16 would round away small lr*grad steps); the frozen
+# reference model is always bf16. Switch this to torch.bfloat16 to also halve
+# the scoring model's weight memory if you still hit OOM, at a small risk to
+# AdaJASA training precision.
+GEMMA1B_SCORING_DTYPE = torch.float32
+
+
 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 = {}
@@ -54,36 +87,52 @@ def load_tokenizer(model_name, for_dataset, cache_dir):
             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]
 
+    # Evaluate the witness statistic in fp32 even when the models run in bf16:
+    # the full-vocab softmax, the probability-weighted variance and the
+    # var-normalised division are precision-sensitive. The upcast buffer is
+    # transient (freed each step), so memory impact is small.
+    logits_ref = logits_ref.float()
+    logits_score = logits_score.float()
+
     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'):
+    def __init__(self, model_name, dataset='xsum', device='cuda', cache_dir='./models', lora_r=None):
         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):
+
+        def load_model(model_name, device, cache_dir, dtype_override=None):
             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))
+            # Gemma-1b's ~256k vocab makes fp32 logits/activations very large;
+            # bf16 ~halves base-model + activation memory at no runtime cost.
+            if 'gemma-1b' in model_name:
+                model_kwargs.update(dict(torch_dtype=torch.bfloat16))
+            # Explicit override (e.g. keep the trainable scoring model in fp32).
+            if dtype_override is not None:
+                model_kwargs.update(dict(torch_dtype=dtype_override))
             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, device=device)
@@ -92,29 +141,26 @@ class ComputeStat(nn.Module):
             model.to(device)
             print(f'DONE ({time.time() - start:.2f}s)')
             return model
-        
-        # load scoring model
+
+        # load scoring model (the trainable one). Keep gemma-1b in fp32 here so
+        # the LoRA witness updates stay precise; the frozen reference below is bf16.
         self.scoring_tokenizer = load_tokenizer(model_name, dataset, cache_dir)
-        scoring_model = load_model(model_name, device, cache_dir)
+        scoring_dtype = GEMMA1B_SCORING_DTYPE if 'gemma-1b' in model_name else None
+        scoring_model = load_model(model_name, device, cache_dir, dtype_override=scoring_dtype)
         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"],
-            )
+            default_r, alpha, dropout = 4, 16, 0.05
         else:
-            self.peft_config = LoraConfig(
-                task_type=TaskType.CAUSAL_LM, 
-                inference_mode=False, 
-                r=8, 
-                lora_alpha=32, 
-                lora_dropout=0.1, 
-            )
+            default_r, alpha, dropout = 8, 32, 0.1
+        self.peft_config = LoraConfig(
+            task_type=TaskType.CAUSAL_LM,
+            inference_mode=False,
+            r=lora_r if lora_r is not None else default_r,
+            lora_alpha=alpha,
+            lora_dropout=dropout,
+            target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
+        )
         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)
@@ -126,14 +172,20 @@ class ComputeStat(nn.Module):
         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}")
-    
+
+        # Optional learned-domain estimator. When present, the special domain
+        # names "estimate" / "softest" route a text through this classifier to
+        # pick (or blend) the null distribution; see `_resolve_domain` /
+        # `_predict_domain_probs` and `compute_p_value` below.
+        self.domain_estimator = None
+
     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'
@@ -144,7 +196,6 @@ class ComputeStat(nn.Module):
     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
 
@@ -154,23 +205,23 @@ class ComputeStat(nn.Module):
         """
         os.makedirs(save_directory, exist_ok=True)
 
-        # 1. 保存 scoring_model (LoRA adapter + 基础模型)
+        # 1. Save the scoring model (LoRA adapter + base model).
         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
+        # 2. Save every null_distr_* buffer.
         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列表）
+
+        # 3. Save config (including the domain list).
         config = {
             "domains": list(null_distrs.keys()),
             "criterion": getattr(self, "criterion", None),
@@ -178,38 +229,47 @@ class ComputeStat(nn.Module):
         with open(os.path.join(save_directory, "config.json"), "w") as f:
             json.dump(config, f)
 
+        # 4. Save the domain estimator, if one has been trained.
+        if not save_null_distr_only and self.domain_estimator is not None:
+            self.domain_estimator.save_pretrained(save_directory)
+
         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.
+        Load the scoring model, reference model, all null_distr buffers, and
+        (if present) the domain estimator.
         """
-        # 1. 初始化类
+        # 1. Construct the class.
         model = cls(*args, **kwargs)
 
-        # 2. 加载 scoring_model
+        # 2. Load the scoring model.
+        # NOTE: pass cache_dir through so the PEFT base model (google/gemma-3-1b-pt,
+        # referenced in adapter_config.json) is resolved from the same cache used by
+        # the constructor above, instead of silently re-downloading into the default
+        # HF cache (~/.cache/huggingface). Pin to a single device instead of
+        # device_map='auto' so this adapter checkpoint can't end up split across
+        # devices from the reference model it sits next to.
         scoring_dir = os.path.join(load_directory, "scoring_model")
         model.scoring_model = AutoPeftModelForCausalLM.from_pretrained(
             scoring_dir,
-            # Same fix as `from_pretrained()` above: pin to model.device
-            # instead of 'auto' so this adapter checkpoint can't end up
-            # split across devices from the reference model it sits next to.
             device_map={"": model.device},
             low_cpu_mem_usage=True,
             use_safetensors=True,
+            cache_dir=kwargs.get("cache_dir"),
             trust_remote_code=True,
         )
 
-        # 3. 加载所有 null_distr
+        # 3. Load every 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. 加载配置信息
+
+        # 4. Load config.
         config_path = os.path.join(load_directory, "config.json")
         if os.path.exists(config_path):
             with open(config_path, "r") as f:
@@ -218,9 +278,42 @@ class ComputeStat(nn.Module):
                 model.criterion = config["criterion"]
             print(f"✅ Loaded config: {config}")
 
+        # 5. Load the domain estimator, if the checkpoint has one.
+        # Pass our already-loaded `scoring_tokenizer` through: the classifier
+        # shares the exact same base-model tokenizer (gemma-1b), so this
+        # skips a second, redundant tokenizer load from disk/hub.
+        #
+        # This is wrapped in its own try/except: the domain estimator is an
+        # *optional* feature (only needed for domain="estimate"/"softest").
+        # A failure here -- e.g. an installed `transformers` version that
+        # doesn't yet map the base model's config to a SequenceClassification
+        # head -- must not take down the whole app; `model` (with manual
+        # domain selection) should still load and serve requests.
+        if os.path.exists(os.path.join(load_directory, DOMAIN_CLF_SUBDIR)):
+            try:
+                model.domain_estimator = DomainClassifier.from_pretrained(
+                    load_directory,
+                    cache_dir=kwargs.get("cache_dir", "./models"),
+                    device=model.device,
+                    tokenizer=model.scoring_tokenizer,
+                )
+            except Exception as e:  # noqa: BLE001 — degrade gracefully, see comment above
+                print(
+                    f"⚠️ Could not load domain estimator from {load_directory}: {e}\n"
+                    f"   Continuing without it -- domain='estimate'/'softest' will be "
+                    f"unavailable, but manual domain selection still works."
+                )
+                model.domain_estimator = None
+
+        # Default to eval mode: a reloaded checkpoint is normally used for
+        # inference, and leaving LoRA dropout active would perturb the witness
+        # statistic away from the calibrated null distribution. Call .train()
+        # explicitly before any further fine-tuning.
+        model.scoring_model.eval()
+
         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,:]
@@ -236,15 +329,15 @@ class ComputeStat(nn.Module):
     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:] 
+        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
@@ -259,33 +352,98 @@ class ComputeStat(nn.Module):
         if not isinstance(null_distr, torch.Tensor):
             null_distr = torch.tensor(null_distr)
 
-        # detach + clone + 移到正确设备
+        # detach + clone + move to the right device
         null_distr = null_distr.detach().clone().to(self.device)
 
-        # 直接覆盖 buffer，避免 delattr 带来的问题
+        # overwrite the buffer directly, to avoid issues with 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
+    def _resolve_domain(self, text, domain: str):
+        """Map the requested domain to a concrete one. The special name
+        ``"estimate"`` predicts a single, most-likely domain from ``text``
+        with the learned estimator (hard routing); any other name is returned
+        unchanged. For the soft/mixture routing used by ``"softest"``, see
+        ``_predict_domain_probs`` and ``compute_p_value_softest`` below."""
+        if domain in (ESTIMATE_DOMAIN, "estimated"):
+            if self.domain_estimator is None:
+                raise ValueError(
+                    "domain='estimate' requested but no domain estimator is "
+                    "loaded. Train one with scripts/train_domain_clf.py (it is "
+                    "saved into the checkpoint), or pass an explicit domain."
+                )
+            texts = [text] if isinstance(text, str) else list(text)
+            return self.domain_estimator.predict(texts)[0]
+        return domain
+
+    def _predict_domain_probs(self, text):
+        """Return ``{domain: probability}`` for ``text``, from the learned
+        domain estimator, restricted to and renormalised over the domains
+        that have a calibrated null distribution.
+
+        This is the *soft* counterpart of ``_resolve_domain``'s ``"estimate"``:
+        instead of collapsing the estimator's output to a single most-likely
+        domain, it keeps the full predicted distribution, so a text that is
+        itself a mixture of domains can be scored against a weighted
+        combination of the candidate null distributions instead of being
+        forced into exactly one of them.
         """
+        if self.domain_estimator is None:
+            raise ValueError(
+                "domain='softest' requested but no domain estimator is "
+                "loaded. Train one with scripts/train_domain_clf.py (it is "
+                "saved into the checkpoint), or pass an explicit domain."
+            )
+        texts = [text] if isinstance(text, str) else list(text)
+        probs = self.domain_estimator.predict_proba(texts)[0]
+
+        available = set(self._null_distr_domains())
+        probs = {d: p for d, p in probs.items() if d in available}
+        z = sum(probs.values())
+        if z <= 0:
+            raise ValueError(
+                "No overlap between the domain estimator's labels and the "
+                "calibrated null distributions; cannot compute a 'softest' "
+                "p-value."
+            )
+        return {d: p / z for d, p in probs.items()}
+
+    def _compute_crit(self, text):
+        """Tokenize ``text`` and compute the AdaJASA witness statistic. Shared
+        by both the hard-domain (``compute_p_value``) and the soft/mixture
+        (``compute_p_value_softest``) p-value paths."""
         tokenized = self.scoring_tokenizer(
-            text, 
-            return_tensors="pt", 
-            padding=True, 
+            text,
+            return_tensors="pt",
+            padding=True,
             return_token_type_ids=False
         ).to(self.device)
-        labels = tokenized.input_ids[:, 1:] 
-        
+        labels = tokenized.input_ids[:, 1:]
+
         with torch.inference_mode():
             crit, _, _ = self.compute_stats(tokenized, labels, training_module=False)
-        
-        # 获取对应domain的null distribution
+        return crit
+
+    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. Pass "estimate" to
+                    let the learned estimator predict a single domain from the
+                    text (hard routing), or "softest" to calibrate against a
+                    weighted mixture of all calibrated domains' null
+                    distributions -- weighted by the estimator's predicted
+                    probabilities.
+        """
+        if domain == SOFTEST_DOMAIN:
+            return self.compute_p_value_softest(text)
+
+        domain = self._resolve_domain(text, domain)
+        crit = self._compute_crit(text)
+
+        # Look up the null distribution for this domain.
         distr_name = f"null_distr_{domain}"
         if not hasattr(self, distr_name):
             raise ValueError(
@@ -297,22 +455,333 @@ class ComputeStat(nn.Module):
 
         return crit, p_value
 
+    def compute_p_value_softest(self, text):
+        """Soft/mixture p-value:
+
+            p-value = (1 + sum_k p_k * count_k) / (1 + sum_k p_k * m_k)
+
+        where, for each calibrated domain k, ``p_k`` is the estimated
+        probability that ``text`` belongs to domain k (from
+        ``_predict_domain_probs``), ``m_k`` is the number of human-written
+        calibration texts collected for domain k, and ``count_k`` is the
+        number of those texts whose statistic falls below the observed
+        statistic S(text).
+        """
+        domain_weights = self._predict_domain_probs(text)
+        crit = self._compute_crit(text)
+
+        numerator = 1.0
+        denominator = 1.0
+        for domain, weight in domain_weights.items():
+            if weight <= 0:
+                continue
+            null_distr = getattr(self, f"null_distr_{domain}")
+            m_k = null_distr.numel()
+            count_k = (m_k - torch.searchsorted(null_distr, crit, right=False)[0]).item()
+            numerator += weight * count_k
+            denominator += weight * m_k
+
+        p_value = torch.tensor(numerator / denominator, device=crit.device)
+        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 _null_distr_domains(self):
+        """Domains with a concretely calibrated null distribution (excludes
+        the pseudo-domain names ``"estimate"`` / ``"softest"``, which are
+        resolved to a concrete domain -- or a weighted mixture of them -- at
+        inference time)."""
+        return [
+            buffer_name.replace("null_distr_", "")
+            for buffer_name in self._buffers.keys()
+            if buffer_name.startswith("null_distr_")
+        ]
+
     def get_available_domains(self):
         """
-        Get list of all available domains with null distributions.
+        Get list of all available domains with null distributions, plus the
+        pseudo-domain names ("estimate", "softest") when a domain estimator
+        is loaded.
         """
-        domains = []
-        for buffer_name in self._buffers.keys():
-            if buffer_name.startswith("null_distr_"):
-                domain = buffer_name.replace("null_distr_", "")
-                domains.append(domain)
+        domains = self._null_distr_domains()
+        if getattr(self, "domain_estimator", None) is not None:
+            domains.append(ESTIMATE_DOMAIN)
+            domains.append(SOFTEST_DOMAIN)
         return domains
+
+
+# Sub-directory (inside the AdaJASA checkpoint directory) that holds the learned
+# domain estimator, so the classifier travels with the null distributions it
+# selects between.
+DOMAIN_CLF_SUBDIR = "domain_clf"
+
+# Special domain name: route a text through the learned estimator instead of
+# assuming the domain is known (oracle). Hard routing -- picks a single,
+# most-likely domain (argmax).
+ESTIMATE_DOMAIN = "estimate"
+
+# Special domain name: soft/mixture routing. Instead of picking a single
+# domain, calibrates against a weighted combination of every calibrated
+# domain's null distribution, weighted by the domain estimator's predicted
+# probabilities -- appropriate when the text may itself be a mixture of
+# domains. See `ComputeStat.compute_p_value_softest`.
+SOFTEST_DOMAIN = "softest"
+
+
+class DomainClassifier(nn.Module):
+    """A LoRA sequence classifier over text *domains*, sharing the gemma-1b base.
+
+    Instead of assuming the test domain is known a priori, we estimate it from
+    the text and let the predicted domain pick (or blend) which pre-computed
+    AdaJASA null distribution is used for the decision.
+
+    It is trained on ``(text, domain_label)`` pairs. Crucially, it never sees
+    the human/machine label and is not part of null-distribution calibration,
+    so it introduces no adaptivity into the p-values -- it only routes a test
+    text to the correct calibration. A separate LoRA adapter on the *same*
+    base model keeps this memory-efficient.
+    """
+
+    def __init__(
+        self,
+        model_name,
+        label_names,
+        device="cuda",
+        cache_dir="./models",
+        lora_r=8,
+        max_length=512,
+        tokenizer=None,
+        _build_base=True,
+    ):
+        super().__init__()
+        self.device = device
+        self.model_name = model_name
+        self.label_names = list(label_names)
+        self.label2id = {name: i for i, name in enumerate(self.label_names)}
+        self.id2label = {i: name for i, name in enumerate(self.label_names)}
+        self.max_length = max_length
+        # The classifier shares its base model with ComputeStat's
+        # scoring/reference models (same `model_name`), so it can reuse an
+        # already-loaded tokenizer (e.g. `ComputeStat.scoring_tokenizer`)
+        # instead of loading -- or saving/reloading -- its own copy.
+        self.tokenizer = tokenizer
+        self.model = None
+
+        # When loading via the classmethod `from_pretrained` below, the PEFT
+        # model is restored from disk, so skip building a fresh base here
+        # (avoids re-downloading / re-initialising the backbone).
+        if not _build_base:
+            return
+
+        model_fullname = get_model_fullname(model_name)
+        if self.tokenizer is None:
+            tok_kwargs = {"padding_side": "right"}
+            self.tokenizer = from_pretrained(AutoTokenizer, model_fullname, tok_kwargs, cache_dir=cache_dir)
+        if self.tokenizer.pad_token_id is None:
+            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+
+        base_kwargs = {
+            "num_labels": len(self.label_names),
+            "id2label": self.id2label,
+            "label2id": self.label2id,
+        }
+        if "gemma-1b" in model_name:
+            base_kwargs["torch_dtype"] = torch.bfloat16
+        base_model = from_pretrained(AutoModelForSequenceClassification, model_fullname, base_kwargs, cache_dir, device=device)
+        base_model.config.pad_token_id = self.tokenizer.pad_token_id
+
+        peft_config = LoraConfig(
+            task_type=TaskType.SEQ_CLS,
+            inference_mode=False,
+            r=lora_r,
+            lora_alpha=lora_r * 4,
+            lora_dropout=0.05,
+            target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
+        )
+        self.model = get_peft_model(base_model, peft_config)
+        self.model.to(device)
+
+        trainable = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
+        total = sum(p.numel() for p in self.model.parameters())
+        print(f"[DomainClassifier] {len(self.label_names)} domains; "
+              f"trainable/total params: {trainable}/{total}={trainable / total:.4f}")
+
+    def fit(self, texts, labels, epochs=3, lr=1e-4, batch_size=8, seed=42):
+        """Train on ``(texts, labels)`` where each label is a domain name."""
+        random.seed(seed)
+        self.model.train()
+        label_ids = [self.label2id[l] for l in labels]
+        optimizer = AdamW(self.model.parameters(), lr=lr)
+        n = len(texts)
+        order = list(range(n))
+        for epoch in range(epochs):
+            random.shuffle(order)
+            total_loss, correct, seen = 0.0, 0, 0
+            for start in range(0, n, batch_size):
+                idx = order[start:start + batch_size]
+                batch_texts = [texts[i] for i in idx]
+                batch_labels = torch.tensor([label_ids[i] for i in idx], device=self.device)
+                enc = self.tokenizer(
+                    batch_texts, return_tensors="pt", padding=True, truncation=True,
+                    max_length=self.max_length, return_token_type_ids=False,
+                ).to(self.device)
+                optimizer.zero_grad()
+                out = self.model(**enc, labels=batch_labels)
+                out.loss.backward()
+                optimizer.step()
+                total_loss += out.loss.item() * len(idx)
+                correct += (out.logits.argmax(dim=-1) == batch_labels).sum().item()
+                seen += len(idx)
+                if (start // batch_size) % 50 == 0:
+                    torch.cuda.empty_cache()
+            print(f"[DomainClassifier] epoch {epoch}: "
+                  f"loss={total_loss / max(seen, 1):.4f} acc={correct / max(seen, 1):.4f}")
+        return self
+
+    @torch.no_grad()
+    def predict(self, texts, batch_size=16):
+        """Return a list of predicted domain *names* for ``texts``."""
+        self.model.eval()
+        preds = []
+        for start in range(0, len(texts), batch_size):
+            batch_texts = texts[start:start + batch_size]
+            enc = self.tokenizer(
+                batch_texts, return_tensors="pt", padding=True, truncation=True,
+                max_length=self.max_length, return_token_type_ids=False,
+            ).to(self.device)
+            logits = self.model(**enc).logits
+            preds.extend(self.id2label[i] for i in logits.argmax(dim=-1).tolist())
+        return preds
+
+    @torch.no_grad()
+    def predict_proba(self, texts, batch_size=16):
+        """Return a list of ``{domain_name: probability}`` dicts (softmax over
+        the classifier's logits), one per text in ``texts``.
+
+        Unlike :meth:`predict` (hard argmax), this keeps the full predicted
+        distribution over domains -- what the soft/mixture p-value
+        (domain="softest") needs: a text that is itself a blend of domains
+        gets calibrated against a weighted combination of null distributions
+        rather than forced into a single one.
+        """
+        self.model.eval()
+        all_probs = []
+        for start in range(0, len(texts), batch_size):
+            batch_texts = texts[start:start + batch_size]
+            enc = self.tokenizer(
+                batch_texts, return_tensors="pt", padding=True, truncation=True,
+                max_length=self.max_length, return_token_type_ids=False,
+            ).to(self.device)
+            logits = self.model(**enc).logits
+            probs = torch.softmax(logits.float(), dim=-1)
+            for row in probs.tolist():
+                all_probs.append({self.id2label[i]: p for i, p in enumerate(row)})
+        return all_probs
+
+    def save_pretrained(self, ckpt_dir):
+        """Save the LoRA adapter and label map under ``ckpt_dir``.
+
+        The tokenizer is deliberately *not* saved here. It's byte-identical
+        to the one `ComputeStat` already loads for its scoring/reference
+        models (same base model), so writing another copy into
+        ``domain_clf/`` would just duplicate ~40MB of vocab files per
+        checkpoint for no benefit -- `from_pretrained` below reloads it from
+        the shared ``cache_dir`` (or reuses a passed-in tokenizer) instead.
+        """
+        save_dir = os.path.join(ckpt_dir, DOMAIN_CLF_SUBDIR)
+        os.makedirs(save_dir, exist_ok=True)
+        self.model.save_pretrained(save_dir, safe_serialization=True)
+        with open(os.path.join(save_dir, "label_names.json"), "w") as f:
+            json.dump({"label_names": self.label_names, "base_model": self.model_name}, f)
+        print(f"✅ Domain classifier saved to {save_dir} (labels: {self.label_names})")
+
+    @classmethod
+    def from_pretrained(cls, ckpt_dir, base_model=None, cache_dir="./models", device="cuda", max_length=512, tokenizer=None):
+        """Load a saved domain classifier from ``ckpt_dir``. ``base_model`` is
+        optional -- it's read from the checkpoint's ``label_names.json`` if
+        omitted.
+
+        Pass ``tokenizer=`` to reuse an already-loaded tokenizer (e.g.
+        ``ComputeStat.scoring_tokenizer``) instead of loading a fresh copy --
+        see the note on `save_pretrained` for why no tokenizer is bundled
+        with this checkpoint in the first place.
+        """
+        save_dir = os.path.join(ckpt_dir, DOMAIN_CLF_SUBDIR)
+        with open(os.path.join(save_dir, "label_names.json")) as f:
+            meta = json.load(f)
+        label_names = meta["label_names"]
+        base_model = base_model or meta.get("base_model", "gemma-1b")
+
+        obj = cls(
+            base_model, label_names, device=device, cache_dir=cache_dir,
+            max_length=max_length, tokenizer=tokenizer, _build_base=False,
+        )
+        if obj.tokenizer is None:
+            model_fullname = get_model_fullname(base_model)
+            tok_kwargs = {"padding_side": "right"}
+            obj.tokenizer = from_pretrained(AutoTokenizer, model_fullname, tok_kwargs, cache_dir=cache_dir)
+        if obj.tokenizer.pad_token_id is None:
+            obj.tokenizer.pad_token_id = obj.tokenizer.eos_token_id
+
+        # Match the training dtype (gemma-1b is trained/saved in bf16). Loaded
+        # without device_map="auto" (we go straight to `.to(device, dtype)`
+        # below) so there's no risk of the offload-split issue that motivates
+        # the device pinning elsewhere in this file.
+        dtype = torch.bfloat16 if "gemma-1b" in base_model else torch.float32
+        obj.model = AutoPeftModelForSequenceClassification.from_pretrained(
+            save_dir,
+            num_labels=len(label_names),
+            torch_dtype=dtype,
+            low_cpu_mem_usage=True,
+            cache_dir=cache_dir,
+        )
+        obj.model.to(device=device, dtype=dtype)
+        obj.model.config.pad_token_id = obj.tokenizer.pad_token_id
+        obj.model.eval()
+        print(f"✅ Domain classifier loaded from {save_dir} (labels: {label_names}, dtype: {dtype})")
+        return obj
+
+
+def train_domain_clf(
+    texts,
+    labels,
+    ckpt_dir,
+    base_model="gemma-1b",
+    cache_dir="./models",
+    device="cuda",
+    epochs=3,
+    lr=1e-4,
+    batch_size=8,
+    lora_r=8,
+    seed=42,
+    tokenizer=None,
+):
+    """Train a LoRA domain estimator on ``(text, domain-label)`` pairs and save
+    its checkpoint into ``ckpt_dir`` -- the *same* directory that holds the
+    AdaJASA null distributions.
+
+    Args:
+        texts:     list[str] of input texts.
+        labels:    list[str] of domain names, aligned with ``texts``.
+        ckpt_dir:  AdaJASA checkpoint directory; the classifier is written to its
+                   ``domain_clf/`` sub-directory.
+        tokenizer: optional, already-loaded tokenizer to reuse (e.g. a
+                   `ComputeStat` instance's `scoring_tokenizer`) instead of
+                   loading a fresh copy of the same base-model tokenizer.
+
+    Returns:
+        The trained :class:`DomainClassifier`.
+    """
+    label_names = sorted(set(labels))
+    clf = DomainClassifier(
+        base_model, label_names, device=device, cache_dir=cache_dir, lora_r=lora_r,
+        tokenizer=tokenizer,
+    )
+    clf.fit(texts, labels, epochs=epochs, lr=lr, batch_size=batch_size, seed=seed)
+    clf.save_pretrained(ckpt_dir)
+    return clf

src/app.py

@@ -29,7 +29,7 @@ if os.environ.get("SPACE_ID"):
 
 import gradio as gr
 
-from FineTune.model import ComputeStat
+from FineTune.model import ComputeStat, ESTIMATE_DOMAIN, SOFTEST_DOMAIN
 from feedback import FeedbackManager
 from stats import StatsManager
 
@@ -79,6 +79,37 @@ DOMAINS = [
     "UserReview",
 ]
 
+# "Domain of the text" is a two-level hierarchical selector built from two
+# cascading gr.Dropdowns, grouped together in one row:
+#   - Level 1 (domain_mode_dropdown): Manual / Soft Estimate / Estimate.
+#   - Level 2 (manual_domain_dropdown): the 9 entries in DOMAINS -- only
+#     meaningful (and only shown) when level 1 is "Manual".
+# "Soft Estimate" calibrates against a probability-weighted blend of every
+# domain's null distribution (FineTune.model.SOFTEST_DOMAIN); "Estimate"
+# calibrates against the single most likely domain (FineTune.model.
+# ESTIMATE_DOMAIN). Both route through the pretrained transformer-based
+# domain classifier bundled in the checkpoint (FineTune/model.py:
+# DomainClassifier).
+DOMAIN_MODES = ["Manual", "Soft Estimate", "Estimate"]
+DEFAULT_DOMAIN_MODE = "Manual"
+DEFAULT_MANUAL_DOMAIN = "General"
+
+DOMAIN_MODE_INFO = {
+    "Manual": (
+        "Pick the domain that best matches your text."
+    ),
+    "Soft Estimate": (
+        "🤖 The estimation is conducted by a pretrained transformer-based classifier, which "
+        "predicts a probability over all supported domains; detection is then calibrated "
+        "against a probability-weighted blend of their null distributions. Useful when the "
+        "text may span more than one domain."
+    ),
+    "Estimate": (
+        "🤖 The estimation is conducted by a pretrained transformer-based classifier, which "
+        "predicts the single most likely domain for your text automatically."
+    ),
+}
+
 FEEDBACK_DATASET_ID = os.environ.get("FEEDBACK_DATASET_ID", f"{ACCOUNT_NAME}/user-feedback")
 
 
@@ -105,6 +136,16 @@ except Exception as e:  # noqa: BLE001 — surfaced in the UI below
     model = None
     model_load_error = str(e)
 
+# `ComputeStat.from_pretrained` degrades gracefully if the bundled domain
+# classifier fails to load (e.g. an installed `transformers` version too old
+# to map the base model to a SequenceClassification head) -- `domain_estimator`
+# is simply left as None. Mirror that here: only offer "Soft Estimate" /
+# "Estimate" in the UI when there's actually a classifier behind them.
+DOMAIN_ESTIMATOR_AVAILABLE = bool(model is not None and model.domain_estimator is not None)
+AVAILABLE_DOMAIN_MODES = (
+    DOMAIN_MODES if DOMAIN_ESTIMATOR_AVAILABLE else [m for m in DOMAIN_MODES if m == "Manual"]
+)
+
 feedback_manager = FeedbackManager(
     dataset_repo_id=FEEDBACK_DATASET_ID,
     hf_token=os.environ.get("HF_TOKEN"),
@@ -159,6 +200,7 @@ INTERPRETATION_TEXT = """
 - **Suggestions for better detection**
     - Provide longer text inputs for more reliable detection results.
     - Select the domain that best matches the content of your text to improve detection accuracy.
+    - Not sure of the domain? Switch **Domain of the text** to *Soft Estimate* or *Estimate* — the estimation is conducted by a pretrained transformer-based classifier that infers it automatically.
 """
 
 FOOTER_TEXT = (
@@ -187,19 +229,55 @@ REFERENCES_BIBTEX = """@article{zhou2026detecting,
 # -----------------------------------------------------------------------
 # Event handlers
 # -----------------------------------------------------------------------
-def run_detection(text: str, domain: str, alpha: float):
-    """Detect button handler: runs inference and refreshes all result widgets."""
+def resolve_domain_selection(domain_mode: str, manual_domain: str) -> str:
+    """Translate the level-1 mode + (level-2, Manual-only) domain pick into
+    the concrete domain string `ComputeStat.compute_p_value` expects."""
+    if domain_mode == "Soft Estimate":
+        return SOFTEST_DOMAIN
+    if domain_mode == "Estimate":
+        return ESTIMATE_DOMAIN
+    return manual_domain
+
+
+def on_domain_mode_change(domain_mode: str):
+    """Level-1 dropdown change handler: show the level-2 domain dropdown only
+    in Manual mode, and swap in the explanatory blurb for whichever mode is
+    now active."""
+    return (
+        gr.update(visible=(domain_mode == "Manual")),
+        DOMAIN_MODE_INFO.get(domain_mode, ""),
+    )
+
+
+def run_detection(text: str, domain_mode: str, manual_domain: str, alpha: float, progress=gr.Progress()):
+    """Detect button handler: runs inference and refreshes all result widgets.
+
+    `progress` is a `gr.Progress` instance (see
+    https://www.gradio.app/guides/progress-bars) -- Gradio recognizes it by
+    its default value and doesn't treat it as a regular input, so it's left
+    out of `detect_btn.click(inputs=...)`. The actual scoring pass
+    (`_run_inference`) is one blocking call with no internal hooks to report
+    from, so the bar advances in stages around it rather than continuously.
+    """
     if not text or not text.strip():
         raise gr.Error("⚠️ Please enter some text before detecting.")
 
+    progress(0, desc="Validating input…")
+    domain = resolve_domain_selection(domain_mode, manual_domain)
+    is_auto_domain = domain in (ESTIMATE_DOMAIN, SOFTEST_DOMAIN)
+
     start_time = time.time()
     try:
+        progress(0.2, desc="Estimating domain…" if is_auto_domain else "Tokenizing text…")
+        progress(0.4, desc="Scoring text against the reference model…")
         crit, p_value = _run_inference(text, domain)
+        progress(0.9, desc="Computing p-value…")
     except gr.Error:
         raise
     except Exception as e:  # noqa: BLE001 — surfaced to the user via gr.Error
         raise gr.Error(f"Detection failed: {e}")
     elapsed_time = time.time() - start_time
+    progress(1.0, desc="Done")
 
     stats_manager.increment_detection()
 
@@ -216,7 +294,8 @@ def run_detection(text: str, domain: str, alpha: float):
         gr.update(value=format_conclusion(p_value, alpha), visible=True),
         gr.update(visible=True),  # interpretation accordion
         gr.update(value=f"⏱️ Processing time: {elapsed_time:.2f} seconds", visible=True),
-        gr.update(visible=True),  # feedback row
+        gr.update(visible=True),  # feedback section (label + buttons) — first reveal happens here
+        gr.update(visible=True),  # feedback buttons row (re-shown in case a prior run hid it)
         gr.update(visible=False),  # feedback thanks message
         detection_state,
     )
@@ -250,8 +329,10 @@ def submit_feedback(feedback_type: str, detection_state: dict | None):
 
 
 def on_load():
+    # Visit count is still tracked (StatsManager persists it alongside the
+    # detection count) but is no longer surfaced in the UI, so this handler
+    # has no outputs.
     stats_manager.increment_visit()
-    return f"{stats_manager.visit_count:,} visits"
 
 
 # -----------------------------------------------------------------------
@@ -300,15 +381,6 @@ CUSTOM_CSS = """
     color: #475569;
     font-size: 0.9rem;
 }
-#stats-chip {
-    position: fixed;
-    top: 3.6rem;
-    right: 1rem;
-    font-size: 0.78rem;
-    color: #9ca3af;
-    z-index: 999;
-    text-align: right;
-}
 #app-footer {
     position: fixed;
     left: 0;
@@ -360,14 +432,30 @@ def build_interface() -> gr.Blocks:
         )
 
         with gr.Row():
-            domain_dropdown = gr.Dropdown(
-                choices=DOMAINS, value="General",
-                label="💡 Domain that matches your text",
-            )
-            detect_btn = gr.Button("🔍 Detect", variant="primary", elem_id="detect-btn")
+            with gr.Row(scale=2):
+                # Level 1 + level 2 of the domain hierarchy, grouped together
+                # in their own row so they read as one control cluster even
+                # though they're two separate gr.Dropdowns.
+                domain_mode_dropdown = gr.Dropdown(
+                    choices=AVAILABLE_DOMAIN_MODES, value=DEFAULT_DOMAIN_MODE,
+                    label="💡 Domain of the text",
+                )
+                manual_domain_dropdown = gr.Dropdown(
+                    choices=DOMAINS, value=DEFAULT_MANUAL_DOMAIN,
+                    label="Domain",
+                    visible=(DEFAULT_DOMAIN_MODE == "Manual"),
+                )
+            detect_btn = gr.Button("🔍 Detect", variant="primary", elem_id="detect-btn", scale=1)
             alpha_slider = gr.Slider(
                 minimum=0.01, maximum=0.2, value=0.05, step=0.005,
-                label="Significance level (α)",
+                label="Significance level (α)", scale=1,
+            )
+
+        domain_mode_info = gr.Markdown(DOMAIN_MODE_INFO[DEFAULT_DOMAIN_MODE])
+        if not DOMAIN_ESTIMATOR_AVAILABLE:
+            gr.Markdown(
+                "_Automatic domain estimation is temporarily unavailable on this deployment "
+                "— please select a domain manually._"
             )
 
         conclusion_md = gr.Markdown(visible=False, latex_delimiters=LATEX_DELIMITERS)
@@ -377,27 +465,35 @@ def build_interface() -> gr.Blocks:
 
         elapsed_caption = gr.Markdown(visible=False)
 
-        gr.HTML(
-            '<div style="margin-top: 0.6rem;"><strong>📝 Result Feedback</strong>: '
-            'Does this detection result meet your expectations? '
-            '<span title="🔒 Your feedback is stored privately and will never be shared with '
-            'third parties. It is used solely to improve detection accuracy.">🔒</span></div>'
-        )
-        with gr.Row(visible=False) as feedback_row:
-            expected_btn = gr.Button("✅ Expected")
-            unexpected_btn = gr.Button("❌ Unexpected")
-        feedback_thanks = gr.Markdown(visible=False)
-
-        stats_chip = gr.Markdown(elem_id="stats-chip")
+        with gr.Column(visible=False) as feedback_section:
+            gr.HTML(
+                '<div style="margin-top: 0.6rem;"><strong>📝 Result Feedback</strong>: '
+                'Does this detection result meet your expectations? '
+                '<span title="🔒 Your feedback is stored privately and will never be shared with '
+                'third parties. It is used solely to improve detection accuracy.">🔒</span></div>'
+            )
+            with gr.Row() as feedback_row:
+                expected_btn = gr.Button("✅ Expected")
+                unexpected_btn = gr.Button("❌ Unexpected")
+            feedback_thanks = gr.Markdown(visible=False)
 
         with gr.Accordion("📚 References", open=False, elem_id="references-box"):
             gr.Markdown(REFERENCES_INTRO)
             gr.Code(value=REFERENCES_BIBTEX, language=None, interactive=False, show_label=False)
 
+        domain_mode_dropdown.change(
+            fn=on_domain_mode_change,
+            inputs=[domain_mode_dropdown],
+            outputs=[manual_domain_dropdown, domain_mode_info],
+        )
+
         detect_btn.click(
             fn=run_detection,
-            inputs=[text_input, domain_dropdown, alpha_slider],
-            outputs=[conclusion_md, interpretation_box, elapsed_caption, feedback_row, feedback_thanks, detection_state],
+            inputs=[text_input, domain_mode_dropdown, manual_domain_dropdown, alpha_slider],
+            outputs=[
+                conclusion_md, interpretation_box, elapsed_caption,
+                feedback_section, feedback_row, feedback_thanks, detection_state,
+            ],
         )
 
         expected_btn.click(
@@ -415,7 +511,9 @@ def build_interface() -> gr.Blocks:
 
         # `demo.load` fires once per browser session (page load) — the closest
         # Gradio equivalent to the Streamlit "count once per session" trick.
-        demo.load(fn=on_load, outputs=[stats_chip])
+        # No outputs: the visit count is tracked (see StatsManager) but no
+        # longer displayed in the UI.
+        demo.load(fn=on_load)
 
     return demo