Curious builder 0 XP earned · 300 to level 2

0 days Finish a lesson to begin

Badge collection

0 of 5 unlocked

small wins to finish your path

Next lesson →

NumPy and pandas basics

Two libraries sit underneath almost all numerical AI work in Python. NumPy gives you fast arrays of numbers — exactly what embeddings, logits, and model weights are. pandas gives you labelled tables — exactly what datasets and evaluation results are. This lesson teaches both from the ground up and connects every feature to a real AI task, including building a tiny semantic search engine.

Last updated: 01 Nov, 2026 · NumPy 2.x · pandas 2.x+

Install and import

Terminal

python -m pip install numpy pandas

By universal convention, NumPy is imported as np and pandas as pd. You will see these aliases in every tutorial, notebook, and codebase.

Python

import numpy as np
import pandas as pd

Part 1 — NumPy

Definition · NumPy

Numerical Python: the foundational library for numeric computing. Its core is the ndarray, and its operations are implemented in optimised C, which makes them 10–100× faster than equivalent Python loops.

Definition · ndarray

An N-dimensional array: a grid of values that all share one data type and are stored contiguously in memory. A 1-D array is a vector, a 2-D array is a matrix, and higher dimensions are often called tensors.

Why does this matter for AI? An embedding is a 1-D array of floats (often 384, 768, 1536, or 3072 numbers). A collection of embeddings is a 2-D matrix with one row per document. A batch of images is a 4-D array. Every one of these is an ndarray or a close cousin like a PyTorch tensor, which deliberately copies NumPy's interface.

Creating arrays

Python

import numpy as np

v = np.array([0.2, 0.5, 0.8])                 # from a Python list
m = np.array([[1, 2, 3], [4, 5, 6]])          # 2-D: list of lists

print(v)
print(m)
print(np.zeros(3), np.ones((2, 2)).tolist())
print(np.arange(0, 10, 3))                    # like range(), as an array
print(np.linspace(0, 1, 5))                   # 5 evenly spaced values

rng = np.random.default_rng(seed=42)          # reproducible random numbers
print(rng.random(3).round(3))

[0.2 0.5 0.8] [[1 2 3] [4 5 6]] [0. 0. 0.] [[1.0, 1.0], [1.0, 1.0]] [0 3 6 9] [0. 0.25 0.5 0.75 1. ] [0.774 0.439 0.859]

Always create random generators with a seed in experiments; it makes results reproducible, which is essential when you compare models or prompts.

Shape, dimensions, and data types

Definition · shape and dtype

shape is a tuple giving the size of each dimension — (2, 3) means 2 rows and 3 columns. dtype is the type shared by every element, such as int64, float32, or float64.

Python

import numpy as np

embeddings = np.zeros((1000, 768), dtype=np.float32)   # 1000 docs, 768 dims

print(embeddings.shape)
print(embeddings.ndim)
print(embeddings.dtype)
print(embeddings.size)                       # total number of elements
print(embeddings.nbytes / 1_000_000, "MB")

(1000, 768) 2 float32 768000 3.072 MB

Embeddings are usually stored as float32 rather than the default float64: it halves memory with no meaningful loss of retrieval quality. At a million documents that is the difference between about 3 GB and 6 GB.

Vectorised operations

Definition · Vectorisation

Applying an operation to a whole array at once instead of writing a Python loop over its elements. NumPy runs the loop internally in compiled code.

Python

import numpy as np

scores = np.array([0.41, 0.88, 0.67, 0.91])

print(scores * 100)            # every element times 100
print(scores + 0.05)           # every element plus 0.05
print(scores >= 0.7)           # element-wise comparison -> booleans
print(np.sqrt(np.array([4, 9, 16])))

a = np.array([1, 2, 3])
b = np.array([10, 20, 30])
print(a + b, a * b)            # element-by-element

[41. 88. 67. 91.] [0.46 0.93 0.72 0.96] [False True False True] [2. 3. 4.] [11 22 33] [10 40 90]

Compare this with the plain-Python version [s * 100 for s in scores]. For four numbers the difference is invisible; for a million embeddings it is the difference between milliseconds and seconds.

Broadcasting

Definition · Broadcasting

NumPy's rule for combining arrays of different shapes: a smaller array is virtually “stretched” along missing or size-1 dimensions to match the larger one, without copying data.

Python

import numpy as np

m = np.array([[1.0, 2.0, 3.0],
              [4.0, 5.0, 6.0]])           # shape (2, 3)

col_means = m.mean(axis=0)                # shape (3,)
print(col_means)
print(m - col_means)                      # (2, 3) - (3,) -> each row minus the means

row_norms = np.linalg.norm(m, axis=1, keepdims=True)   # shape (2, 1)
print((m / row_norms).round(3))           # divide each row by its own length

[2.5 3.5 4.5] [[-1.5 -1.5 -1.5] [ 1.5 1.5 1.5]] [[0.267 0.535 0.802] [0.456 0.57 0.684]]

That last line — dividing each row by its length — is exactly how embeddings are normalised before similarity search. Shapes are compatible when, comparing from the right, each pair of dimensions is equal or one of them is 1.

Aggregations and the axis argument

Aggregations reduce many values to one: sum, mean, min, max, std, argmax. On a 2-D array, axis=0 collapses the rows (one result per column) and axis=1 collapses the columns (one result per row).

Python

import numpy as np

# rows = 3 models, columns = 4 evaluation questions
results = np.array([
    [0.9, 0.7, 0.8, 0.6],
    [0.6, 0.8, 0.7, 0.9],
    [0.5, 0.4, 0.9, 0.8],
])

print(results.mean())               # overall average
print(results.mean(axis=1))         # average per model (row)
print(results.max(axis=0))          # best score per question (column)
print(results.mean(axis=1).argmax())   # index of the best model

0.7166666666666668 [0.75 0.75 0.65] [0.9 0.8 0.9 0.9] 0

argmax returns the position of the largest value, not the value itself. Models 0 and 1 tie at 0.75; argmax returns the first one.

Indexing, slicing, and boolean masks

Python

import numpy as np

m = np.arange(12).reshape(3, 4)
print(m)
print(m[1, 2])          # row 1, column 2
print(m[0])             # first row
print(m[:, 1])          # second column, every row
print(m[:2, -2:])       # first two rows, last two columns

scores = np.array([0.41, 0.88, 0.67, 0.91])
mask = scores > 0.6
print(mask, scores[mask])
print(np.where(scores > 0.6, "keep", "drop"))

[[ 0 1 2 3] [ 4 5 6 7] [ 8 9 10 11]] 6 [0 1 2 3] [1 5 9] [[2 3] [6 7]] [False True True True] [0.88 0.67 0.91] ['drop' 'keep' 'keep' 'keep']

reshape(3, 4) rearranges the 12 values into 3 rows of 4 without copying them. A boolean mask selects exactly the elements where the mask is True — the NumPy equivalent of a filtered list comprehension.

Embeddings and cosine similarity

Definition · Cosine similarity

A measure of how closely two vectors point in the same direction: their dot product divided by the product of their lengths. It ranges from −1 (opposite) through 0 (unrelated) to 1 (same direction). For text embeddings, higher cosine similarity means more similar meaning.

The formula is cos(a, b) = (a · b) / (‖a‖ × ‖b‖), where a · b is the dot product (multiply matching elements, then sum) and ‖a‖ is the vector's length (its L2 norm).

Python

import numpy as np

a = np.array([0.2, 0.5, 0.8])
b = np.array([0.1, 0.6, 0.7])
c = np.array([0.9, -0.2, 0.1])

def cosine(x, y):
    return float(np.dot(x, y) / (np.linalg.norm(x) * np.linalg.norm(y)))

print(round(np.dot(a, b), 3))       # 0.02 + 0.30 + 0.56
print(round(cosine(a, b), 3))       # similar direction
print(round(cosine(a, c), 3))       # different direction

0.88 0.984 0.179

Mini semantic search with NumPy

A vector database does, at heart, exactly this: compare a query vector with every document vector and return the most similar ones. With normalised vectors, cosine similarity is just a dot product, so one matrix multiplication scores every document at once.

Python

import numpy as np

docs = ["RAG basics", "Fine-tuning guide", "Vector databases", "Cooking pasta"]

# toy 4-dimensional "embeddings" (real ones have hundreds of dimensions)
doc_vecs = np.array([
    [0.9, 0.1, 0.4, 0.0],
    [0.2, 0.9, 0.1, 0.0],
    [0.8, 0.0, 0.6, 0.1],
    [0.0, 0.1, 0.0, 0.9],
], dtype=np.float32)
query = np.array([0.85, 0.05, 0.5, 0.0], dtype=np.float32)   # "how does retrieval work?"

# 1. normalise rows to length 1
doc_unit = doc_vecs / np.linalg.norm(doc_vecs, axis=1, keepdims=True)
query_unit = query / np.linalg.norm(query)

# 2. one matrix-vector product = cosine similarity with every doc
sims = doc_unit @ query_unit
print(sims.round(3))

# 3. top-k: sort descending and take the first k positions
k = 2
top = np.argsort(-sims)[:k]
for rank, i in enumerate(top, start=1):
    print(rank, docs[i], round(float(sims[i]), 3))

[0.992 0.289 0.988 0.006] 1 RAG basics 0.992 2 Vector databases 0.988

@ is matrix multiplication, and np.argsort(-sims) returns the indexes that would sort the similarities from highest to lowest. Swap the toy vectors for real embeddings from an API and this is a working retrieval step — fine for up to tens of thousands of documents before you need a dedicated vector index such as FAISS or a vector database.

Watch out

Shapes must line up for @: (4, 4) @ (4,) works, but mixing embeddings from different models — or different dimensions — produces errors or meaningless scores. Always embed queries and documents with the same model.

Saving and loading arrays

Python

import numpy as np

vecs = np.random.default_rng(0).random((3, 4)).astype(np.float32)
np.save("embeddings.npy", vecs)          # fast binary format

loaded = np.load("embeddings.npy")
print(loaded.shape, loaded.dtype, np.array_equal(vecs, loaded))

(3, 4) float32 True

Part 2 — pandas

Definition · pandas

A library for working with labelled, tabular data. It is built on NumPy and provides two main structures: Series and DataFrame.

Definition · Series

A one-dimensional labelled array — a single column of values plus an index that labels each row.

Definition · DataFrame

A two-dimensional table of rows and named columns, where each column is a Series and can have its own type. Think of a spreadsheet or SQL table that you control with code.

In AI projects, pandas is the tool for inspecting datasets before training, analysing evaluation results, comparing prompts or models, and exporting clean data to CSV or JSONL.

Creating a DataFrame

The most common way is from a list of dicts — the same records you learned to build in earlier lessons. Each dict becomes a row and each key becomes a column.

Python

import pandas as pd

rows = [
    {"query": "what is rag", "model": "small", "hit": True, "score": 0.81, "latency_ms": 420},
    {"query": "fine tuning", "model": "small", "hit": False, "score": 0.44, "latency_ms": 390},
    {"query": "what is rag", "model": "large", "hit": True, "score": 0.93, "latency_ms": 910},
    {"query": "fine tuning", "model": "large", "hit": True, "score": 0.77, "latency_ms": 880},
    {"query": "agents", "model": "small", "hit": True, "score": 0.69, "latency_ms": None},
]
df = pd.DataFrame(rows)

print(df)
print(df.shape)
print(list(df.columns))

query model hit score latency_ms 0 what is rag small True 0.81 420.0 1 fine tuning small False 0.44 390.0 2 what is rag large True 0.93 910.0 3 fine tuning large True 0.77 880.0 4 agents small True 0.69 NaN (5, 5) ['query', 'model', 'hit', 'score', 'latency_ms']

The numbers on the left (0–4) are the index. The missing latency became NaN (“Not a Number”), pandas' marker for missing numeric data, and the column turned into floats so it can hold it.

Inspecting data

Method	What it shows
`df.head(n)` / `df.tail(n)`	First / last `n` rows (default 5).
`df.shape`	`(rows, columns)`.
`df.info()`	Column names, types, and non-null counts — the first thing to run on new data.
`df.describe()`	Count, mean, std, min, quartiles, and max of numeric columns.
`df["col"].value_counts()`	How often each value appears.
`df.isna().sum()`	Number of missing values per column.

Python

print(df[["score", "latency_ms"]].describe().round(2))
print(df["model"].value_counts())
print(df.isna().sum()["latency_ms"])

score latency_ms count 5.00 4.00 mean 0.73 650.00 std 0.18 283.43 min 0.44 390.00 25% 0.69 412.50 50% 0.77 650.00 75% 0.81 887.50 max 0.93 910.00 model small 3 large 2 Name: count, dtype: int64 1

Selecting columns and rows

Python

print(df["score"].tolist())                 # one column -> Series
print(df[["query", "score"]].head(2))       # several columns -> DataFrame
print(df.loc[2, "query"])                   # by index label and column name
print(df.iloc[0, 3])                        # by integer position

[0.81, 0.44, 0.93, 0.77, 0.69] query score 0 what is rag 0.81 1 fine tuning 0.44 what is rag 0.81

Accessor	Selects by	Example
`df["col"]`	Column name	`df["score"]`
`df.loc[rows, cols]`	Labels (index values and column names)	`df.loc[df["hit"], "query"]`
`df.iloc[rows, cols]`	Integer positions	`df.iloc[:3, :2]`

Filtering rows

Filtering uses boolean masks, just like NumPy. Combine conditions with & (and), | (or), and ~ (not), and wrap each condition in parentheses.

Python

hits = df[df["hit"]]
print(hits[["query", "model", "score"]])

strong_small = df[(df["model"] == "small") & (df["score"] > 0.6)]
print(strong_small["query"].tolist())

print(df[df["query"].isin(["agents", "fine tuning"])].shape[0])

query model score 0 what is rag small 0.81 2 what is rag large 0.93 3 fine tuning large 0.77 4 agents small 0.69 ['what is rag', 'agents'] 3

Watch out

Use & and |, not and/or, and always parenthesise each condition. df[df.model == "small" and df.score > 0.6] raises “The truth value of a Series is ambiguous”.

Creating and transforming columns

Python

df["score_pct"] = (df["score"] * 100).round().astype(int)
df["grade"] = df["score"].apply(lambda s: "pass" if s >= 0.7 else "fail")
df["query_words"] = df["query"].str.split().str.len()
df["latency_ms"] = df["latency_ms"].fillna(df["latency_ms"].median())

print(df[["query", "score_pct", "grade", "query_words", "latency_ms"]])

query score_pct grade query_words latency_ms 0 what is rag 81 pass 3 420.0 1 fine tuning 44 fail 2 390.0 2 what is rag 93 pass 3 910.0 3 fine tuning 77 pass 2 880.0 4 agents 69 fail 1 650.0

Arithmetic on columns is vectorised, just like NumPy.
.apply(func) runs a Python function on each value — flexible, but slower; prefer vectorised operations for large data.
The .str accessor gives you vectorised string methods: .str.lower(), .str.contains(), .str.len().
fillna() replaces missing values; dropna() removes rows that contain them.

Sorting and grouping

Definition · groupby

Splits the rows into groups by the values of one or more columns, applies an aggregation (mean, count, sum…) to each group, and combines the results into a new table — the “split–apply–combine” pattern, equivalent to SQL's GROUP BY.

Python

print(df.sort_values("score", ascending=False)[["query", "model", "score"]].head(3))

summary = df.groupby("model").agg(
    runs=("query", "count"),
    hit_rate=("hit", "mean"),
    avg_score=("score", "mean"),
    p50_latency=("latency_ms", "median"),
).round(3)
print(summary)

query model score 2 what is rag large 0.93 0 what is rag small 0.81 3 fine tuning large 0.77 runs hit_rate avg_score p50_latency model large 2 1.000 0.850 895.0 small 3 0.667 0.647 420.0

In two lines you now have a model comparison table: the large model hits more often and scores higher, but with roughly double the latency. This exact pattern — group evaluation runs by model or prompt version and aggregate — is how AI teams decide what to ship.

Reading and writing files

Format	Read	Write
CSV	`pd.read_csv("f.csv")`	`df.to_csv("f.csv", index=False)`
JSONL	`pd.read_json("f.jsonl", lines=True)`	`df.to_json("f.jsonl", orient="records", lines=True, force_ascii=False)`
Parquet	`pd.read_parquet("f.parquet")`	`df.to_parquet("f.parquet")` (needs `pyarrow`)
Excel	`pd.read_excel("f.xlsx")`	`df.to_excel("f.xlsx", index=False)` (needs `openpyxl`)

Python

df[["query", "model", "score"]].to_json("results.jsonl", orient="records", lines=True)
print(open("results.jsonl", encoding="utf-8").read().splitlines()[0])

back = pd.read_json("results.jsonl", lines=True)
print(back.shape)

{"query":"what is rag","model":"small","score":0.81} (5, 3)

Note

Pass index=False to to_csv, otherwise pandas writes the index as an unnamed extra column. For large datasets prefer Parquet: it is compressed, keeps column types, and loads far faster than CSV.

NumPy or pandas?

Use NumPy when…	Use pandas when…
Data is purely numeric and same-typed (embeddings, logits, pixel arrays).	Data has named columns of mixed types (text, numbers, booleans, dates).
You need maximum speed for maths: dot products, norms, matrix multiplication.	You need to filter, group, join, and summarise records.
Feeding data into PyTorch, scikit-learn, or FAISS.	Exploring a dataset, analysing eval results, exporting CSV/JSONL.

They work together: df["score"].to_numpy() gives you a NumPy array, and pd.DataFrame(array, columns=[...]) wraps an array in a table.

Try it

Create a 5×8 random matrix of “document embeddings” with np.random.default_rng(7) and a random query vector, then return the top-3 most similar rows by cosine similarity.
Build a DataFrame of 8 evaluation rows with columns prompt_version (v1/v2), correct (bool), and tokens (int).
Group by prompt_version to get accuracy and mean tokens, sort by accuracy, and save the summary to summary.csv.

Key takeaways

A NumPy ndarray is a fast, same-typed grid of numbers; check its shape and dtype constantly.
Vectorised operations and broadcasting replace Python loops and are dramatically faster.
Cosine similarity = dot product of normalised vectors; matrix @ query plus argsort is a complete top-k search.
A pandas DataFrame is a table of named, typed columns built easily from a list of dicts.
Filter with boolean masks (&, |, ~), add columns with vectorised maths, and summarise with groupby().agg().
Read and write CSV, JSONL, and Parquet directly with pandas.