KOS-V5-Retriever / README.md
codybum's picture
Update README.md
776a01c verified
|
Raw History Blame Contribute Delete
12.5 kB
---
license: apache-2.0
library_name: peft
pipeline_tag: feature-extraction
language:
- en
developed_by: University of Kentucky (College of Medicine Office for Research; Center for Clinical and Translational Science) and University of Louisville (Kentucky Center for Digital Innovation)
affiliations:
- name: University of Kentucky, College of Medicine Office for Research
url: https://medicine.uky.edu/sites/research
- name: University of Kentucky, Center for Clinical and Translational Science (CCTS)
url: https://www.ccts.uky.edu/
- name: University of Louisville, Kentucky Center for Digital Innovation
url: https://centers.louisville.edu/kentucky-center-digital-innovation
model_name: KOS-V5-Retriever
model_codename: Catbird
model_type: qwen3
base_model: Kentucky-Open-Science/KOS-V5-Instruct
tags:
- medical
- clinical
- biomedical
- retrieval
- embeddings
- sentence-similarity
- dense-retriever
- llm2vec
- lora
- peft
- from-scratch
- qwen3
- KOS-V5
- university-of-kentucky
- university-of-louisville
- ccts
---
<p align="center">
<img src="catbird_llm_logo.png" alt="Catbird" width="320"/>
</p>
# KOS-V5-Retriever Β· *"Catbird"*
**Developed by**
**University of Kentucky**
- [College of Medicine, Office for Research](https://medicine.uky.edu/sites/research)
- [Center for Clinical and Translational Science (CCTS)](https://www.ccts.uky.edu/)
**University of Louisville**
- [Kentucky Center for Digital Innovation](https://centers.louisville.edu/kentucky-center-digital-innovation)
> πŸ”Ž **This repository is the retrieval / embedding adapter for [KOS-V5-Instruct](https://huggingface.co/Kentucky-Open-Science/KOS-V5-Instruct).**
> It is a **~264 MB LoRA adapter** (PEFT, rank 32), *not* a standalone model. Loaded on top of the **unmodified** KOS-V5-Instruct weights (verified byte-identical β€” `model.safetensors` sha256 matches the published base), it turns the *generator* into a **dense text-retrieval embedding model**: same medical foundation, re-purposed to produce a vector per input for semantic search and RAG. Detach the adapter and you have the original generator back.
**KOS-V5 (codename *Catbird*)** is the fifth-generation Kentucky Open Science model line β€” a **3.72B-parameter medical language model trained from scratch** (not distilled, not pruned, not continued-pretrained from a general base). This adapter shares that foundation; the sections below cover **(1) what the adapter does and how well it retrieves**, then **(2) the base model it is built on**.
> ⚠️ **Research use only.** Provided for research purposes only; not for commercial, clinical, legal, or production-grade use. The user assumes all risks.
---
## 1 Β· Retrieval & embeddings β€” what this adapter does
The base is a **decoder** LLM (it generates text left-to-right). This adapter converts it into a **text encoder** (llm2vec-style) with three inference-time changes plus a small trained delta:
1. **Bidirectional attention** β€” the causal mask is replaced with a padding-only mask, so every token attends to the whole sequence.
2. **Mean pooling** β€” the sequence embedding is the attention-masked mean of the last hidden states, L2-normalized.
3. **LoRA (rank 32)** on all seven linear projections (`q,k,v,o,gate,up,down`), **contrastively trained** so related (query, passage) pairs land close and unrelated pairs far apart. The base weights are frozen; only the LoRA delta is learned.
The adapter does **not** generate text and does **not** alter the base's chat/tool behaviour β€” it is a separate, hot-swappable head for retrieval only.
### Results β€” official BEIR (SciFact)
Scored with the **official `beir` library** (`GenericDataLoader` + `EvaluateRetrieval` + **pytrec_eval** β€” the exact scorer behind the public BEIR/MTEB leaderboard), **zero-shot**: the training data explicitly **excludes SciFact**, verified clean (0 of 300 SciFact test queries appear in training).
| BEIR SciFact (NDCG@10, official) | value |
| :-- | --: |
| **KOS-V5-Retriever (this adapter)** | **0.7007** |
| Recall@10 | 0.8639 |
| BM25 (reference) | 0.665 |
| strong dense retrievers β€” GTR / E5 / BGE (reference) | 0.70–0.76 |
A legitimately **strong dense retriever** β€” it beats BM25 and sits in the strong-dense band, zero-shot, on a base a fraction the size of typical dense-retrieval models.
> **Scope / honesty.** This is **one** BEIR task (SciFact). It has **not yet** been evaluated on other BEIR tasks or MTEB; read 0.70 as a strong single-benchmark result, not a full retrieval profile. Broader evaluation (nfcorpus, fiqa, MTEB-medical) is planned. Contamination on other benchmarks is unchecked; **BEIR SciFact is zero-shot CLEAN.**
### Usage
All four moving parts (base + LoRA + bidirectional patch + mean-pool) must be applied together, or the number above will not reproduce. This snippet **is** the evaluation encoder.
```python
import torch, torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer
from peft import PeftModel
BASE = "Kentucky-Open-Science/KOS-V5-Instruct"
ADAPTER = "Kentucky-Open-Science/KOS-V5-Retriever" # this repo
# 1) bidirectional attention: replace Qwen3's causal mask with a padding-only mask
import transformers.models.qwen3.modeling_qwen3 as Q3
def _bidirectional(*args, **kwargs):
ie = kwargs.get("input_embeds", args[1] if len(args) > 1 else None)
am = kwargs.get("attention_mask", args[2] if len(args) > 2 else None)
if am is None:
return None
dt = ie.dtype if ie is not None else torch.float32
return (1.0 - am[:, None, None, :].to(dt)) * torch.finfo(dt).min
Q3.create_causal_mask = _bidirectional
# 2) base (frozen) + LoRA adapter
tok = AutoTokenizer.from_pretrained(ADAPTER)
base = AutoModel.from_pretrained(BASE, torch_dtype=torch.float32)
model = PeftModel.from_pretrained(base, ADAPTER).cuda().eval()
# 3) mean-pool + L2-normalize
@torch.no_grad()
def encode(texts, max_length=256):
b = tok(texts, padding=True, truncation=True, max_length=max_length, return_tensors="pt").to("cuda")
h = model(input_ids=b["input_ids"], attention_mask=b["attention_mask"]).last_hidden_state
am = b["attention_mask"].unsqueeze(-1).float()
return F.normalize((h * am).sum(1) / am.sum(1).clamp(min=1), dim=1)
docs = encode(["Vitamin D deficiency is associated with increased risk of respiratory infection."])
queries = encode(["Does low vitamin D raise infection risk?"])
print(queries @ docs.T) # cosine similarity; rank candidates by this
```
Retrieval is symmetric β€” use the same `encode` for queries and documents. The adapter's `meta.json` records the required config (`bidirectional: true`, `pooling: mean`).
### Training (this adapter)
- **Base:** `Kentucky-Open-Science/KOS-V5-Instruct` (frozen; byte-identical to the published weights).
- **Adapter:** PEFT LoRA rank 32, targets `q,k,v,o,gate,up,down`, bf16.
- **Objective:** contrastive (in-batch + hard negatives) over public (query, positive, negative) retrieval pairs, **with SciFact excluded** so the evaluation is zero-shot. No dataset content is distributed with this adapter.
### Retrieval limitations
- **Retrieval only.** Produces embeddings; it does **not** generate. For chat / instruction following / tool calling, use the [base model](https://huggingface.co/Kentucky-Open-Science/KOS-V5-Instruct) without the adapter.
- **Single-benchmark evidence** (BEIR SciFact). Generalization is unproven.
- **Requires the exact inference recipe** (bidirectional patch + mean-pool). A plain causal `PeftModel` load will not reproduce the results.
---
## 2 Β· The base model β€” KOS-V5-Instruct
The rest of this card describes the model this adapter is built on. **These are the *base generator's* numbers** (instruction following, tool calling, medical QA); the adapter's own metric is the BEIR retrieval number above.
**A 3.72B-parameter medical language model trained from scratch.** KOS-V5 (codename **Catbird**) holds the **instruction-tuned head** of the line: the [KOS-V5-Base](https://huggingface.co/Kentucky-Open-Science/KOS-V5-Base) pretraining checkpoint taken through SFT and two GRPO reinforcement-learning legs. Unlike the base, it **follows instructions and calls tools**.
**Code name: Catbird.** Native to Kentucky, the Gray Catbird is a songbird famous for its cat-like "meow"; trained from scratch by teams from the University of Kentucky (Cat) and University of Louisville (Bird).
### Core specifications
| Attribute | Detail |
| :--- | :--- |
| **Architecture** | Decoder-only Transformer (`Qwen3ForCausalLM`), Grouped-Query Attention |
| **Parameters** | 3.715 B |
| **Hidden / Layers** | 2560 / 36 |
| **Attention** | 32 query / 8 KV heads (GQA 4:1), head_dim 128, per-head QK-RMSNorm |
| **Feed-forward** | SwiGLU, intermediate 9728 |
| **Vocabulary** | 32,000, custom medical byte-level BPE |
| **Context length** | 32,768 |
| **Position encoding** | RoPE, ΞΈ = 25,000 |
| **Embeddings** | tied |
| **Precision** | bfloat16 (7.43 GB, single shard) |
### The medical foundation
**This is a medical model.** It inherits a base trained on a **54-source medical/biomedical corpus**. The strongest evidence is **bits-per-byte on held-out medical text** (tokenizer-agnostic). In a **17-model pool** β€” including BioMedLM, Meditron-7B, PMC-LLaMA-7B and MedGemma-4B β€” the KOS-V5 base ranks **1 of 17**:
| medical text (BPB, lower is better) | KOS-V5-Base | rank |
| :-- | --: | --: |
| **5-corpus mean, held-out medical text** | **0.4635** | **1 / 17** |
| clinical narratives | **0.4179** | **1 / 17** |
| radiology | **0.5132** | **1 / 17** |
| chest X-ray reports | **0.6688** | **1 / 17** |
| BIOSSES biomedical sentence similarity (Pearson / Spearman) | **0.7097 / 0.7014** | **1 / 17** |
Every comparator was trained on **1.3–153Γ— more data**. See [KOS-V5-Base](https://huggingface.co/Kentucky-Open-Science/KOS-V5-Base) for the full 96-metric evaluation.
### Base generation benchmarks (official suites)
EleutherAI lm-evaluation-harness `0.4.12.dev0` (commit `c1c4bea`), pristine clone, stock tasks; BFCL via the official `bfcl_eval` (FC mode, non-live AST) on vLLM.
| benchmark | KOS-V5-Instruct | KOS-V4-Instruct |
| :-- | --: | --: |
| **IFEval strict-avg** | **72.19** | 61.6 |
| **MMLU (57-subj, 5-shot, `acc`)** | **0.4512** | 0.2782 |
| medical-9 MMLU mean | **0.4915** | 0.2752 |
| PubMedQA / MedQA / MedMCQA | **0.706 / 0.380 / 0.365** | 0.686 / 0.282 / 0.278 |
| BFCL simple / multiple / parallel | **85.0 / 84.0 / 80.5** | 72.8 / 73.0 / 60.5 |
| RGB negative-rejection (official) | **57.33** | β€” |
IFEval is reported as strict-avg = `(prompt-strict + inst-strict) / 2`. On instruction following the base places **first among nine university-built instruct models** and above the original GPT-3.5-turbo generation; its BFCL tool-calling is **above the Qwen3-4B-Instruct-2507 peer**. The peer still leads on parametric knowledge (MMLU 0.7266) and raw IFEval (84.71). Forgetting control: OOD broad-holdout perplexity at **0.99Γ—** the pre-RL base (no measurable forgetting).
> ⚠️ **Read the medical signal from BPB, not the MCQ scores.** KOS models place little probability mass on MCQ answer *letters*; the format, not the knowledge, is the bottleneck. A model ranking **1 of 17** at modelling clinical text while scoring modestly on multiple-choice is exhibiting exactly that gap.
### Data contamination (base)
- **IFEval: CLEAN (verbatim)** β€” 0 exact containments vs the 541 official test prompts.
- **BFCL: CLEAN (verbatim)** β€” 0 exact containments vs 5,437 official prompts.
- **MMLU / PubMedQA / MedQA / MedMCQA: UNCHECKED.**
- **Retrieval (this adapter), BEIR SciFact: CLEAN (zero-shot)** β€” see Β§1.
## Related models
- [**KOS-V5-Instruct**](https://huggingface.co/Kentucky-Open-Science/KOS-V5-Instruct) β€” the base this adapter attaches to (3.72B, medical, instruction + tool-calling).
- [**KOS-V5-Base**](https://huggingface.co/Kentucky-Open-Science/KOS-V5-Base) β€” the from-scratch pretrained foundation (235.2B tokens).
- [**KOS-V4-Instruct**](https://huggingface.co/Kentucky-Open-Science/KOS-V4-Instruct) β€” previous generation.
## Intended use & limitations
Research use only. English only. Not for clinical, commercial, legal, or production-grade use. Retrieval outputs and any downstream results may be wrong or misleading; this model must not be used to make or inform medical decisions.
## Naming
The program is **KOS** (KOS-V1..V6); the V5 series codename is **Catbird**. Earlier internal names are not used.