README.md

---
base_model: meta-llama/Llama-3.2-3B-Instruct
library_name: peft
model_name: entity_model3
tags:
- base_model:adapter:meta-llama/Llama-3.2-3B-Instruct
- lora
- sft
- transformers
- trl
- entity-extraction
license: llama3.2
pipeline_tag: text-generation
---

# entity_model3

A LoRA adapter for [meta-llama/Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct)
that extracts entities from multi-hop questions and labels each one **known** or **unknown**.

An entity is `known` if the question states it outright, and `unknown` if the question refers to it
only by description and it has to be resolved by a downstream lookup. This is intended as the first
stage of a retrieval pipeline over table+text corpora such as OTT-QA and HybridQA.

**This repo contains adapter weights only (~36 MB), not a full model.** You need the base model
as well — see below.

## Requirements

```bash
pip install transformers peft torch
```

The base model is gated. Accept the license at
[meta-llama/Llama-3.2-3B-Instruct](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct),
then authenticate:

```bash
hf auth login
```

Use the **Instruct** checkpoint, not the plain `Llama-3.2-3B` base model. The adapter was trained
on chat-formatted data, and pairing it with the non-instruct base loads without error but produces
degraded output.

## Usage

```python
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel

ADAPTER = "Pranav0511/entity_model3"

tokenizer = AutoTokenizer.from_pretrained(ADAPTER)
base_model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-3B-Instruct",
    torch_dtype=torch.float16,
    device_map="auto",
)
model = PeftModel.from_pretrained(base_model, ADAPTER).eval()

SYSTEM_PROMPT = (
    "Extract entities from the question and classify "
    "each as known or unknown. Return JSON only."
)

def extract_entities(question: str) -> str:
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": question},
    ]
    prompt = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=200, do_sample=False)

    generated = outputs[0][inputs.input_ids.shape[1]:]
    return tokenizer.decode(generated, skip_special_tokens=True).strip()


print(extract_entities(
    "Who was the Conservative Party of Canada candidate of the federal "
    "electoral district that was named in honour of a geographer and "
    "explorer of the Canadian west?"
))
```

The system prompt above is not optional — it is the exact string used in every training example,
and output quality drops sharply without it. Greedy decoding (`do_sample=False`) is recommended
for stable JSON.

## Output format

```json
{
  "entities": [
    {"entity": "geographer and explorer of the Canadian west", "type": "known"},
    {"entity": "federal electoral district", "type": "unknown"},
    {"entity": "Conservative Party of Canada candidate", "type": "unknown"}
  ]
}
```

Generation is not constrained, so parse defensively — slice from the first `{` to the last `}`
and wrap `json.loads` in a try/except rather than trusting the raw string.

## Training

Supervised fine-tuning with TRL's `SFTTrainer` on 3,924 question/entity pairs, with the base model
loaded in 4-bit NF4 (QLoRA) and a bf16 compute dtype.

| | |
|---|---|
| LoRA rank / alpha / dropout | 16 / 32 / 0.05 |
| Target modules | `q_proj`, `k_proj`, `v_proj`, `o_proj` |
| Epochs | 5 |
| Effective batch size | 8 (4 × 2 grad accum) |
| Learning rate | 2e-4 |

### Framework versions

- PEFT 0.16.0
- TRL 0.20.0
- Transformers 4.53.3
- PyTorch 2.6.0+cu124
- Datasets 4.8.5
- Tokenizers 0.21.4

## License

Derived from Llama 3.2 and therefore covered by the
[Llama 3.2 Community License](https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct/blob/main/LICENSE.txt).

## Citation

```bibtex
@misc{vonwerra2022trl,
	title        = {{TRL: Transformer Reinforcement Learning}},
	author       = {Leandro von Werra and Younes Belkada and Lewis Tunstall and Edward Beeching and Tristan Thrush and Nathan Lambert and Shengyi Huang and Kashif Rasul and Quentin Gallou{\'e}dec},
	year         = 2020,
	journal      = {GitHub repository},
	publisher    = {GitHub},
	howpublished = {\url{https://github.com/huggingface/trl}}
}
```