# TinyPEFT

TinyPEFT is a small, pure-PyTorch **Parameter-Efficient Fine-Tuning (PEFT)** engine for fine-tuning large language models. 


You can inject LoRA / Adapters / BitFit / Prompt Tuning into almost any Transformer, train only a few parameters, and optionally export just those PEFT weights.
Models can be loaded locally or via the Hugging Face Transformers library.

---

## Installation

From PyPI (recommended):

```bash
pip install tinypeft
```

Or from source in editable mode (for development):

```bash
git clone https://github.com/kanishkez/TinyPEFT.git
cd TinyPEFT
pip install -e .
```

TinyPEFT depends on:

- `torch>=1.13`
- `transformers>=4.30`
- `safetensors>=0.3.1`

---

## Available Modules

All PEFT modules are designed to work with standard HuggingFace models (`AutoModelForCausalLM`, etc.) and with custom PyTorch models that use linear layers.

- **LoRA (`inject_lora`)**  
  Injects low-rank adapters into linear-like layers.  
  Freezes base weights and adds trainable `A/B` matrices.  
  Works well on GPT‑2–style models and generic `nn.Linear` stacks.

- **Adapters (`inject_adapter`)**  
  Wraps candidate layers with bottleneck adapters.  
  Freezes the original module; only adapter weights are trainable.  
  Targets common MLP and feed-forward layers in Transformers.

- **BitFit (`inject_bitfit`)**  
  Sets all non-bias parameters `requires_grad=False`.  
  Only biases are finetuned.

- **Prompt Tuning (`inject_prompt_tuning`)**  
  Adds trainable soft prompt embeddings in front of the input sequence.  
  Freezes the original token embeddings.

- **Trainer (`PEFTTrainer`)**  
  Very small trainer that only optimizes PEFT parameters.  
  Accepts a model and one or more PEFT configs/layers.

- **QA Finetuning Helper (`QAConfig`, `fine_tune_qa`)**  
  A higher-level helper for question–answer style finetuning.  
  Supports JSONL, CSV, and HuggingFace `datasets` via loader utilities.

All of these are exported from the top-level package:

```python
from tinypeft import (
    inject_lora,
    inject_adapter,
    inject_bitfit,
    inject_prompt_tuning,
    PEFTTrainer,
    QAConfig,
    fine_tune_qa,
)
```

---

## Basic LoRA Usage

### Inject LoRA into a GPT‑2 model

```python
import torch
from tinypeft import inject_lora, PEFTTrainer, load_model_and_tokenizer

model, tokenizer = load_model_and_tokenizer("gpt2")

model, loras = inject_lora(
    model,
    r=8,
    alpha=16,
    target_modules=["c_attn", "c_fc", "c_proj"],
    match_mode="contains",
    dropout=0.1,
)

trainer = PEFTTrainer(model, loras, lr=1e-4)

batch = tokenizer("hello world", return_tensors="pt")
loss = trainer.train_step(batch)
print("Training loss:", loss)
```

This will:

- Freeze all base GPT‑2 parameters.
- Add trainable LoRA weights on the attention and MLP projections.
- Train only those LoRA parameters.

### Saving and Loading LoRA Weights

```python
import torch
from tinypeft import get_lora_state_dict, load_lora_state_dict

# save
state = get_lora_state_dict(model)
torch.save(state, "lora_adapter.pt")

# load into a fresh model
new_model, _ = inject_lora(new_model, r=8, alpha=16)
load_lora_state_dict(new_model, torch.load("lora_adapter.pt"))
```

---

## QA Finetuning with Datasets

TinyPEFT includes a small QA finetuning helper to quickly train a PEFT-augmented model on question–answer pairs.

### 1. JSONL or CSV dataset loaders

Each entry should contain at least `question` and `answer` fields.

```python
from tinypeft import (
    load_qa_from_jsonl,
    load_qa_from_csv,
)

pairs_jsonl = load_qa_from_jsonl("qa.jsonl")
pairs_csv = load_qa_from_csv("qa.csv", question_field="question", answer_field="answer")
```

For HuggingFace datasets:

```python
from tinypeft import load_qa_from_hf_dataset

pairs = load_qa_from_hf_dataset(
    dataset_name="squad",  # or any QA-like dataset
    split="train",
    question_field="question",
    answer_field="answer",
)
```

### 2. Configuring finetuning with `QAConfig`

```python
from tinypeft import QAConfig

config = QAConfig(
    model_name="gpt2",
    peft_type="lora",       # "lora", "adapters", "bitfit", "prompt_tuning"
    lr=1e-4,
    batch_size=4,
    num_epochs=1,
    max_length=256,
    lora_r=8,
    lora_alpha=16,
    lora_dropout=0.1,
)
```

### 3. Running `fine_tune_qa`

```python
from tinypeft import fine_tune_qa

pairs = [
    {"question": "What is TinyPEFT?", "answer": "A minimal PEFT library."},
    {"question": "What does LoRA do?", "answer": "Adds low-rank adapters."},
]

model, peft, tokenizer = fine_tune_qa(pairs, config)
```

Internally this will:

- Load `config.model_name` and its tokenizer.
- Ensure the tokenizer has a `pad_token` (fallbacks to `eos_token` if needed).
- Inject the requested PEFT method (LoRA / Adapters / BitFit / Prompt Tuning).
- Build a PyTorch `DataLoader` over your QA pairs, formatting each example as:

  ```text
  Question: <question>
  Answer: <answer>
  ```

- Run `num_epochs` of training using `PEFTTrainer`.

### 4. Generating after QA finetune (LoRA)

```python
prompt = "Question: What is TinyPEFT?\nAnswer:"
inputs = tokenizer(prompt, return_tensors="pt")
with torch.no_grad():
    out = model.generate(**inputs, max_new_tokens=20, do_sample=False)
print(tokenizer.decode(out[0], skip_special_tokens=True))
```

This uses the same model instance returned by `fine_tune_qa`, with PEFT layers still attached.

---

## Project Structure

```text
tinypeft/
    adapters/        # Adapter injection and wrappers
    bitfit/          # BitFit (bias-only) injection
    loader/          # Model/tokenizer loading helpers
    lora/            # LoRA layers and helpers
    prompt_tuning/   # Soft prompt / prompt tuning wrappers
    trainer/         # PEFTTrainer and QA finetune helper
    utils/           # Introspection and module utils
    composite.py     # Combine multiple PEFT methods
pyproject.toml
README.md
LICENSE
```

---

## Why TinyPEFT?

- **Pure PyTorch** – no extra training framework dependencies.
- **Production-Friendly Minimalism** – small surface area, easy to integrate and debug.
- **Multiple PEFT Methods** – LoRA, Adapters, BitFit, Prompt Tuning, and a simple QA finetune engine.
- **HF-Compatible** – designed for HuggingFace models, but works with vanilla PyTorch modules too.

---

## License

MIT
