README.md

metadata

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

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.

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 · github.com/yashhooda1