# Step 1: Download & filter Common Crawl with Datatrove
from datatrove.pipeline.readers import WarcReader
from datatrove.pipeline.filters import LanguageFilter, GopherQualityFilter
from datatrove.pipeline.writers import JsonlWriter
from datatrove.executor import LocalPipelineExecutor

pipeline = LocalPipelineExecutor(pipeline=[
    WarcReader("s3://commoncrawl/crawl-data/CC-MAIN-2024-10/segments/"),
    LanguageFilter(languages=["en"]),        # Keep English only
    GopherQualityFilter(min_stop_words=2),    # Quality heuristics
    JsonlWriter("output/train_data"),
], tasks=100)
pipeline.run()

# Step 2: Custom scraper for domain-specific data
import requests
from bs4 import BeautifulSoup

def scrape_page(url):
    soup = BeautifulSoup(requests.get(url).text, "html.parser")
    return " ".join(p.get_text() for p in soup.find_all("p"))

import numpy as np
import tiktoken
enc = tiktoken.get_encoding("gpt2")  # 50,257 tokens

def tokenize(doc):
    tokens = [enc.eot_token]            # <|endoftext|> between docs
    tokens.extend(enc.encode_ordinary(doc["text"]))
    return np.array(tokens, dtype=np.uint16)

# Write 100M-token shards (nanoGPT style)
shard, idx, total = np.empty((100_000_000,), dtype=np.uint16), 0, 0
for doc in ds:
    toks = tokenize(doc)
    if idx + len(toks) > len(shard):
        np.save(f"shard_{total:04d}.npy", shard[:idx]); total += 1; idx = 0
    shard[idx:idx+len(toks)] = toks; idx += len(toks)

# CausalSelfAttention with PyTorch 2.0 FlashAttention
class CausalSelfAttention(nn.Module):
    def forward(self, x):
        B, T, C = x.size()
        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
        head_dim = C // self.n_head
        q = q.view(B, T, self.n_head, head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_head, head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_head, head_dim).transpose(1, 2)
        # Fused FlashAttention — no O(T²) attention matrix stored in memory
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.c_proj(y.transpose(1, 2).contiguous().view(B, T, C))

# Weight tying: token embedding = output projection (halves params)
self.transformer.wte.weight = self.lm_head.weight

# Scale-down residual projections to prevent explosion at depth
for name, p in model.named_parameters():
    if 'c_proj' in name:
        std = 0.02 * (2 * config.n_layer) ** -0.5
        nn.init.normal_(p, mean=0.0, std=std)

# Run: torchrun --nproc_per_node=8 train.py
from torch.nn.parallel import DistributedDataParallel as DDP
dist.init_process_group(backend="nccl")
rank = dist.get_rank()

model = GPT(GPTConfig()).to(rank)
model = torch.compile(model)        # PyTorch 2.0: fuse ops, ~2× faster
model = DDP(model, device_ids=[rank])

# Gradient accumulation (simulate large batch on few GPUs)
grad_accum_steps = total_batch_size // (B * T * world_size)
for micro_step in range(grad_accum_steps):
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        logits, loss = model(x, y)
    loss = loss / grad_accum_steps
    loss.backward()
optimizer.step(); optimizer.zero_grad()

# Cosine LR schedule with warmup
def get_lr(step):
    if step < warmup_steps: return max_lr * step / warmup_steps
    if step > max_steps:    return min_lr
    decay = (step - warmup_steps) / (max_steps - warmup_steps)
    return min_lr + 0.5 * (max_lr - min_lr) * (1 + math.cos(math.pi * decay))

from trl import SFTTrainer, DPOTrainer, SFTConfig, DPOConfig
from transformers import AutoModelForCausalLM, AutoTokenizer

model     = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B")

# Stage 1: SFT on instruction data
sft = SFTTrainer(
    model=model,
    train_dataset=sft_dataset,   # {"prompt": ..., "completion": ...}
    args=SFTConfig(max_seq_length=2048, num_train_epochs=3),
)
sft.train()

# Stage 2: DPO on preference data (no reward model needed)
dpo = DPOTrainer(
    model=sft.model,
    train_dataset=dpo_dataset,   # {"prompt", "chosen", "rejected"}
    args=DPOConfig(beta=0.1),    # beta controls KL penalty strength
)
dpo.train()

# One-liner alternative (Axolotl handles everything above)
# axolotl train config.yaml

from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from transformers import BitsAndBytesConfig

# QLoRA: load in 4-bit, train LoRA adapters in bf16
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B", quantization_config=bnb_config
)
model = prepare_model_for_kbit_training(model)

lora_cfg = LoraConfig(
    r=16, lora_alpha=32,               # rank=16, scale=alpha/r=2
    target_modules=["q_proj","v_proj"], # only attention projections
    lora_dropout=0.05, bias="none",
)
model = get_peft_model(model, lora_cfg)
model.print_trainable_parameters()
# trainable params: 6,815,744 || all params: 8,036,564,992 || 0.08%

# Run standard benchmarks via lm-evaluation-harness
lm_eval --model hf \
  --model_args pretrained=./model \
  --tasks hellaswag,arc_challenge,mmlu,truthfulqa \
  --num_fewshot 5 \
  --output_path ./results.json

# Custom eval: measure exact match on your domain
from lm_eval.models.huggingface import HFLM
evaluator = HFLM(pretrained="./model")

def eval_exact_match(model, tokenizer, pairs):
    correct = 0
    for prompt, expected in pairs:
        out = model.generate(**tokenizer(prompt, return_tensors="pt").to("cuda"),
                              max_new_tokens=50)
        pred = tokenizer.decode(out[0], skip_special_tokens=True)
        if expected.lower() in pred.lower(): correct += 1
    return correct / len(pairs)

# vLLM — OpenAI-compatible server with PagedAttention
vllm serve ./merged-model \
  --port 8000 \
  --tensor-parallel-size 2 \
  --served-model-name my-llm \
  --max-model-len 8192

# Convert to GGUF for CPU / edge inference
python llama.cpp/convert_hf_to_gguf.py \
  ./merged-model --outtype q4_k_m \
  --outfile model.gguf
./llama.cpp/llama-server -m model.gguf --port 8080

# Call either server with the OpenAI client
from openai import OpenAI
client = OpenAI(base_url="http://localhost:8000/v1", api_key="-")
resp = client.chat.completions.create(
    model="my-llm",
    messages=[{"role": "user", "content": "Explain attention in one paragraph."}]
)
print(resp.choices[0].message.content)