from transformers import AutoTokenizer

# Choose a model name (e.g., 'bert-base-uncased', 'gpt2', 't5-small')
# model_name = "bert-base-uncased"
model_name = "gpt2"
# model_name = "google-bert/bert-base-uncased" # More explicit name

# Load the tokenizer
try:
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    print(f"Successfully loaded tokenizer for '{model_name}'")
except Exception as e:
    print(f"Error loading tokenizer for '{model_name}': {e}")
    exit()

# Get the entire vocabulary dictionary {token_string: token_id}
vocabulary = tokenizer.get_vocab()

# --- Exploring the Vocabulary ---

# 1. Get the total number of tokens (vocabulary size)
vocab_size = len(vocabulary)
# Alternatively, use the tokenizer's attribute:
vocab_size_attr = tokenizer.vocab_size
print(f"Vocabulary Size: {vocab_size} (from len(vocab))")
print(f"Vocabulary Size: {vocab_size_attr} (from tokenizer.vocab_size)")

# 2. Print the first few items (token: id)
print("\nFirst 10 items in the vocabulary:")
count = 0
for token, token_id in vocabulary.items():
    print(f"'{token}': {token_id}")
    count += 1
    if count >= 10:
        break

# 3. Get a list of all token strings
all_token_strings = list(vocabulary.keys())
print(f"\nFirst 10 token strings: {all_token_strings[:10]}")
print(f"Total number of token strings: {len(all_token_strings)}")

# 4. Get a list of all token IDs
all_token_ids = list(vocabulary.values())
print(f"\nFirst 10 token IDs: {all_token_ids[:10]}")
print(f"Total number of token IDs: {len(all_token_ids)}")
# Note: Token IDs usually range from 0 to vocab_size - 1

# 5. Look up the ID for a specific token (if it exists)
token_to_find = "hello"
if token_to_find in vocabulary:
    print(f"\nID for token '{token_to_find}': {vocabulary[token_to_find]}")
else:
    # For tokenizers like GPT-2/BPE, 'hello' might be tokenized differently (e.g., 'hell', 'o' or 'Ġhello')
    # Let's check for 'Ġhello' which is common in GPT-2/RoBERTa (Ġ indicates start of word/space before)
    gpt2_token = "Ġhello"
    if gpt2_token in vocabulary:
         print(f"\nID for token '{gpt2_token}': {vocabulary[gpt2_token]}")
    else:
        print(f"\nToken '{token_to_find}' (or variants) not found directly in vocabulary keys.")
        # You might need to tokenize the word first to see its components
        tokenized_output = tokenizer.tokenize(token_to_find)
        print(f"'{token_to_find}' tokenizes to: {tokenized_output}")
        for sub_token in tokenized_output:
             if sub_token in vocabulary:
                 print(f"  ID for sub_token '{sub_token}': {vocabulary[sub_token]}")


# 6. Look up the token string for a specific ID
id_to_find = 500
if id_to_find < vocab_size_attr:
    # You can use the tokenizer's conversion method directly
    token_string = tokenizer.convert_ids_to_tokens([id_to_find])[0]
    print(f"\nToken string for ID {id_to_find}: '{token_string}'")

    # Or manually create an inverse mapping (less efficient for single lookups)
    # id_to_token_map = {v: k for k, v in vocabulary.items()}
    # print(f"Token string for ID {id_to_find} (manual lookup): '{id_to_token_map.get(id_to_find, 'ID not found')}'")
else:
     print(f"\nID {id_to_find} is out of range for this vocabulary (size {vocab_size_attr}).")

# 7. Get Special Tokens
print("\nSpecial Tokens Map:", tokenizer.special_tokens_map)
print("All Special Tokens:", tokenizer.all_special_tokens)
print("All Special Token IDs:", tokenizer.all_special_ids)