english_script_model.py

import gc
import os

import pandas as pd
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizer
from utils import update_tokenizer_vocab

MODEL_ID = "swiss-ai/Apertus-8B-Instruct-2509"
SCRIPT_FILE = "token_unicode_scripts.csv"
DEVICE = "cuda:0"
OUTPUT_DIR = "./models/apertus-8b-pruned-latin"

ENGLISH_SCRIPTS = {"Common", "Inherited", "Latin"}


def create_filtered_vocab_model():
    print(f"1. Loading Script Map from {SCRIPT_FILE}...")
    df = pd.read_csv(SCRIPT_FILE)

    if "Unnamed: 0" in df.columns:
        df = df.rename(
            columns={"Unnamed: 0": "token_id", "scripts_common_combined": "script"}
        )

    id_to_script = pd.Series(df.script.values, index=df.token_id).to_dict()

    print("2. Loading Tokenizer...")
    tokenizer: PreTrainedTokenizer = AutoTokenizer.from_pretrained(MODEL_ID, fast=True)
    print(f"The tokenizer has the following class: {tokenizer.__class__}")

    print(f" Active Scripts: {ENGLISH_SCRIPTS}")

    # --- BUILD THE KEEP LIST ---
    keep_indices = df[df["script"].isin(ENGLISH_SCRIPTS)]["token_id"].tolist()

    min_csv_id = df["token_id"].min()
    if min_csv_id > 0:
        print(
            f" Note: CSV starts at {min_csv_id}. Auto-adding IDs 0-{min_csv_id-1} (Special Tokens)."
        )
        keep_indices.extend(range(0, min_csv_id))

    keep_indices = sorted(list(set(keep_indices)))
    new_vocab_size = len(keep_indices)

    print(f" Reduced Vocab Size: {new_vocab_size}")

    index_map = torch.tensor(keep_indices, device=DEVICE)

    # --- BRAIN SURGERY ---
    print(f"\n3. Loading Model: {MODEL_ID}...")
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        dtype=torch.bfloat16,
        device_map=DEVICE,
        trust_remote_code=True,
    )

    # MEASURE STARTING MEMORY
    torch.cuda.empty_cache()
    torch.cuda.synchronize()
    start_mem = torch.cuda.memory_allocated(DEVICE) / (1024**3)

    print("4. Slicing the model...")
    original_head = model.lm_head
    hidden_size = original_head.in_features
    full_weights = original_head.weight.data
    original_input_embeddings = model.model.embed_tokens

    reduced_weights = full_weights[keep_indices, :]
    reduced_embed_weights = original_input_embeddings.weight.data[keep_indices, :]

    del original_head
    del full_weights
    del original_input_embeddings
    del model.lm_head
    del model.model.embed_tokens
    gc.collect()
    torch.cuda.empty_cache()

    # Reset memory stats to see current usage clearly
    torch.cuda.reset_peak_memory_stats(DEVICE)

    # Re-assign the sliced layers
    model.lm_head = torch.nn.Linear(
        hidden_size, new_vocab_size, bias=False, device=DEVICE, dtype=torch.bfloat16
    )
    model.lm_head.weight.data = reduced_weights
    model.config.vocab_size = new_vocab_size

    model.model.embed_tokens = torch.nn.Embedding(
        new_vocab_size, hidden_size, device=DEVICE, dtype=torch.bfloat16
    )
    model.model.embed_tokens.weight.data = reduced_embed_weights

    # Edit tokenizer

    # 3. Update the Python wrapper's config

    original_vocab = tokenizer.get_vocab()
    old_id_to_new_id = {old_id: new_id for new_id, old_id in enumerate(keep_indices)}

    new_vocab = {}
    for token, old_id in original_vocab.items():
        if old_id in old_id_to_new_id:
            new_id = old_id_to_new_id[old_id]
            new_vocab[token] = new_id

    # MEASURE ENDING MEMORY
    torch.cuda.synchronize()
    end_mem = torch.cuda.memory_allocated(DEVICE) / (1024**3)

    savings_gb = start_mem - end_mem
    print(f" Success! RAM Saved: ~{savings_gb:.4f} GB")

    # --- SAVING FOR HUGGING FACE ---
    print(f"\n5. Saving to {OUTPUT_DIR}...")
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    model.save_pretrained(OUTPUT_DIR)
    tokenizer.save_pretrained(OUTPUT_DIR)
    torch.save(index_map, f"{OUTPUT_DIR}/index_map.pt")

    # Update tokenizer vocabulary in tokenizer.json

    update_tokenizer_vocab(
        model_path=OUTPUT_DIR,
        new_vocab=new_vocab,
    )

    print(" All done!")
    
    # Simple test
    print("\n6. Loading pruned model and tokenizer for testing...")

    test_model = AutoModelForCausalLM.from_pretrained(
        OUTPUT_DIR,
        dtype=torch.bfloat16,
        device_map=DEVICE,
        trust_remote_code=True,
    )
    test_tokenizer: PreTrainedTokenizer = AutoTokenizer.from_pretrained(
        OUTPUT_DIR,
        fast=True,
    )
    test_input = "Hello, how are you?"
    inputs = test_tokenizer(test_input, return_tensors="pt").to(DEVICE)
    with torch.no_grad():
        outputs = test_model.generate(**inputs, max_new_tokens=20)
    print(" Test successful! Model output obtained.")
    print("Model response:")
    print(test_tokenizer.decode(outputs[0], skip_special_tokens=True))

if __name__ == "__main__":
    create_filtered_vocab_model()