# GENAI Experiment 3: Text generation: tokenization, embedding, pre-train + fine-tune, sampling (GPT-2 idea)
# ===========================================================================

# # Part 1 — Dataset (typed in code) and tokenization

import numpy as np, torch, torch.nn as nn, torch.nn.functional as F, math
import matplotlib.pyplot as plt
torch.manual_seed(0); np.random.seed(0)

general_text = """the sun rises in the east and sets in the west. people walk to work in the morning.
children play in the park after school. the weather is cool and pleasant today.
we read books to learn new things. the market is busy in the evening. birds sing in the trees.
my friend likes to cook food at home. the train arrives at the station on time. """ * 6

ai_text = """artificial intelligence is the science of making machines intelligent.
machine learning is a branch of artificial intelligence that learns patterns from data.
deep learning uses neural networks with many layers. a neural network learns by updating its weights.
generative models create new data such as text, images and music.
a language model predicts the next token in a sentence. transformers use attention to understand context.
training a model means reducing the loss on the training data. """ * 6

chars = sorted(set(general_text + ai_text))
stoi = {c: i for i, c in enumerate(chars)}; itos = {i: c for c, i in stoi.items()}
encode = lambda s: [stoi[c] for c in s]
decode = lambda ids: "".join(itos[i] for i in ids)
V = len(chars)
print("vocabulary size:", V)
print("TOKENIZATION  'neural' ->", encode("neural"), "->", decode(encode("neural")))

# Embedding is a separate step: it turns token ids into learned vectors.

emb = nn.Embedding(V, 32)
ids = torch.tensor(encode("neural"))
print("EMBEDDING  ids", ids.tolist(), "-> vectors of shape", tuple(emb(ids).shape))

# # Part 2 — Language model (embedding → GRU → next-character scores)

class CharLM(nn.Module):
    def __init__(self, V, E=32, H=128):
        super().__init__()
        self.emb = nn.Embedding(V, E)
        self.rnn = nn.GRU(E, H, batch_first=True)
        self.out = nn.Linear(H, V)
    def forward(self, x, h=None):
        o, h = self.rnn(self.emb(x), h)
        return self.out(o), h

SEQ = 40
def batches(text, n=32):
    data = torch.tensor(encode(text))
    idx = torch.randint(0, len(data) - SEQ - 1, (n,))
    x = torch.stack([data[i:i+SEQ] for i in idx]); y = torch.stack([data[i+1:i+SEQ+1] for i in idx])
    return x, y                                   # y is x shifted by one = causal LM target

def evaluate(model, text):
    model.eval()
    with torch.no_grad():
        data = torch.tensor(encode(text)); x, y = data[:-1].unsqueeze(0), data[1:].unsqueeze(0)
        loss = F.cross_entropy(model(x)[0].view(-1, V), y.view(-1)).item()
    model.train(); return loss, math.exp(loss)

def train(model, text, epochs, steps=60, lr=3e-3, tag=""):
    opt = torch.optim.Adam(model.parameters(), lr=lr); hist = []
    for ep in range(1, epochs + 1):
        for _ in range(steps):
            x, y = batches(text)
            logits, _ = model(x)
            loss = F.cross_entropy(logits.view(-1, V), y.view(-1))
            opt.zero_grad(); loss.backward(); opt.step()
        l, p = evaluate(model, ai_val)
        hist.append((l, p)); print(f"[{tag}] epoch {ep}  loss on AI val {l:.3f}  perplexity {p:.2f}")
    return hist

ai_train = ai_text
ai_val = """a machine learning model learns from data and makes predictions.
neural networks are used in deep learning to recognise images and text.
a generative model can create new images and music from data. """   # unseen sentences, same domain
model = CharLM(V)

# # Part 3 — Pre-train on general text, then fine-tune on AI text (transfer learning)

hist_pre = train(model, general_text, epochs=5, tag="pre-train general")
hist_ft = train(model, ai_train, epochs=6, lr=1e-3, tag="fine-tune AI")

losses = [h[0] for h in hist_pre + hist_ft]; ppl = [h[1] for h in hist_pre + hist_ft]
fig, ax = plt.subplots(1, 2, figsize=(11, 3.5))
ax[0].plot(range(1, len(losses) + 1), losses, "o-"); ax[0].set_title("Loss on AI validation text")
ax[1].plot(range(1, len(ppl) + 1), ppl, "o-", c="C1"); ax[1].set_title("Perplexity on AI validation text")
for a in ax: a.axvline(5.5, ls="--", c="gray"); a.set_xlabel("epoch (left of dashed line = pre-training)")
plt.tight_layout(); plt.show()

# # Part 4 — Generation with temperature, top-k and top-p (written by hand)

def sample(model, prompt, n=120, temperature=1.0, top_k=0, top_p=1.0):
    model.eval(); ids = encode(prompt); h = None
    x = torch.tensor([ids])
    with torch.no_grad():
        for _ in range(n):
            logits, h = model(x, h)
            logits = logits[0, -1] / temperature                      # temperature
            if top_k > 0:                                              # top-k: keep k best
                kth = torch.topk(logits, top_k).values[-1]
                logits[logits < kth] = -float("inf")
            probs = F.softmax(logits, -1)
            if top_p < 1.0:                                            # top-p: smallest set with sum >= p
                sp, si = torch.sort(probs, descending=True)
                keep = torch.cumsum(sp, 0) - sp < top_p
                probs = torch.zeros_like(probs).scatter(0, si[keep], sp[keep]); probs /= probs.sum()
            nxt = torch.multinomial(probs, 1).item()
            ids.append(nxt); x = torch.tensor([[nxt]])
    model.train(); return decode(ids)

torch.manual_seed(1)
prompt = "machine learning"
for name, kw in [("greedy-like T=0.2", dict(temperature=0.2)), ("T=0.8", dict(temperature=0.8)),
                 ("T=1.5", dict(temperature=1.5)), ("top-k=5", dict(top_k=5)), ("top-p=0.9", dict(top_p=0.9))]:
    print(f"[{name}] {sample(model, prompt, **kw)!r}\n")

# Exam explanation: during pre-training on general text the perplexity on AI text gets worse, because
# the model specialises in general sentences. Fine-tuning on AI text drops it sharply (≈151 → ≈7)
# within 2–3 epochs; later epochs raise it again, which is over-fitting on a tiny corpus. Low
# temperature is safe but repetitive, high temperature is creative but full of spelling mistakes.
# Top-k and top-p cut off unlikely characters. Fine-tuning on AI text lowers the perplexity on AI text
# compared with general pre-training alone, which is transfer learning.
#
# *Interactive prompt:* print(sample(model, input("prompt: "), temperature=0.8))
