Adapter trained on 45,000 CaseHOLD questions, 85.92% on test

README.md

@@ -0,0 +1,115 @@
+---
+license: apache-2.0
+base_model: Qwen/Qwen2.5-0.5B
+datasets:
+- coastalcph/lex_glue
+language:
+- en
+tags:
+- lora
+- legal
+- multiple-choice
+- casehold
+model-index:
+- name: casehold-lora
+  results:
+  - task:
+      type: multiple-choice
+      name: Multiple choice QA
+    dataset:
+      name: CaseHOLD (LexGLUE)
+      type: coastalcph/lex_glue
+      config: case_hold
+      split: test
+    metrics:
+    - type: accuracy
+      value: 85.92
+      name: Accuracy
+---
+
+# casehold-lora
+
+A LoRA adapter for [Qwen/Qwen2.5-0.5B](https://huggingface.co/Qwen/Qwen2.5-0.5B), trained on CaseHOLD: given a
+passage from a US court opinion with one cited holding removed, pick that holding from five candidates.
+
+**85.92% on the LexGLUE test split** (3,093 of 3,600).
+
+Code, every number with the command that printed it, and the decisions behind the build:
+[github.com/epochzero-0/casehold-lora](https://github.com/epochzero-0/casehold-lora).
+Demo: [epochzero-0/casehold-lora-demo](https://huggingface.co/spaces/epochzero-0/casehold-lora-demo).
+
+## Results
+
+| Model | CaseHOLD test | Source |
+|---|---|---|
+| TF-IDF + SVM | 22.4 | LexGLUE |
+| Qwen2.5-0.5B, zero-shot | 38.4 (validation) | this work |
+| BERT | 70.8 | LexGLUE |
+| RoBERTa-large | 74.4 | LexGLUE |
+| Legal-BERT | 75.3 | LexGLUE |
+| CaseLaw-BERT | 75.4 | LexGLUE |
+| **Qwen2.5-0.5B + this adapter** | **85.9** | this work |
+
+Validation accuracy by training size: 75.28% at 5,000 questions, 80.33% at 15,000, 87.05% at 45,000.
+
+One seed and no hyperparameter tuning, so there is no error bar. The model was chosen on validation and scored
+on test once. Qwen2.5-0.5B (494M parameters) is larger than BERT and Legal-BERT (about 110M).
+
+## Training
+
+- LoRA written by hand (`lora.py`, about 30 lines), r = 16, alpha = 32, on q/k/v/o of all 24 layers:
+  2,162,688 trained parameters, 0.436% of the model. Checked against `peft` with the same weights: the logits
+  match with a max difference of 0.0.
+- One epoch over the 45,000 train questions, plain PyTorch loop, AdamW, lr 2e-4, 16 questions per step, bf16
+  autocast.
+- 62 minutes (1.04 GPU-hours) on an RTX 4050 Laptop GPU, peak 4.3 GiB of VRAM.
+
+## How it answers
+
+One prompt per question: an instruction, the passage, the five holdings lettered A to E, then `Answer:`. The
+answer is whichever of the tokens ` A` to ` E` has the highest logit at the last position. Prompts over 640
+tokens lose passage text from the left; the holdings are never cut.
+
+## Use
+
+This is not a `peft` checkpoint. `adapter.safetensors` holds the A and B matrices, and `lora.py` in this repo
+wraps the base model with them.
+
+```python
+import json, torch
+from huggingface_hub import hf_hub_download
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+repo = "epochzero-0/casehold-lora"
+for f in ("lora.py", "data.py"):
+    hf_hub_download(repo, f, local_dir=".")
+from data import encode, letter_ids
+from lora import load_adapter
+
+cfg = json.load(open(hf_hub_download(repo, "adapter_config.json")))
+tok = AutoTokenizer.from_pretrained(cfg["base_model"])
+model = AutoModelForCausalLM.from_pretrained(cfg["base_model"], dtype=torch.float32)
+load_adapter(model, hf_hub_download(repo, "adapter.safetensors"), cfg["r"], cfg["alpha"]).eval()
+
+question = {"context": "... (<HOLDING>) ...", "endings": ["holding that ...", "...", "...", "...", "..."]}
+ids = torch.tensor([encode(question, tok, cfg["max_len"])])
+with torch.no_grad():
+    final = model.model(input_ids=ids, use_cache=False).last_hidden_state[0, -1]
+    probs = model.lm_head(final)[letter_ids(tok)].softmax(-1)
+print(dict(zip("ABCDE", probs.tolist())))
+```
+
+## Limits
+
+It picks among five given holdings and nothing else. It is trained on US case law only, it is not legal
+advice, and it has not been tested on holdings written in any other style.
+
+## Licences and citations
+
+Adapter: Apache-2.0, as is Qwen2.5-0.5B. CaseHOLD was released by its authors under Apache-2.0; the LexGLUE
+copy is CC BY 4.0.
+
+- Zheng, Guha, Anderson, Henderson, Ho. *When Does Pretraining Help? Assessing Self-Supervised Learning for
+  Law and the CaseHOLD Dataset.* ICAIL 2021. arXiv:2104.08671
+- Chalkidis et al. *LexGLUE: A Benchmark Dataset for Legal Language Understanding in English.* ACL 2022.
+  arXiv:2110.00976

adapter.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8452f6110b3f44d94d8c6b38129266526630d9be50904db6bd426b7b7f73922b
+size 8670408

adapter_config.json

@@ -0,0 +1,13 @@
+{
+  "base_model": "Qwen/Qwen2.5-0.5B",
+  "r": 16,
+  "alpha": 32,
+  "targets": [
+    "q_proj",
+    "k_proj",
+    "v_proj",
+    "o_proj"
+  ],
+  "max_len": 640,
+  "trained_on": 45000
+}

data.py

@@ -0,0 +1,39 @@
+"""CaseHOLD as one prompt per question: instruction, citing context, holdings A-E, "Answer:"."""
+MODEL = "Qwen/Qwen2.5-0.5B"
+LETTERS = "ABCDE"
+HEAD = (
+    "The passage below is from a US court opinion. <HOLDING> marks a missing holding. "
+    "Pick the holding that fits.\n\n"
+)
+
+
+def load():
+    from datasets import load_dataset  # here, so the demo can import the prompt code without it
+
+    return load_dataset("coastalcph/lex_glue", "case_hold")
+
+
+def tail(ex):
+    """The five holdings and the answer cue. Never truncated."""
+    lines = "\n".join(f"{letter}: {e}" for letter, e in zip(LETTERS, ex["endings"]))
+    return f"\n\n{lines}\nAnswer:"
+
+
+def encode(ex, tok, max_len=None):
+    """Token ids for one question. Over max_len, the context loses tokens from its left end:
+    the citation with <HOLDING> sits at the right end, and the holdings stay whole."""
+    head = tok(HEAD)["input_ids"]
+    ctx = tok(ex["context"])["input_ids"]
+    end = tok(tail(ex))["input_ids"]
+    if max_len is not None:
+        room = max(max_len - len(head) - len(end), 0)
+        if len(ctx) > room:
+            ctx = ctx[len(ctx) - room:]
+    return head + ctx + end
+
+
+def letter_ids(tok):
+    """The token id of ' A' ... ' E', the five possible next tokens after 'Answer:'."""
+    ids = [tok(f" {letter}")["input_ids"] for letter in LETTERS]
+    assert all(len(i) == 1 for i in ids), f"a letter is not one token: {ids}"
+    return [i[0] for i in ids]

lora.py

@@ -0,0 +1,70 @@
+"""LoRA, written out: a frozen linear layer plus a trained low-rank correction."""
+import math
+
+import torch
+import torch.nn as nn
+
+TARGETS = ("q_proj", "k_proj", "v_proj", "o_proj")
+
+
+class LoRALinear(nn.Module):
+    """y = base(x) + (alpha / r) * B A x
+
+    base is the pretrained layer, frozen. A is (r, in), B is (out, r), and only they train.
+    B starts at zero, so at step 0 the correction is zero and the model is the pretrained one.
+    """
+
+    def __init__(self, base: nn.Linear, r: int, alpha: float):
+        super().__init__()
+        self.base = base
+        for p in self.base.parameters():
+            p.requires_grad = False
+        w = base.weight
+        # A and B stay float32 whatever the base is stored in: AdamW's small updates vanish in bf16
+        self.A = nn.Parameter(torch.empty(r, base.in_features, dtype=torch.float32, device=w.device))
+        self.B = nn.Parameter(torch.zeros(base.out_features, r, dtype=torch.float32, device=w.device))
+        nn.init.kaiming_uniform_(self.A, a=math.sqrt(5))
+        self.scale = alpha / r
+
+    def forward(self, x):
+        delta = (x.to(self.A.dtype) @ self.A.T @ self.B.T) * self.scale
+        return self.base(x) + delta.to(x.dtype)
+
+
+def inject(model, r=16, alpha=32, targets=TARGETS):
+    """Freeze every weight, then wrap the attention projections of every layer.
+    Returns {"layers.3.q_proj": LoRALinear, ...}."""
+    for p in model.parameters():
+        p.requires_grad = False
+    wrapped = {}
+    for i, layer in enumerate(model.model.layers):
+        for name in targets:
+            lora = LoRALinear(getattr(layer.self_attn, name), r, alpha)
+            setattr(layer.self_attn, name, lora)
+            wrapped[f"layers.{i}.{name}"] = lora
+    return wrapped
+
+
+def lora_state(model):
+    """Only the trained tensors: this is the whole adapter."""
+    return {k: v.detach().cpu() for k, v in model.state_dict().items() if k.endswith((".A", ".B"))}
+
+
+def load_adapter(model, path, r=16, alpha=32):
+    """Wrap the model and load trained A and B into it."""
+    inject(model, r, alpha)
+    if str(path).endswith(".safetensors"):
+        from safetensors.torch import load_file
+
+        state = load_file(path)
+    else:
+        state = torch.load(path, map_location="cpu")
+    missing, unexpected = model.load_state_dict(state, strict=False)
+    assert not unexpected and not [k for k in missing if k.endswith((".A", ".B"))]
+    return model
+
+
+def count(model):
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    total = sum(p.numel() for p in model.parameters())
+    return trainable, total