Vocabulary Extension: Llama-3 Tokenizer'a 8K TR Token Ekle (Embedding Init Stratejileri)

# === Llama-3 Tokenizer Extension Lab ===
from transformers import AutoTokenizer
from tokenizers import Tokenizer
 
# 1. Base tokenizer'ı yükle
base = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3.1-8B")
print(f"Base vocab size: {base.vocab_size}")  # 128256
 
# 2. TR-spesifik tokenizer'ı yükle (önceki dersten)
tr_tok = Tokenizer.from_file("tr_bpe_50k.json")
tr_vocab = tr_tok.get_vocab()  # {token_str: id}
 
# 3. Aday TR token'ları seç: Llama-3'te olmayan, TR-corpus'ta sık
tr_corpus_sample = open("/data/tr-corpus/oscar-tr.txt").read()[:50_000_000]  # 50MB sample
 
from collections import Counter
# Tokenize sample with TR tokenizer
tr_tokens_in_corpus = [t for t in tr_tok.encode(tr_corpus_sample).tokens]
tr_token_freq = Counter(tr_tokens_in_corpus)
 
# Base tokenizer'ın vocab'ında olmayanları filtre
base_vocab_set = set(base.get_vocab().keys())
new_candidates = []
for tok, freq in tr_token_freq.most_common():
    if tok not in base_vocab_set and freq >= 50:
        new_candidates.append((tok, freq))
    if len(new_candidates) >= 8000:
        break
 
print(f"Aday TR-spesifik token sayısı: {len(new_candidates)}")
 
# 4. Base tokenizer'a ekle
new_tokens = [tok for tok, _ in new_candidates]
n_added = base.add_tokens(new_tokens)
print(f"Eklenen: {n_added}, Yeni vocab: {len(base)}")  # 128256 + 8000 = 136256
 
# 5. Kaydet
base.save_pretrained("llama3-tr-tok-136k")

# === Embedding Init Stratejileri ===
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
 
ext_tok = AutoTokenizer.from_pretrained("llama3-tr-tok-136k")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3.1-8B",
                                              torch_dtype=torch.bfloat16, device_map="cuda")
 
# Mevcut embedding 128256 satır
old_emb = model.get_input_embeddings().weight.data
old_lm_head = model.get_output_embeddings().weight.data
old_vocab = 128256
new_vocab = len(ext_tok)
print(f"old emb shape: {old_emb.shape}, new vocab: {new_vocab}")
 
# 6. Embedding'i resize et
model.resize_token_embeddings(new_vocab)
new_emb = model.get_input_embeddings().weight.data
new_lm_head = model.get_output_embeddings().weight.data
 
# 7. STRATEJİ A — Mean-init (en basit)
mean_vec = old_emb.mean(dim=0, keepdim=True)             # [1, hidden]
new_emb[old_vocab:] = mean_vec.expand(new_vocab - old_vocab, -1)
new_lm_head[old_vocab:] = old_lm_head.mean(dim=0, keepdim=True).expand(new_vocab - old_vocab, -1)
 
# 7b. STRATEJİ B — Byte-decomposition init (en akıllı)
# Yeni TR token "şirket" → base tokenizer'da ["ş", "i", "r", "k", "e", "t"] olarak parçalanabilir
# Yeni embedding'i = ortalama(parça embedding'leri) yap
base = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3.1-8B")
for new_id, tok_str in enumerate(ext_tok.convert_ids_to_tokens(range(old_vocab, new_vocab))):
    # base tokenizer ile yeni token'ı parçala
    decomp_ids = base.encode(tok_str.replace("Ġ", " "), add_special_tokens=False)
    if decomp_ids:
        new_emb[old_vocab + new_id] = old_emb[decomp_ids].mean(dim=0)
        new_lm_head[old_vocab + new_id] = old_lm_head[decomp_ids].mean(dim=0)
 
# 7c. STRATEJİ C — SVD-init (research)
# Old embedding matrix'in low-rank yaklaşımıyla yeni vektör çıkar
U, S, V = torch.svd(old_emb.float())
# Yeni embedding'ler için random combination of top-k singular vectors
torch.manual_seed(42)
k = 32
for new_id in range(old_vocab, new_vocab):
    coeffs = torch.randn(k) * 0.02
    new_emb[new_id] = (U[:, :k] @ (S[:k] * coeffs)).bfloat16()
 
# 8. Save
model.save_pretrained("llama3-tr-emb-136k")
ext_tok.save_pretrained("llama3-tr-emb-136k")