← Lesson
Weighted Averages and Attention
Joeven
Run
Reset
Python loads on first run
import math def dot(u, v): return sum(a * b for a, b in zip(u, v)) def softmax(zs): m = max(zs) exps = [math.exp(z - m) for z in zs] z = sum(exps) return [e / z for e in exps] def attention(query, keys, values): scores = [dot(query, k) for k in keys] weights = softmax(scores) dim = len(values[0]) out = [0.0] * dim for w, val in zip(weights, values): for i, x in enumerate(val): out[i] += w * x return weights, out query = [1.0, 0.0] keys = [ [1.0, 0.0], [0.0, 1.0], ] values = [ [10.0, 0.0], [0.0, 10.0], ] w, out = attention(query, keys, values) print("weights", [round(x, 3) for x in w]) print("output ", [round(x, 3) for x in out]) query2 = [0.0, 1.0] w2, out2 = attention(query2, keys, values) print("weights", [round(x, 3) for x in w2]) print("output ", [round(x, 3) for x in out2])
Run to execute this in your browser. Nothing is sent to a server.