README.md

---
base_model: mistralai/Mistral-7B-Instruct-v0.3
library_name: peft
license: apache-2.0
tags:
  - pyspark
  - data-engineering
  - code-generation
  - qlora
  - lora
  - delta-lake
language:
  - en
pipeline_tag: text-generation
---

# PySpark Coding Assistant — QLoRA (Mistral-7B-Instruct-v0.3)

A LoRA adapter fine-tuned to generate **modern, idiomatic PySpark** for production data engineering tasks — Delta Lake operations, window functions, incremental loading, and data quality checks.

Trained on a synthetic corpus distilled from hand-curated production patterns, with AST-level validation and deprecated-API filtering. Built end-to-end on a free Colab T4.

> **Read the Limitations section before using this.** This is an honest v1 with a known accuracy ceiling, published with its failure cases documented.

---

## Results

| Metric | Value |
|---|---|
| Final training loss | **0.460** |
| Final validation loss | **0.757** |
| Validation perplexity | **1.69** |
| Mean token accuracy | **0.833** |
| Optimizer steps | 33 (3 epochs) |
| Wall-clock training time | 11m 17s |
| Peak VRAM | 6.27 GB / 14.56 GB (T4) |

The ~0.30 train/val gap indicates the model generalized rather than memorizing the seed examples — but see below for what it did and did not learn.

---

## Limitations — read this

On a held-out spot check of three unseen prompts, the model produced **1 correct answer out of 3**.

**What it learned:** the instruction format, idiomatic import style, `spark.table` / Delta read-write conventions, and the general register of production PySpark.

**What it did not learn:** operator semantics. Two representative failures:

- Asked for a **left anti join**, it emitted `how='left_semi'` — the logical inverse — while the accompanying comment described anti-join behavior.
- Asked for a **30-day rolling sum**, it produced invalid window syntax and called `date_sub()` with three arguments (it takes two).

**Root cause:** the corpus is 191 examples and training ran for 33 optimizer steps. That is enough to learn surface form, not enough to learn semantics. Corpus size is the binding constraint, not the architecture, the hyperparameters, or the base model.

**Do not use this to generate production code unreviewed.** It is a demonstration of an end-to-end fine-tuning pipeline, published with its measured accuracy rather than without.

**Planned v2:** scale the corpus to 2,000+ examples, add execution-based validation (run generated code against a real Spark session and keep only what executes correctly), and evaluate against a held-out benchmark rather than spot checks.

---

## Training data

`hoodarunner/pyspark-production-patterns` — 191 validated examples.

Built by **LLM distillation** from 15 hand-written seed examples drawn from real production pipelines (Toast POS ingestion, medallion architecture, franchise analytics), expanded across 18 data engineering topics.

Every generated example passed an automated validator:

- **AST parse** — must be syntactically valid Python
- **Required** — must reference `pyspark`, `spark.`, `DeltaTable`, or `Window`
- **Banned** — rejects `SparkContext`, `sqlContext`, `StreamingContext`, `.rdd`, and any `pandas` / `numpy` / `sklearn` imports in the output

The ban list matters: an earlier attempt filtered a 122k-row public code dataset and yielded only 58 usable rows, most of them Spark 1.x RDD-era API or pandas mislabeled as Spark. **Public code corpora are not a viable source of modern PySpark training data.** Synthetic generation from curated seeds was strictly better.

---

## Training configuration

| | |
|---|---|
| Base model | `mistralai/Mistral-7B-Instruct-v0.3` |
| Quantization | 4-bit NF4, double quant, fp16 compute |
| LoRA rank / alpha / dropout | 16 / 32 / 0.05 |
| Target modules | `q_proj`, `k_proj`, `v_proj`, `o_proj`, `gate_proj`, `up_proj`, `down_proj` |
| Trainable params | 448 tensors (~0.6% of base) |
| Epochs | 3 |
| Effective batch | 16 (batch 2 × grad accum 8) |
| Learning rate | 2e-4, cosine schedule, 3% warmup |
| Optimizer | `paged_adamw_32bit` |
| Max sequence length | 1024 |
| Precision | fp32 (AMP disabled) |
| Hardware | Google Colab, NVIDIA T4 (16GB) |

### T4-specific notes

Three things that are not obvious and cost real time:

1. **Turing has no bf16.** Mistral's config declares `torch_dtype: bfloat16`, so LoRA params are created in bf16 unless you pass `torch_dtype=torch.float16` explicitly. fp16 AMP's `GradScaler` then fails with `_amp_foreach_non_finite_check_and_unscale_cuda not implemented for 'BFloat16'`. Casting trainable params to fp32 was not sufficient — autocast still produced bf16 gradients. **Disabling AMP entirely (`fp16=False, bf16=False`) was the fix.** At rank 16 the fp32 memory cost is affordable.
2. **Do not use `packing=True` without Flash Attention.** T4 cannot run Flash Attention, and packing without it causes cross-contamination between samples in the attention mask. It also collapses a small dataset into very few optimizer steps.
3. **Restart the runtime after any failed model load.** A partially-loaded 7B model is not garbage collected and the next attempt OOMs.

---

## Usage

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

BASE = "mistralai/Mistral-7B-Instruct-v0.3"
ADAPTER = "hoodarunner/pyspark-coding-assistant-lora"

bnb = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.float16,
)

tok = AutoTokenizer.from_pretrained(BASE)
model = AutoModelForCausalLM.from_pretrained(
    BASE, quantization_config=bnb, device_map="auto", torch_dtype=torch.float16
)
model = PeftModel.from_pretrained(model, ADAPTER)
model.eval()

SYSTEM = "You are an expert PySpark engineer. Answer with correct, idiomatic PySpark code."
q = "Write PySpark code to read a Delta table and filter rows where revenue > 10000."

prompt = f"<s>[INST] {SYSTEM}\n\n{q} [/INST]"
ids = tok(prompt, return_tensors="pt").to("cuda")

with torch.no_grad():
    out = model.generate(**ids, max_new_tokens=256, temperature=0.2,
                         do_sample=True, pad_token_id=tok.eos_token_id)

print(tok.decode(out[0][ids["input_ids"].shape[1]:], skip_special_tokens=True))
```

**The prompt format matters.** The adapter was trained on `<s>[INST] {system}\n\n{question} [/INST]`. Deviating from it degrades output quality.

---

## Example output (correct)

**Prompt:** Write PySpark code to read a Delta table and filter rows where revenue > 10000.

```python
from pyspark.sql.functions import col

df = spark.table('silver.pos_transactions')
filtered = df.filter(col('revenue') > 10000)

(filtered.write
    .format('delta')
    .mode('append')
    .save('abfss://silver@lake.dfs.core.windows.net/high_revenue_transactions'))
```

---

## Author

**Yash Hooda** — Data & AI Engineer
[yashhooda.ai](https://yashhooda.ai) · [github.com/yashhooda1](https://github.com/yashhooda1)