# Code examples from: Build Your Own LLM, Lesson 4: Embeddings and Your First Neural Language Model
# https://www.atulvij.com/llm-course-4-embeddings-pytorch
# Review and adapt before running; examples use sample values.

# ---- Example 1 (python) ----
import torch
import torch.nn as nn
import torch.nn.functional as F

text = open("input.txt", encoding="utf-8").read()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
data = torch.tensor([stoi[c] for c in text])
V = len(chars)
n = int(0.9 * len(data))
train_data, val_data = data[:n], data[n:]

# ---- Example 2 (python) ----
class Bigram(nn.Module):
    def __init__(self, V):
        super().__init__()
        self.table = nn.Embedding(V, V)

    def forward(self, idx):
        return self.table(idx)                  # (batch, V) logits

def batch(split, size=256):
    d = train_data if split == "train" else val_data
    ix = torch.randint(0, len(d) - 1, (size,))
    return d[ix], d[ix + 1]                       # inputs, next characters

model = Bigram(V)
opt = torch.optim.AdamW(model.parameters(), lr=1e-2)
for step in range(3001):
    x, y = batch("train")
    loss = F.cross_entropy(model(x), y)
    opt.zero_grad()
    loss.backward()
    opt.step()
    if step % 500 == 0:
        print(step, round(loss.item(), 3))

# ---- Example 3 (python) ----
class EmbedModel(nn.Module):
    def __init__(self, V, d=32):
        super().__init__()
        self.emb = nn.Embedding(V, d)            # each token -> 32 numbers
        self.head = nn.Linear(d, V)              # 32 numbers -> V scores

    def forward(self, idx):
        return self.head(self.emb(idx))

# ---- Example 4 (python) ----
emb_model = EmbedModel(V)
# ... train it with the same loop as above (swap `model` for `emb_model`) ...

E = F.normalize(emb_model.emb.weight.detach(), dim=1)
def similar(c, k=5):
    sims = E @ E[stoi[c]]
    best = sims.argsort(descending=True)[1:k + 1]
    return [(itos[i.item()], round(sims[i].item(), 2)) for i in best]

print(similar("a"))   # often other vowels
print(similar("5"))   # often other digits

# ---- Example 5 (python) ----
@torch.no_grad()
def generate(model, start="T", n=300):
    idx = torch.tensor([stoi[start]])
    out = [start]
    for _ in range(n):
        probs = F.softmax(model(idx)[-1], dim=-1)
        idx = torch.multinomial(probs, 1)
        out.append(itos[idx.item()])
    return "".join(out)

print(generate(emb_model))