apply_script_filter.py

import torch
import pandas as pd
import gc
from transformers import AutoModelForCausalLM, AutoTokenizer

# --- Configuration ---
MODEL_ID = "swiss-ai/Apertus-8B-Instruct-2509"
SCRIPT_FILE = "token_unicode_scripts.csv"
DEVICE = "cuda:0"

# Set this to "Greek", "Cyrillic", etc. if you want to force it.
# Set to None to rely purely on what is detected in the prompt.
FORCE_SCRIPT = None 

def apply_script_filter():
    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 = AutoTokenizer.from_pretrained(MODEL_ID)
    
    # --- DYNAMIC SCRIPT DETECTION ---
    prompt = "The future of AI is" 
    print(f"\nAnalyzing Prompt: '{prompt}'")
    
    prompt_ids = tokenizer.encode(prompt, add_special_tokens=False)
    
    active_scripts = {"Common", "Inherited"}
    
    detected_scripts = set()
    for tid in prompt_ids:
        if tid in id_to_script:
            script = id_to_script[tid]
            detected_scripts.add(script)
            active_scripts.add(script)
    
    if FORCE_SCRIPT:
        print(f"   Forcing Script: {FORCE_SCRIPT}")
        active_scripts.add(FORCE_SCRIPT)
        
    print(f"   Detected in Prompt: {detected_scripts}")
    print(f"   Final Active Scripts: {active_scripts}")
    
    # --- BUILD THE KEEP LIST ---
    keep_indices = df[df['script'].isin(active_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.synchronize()
    start_mem = torch.cuda.max_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
    
    reduced_weights = full_weights[keep_indices, :]
    
    del model.lm_head
    gc.collect()
    torch.cuda.empty_cache()
    
    # Reset memory stats to see current usage clearly
    torch.cuda.reset_peak_memory_stats(DEVICE)
    
    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
    
    # MEASURE ENDING MEMORY
    torch.cuda.synchronize()
    end_mem = torch.cuda.max_memory_allocated(DEVICE) / (1024**3)
    
    # The difference might be tricky because of caching, but comparing peak allocation is usually safe
    savings_gb = start_mem - end_mem
    print(f"   Success! RAM Saved: ~{savings_gb:.4f} GB")

    # --- VERIFICATION ---
    print("-" * 60)
    print("5. Verifying Generation")
    
    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
    curr_ids = inputs.input_ids
    
    for _ in range(20):
        with torch.no_grad():
            outputs = model(curr_ids)
            next_token_logits = outputs.logits[0, -1, :]
            
            new_id = torch.argmax(next_token_logits).item()
            original_id = index_map[new_id].item()
            
            word = tokenizer.decode([original_id])
            print(f"Gen: {word} (Map: {new_id} -> {original_id})")
            
            curr_ids = torch.cat([curr_ids, torch.tensor([[original_id]], device=DEVICE)], dim=1)
            if original_id == tokenizer.eos_token_id:
                break
    
    print(f"\nFull Output: {tokenizer.decode(curr_ids[0])}")

if __name__ == "__main__":
    apply_script_filter()