Fine-tune a local Llama model for PDF field extraction

[marknovak]

Overview

Replace (or improve on) qwen3:30b with a smaller, faster, task-specialized Llama model trained on our own PDF+label pairs. The pipeline already has everything needed to generate training examples — we just need to pair the preprocessed text it produces with the correct JSON outputs from Google Sheets.

---

Phase 1 — Prepare Training Data

Goal: produce a JSONL file of (instruction, output) pairs. Target ≥200 examples; 500+ will give meaningfully better results.

Step 1.1 — Export labels from Google Sheets

Download the Google Sheet as CSV. Each row maps to one PDF (by filename or DOI) with the 9 field values:

filename, species_name, study_location, study_year_range, study_year, study_month, study_day, num_empty, num_nonempty, num_sampled

Leave cells blank (not "N/A") where the field is null — convert to null in JSON later.

Step 1.2 — Extract preprocessed text from each PDF

Run the text extraction on the PDF collection using the same pipeline the model sees at inference time. Call extract_key_sections() from src/extraction/llm_text.py on each PDF's extracted text and save the result to a .txt file named identically to the PDF.

Step 1.3 — Pair text with labels

Write a script (~50 lines) that:

1. Reads the CSV from Step 1.1
2. For each row, reads the matching .txt file from Step 1.2
3. Constructs the JSON output object from the CSV values (blank → null)
4. Writes one JSONL record per pair

Record format (Alpaca-style):

{
  "instruction": "<PRIMARY_PROMPT from src/config.py>",
  "input": "<preprocessed text for this PDF>",
  "output": "{\"species_name\": \"Pygoscelis papua\", \"study_location\": \"...\", ...}"
}

Step 1.4 — Split into train / eval

Reserve 15–20% randomly as a held-out evaluation set.

import random, json

records = [json.loads(l) for l in open("pairs.jsonl")]
random.shuffle(records)
split = int(len(records) * 0.8)

with open("train.jsonl", "w") as f:
    for r in records[:split]: f.write(json.dumps(r) + "\n")

with open("eval.jsonl", "w") as f:
    for r in records[split:]: f.write(json.dumps(r) + "\n")

---

Phase 2 — Choose Model and Hardware

Model recommendation

Model	VRAM (4-bit)	Speed	Recommendation
Llama 3.1 8B	~6 GB	Fast	Best starting point
Llama 3.2 3B	~3 GB	Very fast	If VRAM is limited
Llama 3.1 70B	~40 GB	Slow	Only if high-end hardware

Start with Llama 3.1 8B. A fine-tuned 8B model on this specific 9-field task will likely match or outperform qwen3:30b at general extraction.

Hardware requirements

• Minimum: 8 GB VRAM (RTX 3070 or Apple M-series with 16 GB unified memory)
• Recommended: 16–24 GB VRAM (RTX 3090/4090)
• CPU fallback: Possible but slow — hours per epoch vs. minutes on GPU

---

Phase 3 — Fine-Tune with Unsloth

Unsloth is the fastest local QLoRA trainer with direct Llama 3 support. Reduces VRAM usage ~60% vs. vanilla Hugging Face.

Step 3.1 — Install

python -m venv ~/.venvs/finetune
source ~/.venvs/finetune/bin/activate

pip install "unsloth[cu121-torch240] @ git+https://github.com/unslothai/unsloth.git"
pip install datasets trl

On Apple Silicon:

pip install "unsloth @ git+https://github.com/unslothai/unsloth.git"
pip install datasets trl mlx-lm

Step 3.2 — Training script (finetune.py)

from unsloth import FastLanguageModel
from trl import SFTTrainer
from transformers import TrainingArguments
from datasets import load_dataset

MODEL_NAME  = "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit"
MAX_SEQ_LEN = 8192   # matches Ollama num_ctx in llm_client.py
LORA_RANK   = 16
OUTPUT_DIR  = "./llama-fracfeed-lora"

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name=MODEL_NAME,
    max_seq_length=MAX_SEQ_LEN,
    load_in_4bit=True,
)

model = FastLanguageModel.get_peft_model(
    model,
    r=LORA_RANK,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
    lora_alpha=LORA_RANK * 2,
    lora_dropout=0.05,
    bias="none",
    use_gradient_checkpointing="unsloth",
)

def format_record(example):
    return {"text": tokenizer.apply_chat_template(
        [
            {"role": "system",   "content": example["instruction"]},
            {"role": "user",     "content": example["input"]},
            {"role": "assistant","content": example["output"]},
        ],
        tokenize=False,
        add_generation_prompt=False,
    )}

dataset = load_dataset("json", data_files={"train": "train.jsonl", "eval": "eval.jsonl"})
dataset = dataset.map(format_record)

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset["train"],
    eval_dataset=dataset["eval"],
    dataset_text_field="text",
    max_seq_length=MAX_SEQ_LEN,
    dataset_num_proc=2,
    args=TrainingArguments(
        output_dir=OUTPUT_DIR,
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,
        num_train_epochs=3,
        learning_rate=2e-4,
        warmup_ratio=0.1,
        lr_scheduler_type="cosine",
        fp16=True,
        logging_steps=10,
        eval_strategy="epoch",
        save_strategy="epoch",
        load_best_model_at_end=True,
    ),
)

trainer.train()
model.save_pretrained(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)

Step 3.3 — Run training

python finetune.py

With 300 examples on an RTX 3090: ~15–30 minutes for 3 epochs. If eval loss is still falling after epoch 3, add more epochs.

---

Phase 4 — Evaluate Before Deploying

Minimum bar before deploying: ≥85% per-field exact-match accuracy on fields that appear in ≥50% of eval examples.

Step 4.1 — Run inference on eval set

from unsloth import FastLanguageModel
import json

model, tokenizer = FastLanguageModel.from_pretrained("./llama-fracfeed-lora")
FastLanguageModel.for_inference(model)

results = []
for line in open("eval.jsonl"):
    ex = json.loads(line)
    prompt = tokenizer.apply_chat_template(
        [{"role": "system", "content": ex["instruction"]},
         {"role": "user",   "content": ex["input"]}],
        tokenize=False, add_generation_prompt=True,
    )
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    out = model.generate(**inputs, max_new_tokens=256, temperature=0)
    decoded = tokenizer.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
    results.append({"predicted": decoded, "expected": ex["output"]})

with open("eval_predictions.jsonl", "w") as f:
    for r in results: f.write(json.dumps(r) + "\n")

Step 4.2 — Score field accuracy

import json

fields = ["species_name","study_location","study_year_range","study_year",
          "study_month","study_day","num_empty","num_nonempty","num_sampled"]
correct = {f: 0 for f in fields}
total = 0

for line in open("eval_predictions.jsonl"):
    r = json.loads(line)
    try:
        pred = json.loads(r["predicted"])
        exp  = json.loads(r["expected"])
        for field in correct:
            if str(pred.get(field)) == str(exp.get(field)):
                correct[field] += 1
        total += 1
    except json.JSONDecodeError:
        total += 1

for field, count in correct.items():
    print(f"{field:25s}  {count}/{total}  ({100*count/total:.1f}%)")

---

Phase 5 — Convert to GGUF and Load in Ollama

Step 5.1 — Merge LoRA adapter into base model

from unsloth import FastLanguageModel

model, tokenizer = FastLanguageModel.from_pretrained("./llama-fracfeed-lora")
model.save_pretrained_merged("./llama-fracfeed-merged", tokenizer,
                             save_method="merged_16bit")

Step 5.2 — Convert to GGUF

git clone https://github.com/ggerganov/llama.cpp
cd llama.cpp && make
pip install -r requirements.txt

python convert_hf_to_gguf.py ../llama-fracfeed-merged \
    --outfile ../llama-fracfeed-q4_k_m.gguf \
    --outtype q4_k_m

q4_k_m is the standard 4-bit quantization. Use q8_0 for maximum accuracy if VRAM allows.

Step 5.3 — Create an Ollama Modelfile

FROM /absolute/path/to/llama-fracfeed-q4_k_m.gguf

PARAMETER num_ctx 8192
PARAMETER temperature 0

ollama create llama-fracfeed -f Modelfile
ollama run llama-fracfeed   # smoke test

Step 5.4 — Wire into FracFeedExtractor

In src/config.py, update the default model:

DEFAULT_LLM_MODEL = "llama-fracfeed"   # was "qwen3:30b"

The rest of the pipeline (GBNF schema enforcement, retry logic, Pydantic validation) works unchanged.

---

Phase 6 — Iterate

After deployment, run the batch pipeline on held-out PDFs and compare against Google Sheets ground truth.

Problem	Fix
Model returns malformed JSON	Add examples with noisy PDF text; verify num_ctx=8192 isn't truncating long papers
Specific field consistently wrong	Oversample training examples where that field is populated; add worked examples to the prompt
Model hallucinates values for null fields	Add training examples where the correct output has null for that field
Accuracy regresses on a new species type	Add examples from papers of that type

---

Checklist

• Export Google Sheets → CSV with correct column names
• Run extract_key_sections() on all PDFs → .txt files
• Pair text + labels → train.jsonl / eval.jsonl
• Install Unsloth + TRL
• Run finetune.py (3 epochs)
• Score eval_predictions.jsonl — all fields ≥85%
• Merge adapter → merge_and_save.py
• Convert to GGUF (q4_k_m)
• ollama create llama-fracfeed
• Update DEFAULT_LLM_MODEL in src/config.py
• Run batch pipeline on held-out papers, compare to Sheets