pyspark-coding-assistant-lora

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

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

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://[email protected]/high_revenue_transactions'))

---

Author

Yash Hooda — Data & AI Engineer yashhooda.ai · github.com/yashhooda1