Model Weights and Tokens: Numbers All the Way Down

# Install llama-cpp-python if not already present
try:
    from llama_cpp import Llama
except ImportError:
    %pip install llama-cpp-python
    from llama_cpp import Llama

# Install the gguf reader package
try:
    from gguf import GGUFReader
except ImportError:
    %pip install gguf
    from gguf import GGUFReader

import numpy as np
print("All packages loaded successfully!")

# ── Set these two variables to match your setup ─────────────────────────────
model_path = "/home/jovyan/shared/"           # directory containing .gguf files
model_name = "qwen2-1_5b-instruct-q4_0.gguf" # filename of the model
# ─────────────────────────────────────────────────────────────────────────────

# For local use, uncomment and adjust:
# model_path = "shared-rw/"
#model_path = "/Users/ericvandusen/SmallLM/Models/"

import os
full_model_path = os.path.join(model_path, model_name)
print(f"Looking for: {full_model_path}")
print(f"File exists: {os.path.exists(full_model_path)}")

# See what .gguf files are available
gguf_files = [f for f in os.listdir(model_path) if f.endswith(".gguf")]
print(f"Available .gguf models in {model_path}:")
for f in gguf_files:
    size_mb = os.path.getsize(os.path.join(model_path, f)) / (1024**2)
    print(f"  {f:60s}  {size_mb:7.1f} MB")

# Load the model
# verbose=False silences the C++ loading messages
model = Llama(
    model_path=full_model_path,
    n_ctx=2048,
    verbose=False,
    # n_threads=1, #CPU settings
    # n_gpu_layers=-1 #GPU Settings

)
print(f"Model loaded: {model_name}")
print(f"Vocabulary size: {model.n_vocab()} tokens")
print(f"Context window:  {model.n_ctx()} tokens")

# Economics sentences to tokenize
economics_sentences = [
    "Comparative advantage explains why countries specialize in producing goods they can make at a lower opportunity cost.",
    "When a central bank raises interest rates, borrowing becomes more expensive and inflation tends to fall.",
    "GDP measures the total monetary value of all finished goods and services produced within a country.",
    "A tariff is a tax on imported goods that raises their price for domestic consumers.",
    "Supply and demand determine the equilibrium price in a competitive market.",
]

# Tokenize each sentence and display the result
for sentence in economics_sentences:
    # encode the text to bytes first (llama_cpp expects bytes)
    token_ids = model.tokenize(sentence.encode("utf-8"))
    print(f"Text  : {sentence[:70]}..." if len(sentence) > 70 else f"Text  : {sentence}")
    print(f"Tokens: {token_ids}")
    print(f"Count : {len(token_ids)} tokens for {len(sentence)} characters  "
          f"(ratio: {len(sentence)/len(token_ids):.1f} chars/token)")
    print()

# Pick one sentence to examine closely
example_text = "A tariff is a tax on imported goods that raises their price for domestic consumers."

token_ids = model.tokenize(example_text.encode("utf-8"))

print("=" * 60)
print("TOKENS GOING IN")
print("=" * 60)
print(f"Input text: '{example_text}'")
print(f"\nToken IDs (the integers the model actually sees):")
print(token_ids)
print(f"\nTotal: {len(token_ids)} tokens")

# Now decode each token ID back to text to see what each integer represents
print(f"{'Token ID':>10}  {'Text piece':30}")
print("-" * 45)
for token_id in token_ids:
    # detokenize a single token
    piece_bytes = model.detokenize([token_id])
    piece_str = piece_bytes.decode("utf-8", errors="replace")
    print(f"{token_id:>10}  {repr(piece_str):30}")

# Run inference and capture token-level output
prompt = "Define comparative advantage in one sentence:"

print("=" * 60)
print("TOKENS GOING OUT")
print("=" * 60)
print(f"Prompt: '{prompt}'")
print()

# Use the low-level generate() method to get token IDs
prompt_tokens = model.tokenize(prompt.encode("utf-8"))
print(f"Prompt token IDs ({len(prompt_tokens)} tokens): {prompt_tokens}")
print()

# Generate up to 60 output tokens
output_tokens = []
output_text_pieces = []

for token_id in model.generate(prompt_tokens, top_k=1):  # top_k=1: greedy decoding — always pick the single highest-probability token
    output_tokens.append(token_id)
    piece = model.detokenize([token_id]).decode("utf-8", errors="replace")
    output_text_pieces.append(piece)
    
    # Stop at EOS or after 60 tokens
    if token_id == model.token_eos() or len(output_tokens) >= 60:
        break

print(f"Output token IDs ({len(output_tokens)} tokens):")
print(output_tokens)
print()
print("Decoded output:")
print("".join(output_text_pieces))

# Show output token-by-token in a table
print(f"{'Step':>6}  {'Token ID':>10}  {'Text piece'}")
print("-" * 40)
for i, (tid, piece) in enumerate(zip(output_tokens, output_text_pieces)):
    print(f"{i+1:>6}  {tid:>10}  {repr(piece)}")

# A longer economics passage to chunk
economics_passage = """
Comparative advantage is one of the most important concepts in international trade theory, 
first articulated by David Ricardo in 1817. The principle states that even if one country is 
more efficient at producing all goods than another, both countries can still benefit from 
specialization and trade. The key insight is opportunity cost: a country should specialize 
in producing goods where its relative efficiency advantage is greatest, or its relative 
inefficiency is smallest. For example, if the United States can produce both wheat and 
semiconductors more efficiently than Vietnam, but its advantage in semiconductors is 
proportionally larger, then the US should specialize in semiconductors and import wheat 
from Vietnam. Both countries end up with more of both goods than if each tried to produce 
everything domestically. This principle underpins the argument for free trade and explains 
much of the pattern of international specialization we observe in the global economy.
""".strip()

print(f"Passage length: {len(economics_passage)} characters")
print()
print(economics_passage)

def chunk_text(text, chunk_size=200, overlap=50):
    """
    Split text into overlapping chunks of roughly chunk_size characters.
    overlap: how many characters from the end of one chunk appear at the start of the next.
    """
    chunks = []
    start = 0
    while start < len(text):
        end = min(start + chunk_size, len(text))
        chunks.append(text[start:end])
        if end == len(text):
            break
        start += chunk_size - overlap
    return chunks

chunks = chunk_text(economics_passage, chunk_size=200, overlap=40)

for i, chunk in enumerate(chunks):
    print(f"── Chunk {i+1} ({len(chunk)} chars) ──────────────────────────")
    print(chunk)
    print()

for i, chunk in enumerate(chunks):
    token_ids = model.tokenize(chunk.encode("utf-8"))
    print(f"── Chunk {i+1} ────────────────────────────────────────")
    print(f"  Text  ({len(chunk):3d} chars):  {chunk[:80]}..." if len(chunk) > 80 else 
          f"  Text  ({len(chunk):3d} chars):  {chunk}")
    print(f"  Token IDs ({len(token_ids):2d} tokens): {token_ids}")
    print(f"  Chars-per-token ratio: {len(chunk)/len(token_ids):.2f}")
    print()

# Side-by-side comparison: raw chunk vs tokenized chunk (first chunk only)
chunk = chunks[0]
token_ids = model.tokenize(chunk.encode("utf-8"))

print("RAW TEXT CHUNK:")
print("-" * 60)
print(chunk)
print()

print("TOKENIZED CHUNK (token IDs):")
print("-" * 60)
print(token_ids)
print()

print("TOKEN → TEXT MAPPING (first 20 tokens):")
print("-" * 60)
print(f"{'ID':>8}  {'Text piece'}")
for tid in token_ids[:20]:
    piece = model.detokenize([tid]).decode("utf-8", errors="replace")
    print(f"{tid:>8}  {repr(piece)}")
if len(token_ids) > 20:
    print(f"    ... ({len(token_ids)-20} more tokens)")

from gguf import GGUFReader

# Open the GGUF file without loading it into GPU/CPU as a model
reader = GGUFReader(full_model_path)

print(f"GGUF file: {full_model_path}")
print(f"File size: {os.path.getsize(full_model_path) / (1024**2):.1f} MB")
print()
print(f"Number of metadata key-value pairs: {len(reader.fields)}")
print(f"Number of weight tensors:           {len(reader.tensors)}")

    from gguf import GGUFValueType
    # Print all metadata fields
    print("=" * 70)
    print("GGUF METADATA (key-value pairs)")
    print("=" * 70)

    for name, field in reader.fields.items():
        # field.parts contains the raw data; field.data gives usable values
        try:
            # field.parts[field.data[0]] gets the actual value
            if field.types[0] == GGUFValueType.STRING:
                value = str(bytes(field.parts[-1]), encoding='utf-8')
            elif len(field.data) == 1:
                value = field.parts[field.data[0]][0]
            else:
                value = f"[array of {len(field.data)} items]"

            # Skip huge tokenizer arrays for now (we'll look at them separately)
            if "token" in name.lower() and isinstance(value, str) and value.startswith("["):
                value = "[vocabulary array — shown below]"
        
            print(f"  {name:<50} = {str(value)[:60]}")
        except Exception as e:
            value = f"<unreadable: {e}>"
        

        # Skip huge tokenizer arrays for now (we'll look at them separately)
        if "token" in name.lower() and isinstance(value, str) and value.startswith("["):
            value = "[vocabulary array — shown below]"
        
        print(f"  {name:<50} = {str(value)[:60]}")

# Highlight the architecture-defining fields
architecture_keys = [
    "general.architecture",
    "general.name",
    "general.parameter_count",
    "general.quantization_version",
]

# Collect keys that mention important dimensions
dimension_keywords = ["hidden", "head", "layer", "context", "embed", "ff", "feed",
                      "attention", "block", "n_head", "n_layer", "n_ctx", "n_embd"]

print("=" * 70)
print("KEY ARCHITECTURAL PARAMETERS")
print("=" * 70)

for name, field in reader.fields.items():
    if any(kw in name.lower() for kw in dimension_keywords) or name in architecture_keys:
        try:
            # field.parts[field.data[0]] gets the actual value
            if field.types[0] == GGUFValueType.STRING:
                value = str(bytes(field.parts[-1]), encoding='utf-8')
            elif len(field.data) == 1:
                value = field.parts[field.data[0]][0]
            else:
                value = f"[array of {len(field.data)} items]"
        except Exception:
            value = "<unreadable>"
        print(f"  {name:<50} = {value}")

# Find the tokenizer vocabulary field
vocab_field = None
for name, field in reader.fields.items():
    if "tokens" in name.lower() and "tokenizer" in name.lower():
        vocab_field = (name, field)
        break

if vocab_field is None:
    print("Vocabulary field not found — trying alternative names...")
    for name, field in reader.fields.items():
        if "vocab" in name.lower() or "token" in name.lower():
            print(f"  Found: {name} ({len(field.data)} entries)")
else:
    name, field = vocab_field
    print(f"Vocabulary field: '{name}'")
    print(f"Total vocabulary size: {len(field.data)} tokens")
    print()
    
    # Show first 20 tokens
    print("First 20 tokens (low IDs = most common in training data):")
    print(f"{'Token ID':>10}  {'Token text'}")
    print("-" * 35)
    for i in range(min(20, len(field.data))):
        token_bytes = field.data[i]
        if isinstance(token_bytes, (bytes, bytearray, memoryview)):
            token_text = bytes(token_bytes).decode("utf-8", errors="replace")
        else:
            token_text = str(token_bytes)
        print(f"{i:>10}  {repr(token_text)}")
    
    print()
    print("Last 10 tokens (high IDs = rare or special tokens):")
    print(f"{'Token ID':>10}  {'Token text'}")
    print("-" * 35)
    total = len(field.data)
    for i in range(total - 10, total):
        token_bytes = field.data[i]
        if isinstance(token_bytes, (bytes, bytearray, memoryview)):
            token_text = bytes(token_bytes).decode("utf-8", errors="replace")
        else:
            token_text = str(token_bytes)
        print(f"{i:>10}  {repr(token_text)}")

print("=" * 80)
print("WEIGHT TENSORS IN THE GGUF FILE")
print("=" * 80)
print(f"{'Tensor name':<50} {'Shape':<25} {'Type':<10}")
print("-" * 85)

for tensor in reader.tensors:
    shape_str = str(list(tensor.shape))
    print(f"{tensor.name:<50} {shape_str:<25} {str(tensor.tensor_type.name):<10}")

print()
print(f"Total: {len(reader.tensors)} tensors")

# Count tensor types to understand the model structure
from collections import Counter

# Categorize tensors by their role
categories = {
    "embedding": [],
    "attention_q": [],
    "attention_k": [],
    "attention_v": [],
    "attention_output": [],
    "feed_forward": [],
    "normalization": [],
    "output": [],
    "other": []
}

for tensor in reader.tensors:
    n = tensor.name
    if "token_embd" in n:
        categories["embedding"].append(n)
    elif "attn_q" in n:
        categories["attention_q"].append(n)
    elif "attn_k" in n:
        categories["attention_k"].append(n)
    elif "attn_v" in n:
        categories["attention_v"].append(n)
    elif "attn_output" in n or "attn_out" in n:
        categories["attention_output"].append(n)
    elif "ffn" in n:
        categories["feed_forward"].append(n)
    elif "norm" in n:
        categories["normalization"].append(n)
    elif n in ("output.weight", "output_norm.weight"):
        categories["output"].append(n)
    else:
        categories["other"].append(n)

print("Tensor breakdown by role:")
print(f"{'Role':<25} {'Count':>6}")
print("-" * 33)
for role, tensors in categories.items():
    if tensors:
        print(f"{role:<25} {len(tensors):>6}")
print("-" * 33)
print(f"{'TOTAL':<25} {len(reader.tensors):>6}")

# Find the token embedding tensor
embd_tensor = None
for tensor in reader.tensors:
    if tensor.name == "token_embd.weight":
        embd_tensor = tensor
        break

if embd_tensor is not None:
    print(f"Tensor name:  {embd_tensor.name}")
    print(f"Shape:        {list(embd_tensor.shape)}")
    print(f"  → {embd_tensor.shape[0]} tokens in vocabulary")
    print(f"  → {embd_tensor.shape[1]} floats per token embedding")
    print(f"Storage type: {embd_tensor.tensor_type.name}  (quantized to save space)")
    
    vocab_size   = embd_tensor.shape[0]
    embd_dim     = embd_tensor.shape[1]
    full_fp32_mb = vocab_size * embd_dim * 4 / (1024**2)
    print(f"\nIf stored as full float32: {full_fp32_mb:.1f} MB")
    print(f"(Quantization compresses this significantly)")
else:
    print("Embedding tensor not found — check tensor names above.")

# Read the raw quantized data from the embedding tensor
# GGUFReader stores the raw bytes; we dequantize to float32 for display
if embd_tensor is not None:
    # Convert to float32 numpy array
    raw_data = embd_tensor.data          # numpy array (possibly quantized uint8)
    
    print(f"Raw data shape (quantized storage): {raw_data.shape}")
    print(f"Raw data dtype: {raw_data.dtype}")
    print()
    print("First 20 raw quantized bytes:")
    print(raw_data.flat[:20])
    print()
    print("(These bytes encode groups of 32 float values in a compact quantized format.)")

# Attempt to dequantize using numpy (manual Q4_0 dequantization)
# Q4_0 format: each block of 32 weights = 1 fp16 scale + 16 bytes of 4-bit ints
# Block size in bytes: 2 (scale) + 16 (4-bit data) = 18 bytes per 32 weights

if embd_tensor is not None and str(embd_tensor.tensor_type.name) == "Q4_0":
    raw = embd_tensor.data   # shape: (n_bytes,) uint8
    
    # Each Q4_0 block = 18 bytes = 2 bytes (fp16 scale) + 16 bytes (4-bit * 32 weights)
    block_size_bytes = 18
    weights_per_block = 32
    n_blocks = len(raw) // block_size_bytes
    
    print(f"Dequantizing Q4_0 embedding matrix...")
    print(f"  Total blocks: {n_blocks}")
    print(f"  Total weights: {n_blocks * weights_per_block:,}")
    print()
    
    # Dequantize first few blocks so we can display the actual float values
    blocks_to_show = min(4, n_blocks)  # just first 4 blocks = 128 weights
    decoded_weights = []
    
    for b in range(blocks_to_show):
        block = raw[b * block_size_bytes : (b + 1) * block_size_bytes]
        # Extract fp16 scale (first 2 bytes)
        scale = np.frombuffer(block[:2], dtype=np.float16)[0].astype(np.float32)
        # Extract 4-bit quantized values (remaining 16 bytes = 32 values)
        quant_bytes = block[2:]
        lo = (quant_bytes & 0x0F).astype(np.int8)
        hi = (quant_bytes >> 4).astype(np.int8)
        # Values are stored as signed 4-bit: range -8..7
        nibbles = np.empty(32, dtype=np.int8)
        nibbles[0::2] = lo - 8
        nibbles[1::2] = hi - 8
        dequant = nibbles.astype(np.float32) * scale
        decoded_weights.append((scale, dequant))
    
    # Display block 0 (first 32 weights of the first token's embedding)
    scale0, weights0 = decoded_weights[0]
    print(f"First Q4_0 block of the embedding matrix:")
    print(f"  Scale factor: {scale0:.6f}")
    print(f"  32 dequantized float values:")
    print(np.round(weights0, 5))

elif embd_tensor is not None:
    ttype = str(embd_tensor.tensor_type.name)
    print(f"Tensor is stored as {ttype} — showing raw data sample:")
    print(embd_tensor.data[:64])

# Find a normalization tensor (usually stored as F32 or F16 — easy to read)
norm_tensor = None
for tensor in reader.tensors:
    if "norm" in tensor.name and "blk.0" in tensor.name:
        norm_tensor = tensor
        break

if norm_tensor is not None:
    print(f"Tensor: {norm_tensor.name}")
    print(f"Shape:  {list(norm_tensor.shape)}")
    print(f"Type:   {norm_tensor.tensor_type.name}")
    print()
    
    norm_data = norm_tensor.data
    print(f"Data dtype: {norm_data.dtype}")
    
    # Convert to float32 for display
    floats = norm_data.astype(np.float32)
    print(f"First 20 weight values:")
    print(np.round(floats[:20], 5))
    print()
    print(f"Statistics: min={floats.min():.5f}, max={floats.max():.5f}, "
          f"mean={floats.mean():.5f}, std={floats.std():.5f}")
    print()
    print("Note: normalization weights are usually close to 1.0 (they start at 1.0 during training)")
else:
    # Fallback: show any tensor with F32 data
    for tensor in reader.tensors:
        if "norm" in tensor.name:
            print(f"Found norm tensor: {tensor.name}")
            print(f"  Shape: {list(tensor.shape)}, Type: {tensor.tensor_type.name}")
            data = tensor.data.astype(np.float32)
            print(f"  First 10 values: {np.round(data[:10], 5)}")
            break

# Load the model in embedding mode to extract token vectors
# Note: we need embedding=True to get per-token vectors
embed_model = Llama(
    model_path=full_model_path,
    embedding=True,
    n_ctx=128,
    verbose=False
)
print("Embedding model ready.")

# Economics words to compare
economics_words = [
    "inflation",
    "deflation",
    "tariff",
    "trade",
    "GDP",
    "recession",
    "growth",
    "supply",
    "demand",
    "price",
]


def collapse_embedding(vec: np.ndarray) -> np.ndarray:
    """Return a single vector even if llama-cpp gives per-token rows."""
    if vec.ndim == 1:
        return vec
    if vec.ndim == 2:
        return vec.mean(axis=0)
    raise ValueError(f"Unexpected embedding shape: {vec.shape}")


# Get embeddings for each word
word_embeddings = {}
for word in economics_words:
    emb = embed_model.embed(word)
    emb_array = np.array(emb, dtype=np.float32)
    word_embeddings[word] = collapse_embedding(emb_array)

# Print the shape of one embedding
sample_word = economics_words[0]
print(f"Embedding for '{sample_word}':")
print(f"  Shape: {word_embeddings[sample_word].shape}")
print(f"  First 10 values: {np.round(word_embeddings[sample_word][:10], 4)}")
print(f"  Min: {word_embeddings[sample_word].min():.4f}")
print(f"  Max: {word_embeddings[sample_word].max():.4f}")
print(f"  Norm (length): {np.linalg.norm(word_embeddings[sample_word]):.4f}")

def cosine_similarity(v1, v2):
    """Cosine similarity between two vectors. Range: -1 (opposite) to +1 (same direction)."""
    n1, n2 = np.linalg.norm(v1), np.linalg.norm(v2)
    if n1 == 0 or n2 == 0:
        return 0.0
    return float(np.dot(v1, v2) / (n1 * n2))

# Build a similarity matrix
words = economics_words
n = len(words)
sim_matrix = np.zeros((n, n))

for i, w1 in enumerate(words):
    for j, w2 in enumerate(words):
        sim_matrix[i, j] = cosine_similarity(word_embeddings[w1], word_embeddings[w2])

# Print the similarity matrix
print("Cosine Similarity Matrix for Economics Words")
print("(1.0 = identical direction, 0.0 = unrelated, -1.0 = opposite)")
print()
print(f"{'':15}", end="")
for w in words:
    print(f"{w:>10}", end="")
print()
print("-" * (15 + 10 * n))
for i, w1 in enumerate(words):
    print(f"{w1:<15}", end="")
    for j in range(n):
        v = sim_matrix[i, j]
        print(f"{v:>10.3f}", end="")
    print()

# Find the most similar pairs (excluding self-similarity)
print("Most Similar Word Pairs:")
print("-" * 40)

pairs = []
for i in range(n):
    for j in range(i + 1, n):
        pairs.append((sim_matrix[i, j], words[i], words[j]))

pairs.sort(reverse=True)
for sim, w1, w2 in pairs[:8]:
    print(f"  {w1:12} ↔ {w2:12}  similarity = {sim:.4f}")

print()
print("Least Similar Word Pairs:")
print("-" * 40)
for sim, w1, w2 in pairs[-5:]:
    print(f"  {w1:12} ↔ {w2:12}  similarity = {sim:.4f}")

# Economics analogy: which word completes the relationship?
# Example: inflation is to price as recession is to ???

def find_closest(query_vec, candidates, word_embeddings, exclude=None):
    """Find the word in candidates whose embedding is closest to query_vec."""
    exclude = exclude or []
    best_sim, best_word = -np.inf, None
    for word in candidates:
        if word in exclude:
            continue
        sim = cosine_similarity(query_vec, word_embeddings[word])
        if sim > best_sim:
            best_sim, best_word = sim, word
    return best_word, best_sim

# Extended word list for analogy search
analogy_words = [
    "inflation", "deflation", "tariff", "trade", "GDP", "recession",
    "growth", "supply", "demand", "price", "interest", "bank",
    "export", "import", "tax", "subsidy", "unemployment", "wages"
]

# Get embeddings for the extended list
for word in analogy_words:
    if word not in word_embeddings:
        emb = embed_model.embed(word)
        word_embeddings[word] = np.array(emb, dtype=np.float32).flatten()

# Test: "export" - "trade" + "tax" ≈ ???  (export is to trade as tax is to ?)
analogies = [
    ("export",    "trade",    "tariff",   "import"),     # A:B as C:D
    ("inflation", "price",    "recession", "growth"),
    ("supply",    "producer", "demand",   "consumer"),
]

print("Embedding Arithmetic: A is to B as C is to ???")
print("Formula: vec(A) - vec(B) + vec(C) → find closest word")
print("-" * 60)

for A, B, C, expected in analogies:
    if all(w in word_embeddings for w in [A, B, C]):
        query = word_embeddings[A] - word_embeddings[B] + word_embeddings[C]
        result, sim = find_closest(query, analogy_words, word_embeddings, exclude=[A, B, C])
        print(f"  {A:12} - {B:12} + {C:12} ≈ {result:12}  (similarity: {sim:.3f})")
        print(f"     Expected: {expected}")
        print()
    else:
        missing = [w for w in [A, B, C] if w not in word_embeddings]
        print(f"  Skipped (missing embeddings for: {missing})")

input_text = "Lower interest rates stimulate investment and economic growth."

print("=" * 70)
print("STEP 1: RAW TEXT")
print("=" * 70)
print(f"  '{input_text}'")
print(f"  Length: {len(input_text)} characters")

print()
print("=" * 70)
print("STEP 2: TOKENIZATION (text → integers)")
print("=" * 70)
token_ids = model.tokenize(input_text.encode("utf-8"))
print(f"  Token IDs: {token_ids}")
print(f"  Count: {len(token_ids)} tokens")

print()
print("=" * 70)
print("STEP 3: TOKEN → TEXT MAPPING")
print("=" * 70)
pieces = [model.detokenize([tid]).decode("utf-8", errors="replace") for tid in token_ids]
print(f"  {'ID':>8}  {'Text piece'}")
for tid, piece in zip(token_ids, pieces):
    print(f"  {tid:>8}  {repr(piece)}")

print()
print("=" * 70)
print("STEP 4: EMBEDDING LOOKUP (each integer → vector of floats)")
print("=" * 70)
print("  (Getting embeddings for each token piece...)")
for tid, piece in zip(token_ids[:5], pieces[:5]):  # show first 5 to keep output manageable
    emb = embed_model.embed(piece.strip() or piece)
    emb_arr = np.array(emb, dtype=np.float32)
    print(f"  Token {tid:>6} {repr(piece):>20}  →  vector shape {emb_arr.shape}")
    print(f"    First 8 values: {np.round(emb_arr[:8], 4)}")
    print(f"    Norm: {np.linalg.norm(emb_arr):.4f}")
if len(token_ids) > 5:
    print(f"  ... ({len(token_ids) - 5} more tokens)")

print()
print("=" * 70)
print("STEP 5: TRANSFORMER LAYERS")
print("=" * 70)
n_layers = sum(1 for t in reader.tensors if "blk." in t.name and "attn_q.weight" in t.name)
print(f"  The model runs {n_layers} transformer blocks in sequence.")
print(f"  Each block applies: attention (Q/K/V) + feed-forward network + layer norm")
print(f"  This is {n_layers * 4}+ matrix multiplications per token per forward pass.")

print()
print("=" * 70)
print("STEP 6: GENERATE NEXT TOKEN")
print("=" * 70)
prompt_tokens = model.tokenize(input_text.encode("utf-8"))
next_tok = next(model.generate(prompt_tokens, top_k=1))
next_text = model.detokenize([next_tok]).decode("utf-8", errors="replace")
print(f"  Most probable next token ID: {next_tok}")
print(f"  Decoded: {repr(next_text)}")
print(f"  Continued sentence: '{input_text}{next_text}...'") 

# Clean up: close model handles to free memory
del model
del embed_model
del reader
print("Models and file handles released.")