entity_model3

A LoRA adapter for 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

pip install transformers peft torch

The base model is gated. Accept the license at meta-llama/Llama-3.2-3B-Instruct, then authenticate:

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

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

{
  "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.

Citation

@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}}
}