import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(
            nn.Linear(128, 256),
            nn.LayerNorm(256),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(256, 10)
        )
    def forward(self, x):
        return self.layers(x)

for epoch in range(epochs):
    model.train()
    for xb, yb in train_loader:
        xb, yb = xb.to(device), yb.to(device)
        pred = model(xb)
        loss = criterion(pred, yb)
        opt.zero_grad()
        loss.backward()
        nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        scheduler.step()

    model.eval()
    with torch.no_grad():
        val_loss = sum(criterion(model(xb.to(device)), yb.to(device))
                       for xb, yb in val_loader)

opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)

torch.save({
    "epoch":      epoch,
    "model":      model.state_dict(),
    "optimizer":  opt.state_dict(),
    "scheduler": sched.state_dict(),
    "loss":       loss.item(),
}, "checkpoint.pt")

ckpt = torch.load("checkpoint.pt", map_location=device)
model.load_state_dict(ckpt["model"])

nn.Dropout(0.3)
nn.Dropout2d(0.2)
nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)

import torch, torch.nn.functional as F

def scaled_dot_attention(Q, K, V, mask=None):
    d_k = Q.size(-1)
    scores = (Q @ K.transpose(-2, -1)) / d_k ** 0.5
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    attn = F.softmax(scores, dim=-1)
    return attn @ V, attn

from transformers import pipeline

tasks = {
  "sentiment":    pipeline("sentiment-analysis"),
  "ner":           pipeline("ner", aggregation_strategy="simple"),
  "summarize":     pipeline("summarization", model="facebook/bart-large-cnn"),
  "qa":            pipeline("question-answering"),
  "translate":     pipeline("translation_en_to_fr"),
  "zero-shot":     pipeline("zero-shot-classification"),
  "fill-mask":     pipeline("fill-mask", model="bert-base-uncased"),
  "text-gen":      pipeline("text-generation", model="gpt2"),
}

from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer, Trainer, TrainingArguments
)
model = AutoModelForSequenceClassification.from_pretrained(
    "bert-base-uncased", num_labels=3
)
args = TrainingArguments(
    output_dir="./out", num_train_epochs=3,
    per_device_train_batch_size=16,
    learning_rate=2e-5, weight_decay=0.01,
    evaluation_strategy="epoch", load_best_model_at_end=True
)
trainer = Trainer(model=model, args=args,
                  train_dataset=train_ds, eval_dataset=val_ds)
trainer.train()

from peft import get_peft_model, LoraConfig, TaskType

config = LoraConfig(
    task_type=TaskType.CAUSAL_LM,
    r=16, lora_alpha=32,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05,
    bias="none"
)
model = get_peft_model(base_model, config)
model.print_trainable_parameters()
# trainable: 0.1% of total params

from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain.chains import RetrievalQA

embeddings = OpenAIEmbeddings()
db = FAISS.from_documents(docs, embeddings)
retriever = db.as_retriever(search_kwargs={"k": 4})
chain = RetrievalQA.from_chain_type(
    llm=ChatOpenAI(model="gpt-4o-mini"),
    retriever=retriever
)
result = chain.invoke({"query": question})

from torchvision import transforms as T

train_tfm = T.Compose([
    T.RandomResizedCrop(224),
    T.RandomHorizontalFlip(),
    T.ColorJitter(0.4, 0.4, 0.4, 0.1),
    T.RandomGrayscale(p=0.2),
    T.ToTensor(),
    T.Normalize([0.485,0.456,0.406],
                [0.229,0.224,0.225]),
])
val_tfm = T.Compose([
    T.Resize(256), T.CenterCrop(224),
    T.ToTensor(),
    T.Normalize([0.485,0.456,0.406],
                [0.229,0.224,0.225])
])

from torchvision.models import resnet50, ResNet50_Weights

model = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)

for param in model.parameters():
    param.requires_grad = False

model.fc = nn.Sequential(
    nn.Linear(model.fc.in_features, 256),
    nn.ReLU(), nn.Dropout(0.3),
    nn.Linear(256, num_classes)
)
opt = torch.optim.Adam(model.fc.parameters(), lr=1e-3)

import gymnasium as gym

env = gym.make("CartPole-v1", render_mode="rgb_array")
obs, info = env.reset(seed=42)

for _ in range(500):
    action = env.action_space.sample()
    obs, reward, terminated, truncated, info = env.step(action)
    if terminated or truncated:
        obs, info = env.reset()

env.close()

from stable_baselines3 import PPO, SAC
from stable_baselines3.common.env_util import make_vec_env

env = make_vec_env("LunarLander-v2", n_envs=4)

model = PPO("MlpPolicy", env,
            n_steps=2048, batch_size=64,
            learning_rate=3e-4, verbose=1)
model.learn(total_timesteps=500_000)
model.save("ppo_lunar")

import numpy as np

A = np.random.randn(4, 4)
U, S, Vt = np.linalg.svd(A)
eigenvals = np.linalg.eigvals(A)
rank = np.linalg.matrix_rank(A)

x = np.linalg.solve(A, b)
inv_A = np.linalg.inv(A)
det_A = np.linalg.det(A)
norm_A = np.linalg.norm(A, ord="fro")

from accelerate import Accelerator

accelerator = Accelerator(mixed_precision="bf16")
model, optimizer, train_loader = accelerator.prepare(
    model, optimizer, train_loader
)
with accelerator.autocast():
    outputs = model(batch)
    loss = criterion(outputs, targets)
accelerator.backward(loss)
optimizer.step()

import wandb

wandb.init(project="my-model", config={
    "lr": 3e-4, "epochs": 20,
    "batch_size": 64, "arch": "resnet50"
})

wandb.log({"loss": loss, "acc": acc, "epoch": epoch})
wandb.watch(model, log="all", log_freq=100)
wandb.finish()

import faiss, numpy as np

d = 768
index = faiss.IndexFlatIP(d)
faiss.normalize_L2(vectors)
index.add(vectors)

query = embed_text("What is attention?")
faiss.normalize_L2(query)
D, I = index.search(query, k=5)
print(f"Top-5 ids: {I[0]}, scores: {D[0]}")

torch.cuda.empty_cache()
torch.backends.cuda.matmul.allow_tf32 = True

model = model.half()

with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
    output = model(input)

torch.utils.checkpoint.checkpoint(fn, *args)