import random

text = "hello world"
vocab = sorted(set(text))

print(vocab)

char_to_id = {}

for token_id, character in enumerate(vocab):
    char_to_id[character] = token_id

print(char_to_id)

def encode(text):
    token_ids = []

    for character in text:
        token_ids.append(char_to_id[character])

    return token_ids

print(encode("hello"))

def decode(token_ids):
    characters = []

    for token_id in token_ids:
        characters.append(vocab[token_id])

    return "".join(characters)

print(decode([3, 2, 4, 4, 5]))

original = "hello world"
reconstructed = decode(encode(original))

assert reconstructed == original
print("Round trip passed!")

tokens = encode("hello world")

inputs = tokens[:-1]
targets = tokens[1:]

for current_id, next_id in zip(inputs, targets):
    current_character = vocab[current_id]
    next_character = vocab[next_id]

    print(repr(current_character), "->", repr(next_character))

counts = {}

for current_id, next_id in zip(inputs, targets):
    if current_id not in counts:
        counts[current_id] = {}

    following = counts[current_id]

    if next_id not in following:
        following[next_id] = 0

    following[next_id] += 1

l_id = char_to_id["l"]

for next_id, count in counts[l_id].items():
    print(repr(vocab[next_id]), count)

probabilities = {}

for current_id, following in counts.items():
    total = sum(following.values())
    probabilities[current_id] = {}

    for next_id, count in following.items():
        probabilities[current_id][next_id] = count / total

for next_id, probability in probabilities[l_id].items():
    print(repr(vocab[next_id]), probability)

def generate(start_character, max_new_tokens):
    current_id = char_to_id[start_character]
    genereated_ids = [current_id]

    for _ in range(max_new_tokens):
        if current_id not in probabilities:
            break

        following = probabilities[current_id]

        next_id = random.choices(
            population=list(following.keys()),
            weights=list(following.values()),
            k=1,
        )[0]

        genereated_ids.append(next_id)
        current_id = next_id

    return decode(genereated_ids)
print(generate("h", 30))