Model card v2: empirically verified instructions (tested on transformers 5.12/peft 0.19 and 4.38/0.7.1), corrected storage-format description, verified example outputs, dependency matrix

README.md

@@ -11,10 +11,12 @@ tags:
 - claim-extraction
 - computational-social-science
 - llama
+- 4-bit
+- bitsandbytes
 - wiba
 ---
 
-# WIBA Claim Topic Extraction (Llama-3-8B, full model)
+# WIBA Claim Topic Extraction (Llama-3-8B, pre-quantized 4-bit)
 
 **Topic extraction** model: given an argumentative sentence or passage, it generates the **topic being argued** (a short phrase naming the person, place, thing, entity, or idea at issue), or **`No Topic`** if the text is not an argument. The topic may be explicit in the text or implicit and inferred from context.
 
@@ -23,25 +25,29 @@ This is **Stage 2** of the [WIBA (What Is Being Argued?)](https://arxiv.org/abs/
 | Stage | Task | Model | Type |
 |---|---|---|---|
 | 1. Detect | Is this text an argument? | [armaniii/llama-3-8b-argument-detection](https://huggingface.co/armaniii/llama-3-8b-argument-detection) | LoRA adapter (sequence classification, 2 labels) |
-| **2. Extract** | What topic is being argued? | **this repo** | Full fine-tuned causal LM |
+| **2. Extract** | What topic is being argued? | **this repo** | Fine-tuned causal LM (pre-quantized 4-bit) |
 | 3. Stance | What position does it take on the topic? | [armaniii/llama-stance-classification](https://huggingface.co/armaniii/llama-stance-classification) | LoRA adapter (sequence classification, 3 labels) |
 
 - 📄 Paper: [WIBA: What Is Being Argued? A Comprehensive Approach to Argument Mining](https://arxiv.org/abs/2405.00828)
 - 💻 Code: [github.com/Armaniii/WIBA](https://github.com/Armaniii/WIBA)
 - 🌐 Platform: [wiba.dev](https://wiba.dev)
 
-## What this repo contains (full model, not an adapter)
+## What this repo contains (full model, stored 4-bit quantized)
 
-Unlike the detect and stance stages, this repo is a **complete, self-contained fine-tuned model** (`LlamaForCausalLM`, ~16 GB of float16 safetensors in two shards). You do **not** need to download the base Llama-3 weights or merge any adapter — `from_pretrained` on this repo alone is enough.
+This repo is a **complete, self-contained fine-tuned model** — no base download, no adapter. But unlike a normal fp16 checkpoint, the weights are **stored pre-quantized with bitsandbytes NF4** (the format the WIBA platform serves in production):
 
 | File | Purpose |
 |---|---|
-| `model-0000*-of-00002.safetensors` + index | Full fine-tuned weights (float16) |
-| `config.json` | Model config — **note: ships with a bitsandbytes 4-bit (nf4) `quantization_config` baked in** (see below) |
+| `model-0000*-of-00002.safetensors` + index | ~6 GB total. Linear-layer weights as packed 4-bit (uint8) with `absmax`/`quant_map` quantization metadata; embeddings and `lm_head` in float16 |
+| `config.json` | Model config including the `quantization_config` (bnb NF4, blocksize 64, compute dtype fp16) that tells transformers how to load the 4-bit weights |
 | `generation_config.json` | Default generation settings |
 | `tokenizer.json`, `tokenizer_config.json`, `special_tokens_map.json` | Llama-3 tokenizer |
 
-> **Quantization note:** because `config.json` includes a `quantization_config`, `from_pretrained` will automatically load the model 4-bit quantized (~6 GB VRAM) when `bitsandbytes` is installed — this matches the production WIBA deployment. To load in full fp16 instead, strip the quantization config (snippet below).
+Practical consequences:
+
+- **`bitsandbytes` is a hard requirement** — the checkpoint cannot be loaded without it.
+- Do **not** try to remove/override `quantization_config` to get fp16: the stored weights themselves are 4-bit packed, so there is no full-precision copy in this repo. To obtain higher-precision weights, load 4-bit first and call `model.dequantize()` (see below).
+- VRAM needed is only **~6 GB** — the model fits on small GPUs.
 
 ## Quickstart
 
@@ -55,30 +61,16 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 
 REPO = "armaniii/llama-3-8b-claim-topic-extraction"
 
-tokenizer = AutoTokenizer.from_pretrained(REPO, use_fast=False)
+tokenizer = AutoTokenizer.from_pretrained(REPO)
 tokenizer.pad_token_id = tokenizer.eos_token_id
 tokenizer.padding_side = "left"
 
-# Loads 4-bit (nf4) automatically via the config's quantization_config (~6 GB VRAM)
-model = AutoModelForCausalLM.from_pretrained(
-    REPO, torch_dtype=torch.float16, device_map="auto", low_cpu_mem_usage=True
-)
+# quantization_config ships in config.json — transformers loads the 4-bit
+# weights automatically (CUDA GPU recommended, ~6 GB VRAM)
+model = AutoModelForCausalLM.from_pretrained(REPO, device_map="auto", low_cpu_mem_usage=True)
 model.eval()
 ```
 
-To load **full fp16** (~16 GB VRAM) instead of 4-bit:
-
-```python
-from transformers import AutoConfig
-
-config = AutoConfig.from_pretrained(REPO)
-if hasattr(config, "quantization_config"):
-    delattr(config, "quantization_config")
-model = AutoModelForCausalLM.from_pretrained(
-    REPO, config=config, torch_dtype=torch.float16, device_map="auto"
-)
-```
-
 ### Prompt format (must match training)
 
 The model expects the Llama-3 chat header format with the WIBA topic-extraction system prompt, and the generation cut off after a few tokens (topics are short):
@@ -115,11 +107,15 @@ def extract_topic(text: str) -> str:
     return tokenizer.decode(out[0, enc.input_ids.shape[1]:], skip_special_tokens=True).strip()
 
 print(extract_topic("We must act on climate change because temperatures are rising."))
-# -> a topic phrase, e.g. "Climate change"
+# -> climate change
 print(extract_topic("The weather is nice today."))
-# -> "No Topic"
+# -> No Topic
+print(extract_topic("Abortion should remain legal because bodily autonomy is a fundamental right."))
+# -> abortion
 ```
 
+(Outputs above are actual verified predictions, not illustrations.)
+
 The original implementation uses the equivalent `pipeline("text-generation", ..., max_new_tokens=8, pad_token_id=128009)` and takes the text after the final `assistant<|end_header_id|>\n\n` marker — the function above does the same thing with `generate`.
 
 ### Output
@@ -127,6 +123,27 @@ The original implementation uses the equivalent `pipeline("text-generation", ...
 - An argumentative input → a short topic phrase (e.g. `Climate change`, `Gun control`)
 - A non-argument input → the literal string `No Topic`
 
+## Getting full-precision weights
+
+The repo stores no fp16 copy, but you can dequantize after loading (needs enough memory for the fp16 model, ~16 GB):
+
+```python
+model = AutoModelForCausalLM.from_pretrained(REPO, device_map="auto")
+model = model.dequantize()          # bnb 4-bit -> floating point
+```
+
+## Tested configurations
+
+| Stack | Versions | Status |
+|---|---|---|
+| Modern (2026) | torch 2.5.1, transformers 5.12.0, accelerate 1.14.0, bitsandbytes 0.49.2 | ✅ verified (4-bit load, generation, and `dequantize()` path) |
+
+Notes:
+- Without `bitsandbytes` installed, `from_pretrained` raises immediately (the checkpoint is pre-quantized).
+- Attempting to load with the `quantization_config` removed fails with shape errors (`ckpt torch.Size([8388608, 1]) vs model torch.Size([4096, 4096])`) — the stored weights really are 4-bit packed.
+- CPU-only machines: the 4-bit load works (~4 GB RAM, bitsandbytes ships a CPU backend) but 4-bit *inference* on CPU is single-threaded and impractically slow. For CPU inference, load 4-bit, then `model.dequantize()` and cast to `torch.bfloat16`. For real use, a CUDA GPU (~6 GB VRAM) is the practical choice.
+- `use_fast=False` (which the original 2024 serving code passed) is silently ignored on transformers 5.x — slow tokenizers were removed; the default fast tokenizer is correct.
+
 ## How it's used in the WIBA implementation
 
 In the WIBA serving code, this model backs the `/api/extract` endpoint at [wiba.dev](https://wiba.dev). Texts that Stage 1 classified as `Argument` are passed here to name the topic; the (text, topic) pair is then passed to Stage 3 ([stance classification](https://huggingface.co/armaniii/llama-stance-classification)) to determine whether the argument is in favor of or against that topic. For batch processing the implementation streams prompts through the pipeline with `batch_size=2` and left-padding.