The Python you need for AI

Python offers a rich ecosystem of ML libraries (NumPy, pandas, scikit-learn, PyTorch, TensorFlow), clean syntax that enables rapid prototyping, strong community support, and excellent interoperability with C/C++ for performance-critical components. Its interpreted nature allows interactive development through Jupyter notebooks, which is ideal for experimentation.

Python lists are dynamic, can hold mixed types, and support general-purpose operations. NumPy arrays are fixed-type, contiguous in memory, and support vectorised operations that are 10-100x faster for numerical computation. Use lists for general data structures and NumPy arrays for any numerical or matrix computation in ML pipelines.

Use NumPy: def cosine_similarity(a, b): return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)). This computes the dot product of the vectors divided by the product of their magnitudes. The result ranges from -1 (opposite) to 1 (identical direction). Add epsilon to the denominator to prevent division by zero.

Use virtual environments (venv, conda, or poetry) to isolate project dependencies. Pin exact versions in requirements.txt or pyproject.toml for reproducibility. Use conda for complex native dependencies like CUDA. Docker containers ensure full environment reproducibility. Tools like pip-tools or poetry provide dependency resolution and lock files.

Type hints annotate function parameters and return types (e.g., def predict(input: np.ndarray) -> np.ndarray). They improve code readability, enable static analysis tools like mypy to catch bugs before runtime, and provide better IDE support. They are critical for ML production code where data type mismatches cause silent errors.

First understand the missing data pattern (MCAR, MAR, MNAR) using df.isnull().sum() and missingno library. Options include dropping rows/columns with too many missing values, imputing with mean/median/mode, using advanced imputation like KNN or iterative imputer, or using algorithms that handle NaN natively like XGBoost. Document your strategy for reproducibility.

Start with exploratory data analysis and feature importance analysis. Reduce dimensionality using PCA or feature selection methods. Split data into train/validation/test sets. Begin with a simple baseline like logistic regression, then try gradient boosting (XGBoost/LightGBM) for structured data. Use cross-validation for hyperparameter tuning and evaluate with appropriate metrics.

Broadcasting is NumPy's mechanism for performing operations on arrays of different shapes. It automatically expands the smaller array to match the larger one without copying data. Rules require dimensions to be either equal or one of them to be 1, aligning from the trailing dimension. This enables concise and efficient vectorised code.

Use chunked reading with pd.read_csv(chunksize=...), specify dtypes to reduce memory, use category type for low-cardinality columns, and select only needed columns. For larger scale, consider Dask, Polars, or Vaex for out-of-core computation. Profile memory with df.info(memory_usage='deep') and optimise column types aggressively.

The Global Interpreter Lock (GIL) prevents multiple threads from executing Python bytecode simultaneously. For CPU-bound ML code, it limits true parallelism. However, most ML libraries (NumPy, PyTorch, TensorFlow) release the GIL during computation by running native C/C++/CUDA code. For true Python parallelism, use multiprocessing instead of threading.

PyTorch uses eager execution by default, making it intuitive for debugging and research with a Pythonic API. TensorFlow uses graph-based execution for better production optimisation and has a more mature serving ecosystem (TF Serving, TFLite). PyTorch dominates research and is increasingly popular in production; TensorFlow is still strong in mobile and edge deployment.

Decorators are functions that modify the behaviour of other functions, applied using the @decorator syntax. In ML codebases, they are used for logging (tracking experiments), timing (measuring function performance), caching (memoising expensive computations), and in frameworks like PyTorch for torch.no_grad() and custom autograd functions.

Generators are functions that yield items one at a time using the yield keyword, maintaining state between calls. They produce items lazily, consuming minimal memory regardless of dataset size. They are essential for data loading pipelines where the entire dataset cannot fit in memory, and PyTorch DataLoader uses generators internally.

Subclass torch.utils.data.Dataset, implement __len__ and __getitem__ methods. In __getitem__, load the image, apply transforms (resize, normalise, augment), and return the tensor and label. Wrap it in a DataLoader with batch_size, shuffle=True for training, num_workers for parallel loading, and pin_memory=True for GPU transfer speed.

Shallow copy creates a new object but references the same nested objects; deep copy recursively copies everything. In ML, this matters when modifying model configurations or data preprocessing pipelines. Accidentally modifying shared references can corrupt data or cause subtle training bugs. Use copy.deepcopy() for model state dicts and configuration objects.

Use cProfile or line_profiler to identify bottlenecks. Replace Python loops with vectorised NumPy/pandas operations. Use Cython or Numba for CPU-intensive functions. For GPU code, use torch.cuda.Event for timing and nvidia-smi for utilisation monitoring. Memory profiling with memory_profiler helps find leaks in data pipelines.

groupby splits a DataFrame by categorical columns, and apply runs a function on each group. For feature engineering, compute group-level statistics like mean, count, or std as features. Example: df.groupby('customer_id')['amount'].transform('mean') adds per-customer average amount as a feature. Use agg for multiple statistics at once.

Multithreading shares memory but is limited by the GIL for CPU tasks, suitable for IO-bound work like data downloading. Multiprocessing creates separate processes with their own memory, bypassing the GIL for true parallelism on CPU-bound tasks. For ML data preprocessing, multiprocessing is preferred; for async API calls, threading or asyncio suffices.

Test each transformation function with known inputs and expected outputs. Test edge cases like empty DataFrames, missing values, and unexpected types. Use pytest fixtures for reusable test data. Test the full pipeline end-to-end with a small sample. Verify output shapes, dtypes, and value ranges. Use hypothesis library for property-based testing of transformations.

asyncio provides an event loop for concurrent IO-bound operations using async/await syntax. In AI applications, use it for making concurrent API calls to LLM providers, handling multiple WebSocket connections for real-time inference, or orchestrating parallel data fetches in a RAG pipeline. It does not help with CPU-bound computation.

Polars is a DataFrame library written in Rust that offers lazy evaluation, multi-threaded execution, and memory-efficient processing. It can be 10-100x faster than pandas on large datasets. Unlike pandas, it uses Apache Arrow columnar format and supports lazy query optimisation. It is increasingly adopted in production ML pipelines where pandas becomes a bottleneck.

__init__ defines the layers and parameters of the module. forward defines the computation graph for the forward pass. __call__ is inherited from nn.Module and wraps forward with additional logic like hooks and autograd tracking. You should always call the module directly (model(x)) rather than model.forward(x) to ensure hooks execute properly.

Pickle serialises Python objects to byte streams for storage and loading. It is commonly used to save scikit-learn models and data. However, pickle can execute arbitrary code during deserialisation, making it a security risk for untrusted sources. Use safer alternatives like ONNX, safetensors, or joblib with known-source files. Never unpickle data from untrusted origins.

Sklearn Pipelines chain preprocessing and model steps so that transformations are fit only on training data and applied consistently to test data. Use Pipeline with ColumnTransformer for feature-specific preprocessing. This prevents leakage from fitting scalers or encoders on the full dataset. Always call pipeline.fit(X_train) and pipeline.predict(X_test) rather than transforming separately.

Context managers implement __enter__ and __exit__ methods and are used with the 'with' statement for resource management. In ML, they are used for managing GPU memory (torch.no_grad(), torch.cuda.amp.autocast()), file handles for large data loading, database connections, and experiment tracking contexts like mlflow.start_run().

Use FastAPI or Flask to create endpoints. Load the model at startup (not per request). Define Pydantic models for request/response validation. Implement async handlers for IO-bound operations. Add health checks, logging, and error handling. Use gunicorn/uvicorn with multiple workers for production. Containerise with Docker and add rate limiting.

Dataclasses provide a decorator to auto-generate __init__, __repr__, and comparison methods for data-holding classes. Pydantic models add runtime validation, serialisation, and type coercion. Use dataclasses for internal data structures like model configs. Use Pydantic for API request/response schemas and configuration files where validation of external inputs is critical.

Load the base model with appropriate quantisation config (BitsAndBytesConfig for QLoRA). Apply LoRA config specifying target modules, rank, and alpha. Prepare the dataset in the instruction format using a tokeniser. Use SFTTrainer from trl library with training arguments (learning rate, epochs, batch size). Monitor training loss and validation metrics. Save and merge LoRA weights after training.

Use the merge_and_unload() method from the PEFT library, which adds the LoRA weight deltas (A * B scaled by alpha/r) to the corresponding base model weights. The merged model has the same architecture as the base model with no inference overhead. Save the merged model for deployment. Verify that the merged model produces identical outputs to the LoRA model before deployment.

Profile GPU utilisation with nvidia-smi and PyTorch profiler. Increase batch size to fill GPU memory. Use mixed precision training (FP16/BF16) for higher throughput. Optimise data loading with multiple workers and prefetching. Enable gradient checkpointing to trade compute for memory. Use DeepSpeed or FSDP for multi-GPU efficiency. Reduce CPU-GPU data transfer bottlenecks.

Use PyTorch's DistributedDataParallel (DDP) with NCCL backend for inter-GPU communication. Set up process groups with one process per GPU. Use DistributedSampler for data partitioning. Launch with torchrun specifying nproc_per_node. Enable mixed precision with AMP. Ensure batch size is per-GPU and adjust learning rate for effective global batch size. Profile to verify GPU utilisation.

Set accumulation_steps to the desired factor. In the loop, divide loss by accumulation_steps before backward(). Call optimizer.step() and optimizer.zero_grad() only every accumulation_steps iterations. This effectively increases batch size without requiring proportionally more GPU memory, useful when training large models on limited hardware.

Use Apache Kafka or Redis Streams for message ingestion, asyncio or Faust for stream processing, and a FastAPI/gRPC endpoint for inference. Implement windowed feature computation with a sliding buffer. Use connection pooling for database lookups and batch inference requests when possible. Monitor latency and throughput with Prometheus metrics.

Mixed precision uses FP16 for most operations and FP32 for numerically sensitive operations, reducing memory usage by nearly half and speeding up training on modern GPUs. In PyTorch, use torch.cuda.amp.autocast() for the forward pass and torch.cuda.amp.GradScaler for gradient scaling to prevent underflow. This typically yields 1.5-2x speedup.

Use torch.nn.parallel.DistributedDataParallel (DDP) for data parallelism. Initialise process groups with torch.distributed.init_process_group. Wrap the model in DDP and use DistributedSampler for the DataLoader. Launch with torchrun or torch.distributed.launch. For model parallelism across GPUs, use pipeline parallelism or tensor parallelism with libraries like DeepSpeed or FSDP.

Give the agent tools for data loading (CSV, SQL, APIs), Python code execution for analysis, chart generation, and report writing. Implement a sandboxed code execution environment. Include data profiling as a first step. Handle large datasets with sampling strategies. Add validation tools to verify statistical conclusions. Include guardrails against executing harmful code.

Provide tools for file read/write, code execution in a sandbox, test running, and linting. Implement an iterative loop: generate code, run tests, analyse failures, fix issues. Use the LLM to reason about error messages and fix bugs. Include static analysis tools. Sandbox execution for security. Limit iterations and set timeouts. Log all generated code for review.

DeepSpeed is a deep learning optimisation library from Microsoft that enables training models with billions of parameters. Its ZeRO (Zero Redundancy Optimiser) partitions model states across GPUs to reduce memory per GPU. It supports 3D parallelism (data, tensor, pipeline), mixed precision, gradient checkpointing, and offloading to CPU/NVMe. It can train trillion-parameter models on clusters of GPUs.