penkovsky_sa/llm-arch-research
1
0
Fork 0
mirror of https://github.com/pese-git/llm-arch-research.git synced 2026-01-23 13:00:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Рефакторинг: единообразие оформления кода (пробелы, кавычки, пустые строки), без изменения логики по всему проекту.

Browse Source
This commit is contained in:
Sergey Penkovsky
2025-10-06 22:57:19 +03:00
parent 332cad6159
commit 712278e33c
49 changed files with 2324 additions and 2004 deletions
Download Patch File Download Diff File Expand all files Collapse all files

209
experiments/hf_integration/generate_with_hf_tools.py
View File

@@ -14,54 +14,50 @@ sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from hf_proxy import HFAdapter, HFTokenizerAdapter, create_hf_pipeline
from shared.configs import (
TEST_PROMPTS, GENERATION_CONFIG, PATHS
)
from shared.data import (
print_experiment_info, ensure_directories, ExperimentLogger
)
from shared.configs import TEST_PROMPTS, GENERATION_CONFIG, PATHS
from shared.data import print_experiment_info, ensure_directories, ExperimentLogger
def load_hf_model_and_tokenizer() -> tuple:
"""
Загружает модель и токенизатор в формате HuggingFace.
Returns:
tuple: (hf_model, hf_tokenizer, model_config)
"""
# Используем упрощенную версию модели
model_path = "checkpoints/hf_simple_trained"
tokenizer_path = "checkpoints/hf_simple_tokenizer"
# Проверяем существование файлов
if not os.path.exists(model_path):
raise FileNotFoundError(
f"Модель не найдена: {model_path}\n"
f"Сначала обучите модель: uv run python experiments/hf_integration/simple_hf_training.py"
)
if not os.path.exists(tokenizer_path):
raise FileNotFoundError(
f"Токенизатор не найден: {tokenizer_path}"
)
raise FileNotFoundError(f"Токенизатор не найден: {tokenizer_path}")
# Загружаем адаптированный токенизатор
print("🔧 Загрузка адаптированного токенизатора...")
hf_tokenizer = HFTokenizerAdapter.from_pretrained(tokenizer_path)
print(f"✅ Токенизатор загружен (vocab_size={hf_tokenizer.vocab_size})")
# Загружаем конфигурацию модели
import json
config_path = os.path.join(model_path, "config.json")
with open(config_path, 'r', encoding='utf-8') as f:
with open(config_path, "r", encoding="utf-8") as f:
model_config = json.load(f)
# Загружаем модель через HFAdapter с правильной конфигурацией
print("🔧 Загрузка адаптированной модели...")
model_bin_path = os.path.join(model_path, "pytorch_model.bin")
# Создаем конфигурацию из сохраненного config.json
from hf_proxy import HFAdapterConfig
hf_config = HFAdapterConfig(
vocab_size=model_config["vocab_size"],
hidden_size=model_config["hidden_size"],
@@ -69,26 +65,28 @@ def load_hf_model_and_tokenizer() -> tuple:
num_attention_heads=model_config["num_attention_heads"],
max_position_embeddings=model_config["max_position_embeddings"],
hidden_dropout_prob=model_config.get("hidden_dropout_prob", 0.1),
attention_probs_dropout_prob=model_config.get("attention_probs_dropout_prob", 0.1),
attention_probs_dropout_prob=model_config.get(
"attention_probs_dropout_prob", 0.1
),
)
hf_model = HFAdapter.from_pretrained(model_bin_path, hf_config=hf_config)
hf_model.eval()
print("✅ Модель загружена")
return hf_model, hf_tokenizer, model_config
def test_hf_pipeline(hf_model, hf_tokenizer):
"""
Тестирует создание HuggingFace pipeline.
Args:
hf_model: Адаптированная модель
hf_tokenizer: Адаптированный токенизатор
"""
print("\n🧪 Тестирование HuggingFace pipeline...")
try:
# Создаем pipeline
pipe = create_hf_pipeline(
@@ -97,23 +95,23 @@ def test_hf_pipeline(hf_model, hf_tokenizer):
device="cpu",
max_length=50,
do_sample=True,
temperature=0.7
temperature=0.7,
)
print("✅ HuggingFace pipeline создан")
# Тестируем pipeline
test_prompts = TEST_PROMPTS[:3]
for prompt in test_prompts:
print(f"\n🔤 Промпт: '{prompt}'")
try:
result = pipe(prompt, max_new_tokens=20)
print(f"🎯 Результат: {result[0]['generated_text']}")
except Exception as e:
print(f"❌ Ошибка в pipeline: {e}")
except Exception as e:
print(f"❌ Ошибка создания pipeline: {e}")
@@ -121,47 +119,49 @@ def test_hf_pipeline(hf_model, hf_tokenizer):
def generate_with_hf_model(hf_model, hf_tokenizer, prompt: str, config: dict) -> str:
"""
Генерирует текст через адаптированную модель HF.
Args:
hf_model: Адаптированная модель
hf_tokenizer: Адаптированный токенизатор
prompt: Входной текст
config: Конфигурация генерации
Returns:
str: Сгенерированный текст
"""
print(f"🔤 Промпт: '{prompt}'")
print(f"📊 Параметры: max_tokens={config['max_new_tokens']}, "
f"temp={config['temperature']}, sample={config['do_sample']}")
print(
f"📊 Параметры: max_tokens={config['max_new_tokens']}, "
f"temp={config['temperature']}, sample={config['do_sample']}"
)
# Кодируем через адаптированный токенизатор
inputs = hf_tokenizer(prompt, return_tensors="pt")
print(f"🎯 Токены промпта: {inputs['input_ids'].tolist()[0]}")
print("🔄 Генерация через HF адаптер...")
# Генерируем через адаптированную модель
with torch.no_grad():
generated_ids = hf_model.generate(
input_ids=inputs['input_ids'],
input_ids=inputs["input_ids"],
max_new_tokens=config["max_new_tokens"],
do_sample=config["do_sample"],
temperature=config["temperature"],
top_k=config["top_k"],
top_p=config["top_p"]
top_p=config["top_p"],
)
# Декодируем через адаптированный токенизатор
generated_text = hf_tokenizer.decode(generated_ids[0], skip_special_tokens=True)
return generated_text
def test_different_hf_strategies(hf_model, hf_tokenizer, prompt: str):
"""
Тестирует разные стратегии генерации через HF интерфейс.
Args:
hf_model: Адаптированная модель
hf_tokenizer: Адаптированный токенизатор
@@ -169,32 +169,38 @@ def test_different_hf_strategies(hf_model, hf_tokenizer, prompt: str):
"""
print(f"\n🎭 Сравнение стратегий генерации через HF для промпта: '{prompt}'")
print("=" * 70)
strategies = [
{"name": "🎯 Жадный поиск", "do_sample": False, "temperature": 1.0},
{"name": "🎲 Вероятностная (temp=0.7)", "do_sample": True, "temperature": 0.7},
{"name": "🔥 Случайная (temp=1.2)", "do_sample": True, "temperature": 1.2},
{"name": "❄️ Детерминированная (temp=0.3)", "do_sample": True, "temperature": 0.3},
{
"name": "❄️ Детерминированная (temp=0.3)",
"do_sample": True,
"temperature": 0.3,
},
]
for strategy in strategies:
print(f"\n{strategy['name']}:")
try:
config = GENERATION_CONFIG.copy()
config.update({
"do_sample": strategy["do_sample"],
"temperature": strategy["temperature"],
"max_new_tokens": 20
})
config.update(
{
"do_sample": strategy["do_sample"],
"temperature": strategy["temperature"],
"max_new_tokens": 20,
}
)
generated = generate_with_hf_model(hf_model, hf_tokenizer, prompt, config)
# Выделяем сгенерированную часть
generated_part = generated[len(prompt):]
generated_part = generated[len(prompt) :]
print(f" 📤 Промпт: '{prompt}'")
print(f" 🎯 Сгенерировано: '{generated_part}'")
print(f" 📄 Полный текст: '{generated}'")
except Exception as e:
print(f" ❌ Ошибка: {e}")
@@ -202,30 +208,30 @@ def test_different_hf_strategies(hf_model, hf_tokenizer, prompt: str):
def analyze_hf_tokenization(hf_tokenizer, texts: list):
"""
Анализирует токенизацию через адаптированный токенизатор.
Args:
hf_tokenizer: Адаптированный токенизатор
texts: Список текстов для анализа
"""
print(f"\n🔍 Анализ токенизации через HF адаптер:")
print("=" * 60)
for i, text in enumerate(texts):
print(f"\nТекст {i+1}: '{text}'")
# Токенизация через адаптер
inputs = hf_tokenizer(text, return_tensors="pt")
tokens = inputs['input_ids'].tolist()[0]
tokens = inputs["input_ids"].tolist()[0]
token_strings = hf_tokenizer.tokenize(text)
print(f" Токены (ID): {tokens}")
print(f" Токены (текст): {token_strings}")
print(f" Количество токенов: {len(tokens)}")
# Декодирование обратно
decoded = hf_tokenizer.decode(tokens)
print(f" Декодированный: '{decoded}'")
if text == decoded:
print(f" ✅ Декодирование корректно")
else:
@@ -235,51 +241,55 @@ def analyze_hf_tokenization(hf_tokenizer, texts: list):
def interactive_hf_generation(hf_model, hf_tokenizer):
"""
Режим интерактивной генерации через HF интерфейс.
Args:
hf_model: Адаптированная модель
hf_tokenizer: Адаптированный токенизатор
"""
print(f"\n💬 Интерактивная генерация через HF (для выхода введите 'exit')")
print("-" * 60)
while True:
try:
user_input = input("\n🔤 Введите промпт: ").strip()
if user_input.lower() in ['exit', 'quit', 'выход']:
if user_input.lower() in ["exit", "quit", "выход"]:
break
if not user_input:
continue
# Запрашиваем параметры
try:
max_tokens = int(input("📏 Макс. токенов [50]: ") or "50")
temperature = float(input("🌡️ Температура [0.7]: ") or "0.7")
do_sample_input = input("🎲 Сэмплирование (y/n) [y]: ").lower()
do_sample = do_sample_input != 'n'
do_sample = do_sample_input != "n"
except:
max_tokens = 50
temperature = 0.7
do_sample = True
print("⚠️ Использую параметры по умолчанию")
config = GENERATION_CONFIG.copy()
config.update({
"max_new_tokens": max_tokens,
"temperature": temperature,
"do_sample": do_sample
})
generated = generate_with_hf_model(hf_model, hf_tokenizer, user_input, config)
generated_part = generated[len(user_input):]
config.update(
{
"max_new_tokens": max_tokens,
"temperature": temperature,
"do_sample": do_sample,
}
)
generated = generate_with_hf_model(
hf_model, hf_tokenizer, user_input, config
)
generated_part = generated[len(user_input) :]
print(f"\n🎯 Результат:")
print(f" 📤 Промпт: '{user_input}'")
print(f" 🎯 Сгенерировано: '{generated_part}'")
print(f" 📄 Полный текст: '{generated}'")
except KeyboardInterrupt:
print("\n👋 Завершение работы...")
break
@@ -295,76 +305,79 @@ def main():
"model": "GPT через HFAdapter",
"tokenizer": "BPE через HFTokenizerAdapter",
"инструменты": "HuggingFace pipeline & генерация",
"стратегия": "интеграция с HF экосистемой"
"стратегия": "интеграция с HF экосистемой",
}
print_experiment_info(experiment_name, experiment_config)
ensure_directories()
logger = ExperimentLogger(experiment_name)
try:
# Загружаем модель и токенизатор в HF формате
hf_model, hf_tokenizer, model_config = load_hf_model_and_tokenizer()
# === Анализ токенизации ===
analysis_texts = [
"Искусственный интеллект",
"Нейронные сети",
"Машинное обучение"
"Нейронные сети",
"Машинное обучение",
]
analyze_hf_tokenization(hf_tokenizer, analysis_texts)
# === Тестирование HF pipeline ===
test_hf_pipeline(hf_model, hf_tokenizer)
# === Генерация с разными промптами ===
print(f"\n🎯 Генерация текста через HF адаптер")
print("=" * 60)
for i, prompt in enumerate(TEST_PROMPTS):
print(f"\n📝 Пример {i+1}/{len(TEST_PROMPTS)}")
print("-" * 40)
try:
generated = generate_with_hf_model(hf_model, hf_tokenizer, prompt, GENERATION_CONFIG)
generated = generate_with_hf_model(
hf_model, hf_tokenizer, prompt, GENERATION_CONFIG
)
# Выделяем сгенерированную часть
generated_part = generated[len(prompt):]
generated_part = generated[len(prompt) :]
print(f"📤 Промпт: '{prompt}'")
print(f"🎯 Сгенерировано: '{generated_part}'")
print(f"📄 Полный текст: '{generated}'")
print(f"📏 Длина: {len(generated)} символов")
# Логируем успешную генерацию
logger.log_metric(f"hf_generation_length_{i}", len(generated))
except Exception as e:
print(f"❌ Ошибка при генерации: {e}")
continue
# === Сравнение стратегий генерации ===
test_prompt = "Искусственный"
test_different_hf_strategies(hf_model, hf_tokenizer, test_prompt)
# === Интерактивная генерация ===
interactive_hf_generation(hf_model, hf_tokenizer)
# === Сохранение результатов ===
logger.save_logs("checkpoints/hf_integration_generation_logs.json")
print(f"\n🎉 Эксперимент с HF интеграцией завершен успешно!")
print(f"\n📚 Достигнутая интеграция:")
print(f" ✅ Загрузка модели и токенизатора в HF формате")
print(f" ✅ Использование HF pipeline")
print(f" ✅ Генерация через стандартные HF интерфейсы")
print(f" ✅ Совместимость с HF экосистемой")
except FileNotFoundError as e:
print(f"❌ {e}")
except Exception as e:
print(f"❌ Ошибка в эксперименте: {e}")
import traceback
traceback.print_exc()

178
experiments/hf_integration/simple_hf_training.py
View File

@@ -19,141 +19,139 @@ from llm.tokenizers import BPETokenizer
from hf_proxy import HFAdapter, HFTokenizerAdapter
from shared.configs import (
TRAIN_TEXTS, BASE_GPT_CONFIG, BPE_CONFIG,
TRAINING_CONFIG, PATHS, TEST_PROMPTS
TRAIN_TEXTS,
BASE_GPT_CONFIG,
BPE_CONFIG,
TRAINING_CONFIG,
PATHS,
TEST_PROMPTS,
)
def create_dataset(hf_tokenizer, texts, max_length=128):
"""
Создает простой датасет для обучения.
Args:
hf_tokenizer: Адаптированный токенизатор
texts: Список текстов
max_length: Максимальная длина последовательности
Returns:
list: Список тензоров input_ids
"""
dataset = []
for text in texts:
# Токенизируем текст
inputs = hf_tokenizer(
text,
max_length=max_length,
text,
max_length=max_length,
truncation=True,
padding=False,
return_tensors="pt"
return_tensors="pt",
)
input_ids = inputs['input_ids'][0]
input_ids = inputs["input_ids"][0]
# Создаем метки для языкового моделирования
labels = input_ids.clone()
dataset.append({
'input_ids': input_ids,
'labels': labels
})
dataset.append({"input_ids": input_ids, "labels": labels})
return dataset
def manual_training_loop(hf_model, hf_tokenizer, train_texts, val_texts, config):
"""
Ручной цикл обучения без использования Trainer.
Args:
hf_model: Адаптированная модель
hf_tokenizer: Адаптированный токенизатор
train_texts: Тексты для обучения
val_texts: Тексты для валидации
config: Конфигурация обучения
Returns:
dict: Результаты обучения
"""
print("🎯 Запуск ручного обучения...")
# Создаем датасеты
train_dataset = create_dataset(hf_tokenizer, train_texts)
val_dataset = create_dataset(hf_tokenizer, val_texts)
print(f"📊 Данные: {len(train_dataset)} train, {len(val_dataset)} validation")
# Оптимизатор
optimizer = torch.optim.AdamW(
hf_model.parameters(),
lr=config["learning_rate"]
)
optimizer = torch.optim.AdamW(hf_model.parameters(), lr=config["learning_rate"])
# Функция потерь
loss_fn = nn.CrossEntropyLoss()
# Обучение
hf_model.train()
train_losses = []
val_losses = []
for epoch in range(config["num_epochs"]):
print(f"\n📅 Эпоха {epoch + 1}/{config['num_epochs']}")
# Обучение
epoch_train_loss = 0
for i, batch in enumerate(train_dataset):
optimizer.zero_grad()
input_ids = batch['input_ids'].unsqueeze(0) # [1, seq_len]
labels = batch['labels'].unsqueeze(0) # [1, seq_len]
input_ids = batch["input_ids"].unsqueeze(0) # [1, seq_len]
labels = batch["labels"].unsqueeze(0) # [1, seq_len]
# Forward pass
outputs = hf_model(input_ids=input_ids, labels=labels)
loss = outputs.loss
# Backward pass
loss.backward()
optimizer.step()
epoch_train_loss += loss.item()
if i % 5 == 0:
print(f" Batch {i}/{len(train_dataset)}: loss = {loss.item():.4f}")
avg_train_loss = epoch_train_loss / len(train_dataset)
train_losses.append(avg_train_loss)
print(f" 📊 Средняя train loss: {avg_train_loss:.4f}")
# Валидация
hf_model.eval()
epoch_val_loss = 0
with torch.no_grad():
for batch in val_dataset:
input_ids = batch['input_ids'].unsqueeze(0)
labels = batch['labels'].unsqueeze(0)
input_ids = batch["input_ids"].unsqueeze(0)
labels = batch["labels"].unsqueeze(0)
outputs = hf_model(input_ids=input_ids, labels=labels)
epoch_val_loss += outputs.loss.item()
avg_val_loss = epoch_val_loss / len(val_dataset)
val_losses.append(avg_val_loss)
print(f" 📊 Средняя val loss: {avg_val_loss:.4f}")
hf_model.train()
return {
'train_losses': train_losses,
'val_losses': val_losses,
'final_train_loss': train_losses[-1],
'final_val_loss': val_losses[-1]
"train_losses": train_losses,
"val_losses": val_losses,
"final_train_loss": train_losses[-1],
"final_val_loss": val_losses[-1],
}
def test_generation_after_training(hf_model, hf_tokenizer, test_prompts):
"""
Тестирует генерацию после обучения.
Args:
hf_model: Обученная модель
hf_tokenizer: Токенизатор
@@ -161,24 +159,24 @@ def test_generation_after_training(hf_model, hf_tokenizer, test_prompts):
"""
print("\n🧪 Тестирование генерации после обучения...")
hf_model.eval()
for prompt in test_prompts[:3]:
print(f"\n🔤 Промпт: '{prompt}'")
try:
inputs = hf_tokenizer(prompt, return_tensors="pt")
with torch.no_grad():
generated = hf_model.generate(
input_ids=inputs['input_ids'],
input_ids=inputs["input_ids"],
max_new_tokens=20,
do_sample=True,
temperature=0.8
temperature=0.8,
)
generated_text = hf_tokenizer.decode(generated[0], skip_special_tokens=True)
print(f"🎯 Результат: '{generated_text}'")
except Exception as e:
print(f"❌ Ошибка генерации: {e}")
@@ -188,96 +186,102 @@ def main():
print("=" * 60)
print("🚀 УПРОЩЕННОЕ ОБУЧЕНИЕ GPT С HF-PROXY")
print("=" * 60)
try:
# === Подготовка данных ===
print("🔧 Подготовка данных...")
train_texts = TRAIN_TEXTS[:10] # Используем меньше данных для быстрого тестирования
train_texts = TRAIN_TEXTS[
:10
] # Используем меньше данных для быстрого тестирования
val_texts = TRAIN_TEXTS[10:12]
print(f"📊 Данные: {len(train_texts)} train, {len(val_texts)} validation")
# === Подготовка токенизатора ===
print("🔧 Подготовка токенизатора...")
llm_tokenizer = BPETokenizer()
llm_tokenizer.train(
texts=train_texts,
vocab_size=BPE_CONFIG["vocab_size"],
special_tokens=BPE_CONFIG["special_tokens"]
special_tokens=BPE_CONFIG["special_tokens"],
)
hf_tokenizer = HFTokenizerAdapter(llm_tokenizer)
print(f"✅ Токенизатор создан (vocab_size={hf_tokenizer.vocab_size})")
# === Подготовка модели ===
print("🔧 Подготовка модели...")
model_config = BASE_GPT_CONFIG.copy()
model_config["vocab_size"] = hf_tokenizer.vocab_size
llm_model = GPT(model_config)
hf_model = HFAdapter.from_llm_model(llm_model)
print(f"✅ Модель создана")
# === Тестирование до обучения ===
print("\n🧪 Тестирование до обучения...")
test_generation_after_training(hf_model, hf_tokenizer, TEST_PROMPTS)
# === Обучение ===
print(f"\n🎯 Обучение модели...")
training_config = {
"learning_rate": TRAINING_CONFIG["learning_rate"],
"num_epochs": 2, # Меньше эпох для быстрого тестирования
"batch_size": TRAINING_CONFIG["batch_size"]
"batch_size": TRAINING_CONFIG["batch_size"],
}
results = manual_training_loop(
hf_model, hf_tokenizer, train_texts, val_texts, training_config
)
print(f"\n📊 Результаты обучения:")
print(f" Final train loss: {results['final_train_loss']:.4f}")
print(f" Final val loss: {results['final_val_loss']:.4f}")
# === Тестирование после обучения ===
print("\n🧪 Тестирование после обучения...")
test_generation_after_training(hf_model, hf_tokenizer, TEST_PROMPTS)
# === Сохранение модели ===
print(f"\n💾 Сохранение модели...")
# Создаем директории
os.makedirs("checkpoints/hf_simple_trained", exist_ok=True)
os.makedirs("checkpoints/hf_simple_tokenizer", exist_ok=True)
# Сохраняем токенизатор
hf_tokenizer.save_pretrained("checkpoints/hf_simple_tokenizer")
print("✅ Токенизатор сохранен")
# Сохраняем модель
HFAdapter.save_pretrained(
hf_model,
"checkpoints/hf_simple_trained",
tokenizer=hf_tokenizer
hf_model, "checkpoints/hf_simple_trained", tokenizer=hf_tokenizer
)
print("✅ Модель сохранена")
# Сохраняем результаты
results_path = "checkpoints/simple_training_results.json"
with open(results_path, 'w', encoding='utf-8') as f:
json.dump({
'training_config': training_config,
'model_config': model_config,
'results': results
}, f, indent=2, ensure_ascii=False)
with open(results_path, "w", encoding="utf-8") as f:
json.dump(
{
"training_config": training_config,
"model_config": model_config,
"results": results,
},
f,
indent=2,
ensure_ascii=False,
)
print(f"✅ Результаты сохранены в {results_path}")
print(f"\n🎉 Упрощенное обучение завершено успешно!")
print(f"\n💡 Для использования обученной модели:")
print(f" uv run python experiments/hf_integration/generate_with_hf_tools.py")
except Exception as e:
print(f"❌ Ошибка в эксперименте: {e}")
import traceback
traceback.print_exc()

98
experiments/hf_integration/test_hf_proxy.py
View File

@@ -16,158 +16,163 @@ from llm.tokenizers import BPETokenizer
from hf_proxy import HFAdapter, HFTokenizerAdapter
from shared.configs import (
TRAIN_TEXTS, BASE_GPT_CONFIG, BPE_CONFIG,
TEST_PROMPTS, GENERATION_CONFIG
TRAIN_TEXTS,
BASE_GPT_CONFIG,
BPE_CONFIG,
TEST_PROMPTS,
GENERATION_CONFIG,
)
def test_basic_hf_integration():
"""Тестирует базовую интеграцию hf-proxy."""
print("🧪 Тестирование базовой интеграции hf-proxy...")
# === Подготовка токенизатора ===
print("1. Подготовка токенизатора...")
llm_tokenizer = BPETokenizer()
llm_tokenizer.train(
texts=TRAIN_TEXTS,
vocab_size=BPE_CONFIG["vocab_size"],
special_tokens=BPE_CONFIG["special_tokens"]
special_tokens=BPE_CONFIG["special_tokens"],
)
hf_tokenizer = HFTokenizerAdapter(llm_tokenizer)
print(f" ✅ Токенизатор создан (vocab_size={hf_tokenizer.vocab_size})")
# === Подготовка модели ===
print("2. Подготовка модели...")
model_config = BASE_GPT_CONFIG.copy()
model_config["vocab_size"] = hf_tokenizer.vocab_size
llm_model = GPT(model_config)
hf_model = HFAdapter.from_llm_model(llm_model)
print(f" ✅ Модель создана")
# === Тестирование токенизации ===
print("3. Тестирование токенизации...")
test_texts = ["Искусственный интеллект", "Нейронные сети"]
for text in test_texts:
print(f" 📝 Текст: '{text}'")
# Оригинальный токенизатор
original_tokens = llm_tokenizer.encode(text)
print(f" Оригинальный: {len(original_tokens)} токенов")
# HF адаптер
hf_inputs = hf_tokenizer(text, return_tensors="pt")
print(f" HF адаптер: {hf_inputs['input_ids'].shape}")
# Декодирование
decoded = hf_tokenizer.decode(hf_inputs['input_ids'][0])
decoded = hf_tokenizer.decode(hf_inputs["input_ids"][0])
print(f" Декодированный: '{decoded}'")
# === Тестирование forward pass ===
print("4. Тестирование forward pass...")
for text in test_texts:
hf_inputs = hf_tokenizer(text, return_tensors="pt")
with torch.no_grad():
outputs = hf_model(**hf_inputs)
print(f" 📝 '{text}' -> logits: {outputs.logits.shape}")
# === Тестирование генерации ===
print("5. Тестирование генерации...")
hf_model.eval()
for prompt in TEST_PROMPTS[:3]:
print(f" 🔤 Промпт: '{prompt}'")
try:
inputs = hf_tokenizer(prompt, return_tensors="pt")
with torch.no_grad():
generated = hf_model.generate(
input_ids=inputs['input_ids'],
input_ids=inputs["input_ids"],
max_new_tokens=10,
do_sample=True,
temperature=0.8
temperature=0.8,
)
generated_text = hf_tokenizer.decode(generated[0], skip_special_tokens=True)
print(f" 🎯 Результат: '{generated_text}'")
except Exception as e:
print(f" ❌ Ошибка: {e}")
# === Тестирование сохранения/загрузки ===
print("6. Тестирование сохранения/загрузки...")
try:
# Сохраняем токенизатор
hf_tokenizer.save_pretrained("test_save/tokenizer")
print(" ✅ Токенизатор сохранен")
# Сохраняем модель
HFAdapter.save_pretrained(hf_model, "test_save/model", tokenizer=hf_tokenizer)
print(" ✅ Модель сохранена")
# Загружаем токенизатор
loaded_tokenizer = HFTokenizerAdapter.from_pretrained("test_save/tokenizer")
print(f" ✅ Токенизатор загружен (vocab_size={loaded_tokenizer.vocab_size})")
# Загружаем модель
model_path = os.path.join("test_save/model", "pytorch_model.bin")
loaded_model = HFAdapter.from_pretrained(model_path)
print(" ✅ Модель загружена")
# Проверяем работоспособность загруженной модели
test_input = hf_tokenizer("Тест", return_tensors="pt")
with torch.no_grad():
loaded_outputs = loaded_model(**test_input)
print(f" ✅ Загруженная модель работает (logits: {loaded_outputs.logits.shape})")
print(
f" ✅ Загруженная модель работает (logits: {loaded_outputs.logits.shape})"
)
except Exception as e:
print(f" ❌ Ошибка сохранения/загрузки: {e}")
print("\n🎉 Базовое тестирование hf-proxy завершено!")
def test_hf_tokenizer_methods():
"""Тестирует различные методы HF токенизатора."""
print("\n🧪 Тестирование методов HF токенизатора...")
# Создаем токенизатор
llm_tokenizer = BPETokenizer()
llm_tokenizer.train(
texts=TRAIN_TEXTS[:5],
vocab_size=500,
special_tokens=BPE_CONFIG["special_tokens"]
special_tokens=BPE_CONFIG["special_tokens"],
)
hf_tokenizer = HFTokenizerAdapter(llm_tokenizer)
test_text = "Искусственный интеллект и машинное обучение"
# Тестируем разные методы
print("1. Метод __call__:")
result = hf_tokenizer(test_text, return_tensors="pt")
print(f" Результат: {result}")
print("2. Метод encode:")
encoded = hf_tokenizer.encode(test_text)
print(f" Закодировано: {encoded}")
print("3. Метод decode:")
decoded = hf_tokenizer.decode(encoded)
print(f" Декодировано: '{decoded}'")
print("4. Метод tokenize:")
tokens = hf_tokenizer.tokenize(test_text)
print(f" Токены: {tokens}")
print("5. Метод get_vocab:")
vocab = hf_tokenizer.get_vocab()
print(f" Размер словаря: {len(vocab)}")
print("✅ Все методы токенизатора работают!")
@@ -176,14 +181,14 @@ def main():
print("=" * 60)
print("🧪 ТЕСТИРОВАНИЕ HF-PROXY")
print("=" * 60)
try:
# Тестируем базовую интеграцию
test_basic_hf_integration()
# Тестируем методы токенизатора
test_hf_tokenizer_methods()
print("\n" + "=" * 60)
print("🎉 ВСЕ ТЕСТЫ ПРОЙДЕНЫ УСПЕШНО!")
print("=" * 60)
@@ -195,10 +200,11 @@ def main():
print(" ✅ Генерация текста")
print(" ✅ Сохранение и загрузка моделей")
print(" ✅ Все методы HF токенизатора")
except Exception as e:
print(f"\n❌ Ошибка в тестировании: {e}")
import traceback
traceback.print_exc()

144
experiments/hf_integration/train_with_hf_trainer.py
View File

@@ -17,28 +17,34 @@ from llm.tokenizers import BPETokenizer
from hf_proxy import HFAdapter, HFTokenizerAdapter
from shared.configs import (
TRAIN_TEXTS, BASE_GPT_CONFIG, BPE_CONFIG,
TRAINING_CONFIG, PATHS, TEST_PROMPTS
TRAIN_TEXTS,
BASE_GPT_CONFIG,
BPE_CONFIG,
TRAINING_CONFIG,
PATHS,
TEST_PROMPTS,
)
from shared.data import (
load_training_data, ensure_directories,
print_experiment_info, ExperimentLogger
load_training_data,
ensure_directories,
print_experiment_info,
ExperimentLogger,
)
def setup_hf_training():
"""
Настраивает окружение для обучения через HuggingFace Trainer.
Returns:
tuple: (hf_model, hf_tokenizer, llm_tokenizer, model_config)
"""
print("🔧 Настройка HuggingFace обучения...")
# === Подготовка данных ===
train_texts, val_texts = load_training_data()
print(f"📊 Данные: {len(train_texts)} train, {len(val_texts)} validation")
# === Обучение/загрузка токенизатора ===
if os.path.exists(PATHS["bpe_tokenizer"]):
print("📝 Загрузка BPE токенизатора...")
@@ -50,55 +56,55 @@ def setup_hf_training():
llm_tokenizer.train(
texts=TRAIN_TEXTS,
vocab_size=BPE_CONFIG["vocab_size"],
special_tokens=BPE_CONFIG["special_tokens"]
special_tokens=BPE_CONFIG["special_tokens"],
)
llm_tokenizer.save(PATHS["bpe_tokenizer"])
print(f"✅ Токенизатор обучен и сохранен")
# === Создание адаптера токенизатора ===
print("🔧 Создание адаптера HuggingFace для токенизатора...")
hf_tokenizer = HFTokenizerAdapter(llm_tokenizer)
print(f"✅ Адаптер токенизатора создан")
# === Инициализация модели ===
model_config = BASE_GPT_CONFIG.copy()
model_config["vocab_size"] = llm_tokenizer.get_vocab_size()
print("🔧 Создание GPT модели...")
llm_model = GPT(model_config)
# === Создание адаптера модели ===
print("🔧 Создание адаптера HuggingFace для модели...")
hf_model = HFAdapter.from_llm_model(llm_model)
print(f"✅ Адаптер модели создан")
return hf_model, hf_tokenizer, llm_tokenizer, model_config, train_texts, val_texts
def test_hf_integration(hf_model, hf_tokenizer, llm_tokenizer):
"""
Тестирует интеграцию с HuggingFace инструментами.
Args:
hf_model: Адаптированная модель
hf_tokenizer: Адаптированный токенизатор
llm_tokenizer: Оригинальный токенизатор
"""
print("\n🧪 Тестирование интеграции с HuggingFace...")
test_texts = ["Искусственный интеллект", "Нейронные сети"]
for text in test_texts:
print(f"\n🔤 Текст: '{text}'")
# Тестируем адаптированный токенизатор
hf_inputs = hf_tokenizer(text, return_tensors="pt")
print(f" HF токенизатор: {hf_inputs['input_ids'].shape}")
# Тестируем оригинальный токенизатор для сравнения
original_tokens = llm_tokenizer.encode(text)
print(f" Оригинальный токенизатор: {len(original_tokens)} токенов")
# Тестируем forward pass через адаптированную модель
try:
with torch.no_grad():
@@ -114,28 +120,35 @@ def main():
experiment_name = "Обучение GPT через HF Trainer (с hf-proxy)"
experiment_config = {
"model": "GPT через HFAdapter",
"tokenizer": "BPE через HFTokenizerAdapter",
"tokenizer": "BPE через HFTokenizerAdapter",
"trainer": "HuggingFace Trainer",
"vocab_size": BPE_CONFIG["vocab_size"],
"training_epochs": TRAINING_CONFIG["num_epochs"]
"training_epochs": TRAINING_CONFIG["num_epochs"],
}
print_experiment_info(experiment_name, experiment_config)
ensure_directories()
logger = ExperimentLogger(experiment_name)
try:
# Настраиваем обучение
hf_model, hf_tokenizer, llm_tokenizer, model_config, train_texts, val_texts = setup_hf_training()
(
hf_model,
hf_tokenizer,
llm_tokenizer,
model_config,
train_texts,
val_texts,
) = setup_hf_training()
# Тестируем интеграцию
test_hf_integration(hf_model, hf_tokenizer, llm_tokenizer)
# === Подготовка датасетов HuggingFace ===
print(f"\n📊 Подготовка датасетов HuggingFace...")
from datasets import Dataset
def tokenize_function(examples):
"""Функция токенизации для HF datasets."""
# Используем адаптированный токенизатор
@@ -147,11 +160,11 @@ def main():
)
tokenized["labels"] = tokenized["input_ids"].copy()
return tokenized
# Создаем датасеты
train_dataset = Dataset.from_dict({"text": train_texts})
val_dataset = Dataset.from_dict({"text": val_texts})
# Токенизируем
train_dataset = train_dataset.map(
tokenize_function,
@@ -163,26 +176,26 @@ def main():
batched=True,
remove_columns=val_dataset.column_names,
)
print(f" Train датасет: {len(train_dataset)} примеров")
print(f" Validation датасет: {len(val_dataset)} примеров")
# === Настройка HuggingFace Trainer ===
print(f"\n🔧 Настройка HuggingFace Trainer...")
from transformers import (
Trainer,
Trainer,
TrainingArguments,
DataCollatorForLanguageModeling
DataCollatorForLanguageModeling,
)
# Data collator для языкового моделирования
data_collator = DataCollatorForLanguageModeling(
tokenizer=hf_tokenizer,
mlm=False,
pad_to_multiple_of=8,
)
# Аргументы обучения
training_args = TrainingArguments(
output_dir=PATHS["hf_model"],
@@ -204,7 +217,7 @@ def main():
dataloader_pin_memory=False,
report_to=None,
)
# Создаем Trainer
trainer = Trainer(
model=hf_model,
@@ -213,84 +226,87 @@ def main():
eval_dataset=val_dataset,
data_collator=data_collator,
)
print("✅ HuggingFace Trainer настроен")
# === Запуск обучения ===
print(f"\n🎯 Запуск обучения через HuggingFace Trainer...")
train_result = trainer.train()
# Сохраняем лучшую модель
trainer.save_model()
hf_tokenizer.save_pretrained(PATHS["hf_model"])
print("✅ Обучение завершено успешно!")
print(f"📊 Final train loss: {train_result.metrics['train_loss']:.4f}")
if "eval_loss" in train_result.metrics:
print(f"📊 Final eval loss: {train_result.metrics['eval_loss']:.4f}")
# === Сохранение через hf-proxy ===
print(f"\n💾 Сохранение через hf-proxy...")
from hf_proxy import convert_to_hf_format
# Сохраняем токенизатор в HF формате
hf_tokenizer_dir = PATHS["hf_tokenizer"]
hf_tokenizer.save_pretrained(hf_tokenizer_dir)
# Сохраняем модель через hf-proxy
hf_proxy_dir = PATHS["hf_proxy_model"]
HFAdapter.save_pretrained(hf_model, hf_proxy_dir, tokenizer=hf_tokenizer)
print(f"✅ Модель сохранена в HF формате:")
print(f" - {PATHS['hf_model']}: стандартный HF формат")
print(f" - {hf_proxy_dir}: через hf-proxy")
print(f" - {hf_tokenizer_dir}: токенизатор в HF формате")
# === Тестирование генерации ===
print(f"\n🧪 Тестирование генерации после обучения...")
hf_model.eval()
for prompt in TEST_PROMPTS[:3]:
print(f"\n🔤 Промпт: '{prompt}'")
try:
inputs = hf_tokenizer(prompt, return_tensors="pt")
with torch.no_grad():
generated = hf_model.generate(
input_ids=inputs['input_ids'],
input_ids=inputs["input_ids"],
max_new_tokens=20,
do_sample=True,
temperature=0.8
temperature=0.8,
)
generated_text = hf_tokenizer.decode(generated[0], skip_special_tokens=True)
generated_text = hf_tokenizer.decode(
generated[0], skip_special_tokens=True
)
print(f"🎯 Результат: '{generated_text}'")
except Exception as e:
print(f"❌ Ошибка генерации: {e}")
# === Сохранение результатов ===
results = {
"experiment": experiment_name,
"model_config": model_config,
"training_config": TRAINING_CONFIG,
"final_loss": train_result.metrics.get('train_loss', 'N/A'),
"eval_loss": train_result.metrics.get('eval_loss', 'N/A')
"final_loss": train_result.metrics.get("train_loss", "N/A"),
"eval_loss": train_result.metrics.get("eval_loss", "N/A"),
}
logger.save_logs("checkpoints/hf_integration_training_logs.json")
print(f"\n🎉 Эксперимент с HF интеграцией завершен успешно!")
print(f"\n💡 Для использования обученной модели:")
print(f" uv run python experiments/hf_integration/generate_with_hf_tools.py")
except Exception as e:
print(f"❌ Ошибка в эксперименте: {e}")
import traceback
traceback.print_exc()

171
experiments/llm_only/generate_gpt2_bpe.py
View File

@@ -15,18 +15,14 @@ sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from llm.models.gpt import GPT2
from llm.tokenizers import BPETokenizer
from shared.configs import (
BASE_GPT_CONFIG, TEST_PROMPTS, GENERATION_CONFIG, PATHS
)
from shared.data import (
print_experiment_info, ensure_directories, ExperimentLogger
)
from shared.configs import BASE_GPT_CONFIG, TEST_PROMPTS, GENERATION_CONFIG, PATHS
from shared.data import print_experiment_info, ensure_directories, ExperimentLogger
def load_model_and_tokenizer() -> tuple:
"""
Загружает обученную модель и токенизатор.
Returns:
tuple: (модель, токенизатор, конфигурация)
"""
@@ -36,62 +32,60 @@ def load_model_and_tokenizer() -> tuple:
f"Модель не найдена: {PATHS['gpt_bpe_model']}\n"
f"Сначала обучите модель: uv run python experiments/llm_only/train_gpt_bpe.py"
)
if not os.path.exists(PATHS["bpe_tokenizer"]):
raise FileNotFoundError(
f"Токенизатор не найден: {PATHS['bpe_tokenizer']}"
)
raise FileNotFoundError(f"Токенизатор не найден: {PATHS['bpe_tokenizer']}")
# Загружаем конфигурацию модели
import json
with open(PATHS["gpt_bpe_config"], 'r', encoding='utf-8') as f:
with open(PATHS["gpt_bpe_config"], "r", encoding="utf-8") as f:
model_config = json.load(f)
# Загружаем токенизатор
print("🔧 Загрузка BPE токенизатора...")
tokenizer = BPETokenizer.load(PATHS["bpe_tokenizer"])
print(f"✅ Токенизатор загружен (vocab_size={tokenizer.get_vocab_size()})")
# Загружаем модель
print("🔧 Загрузка GPT2 модели...")
model = GPT2(model_config)
model.load_state_dict(torch.load(PATHS["gpt_bpe_model"], map_location='cpu'))
model.load_state_dict(torch.load(PATHS["gpt_bpe_model"], map_location="cpu"))
model.eval()
print("✅ Модель загружена")
return model, tokenizer, model_config
def generate_text(
model: GPT2,
tokenizer: BPETokenizer,
prompt: str,
config: dict
model: GPT2, tokenizer: BPETokenizer, prompt: str, config: dict
) -> str:
"""
Генерирует текст на основе промпта.
Args:
model: Обученная GPT модель
tokenizer: BPE токенизатор
prompt: Входной текст
config: Конфигурация генерации
Returns:
str: Сгенерированный текст
"""
print(f"🔤 Промпт: '{prompt}'")
print(f"📊 Параметры: max_tokens={config['max_new_tokens']}, "
f"temp={config['temperature']}, sample={config['do_sample']}")
print(
f"📊 Параметры: max_tokens={config['max_new_tokens']}, "
f"temp={config['temperature']}, sample={config['do_sample']}"
)
# Кодируем промпт
input_ids = tokenizer.encode(prompt, add_special_tokens=False)
input_tensor = torch.tensor([input_ids], dtype=torch.long)
print(f"🎯 Токены промпта: {input_ids}")
print(f"🎯 Токены (текст): {tokenizer.tokenize(prompt)}")
print("🔄 Генерация...")
# Генерируем текст
with torch.no_grad():
generated_ids = model.generate(
@@ -100,19 +94,19 @@ def generate_text(
do_sample=config["do_sample"],
temperature=config["temperature"],
top_k=config["top_k"],
top_p=config["top_p"]
top_p=config["top_p"],
)
# Декодируем результат
generated_text = tokenizer.decode(generated_ids[0].tolist())
return generated_text
def test_different_strategies(model: GPT2, tokenizer: BPETokenizer, prompt: str):
"""
Тестирует разные стратегии генерации на одном промпте.
Args:
model: Обученная модель
tokenizer: BPE токенизатор
@@ -120,32 +114,38 @@ def test_different_strategies(model: GPT2, tokenizer: BPETokenizer, prompt: str)
"""
print(f"\n🎭 Сравнение стратегий генерации для промпта: '{prompt}'")
print("=" * 60)
strategies = [
{"name": "🎯 Жадный поиск", "do_sample": False, "temperature": 1.0},
{"name": "🎲 Вероятностная (temp=0.7)", "do_sample": True, "temperature": 0.7},
{"name": "🔥 Случайная (temp=1.2)", "do_sample": True, "temperature": 1.2},
{"name": "❄️ Детерминированная (temp=0.3)", "do_sample": True, "temperature": 0.3},
{
"name": "❄️ Детерминированная (temp=0.3)",
"do_sample": True,
"temperature": 0.3,
},
]
for strategy in strategies:
print(f"\n{strategy['name']}:")
try:
config = GENERATION_CONFIG.copy()
config.update({
"do_sample": strategy["do_sample"],
"temperature": strategy["temperature"],
"max_new_tokens": 20
})
config.update(
{
"do_sample": strategy["do_sample"],
"temperature": strategy["temperature"],
"max_new_tokens": 20,
}
)
generated = generate_text(model, tokenizer, prompt, config)
# Выделяем сгенерированную часть
generated_part = generated[len(prompt):]
generated_part = generated[len(prompt) :]
print(f" 📤 Промпт: '{prompt}'")
print(f" 🎯 Сгенерировано: '{generated_part}'")
print(f" 📄 Полный текст: '{generated}'")
except Exception as e:
print(f" ❌ Ошибка: {e}")
@@ -153,26 +153,26 @@ def test_different_strategies(model: GPT2, tokenizer: BPETokenizer, prompt: str)
def analyze_tokenization(tokenizer: BPETokenizer, texts: list):
"""
Анализирует токенизацию различных текстов.
Args:
tokenizer: BPE токенизатор
texts: Список текстов для анализа
"""
print(f"\n🔍 Анализ токенизации BPE:")
print("=" * 50)
for i, text in enumerate(texts):
print(f"\nТекст {i+1}: '{text}'")
# Токенизация
tokens = tokenizer.encode(text, add_special_tokens=False)
token_strings = tokenizer.tokenize(text)
print(f" Токены (ID): {tokens}")
print(f" Токены (текст): {token_strings}")
print(f" Количество токенов: {len(tokens)}")
print(f" Эффективность: {len(text)} символов → {len(tokens)} токенов")
# Декодирование обратно
decoded = tokenizer.decode(tokens)
if text == decoded:
@@ -184,51 +184,53 @@ def analyze_tokenization(tokenizer: BPETokenizer, texts: list):
def interactive_generation(model: GPT2, tokenizer: BPETokenizer):
"""
Режим интерактивной генерации.
Args:
model: Обученная модель
tokenizer: BPE токенизатор
"""
print(f"\n💬 Интерактивная генерация (для выхода введите 'exit')")
print("-" * 50)
while True:
try:
user_input = input("\n🔤 Введите промпт: ").strip()
if user_input.lower() in ['exit', 'quit', 'выход']:
if user_input.lower() in ["exit", "quit", "выход"]:
break
if not user_input:
continue
# Запрашиваем параметры
try:
max_tokens = int(input("📏 Макс. токенов [50]: ") or "50")
temperature = float(input("🌡️ Температура [0.7]: ") or "0.7")
do_sample_input = input("🎲 Сэмплирование (y/n) [y]: ").lower()
do_sample = do_sample_input != 'n'
do_sample = do_sample_input != "n"
except:
max_tokens = 50
temperature = 0.7
do_sample = True
print("⚠️ Использую параметры по умолчанию")
config = GENERATION_CONFIG.copy()
config.update({
"max_new_tokens": max_tokens,
"temperature": temperature,
"do_sample": do_sample
})
config.update(
{
"max_new_tokens": max_tokens,
"temperature": temperature,
"do_sample": do_sample,
}
)
generated = generate_text(model, tokenizer, user_input, config)
generated_part = generated[len(user_input):]
generated_part = generated[len(user_input) :]
print(f"\n🎯 Результат:")
print(f" 📤 Промпт: '{user_input}'")
print(f" 🎯 Сгенерировано: '{generated_part}'")
print(f" 📄 Полный текст: '{generated}'")
except KeyboardInterrupt:
print("\n👋 Завершение работы...")
break
@@ -244,68 +246,69 @@ def main():
"model": "GPT2 с BPE токенизатором",
"стратегия": "автономная генерация",
"вход": "промпты",
"выход": "сгенерированный текст"
"выход": "сгенерированный текст",
}
print_experiment_info(experiment_name, experiment_config)
ensure_directories()
logger = ExperimentLogger(experiment_name)
try:
# Загружаем модель и токенизатор
model, tokenizer, model_config = load_model_and_tokenizer()
# === Анализ токенизации ===
analysis_texts = [
"Искусственный интеллект",
"Нейронные сети",
"Машинное обучение",
"Машинное обучение",
]
analyze_tokenization(tokenizer, analysis_texts)
# === Генерация с разными промптами ===
print(f"\n🎯 Генерация текста с разными промптами")
print("=" * 60)
for i, prompt in enumerate(TEST_PROMPTS):
print(f"\n📝 Пример {i+1}/{len(TEST_PROMPTS)}")
print("-" * 40)
try:
generated = generate_text(model, tokenizer, prompt, GENERATION_CONFIG)
# Выделяем сгенерированную часть
generated_part = generated[len(prompt):]
generated_part = generated[len(prompt) :]
print(f"📤 Промпт: '{prompt}'")
print(f"🎯 Сгенерировано: '{generated_part}'")
print(f"📄 Полный текст: '{generated}'")
print(f"📏 Длина: {len(generated)} символов")
# Логируем успешную генерацию
logger.log_metric(f"generation_length_{i}", len(generated))
except Exception as e:
print(f"❌ Ошибка при генерации: {e}")
continue
# === Сравнение стратегий генерации ===
test_prompt = "Искусственный"
test_different_strategies(model, tokenizer, test_prompt)
# === Интерактивная генерация ===
interactive_generation(model, tokenizer)
# === Сохранение результатов ===
logger.save_logs("checkpoints/llm_only_generation_logs.json")
print(f"\n🎉 Эксперимент генерации завершен успешно!")
except FileNotFoundError as e:
print(f"❌ {e}")
except Exception as e:
print(f"❌ Ошибка в эксперименте: {e}")
import traceback
traceback.print_exc()

171
experiments/llm_only/generate_gpt_bpe.py
View File

@@ -15,18 +15,14 @@ sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from llm.models.gpt import GPT
from llm.tokenizers import BPETokenizer
from shared.configs import (
BASE_GPT_CONFIG, TEST_PROMPTS, GENERATION_CONFIG, PATHS
)
from shared.data import (
print_experiment_info, ensure_directories, ExperimentLogger
)
from shared.configs import BASE_GPT_CONFIG, TEST_PROMPTS, GENERATION_CONFIG, PATHS
from shared.data import print_experiment_info, ensure_directories, ExperimentLogger
def load_model_and_tokenizer() -> tuple:
"""
Загружает обученную модель и токенизатор.
Returns:
tuple: (модель, токенизатор, конфигурация)
"""
@@ -36,62 +32,60 @@ def load_model_and_tokenizer() -> tuple:
f"Модель не найдена: {PATHS['gpt_bpe_model']}\n"
f"Сначала обучите модель: uv run python experiments/llm_only/train_gpt_bpe.py"
)
if not os.path.exists(PATHS["bpe_tokenizer"]):
raise FileNotFoundError(
f"Токенизатор не найден: {PATHS['bpe_tokenizer']}"
)
raise FileNotFoundError(f"Токенизатор не найден: {PATHS['bpe_tokenizer']}")
# Загружаем конфигурацию модели
import json
with open(PATHS["gpt_bpe_config"], 'r', encoding='utf-8') as f:
with open(PATHS["gpt_bpe_config"], "r", encoding="utf-8") as f:
model_config = json.load(f)
# Загружаем токенизатор
print("🔧 Загрузка BPE токенизатора...")
tokenizer = BPETokenizer.load(PATHS["bpe_tokenizer"])
print(f"✅ Токенизатор загружен (vocab_size={tokenizer.get_vocab_size()})")
# Загружаем модель
print("🔧 Загрузка GPT модели...")
model = GPT(model_config)
model.load_state_dict(torch.load(PATHS["gpt_bpe_model"], map_location='cpu'))
model.load_state_dict(torch.load(PATHS["gpt_bpe_model"], map_location="cpu"))
model.eval()
print("✅ Модель загружена")
return model, tokenizer, model_config
def generate_text(
model: GPT,
tokenizer: BPETokenizer,
prompt: str,
config: dict
model: GPT, tokenizer: BPETokenizer, prompt: str, config: dict
) -> str:
"""
Генерирует текст на основе промпта.
Args:
model: Обученная GPT модель
tokenizer: BPE токенизатор
prompt: Входной текст
config: Конфигурация генерации
Returns:
str: Сгенерированный текст
"""
print(f"🔤 Промпт: '{prompt}'")
print(f"📊 Параметры: max_tokens={config['max_new_tokens']}, "
f"temp={config['temperature']}, sample={config['do_sample']}")
print(
f"📊 Параметры: max_tokens={config['max_new_tokens']}, "
f"temp={config['temperature']}, sample={config['do_sample']}"
)
# Кодируем промпт
input_ids = tokenizer.encode(prompt, add_special_tokens=False)
input_tensor = torch.tensor([input_ids], dtype=torch.long)
print(f"🎯 Токены промпта: {input_ids}")
print(f"🎯 Токены (текст): {tokenizer.tokenize(prompt)}")
print("🔄 Генерация...")
# Генерируем текст
with torch.no_grad():
generated_ids = model.generate(
@@ -100,19 +94,19 @@ def generate_text(
do_sample=config["do_sample"],
temperature=config["temperature"],
top_k=config["top_k"],
top_p=config["top_p"]
top_p=config["top_p"],
)
# Декодируем результат
generated_text = tokenizer.decode(generated_ids[0].tolist())
return generated_text
def test_different_strategies(model: GPT, tokenizer: BPETokenizer, prompt: str):
"""
Тестирует разные стратегии генерации на одном промпте.
Args:
model: Обученная модель
tokenizer: BPE токенизатор
@@ -120,32 +114,38 @@ def test_different_strategies(model: GPT, tokenizer: BPETokenizer, prompt: str):
"""
print(f"\n🎭 Сравнение стратегий генерации для промпта: '{prompt}'")
print("=" * 60)
strategies = [
{"name": "🎯 Жадный поиск", "do_sample": False, "temperature": 1.0},
{"name": "🎲 Вероятностная (temp=0.7)", "do_sample": True, "temperature": 0.7},
{"name": "🔥 Случайная (temp=1.2)", "do_sample": True, "temperature": 1.2},
{"name": "❄️ Детерминированная (temp=0.3)", "do_sample": True, "temperature": 0.3},
{
"name": "❄️ Детерминированная (temp=0.3)",
"do_sample": True,
"temperature": 0.3,
},
]
for strategy in strategies:
print(f"\n{strategy['name']}:")
try:
config = GENERATION_CONFIG.copy()
config.update({
"do_sample": strategy["do_sample"],
"temperature": strategy["temperature"],
"max_new_tokens": 20
})
config.update(
{
"do_sample": strategy["do_sample"],
"temperature": strategy["temperature"],
"max_new_tokens": 20,
}
)
generated = generate_text(model, tokenizer, prompt, config)
# Выделяем сгенерированную часть
generated_part = generated[len(prompt):]
generated_part = generated[len(prompt) :]
print(f" 📤 Промпт: '{prompt}'")
print(f" 🎯 Сгенерировано: '{generated_part}'")
print(f" 📄 Полный текст: '{generated}'")
except Exception as e:
print(f" ❌ Ошибка: {e}")
@@ -153,26 +153,26 @@ def test_different_strategies(model: GPT, tokenizer: BPETokenizer, prompt: str):
def analyze_tokenization(tokenizer: BPETokenizer, texts: list):
"""
Анализирует токенизацию различных текстов.
Args:
tokenizer: BPE токенизатор
texts: Список текстов для анализа
"""
print(f"\n🔍 Анализ токенизации BPE:")
print("=" * 50)
for i, text in enumerate(texts):
print(f"\nТекст {i+1}: '{text}'")
# Токенизация
tokens = tokenizer.encode(text, add_special_tokens=False)
token_strings = tokenizer.tokenize(text)
print(f" Токены (ID): {tokens}")
print(f" Токены (текст): {token_strings}")
print(f" Количество токенов: {len(tokens)}")
print(f" Эффективность: {len(text)} символов → {len(tokens)} токенов")
# Декодирование обратно
decoded = tokenizer.decode(tokens)
if text == decoded:
@@ -184,51 +184,53 @@ def analyze_tokenization(tokenizer: BPETokenizer, texts: list):
def interactive_generation(model: GPT, tokenizer: BPETokenizer):
"""
Режим интерактивной генерации.
Args:
model: Обученная модель
tokenizer: BPE токенизатор
"""
print(f"\n💬 Интерактивная генерация (для выхода введите 'exit')")
print("-" * 50)
while True:
try:
user_input = input("\n🔤 Введите промпт: ").strip()
if user_input.lower() in ['exit', 'quit', 'выход']:
if user_input.lower() in ["exit", "quit", "выход"]:
break
if not user_input:
continue
# Запрашиваем параметры
try:
max_tokens = int(input("📏 Макс. токенов [50]: ") or "50")
temperature = float(input("🌡️ Температура [0.7]: ") or "0.7")
do_sample_input = input("🎲 Сэмплирование (y/n) [y]: ").lower()
do_sample = do_sample_input != 'n'
do_sample = do_sample_input != "n"
except:
max_tokens = 50
temperature = 0.7
do_sample = True
print("⚠️ Использую параметры по умолчанию")
config = GENERATION_CONFIG.copy()
config.update({
"max_new_tokens": max_tokens,
"temperature": temperature,
"do_sample": do_sample
})
config.update(
{
"max_new_tokens": max_tokens,
"temperature": temperature,
"do_sample": do_sample,
}
)
generated = generate_text(model, tokenizer, user_input, config)
generated_part = generated[len(user_input):]
generated_part = generated[len(user_input) :]
print(f"\n🎯 Результат:")
print(f" 📤 Промпт: '{user_input}'")
print(f" 🎯 Сгенерировано: '{generated_part}'")
print(f" 📄 Полный текст: '{generated}'")
except KeyboardInterrupt:
print("\n👋 Завершение работы...")
break
@@ -244,68 +246,69 @@ def main():
"model": "GPT с BPE токенизатором",
"стратегия": "автономная генерация",
"вход": "промпты",
"выход": "сгенерированный текст"
"выход": "сгенерированный текст",
}
print_experiment_info(experiment_name, experiment_config)
ensure_directories()
logger = ExperimentLogger(experiment_name)
try:
# Загружаем модель и токенизатор
model, tokenizer, model_config = load_model_and_tokenizer()
# === Анализ токенизации ===
analysis_texts = [
"Искусственный интеллект",
"Нейронные сети",
"Машинное обучение",
"Машинное обучение",
]
analyze_tokenization(tokenizer, analysis_texts)
# === Генерация с разными промптами ===
print(f"\n🎯 Генерация текста с разными промптами")
print("=" * 60)
for i, prompt in enumerate(TEST_PROMPTS):
print(f"\n📝 Пример {i+1}/{len(TEST_PROMPTS)}")
print("-" * 40)
try:
generated = generate_text(model, tokenizer, prompt, GENERATION_CONFIG)
# Выделяем сгенерированную часть
generated_part = generated[len(prompt):]
generated_part = generated[len(prompt) :]
print(f"📤 Промпт: '{prompt}'")
print(f"🎯 Сгенерировано: '{generated_part}'")
print(f"📄 Полный текст: '{generated}'")
print(f"📏 Длина: {len(generated)} символов")
# Логируем успешную генерацию
logger.log_metric(f"generation_length_{i}", len(generated))
except Exception as e:
print(f"❌ Ошибка при генерации: {e}")
continue
# === Сравнение стратегий генерации ===
test_prompt = "Искусственный"
test_different_strategies(model, tokenizer, test_prompt)
# === Интерактивная генерация ===
interactive_generation(model, tokenizer)
# === Сохранение результатов ===
logger.save_logs("checkpoints/llm_only_generation_logs.json")
print(f"\n🎉 Эксперимент генерации завершен успешно!")
except FileNotFoundError as e:
print(f"❌ {e}")
except Exception as e:
print(f"❌ Ошибка в эксперименте: {e}")
import traceback
traceback.print_exc()

112
experiments/llm_only/train_gpt2_bpe.py
View File

@@ -18,71 +18,77 @@ from llm.training.dataset import TextDataset
from llm.training.trainer import Trainer
from shared.configs import (
TRAIN_TEXTS, BASE_GPT_CONFIG, BPE_CONFIG,
TRAINING_CONFIG, PATHS, TEST_PROMPTS
TRAIN_TEXTS,
BASE_GPT_CONFIG,
BPE_CONFIG,
TRAINING_CONFIG,
PATHS,
TEST_PROMPTS,
)
from shared.data import (
load_training_data, ensure_directories,
print_experiment_info, ExperimentLogger
load_training_data,
ensure_directories,
print_experiment_info,
ExperimentLogger,
)
def train_bpe_tokenizer(texts: list, config: dict) -> BPETokenizer:
"""
Обучает BPE токенизатор на текстах.
Args:
texts: Список текстов для обучения
config: Конфигурация токенизатора
Returns:
BPETokenizer: Обученный токенизатор
"""
print("🔧 Обучение BPE токенизатора...")
tokenizer = BPETokenizer()
tokenizer.train(
texts=texts,
vocab_size=config["vocab_size"],
special_tokens=config["special_tokens"]
special_tokens=config["special_tokens"],
)
# Сохраняем токенизатор
os.makedirs(os.path.dirname(PATHS["bpe_tokenizer"]), exist_ok=True)
tokenizer.save(PATHS["bpe_tokenizer"])
print(f"✅ BPE токенизатор обучен и сохранен: {PATHS['bpe_tokenizer']}")
print(f"📊 Размер словаря: {tokenizer.get_vocab_size()}")
return tokenizer
def test_tokenizer(tokenizer: BPETokenizer, texts: list):
"""
Тестирует токенизатор на примерах.
Args:
tokenizer: Обученный токенизатор
texts: Список тестовых текстов
"""
print("\n🧪 Тестирование токенизатора:")
for i, text in enumerate(texts[:3]):
print(f"\nПример {i+1}:")
print(f" Исходный текст: '{text}'")
# Кодирование
tokens = tokenizer.encode(text)
token_strings = tokenizer.tokenize(text)
print(f" Токены (ID): {tokens}")
print(f" Токены (текст): {token_strings}")
print(f" Количество токенов: {len(tokens)}")
# Декодирование
decoded = tokenizer.decode(tokens)
print(f" Декодированный: '{decoded}'")
if text == decoded:
print(" ✅ Кодирование/декодирование корректно")
else:
@@ -95,22 +101,22 @@ def main():
experiment_name = "Обучение GPT2 с BPE токенизатором (только llm)"
experiment_config = {
"model": "GPT2",
"tokenizer": "BPE",
"tokenizer": "BPE",
"vocab_size": BPE_CONFIG["vocab_size"],
"training_epochs": TRAINING_CONFIG["num_epochs"],
"batch_size": TRAINING_CONFIG["batch_size"],
"learning_rate": TRAINING_CONFIG["learning_rate"]
"learning_rate": TRAINING_CONFIG["learning_rate"],
}
print_experiment_info(experiment_name, experiment_config)
ensure_directories()
logger = ExperimentLogger(experiment_name)
try:
# === Подготовка данных ===
train_texts, val_texts = load_training_data()
print(f"📊 Данные: {len(train_texts)} train, {len(val_texts)} validation")
# === Обучение токенизатора ===
if os.path.exists(PATHS["bpe_tokenizer"]):
print("📝 Загрузка предварительно обученного токенизатора...")
@@ -118,112 +124,112 @@ def main():
print(f"✅ Токенизатор загружен (vocab_size={tokenizer.get_vocab_size()})")
else:
tokenizer = train_bpe_tokenizer(TRAIN_TEXTS, BPE_CONFIG)
# Тестируем токенизатор
test_tokenizer(tokenizer, TEST_PROMPTS[:3])
# === Инициализация модели ===
model_config = BASE_GPT_CONFIG.copy()
model_config["vocab_size"] = tokenizer.get_vocab_size()
print(f"\n🔧 Инициализация GPT2 модели...")
print(f" Размер словаря: {model_config['vocab_size']}")
print(f" Размер эмбеддингов: {model_config['embed_dim']}")
print(f" Количество слоев: {model_config['num_layers']}")
print(f" Количество голов внимания: {model_config['num_heads']}")
model = GPT2(model_config)
# === Подготовка датасета ===
print(f"\n📊 Подготовка датасета...")
train_dataset = TextDataset(
train_texts,
tokenizer,
block_size=model_config["max_position_embeddings"]
train_texts, tokenizer, block_size=model_config["max_position_embeddings"]
)
print(f" Размер train датасета: {len(train_dataset)} примеров")
# === Обучение модели ===
print(f"\n🎯 Начало обучения GPT2 модели...")
trainer = Trainer(
model=model,
train_dataset=train_dataset,
lr=TRAINING_CONFIG["learning_rate"],
batch_size=TRAINING_CONFIG["batch_size"],
num_epochs=TRAINING_CONFIG["num_epochs"],
warmup_steps=TRAINING_CONFIG["warmup_steps"]
warmup_steps=TRAINING_CONFIG["warmup_steps"],
)
# Запускаем обучение
trainer.train()
# === Сохранение модели ===
print(f"\n💾 Сохранение модели...")
os.makedirs(os.path.dirname(PATHS["gpt_bpe_model"]), exist_ok=True)
# Сохраняем модель
torch.save(model.state_dict(), PATHS["gpt_bpe_model"])
# Сохраняем конфигурацию
import json
with open(PATHS["gpt_bpe_config"], 'w', encoding='utf-8') as f:
with open(PATHS["gpt_bpe_config"], "w", encoding="utf-8") as f:
json.dump(model_config, f, indent=2, ensure_ascii=False)
print(f"✅ Модель сохранена:")
print(f" - {PATHS['gpt_bpe_model']}: веса модели")
print(f" - {PATHS['gpt_bpe_config']}: конфигурация модели")
print(f" - {PATHS['bpe_tokenizer']}: токенизатор")
# === Тестирование генерации ===
print(f"\n🧪 Тестирование генерации текста...")
model.eval()
for prompt in TEST_PROMPTS[:3]:
print(f"\n🔤 Промпт: '{prompt}'")
try:
# Кодируем промпт
input_ids = tokenizer.encode(prompt, add_special_tokens=False)
input_tensor = torch.tensor([input_ids], dtype=torch.long)
# Генерируем текст
with torch.no_grad():
generated_ids = model.generate(
x=input_tensor,
max_new_tokens=20,
do_sample=True,
temperature=0.8
temperature=0.8,
)
# Декодируем результат
generated_text = tokenizer.decode(generated_ids[0].tolist())
generated_part = generated_text[len(prompt):]
generated_part = generated_text[len(prompt) :]
print(f"🎯 Сгенерировано: '{generated_part}'")
print(f"📄 Полный текст: '{generated_text}'")
except Exception as e:
print(f"❌ Ошибка генерации: {e}")
# === Сохранение результатов ===
results = {
"experiment": experiment_name,
"model_config": model_config,
"training_config": TRAINING_CONFIG,
"tokenizer_vocab_size": tokenizer.get_vocab_size(),
"final_loss": "см. логи обучения" # В реальном эксперименте можно сохранить final loss
"final_loss": "см. логи обучения", # В реальном эксперименте можно сохранить final loss
}
logger.save_logs("checkpoints/llm_only_training_logs.json")
print(f"\n🎉 Эксперимент завершен успешно!")
print(f"\n💡 Для использования обученной модели:")
print(f" uv run python experiments/llm_only/generate_gpt_bpe.py")
except Exception as e:
print(f"❌ Ошибка в эксперименте: {e}")
import traceback
traceback.print_exc()

112
experiments/llm_only/train_gpt_bpe.py
View File

@@ -18,71 +18,77 @@ from llm.training.dataset import TextDataset
from llm.training.trainer import Trainer
from shared.configs import (
TRAIN_TEXTS, BASE_GPT_CONFIG, BPE_CONFIG,
TRAINING_CONFIG, PATHS, TEST_PROMPTS
TRAIN_TEXTS,
BASE_GPT_CONFIG,
BPE_CONFIG,
TRAINING_CONFIG,
PATHS,
TEST_PROMPTS,
)
from shared.data import (
load_training_data, ensure_directories,
print_experiment_info, ExperimentLogger
load_training_data,
ensure_directories,
print_experiment_info,
ExperimentLogger,
)
def train_bpe_tokenizer(texts: list, config: dict) -> BPETokenizer:
"""
Обучает BPE токенизатор на текстах.
Args:
texts: Список текстов для обучения
config: Конфигурация токенизатора
Returns:
BPETokenizer: Обученный токенизатор
"""
print("🔧 Обучение BPE токенизатора...")
tokenizer = BPETokenizer()
tokenizer.train(
texts=texts,
vocab_size=config["vocab_size"],
special_tokens=config["special_tokens"]
special_tokens=config["special_tokens"],
)
# Сохраняем токенизатор
os.makedirs(os.path.dirname(PATHS["bpe_tokenizer"]), exist_ok=True)
tokenizer.save(PATHS["bpe_tokenizer"])
print(f"✅ BPE токенизатор обучен и сохранен: {PATHS['bpe_tokenizer']}")
print(f"📊 Размер словаря: {tokenizer.get_vocab_size()}")
return tokenizer
def test_tokenizer(tokenizer: BPETokenizer, texts: list):
"""
Тестирует токенизатор на примерах.
Args:
tokenizer: Обученный токенизатор
texts: Список тестовых текстов
"""
print("\n🧪 Тестирование токенизатора:")
for i, text in enumerate(texts[:3]):
print(f"\nПример {i+1}:")
print(f" Исходный текст: '{text}'")
# Кодирование
tokens = tokenizer.encode(text)
token_strings = tokenizer.tokenize(text)
print(f" Токены (ID): {tokens}")
print(f" Токены (текст): {token_strings}")
print(f" Количество токенов: {len(tokens)}")
# Декодирование
decoded = tokenizer.decode(tokens)
print(f" Декодированный: '{decoded}'")
if text == decoded:
print(" ✅ Кодирование/декодирование корректно")
else:
@@ -95,22 +101,22 @@ def main():
experiment_name = "Обучение GPT с BPE токенизатором (только llm)"
experiment_config = {
"model": "GPT",
"tokenizer": "BPE",
"tokenizer": "BPE",
"vocab_size": BPE_CONFIG["vocab_size"],
"training_epochs": TRAINING_CONFIG["num_epochs"],
"batch_size": TRAINING_CONFIG["batch_size"],
"learning_rate": TRAINING_CONFIG["learning_rate"]
"learning_rate": TRAINING_CONFIG["learning_rate"],
}
print_experiment_info(experiment_name, experiment_config)
ensure_directories()
logger = ExperimentLogger(experiment_name)
try:
# === Подготовка данных ===
train_texts, val_texts = load_training_data()
print(f"📊 Данные: {len(train_texts)} train, {len(val_texts)} validation")
# === Обучение токенизатора ===
if os.path.exists(PATHS["bpe_tokenizer"]):
print("📝 Загрузка предварительно обученного токенизатора...")
@@ -118,112 +124,112 @@ def main():
print(f"✅ Токенизатор загружен (vocab_size={tokenizer.get_vocab_size()})")
else:
tokenizer = train_bpe_tokenizer(TRAIN_TEXTS, BPE_CONFIG)
# Тестируем токенизатор
test_tokenizer(tokenizer, TEST_PROMPTS[:3])
# === Инициализация модели ===
model_config = BASE_GPT_CONFIG.copy()
model_config["vocab_size"] = tokenizer.get_vocab_size()
print(f"\n🔧 Инициализация GPT модели...")
print(f" Размер словаря: {model_config['vocab_size']}")
print(f" Размер эмбеддингов: {model_config['embed_dim']}")
print(f" Количество слоев: {model_config['num_layers']}")
print(f" Количество голов внимания: {model_config['num_heads']}")
model = GPT(model_config)
# === Подготовка датасета ===
print(f"\n📊 Подготовка датасета...")
train_dataset = TextDataset(
train_texts,
tokenizer,
block_size=model_config["max_position_embeddings"]
train_texts, tokenizer, block_size=model_config["max_position_embeddings"]
)
print(f" Размер train датасета: {len(train_dataset)} примеров")
# === Обучение модели ===
print(f"\n🎯 Начало обучения GPT модели...")
trainer = Trainer(
model=model,
train_dataset=train_dataset,
lr=TRAINING_CONFIG["learning_rate"],
batch_size=TRAINING_CONFIG["batch_size"],
num_epochs=TRAINING_CONFIG["num_epochs"],
warmup_steps=TRAINING_CONFIG["warmup_steps"]
warmup_steps=TRAINING_CONFIG["warmup_steps"],
)
# Запускаем обучение
trainer.train()
# === Сохранение модели ===
print(f"\n💾 Сохранение модели...")
os.makedirs(os.path.dirname(PATHS["gpt_bpe_model"]), exist_ok=True)
# Сохраняем модель
torch.save(model.state_dict(), PATHS["gpt_bpe_model"])
# Сохраняем конфигурацию
import json
with open(PATHS["gpt_bpe_config"], 'w', encoding='utf-8') as f:
with open(PATHS["gpt_bpe_config"], "w", encoding="utf-8") as f:
json.dump(model_config, f, indent=2, ensure_ascii=False)
print(f"✅ Модель сохранена:")
print(f" - {PATHS['gpt_bpe_model']}: веса модели")
print(f" - {PATHS['gpt_bpe_config']}: конфигурация модели")
print(f" - {PATHS['bpe_tokenizer']}: токенизатор")
# === Тестирование генерации ===
print(f"\n🧪 Тестирование генерации текста...")
model.eval()
for prompt in TEST_PROMPTS[:3]:
print(f"\n🔤 Промпт: '{prompt}'")
try:
# Кодируем промпт
input_ids = tokenizer.encode(prompt, add_special_tokens=False)
input_tensor = torch.tensor([input_ids], dtype=torch.long)
# Генерируем текст
with torch.no_grad():
generated_ids = model.generate(
x=input_tensor,
max_new_tokens=20,
do_sample=True,
temperature=0.8
temperature=0.8,
)
# Декодируем результат
generated_text = tokenizer.decode(generated_ids[0].tolist())
generated_part = generated_text[len(prompt):]
generated_part = generated_text[len(prompt) :]
print(f"🎯 Сгенерировано: '{generated_part}'")
print(f"📄 Полный текст: '{generated_text}'")
except Exception as e:
print(f"❌ Ошибка генерации: {e}")
# === Сохранение результатов ===
results = {
"experiment": experiment_name,
"model_config": model_config,
"training_config": TRAINING_CONFIG,
"tokenizer_vocab_size": tokenizer.get_vocab_size(),
"final_loss": "см. логи обучения" # В реальном эксперименте можно сохранить final loss
"final_loss": "см. логи обучения", # В реальном эксперименте можно сохранить final loss
}
logger.save_logs("checkpoints/llm_only_training_logs.json")
print(f"\n🎉 Эксперимент завершен успешно!")
print(f"\n💡 Для использования обученной модели:")
print(f" uv run python experiments/llm_only/generate_gpt_bpe.py")
except Exception as e:
print(f"❌ Ошибка в эксперименте: {e}")
import traceback
traceback.print_exc()

112
experiments/llm_only/train_llama_bpe.py
View File

@@ -18,71 +18,77 @@ from llm.training.dataset import TextDataset
from llm.training.trainer import Trainer
from shared.configs import (
TRAIN_TEXTS, BASE_GPT_CONFIG, BPE_CONFIG,
TRAINING_CONFIG, PATHS, TEST_PROMPTS
TRAIN_TEXTS,
BASE_GPT_CONFIG,
BPE_CONFIG,
TRAINING_CONFIG,
PATHS,
TEST_PROMPTS,
)
from shared.data import (
load_training_data, ensure_directories,
print_experiment_info, ExperimentLogger
load_training_data,
ensure_directories,
print_experiment_info,
ExperimentLogger,
)
def train_bpe_tokenizer(texts: list, config: dict) -> BPETokenizer:
"""
Обучает BPE токенизатор на текстах.
Args:
texts: Список текстов для обучения
config: Конфигурация токенизатора
Returns:
BPETokenizer: Обученный токенизатор
"""
print("🔧 Обучение BPE токенизатора...")
tokenizer = BPETokenizer()
tokenizer.train(
texts=texts,
vocab_size=config["vocab_size"],
special_tokens=config["special_tokens"]
special_tokens=config["special_tokens"],
)
# Сохраняем токенизатор
os.makedirs(os.path.dirname(PATHS["bpe_tokenizer"]), exist_ok=True)
tokenizer.save(PATHS["bpe_tokenizer"])
print(f"✅ BPE токенизатор обучен и сохранен: {PATHS['bpe_tokenizer']}")
print(f"📊 Размер словаря: {tokenizer.get_vocab_size()}")
return tokenizer
def test_tokenizer(tokenizer: BPETokenizer, texts: list):
"""
Тестирует токенизатор на примерах.
Args:
tokenizer: Обученный токенизатор
texts: Список тестовых текстов
"""
print("\n🧪 Тестирование токенизатора:")
for i, text in enumerate(texts[:3]):
print(f"\nПример {i+1}:")
print(f" Исходный текст: '{text}'")
# Кодирование
tokens = tokenizer.encode(text)
token_strings = tokenizer.tokenize(text)
print(f" Токены (ID): {tokens}")
print(f" Токены (текст): {token_strings}")
print(f" Количество токенов: {len(tokens)}")
# Декодирование
decoded = tokenizer.decode(tokens)
print(f" Декодированный: '{decoded}'")
if text == decoded:
print(" ✅ Кодирование/декодирование корректно")
else:
@@ -95,22 +101,22 @@ def main():
experiment_name = "Обучение Llama с BPE токенизатором (только llm)"
experiment_config = {
"model": "Llama",
"tokenizer": "BPE",
"tokenizer": "BPE",
"vocab_size": BPE_CONFIG["vocab_size"],
"training_epochs": TRAINING_CONFIG["num_epochs"],
"batch_size": TRAINING_CONFIG["batch_size"],
"learning_rate": TRAINING_CONFIG["learning_rate"]
"learning_rate": TRAINING_CONFIG["learning_rate"],
}
print_experiment_info(experiment_name, experiment_config)
ensure_directories()
logger = ExperimentLogger(experiment_name)
try:
# === Подготовка данных ===
train_texts, val_texts = load_training_data()
print(f"📊 Данные: {len(train_texts)} train, {len(val_texts)} validation")
# === Обучение токенизатора ===
if os.path.exists(PATHS["bpe_tokenizer"]):
print("📝 Загрузка предварительно обученного токенизатора...")
@@ -118,112 +124,112 @@ def main():
print(f"✅ Токенизатор загружен (vocab_size={tokenizer.get_vocab_size()})")
else:
tokenizer = train_bpe_tokenizer(TRAIN_TEXTS, BPE_CONFIG)
# Тестируем токенизатор
test_tokenizer(tokenizer, TEST_PROMPTS[:3])
# === Инициализация модели ===
model_config = BASE_GPT_CONFIG.copy()
model_config["vocab_size"] = tokenizer.get_vocab_size()
print(f"\n🔧 Инициализация Llama модели...")
print(f" Размер словаря: {model_config['vocab_size']}")
print(f" Размер эмбеддингов: {model_config['embed_dim']}")
print(f" Количество слоев: {model_config['num_layers']}")
print(f" Количество голов внимания: {model_config['num_heads']}")
model = Llama(model_config)
# === Подготовка датасета ===
print(f"\n📊 Подготовка датасета...")
train_dataset = TextDataset(
train_texts,
tokenizer,
block_size=model_config["max_position_embeddings"]
train_texts, tokenizer, block_size=model_config["max_position_embeddings"]
)
print(f" Размер train датасета: {len(train_dataset)} примеров")
# === Обучение модели ===
print(f"\n🎯 Начало обучения Llama модели...")
trainer = Trainer(
model=model,
train_dataset=train_dataset,
lr=TRAINING_CONFIG["learning_rate"],
batch_size=TRAINING_CONFIG["batch_size"],
num_epochs=TRAINING_CONFIG["num_epochs"],
warmup_steps=TRAINING_CONFIG["warmup_steps"]
warmup_steps=TRAINING_CONFIG["warmup_steps"],
)
# Запускаем обучение
trainer.train()
# === Сохранение модели ===
print(f"\n💾 Сохранение модели...")
os.makedirs(os.path.dirname(PATHS["gpt_bpe_model"]), exist_ok=True)
# Сохраняем модель
torch.save(model.state_dict(), PATHS["gpt_bpe_model"])
# Сохраняем конфигурацию
import json
with open(PATHS["gpt_bpe_config"], 'w', encoding='utf-8') as f:
with open(PATHS["gpt_bpe_config"], "w", encoding="utf-8") as f:
json.dump(model_config, f, indent=2, ensure_ascii=False)
print(f"✅ Модель сохранена:")
print(f" - {PATHS['gpt_bpe_model']}: веса модели")
print(f" - {PATHS['gpt_bpe_config']}: конфигурация модели")
print(f" - {PATHS['bpe_tokenizer']}: токенизатор")
# === Тестирование генерации ===
print(f"\n🧪 Тестирование генерации текста...")
model.eval()
for prompt in TEST_PROMPTS[:3]:
print(f"\n🔤 Промпт: '{prompt}'")
try:
# Кодируем промпт
input_ids = tokenizer.encode(prompt, add_special_tokens=False)
input_tensor = torch.tensor([input_ids], dtype=torch.long)
# Генерируем текст
with torch.no_grad():
generated_ids = model.generate(
x=input_tensor,
max_new_tokens=20,
do_sample=True,
temperature=0.8
temperature=0.8,
)
# Декодируем результат
generated_text = tokenizer.decode(generated_ids[0].tolist())
generated_part = generated_text[len(prompt):]
generated_part = generated_text[len(prompt) :]
print(f"🎯 Сгенерировано: '{generated_part}'")
print(f"📄 Полный текст: '{generated_text}'")
except Exception as e:
print(f"❌ Ошибка генерации: {e}")
# === Сохранение результатов ===
results = {
"experiment": experiment_name,
"model_config": model_config,
"training_config": TRAINING_CONFIG,
"tokenizer_vocab_size": tokenizer.get_vocab_size(),
"final_loss": "см. логи обучения" # В реальном эксперименте можно сохранить final loss
"final_loss": "см. логи обучения", # В реальном эксперименте можно сохранить final loss
}
logger.save_logs("checkpoints/llm_only_training_logs.json")
print(f"\n🎉 Эксперимент завершен успешно!")
print(f"\n💡 Для использования обученной модели:")
print(f" uv run python experiments/llm_only/generate_gpt_bpe.py")
except Exception as e:
print(f"❌ Ошибка в эксперименте: {e}")
import traceback
traceback.print_exc()

14
experiments/shared/configs.py
View File

@@ -30,7 +30,7 @@ BASE_GPT_CONFIG = {
"num_heads": 4,
"num_layers": 4,
"max_position_embeddings": 128,
"dropout": 0.1
"dropout": 0.1,
}
# Конфигурация для маленькой модели (быстрое тестирование)
@@ -40,7 +40,7 @@ SMALL_GPT_CONFIG = {
"num_heads": 2,
"num_layers": 2,
"max_position_embeddings": 64,
"dropout": 0.1
"dropout": 0.1,
}
# Конфигурация для большой модели (качественное обучение)
@@ -50,13 +50,13 @@ LARGE_GPT_CONFIG = {
"num_heads": 8,
"num_layers": 6,
"max_position_embeddings": 256,
"dropout": 0.1
"dropout": 0.1,
}
# === Конфигурации токенизатора ===
BPE_CONFIG = {
"vocab_size": 1000,
"special_tokens": ["<pad>", "<unk>", "<bos>", "<eos>"]
"special_tokens": ["<pad>", "<unk>", "<bos>", "<eos>"],
}
# === Конфигурации обучения ===
@@ -65,7 +65,7 @@ TRAINING_CONFIG = {
"batch_size": 2,
"num_epochs": 3,
"warmup_steps": 50,
"gradient_clip": 1.0
"gradient_clip": 1.0,
}
# === Конфигурации генерации ===
@@ -74,7 +74,7 @@ GENERATION_CONFIG = {
"temperature": 0.7,
"do_sample": True,
"top_k": None,
"top_p": None
"top_p": None,
}
# === Пути для сохранения ===
@@ -84,7 +84,7 @@ PATHS = {
"gpt_bpe_config": "checkpoints/gpt-bpe/config.json",
"hf_tokenizer": "checkpoints/hf-bpe-tokenizer",
"hf_model": "checkpoints/hf-trained",
"hf_proxy_model": "checkpoints/hf-trained-proxy"
"hf_proxy_model": "checkpoints/hf-trained-proxy",
}
# === Тестовые промпты ===

80
experiments/shared/data.py
View File

@@ -10,17 +10,17 @@ from .configs import TRAIN_TEXTS, PATHS
def load_training_data(split_ratio: float = 0.8) -> Tuple[List[str], List[str]]:
"""
Загружает данные для обучения и разделяет на train/validation.
Args:
split_ratio: Доля данных для обучения
Returns:
Tuple: (train_texts, val_texts)
"""
train_size = int(len(TRAIN_TEXTS) * split_ratio)
train_data = TRAIN_TEXTS[:train_size]
val_data = TRAIN_TEXTS[train_size:]
return train_data, val_data
@@ -28,13 +28,13 @@ def ensure_directories():
"""Создает необходимые директории если они не существуют."""
directories = [
"checkpoints",
"checkpoints/gpt-bpe",
"checkpoints/gpt-bpe",
"checkpoints/hf-bpe-tokenizer",
"checkpoints/hf-trained",
"checkpoints/hf-trained-proxy",
"logs"
"logs",
]
for directory in directories:
os.makedirs(directory, exist_ok=True)
@@ -42,33 +42,34 @@ def ensure_directories():
def get_model_paths(experiment_type: str = "llm_only") -> dict:
"""
Возвращает пути для конкретного типа эксперимента.
Args:
experiment_type: Тип эксперимента ('llm_only' или 'hf_integration')
Returns:
dict: Словарь с путями
"""
base_paths = PATHS.copy()
if experiment_type == "hf_integration":
base_paths.update({
"model": base_paths["hf_model"],
"tokenizer": base_paths["hf_tokenizer"]
})
base_paths.update(
{"model": base_paths["hf_model"], "tokenizer": base_paths["hf_tokenizer"]}
)
else: # llm_only
base_paths.update({
"model": base_paths["gpt_bpe_model"],
"tokenizer": base_paths["bpe_tokenizer"]
})
base_paths.update(
{
"model": base_paths["gpt_bpe_model"],
"tokenizer": base_paths["bpe_tokenizer"],
}
)
return base_paths
def print_experiment_info(experiment_name: str, config: dict):
"""
Выводит информацию о запускаемом эксперименте.
Args:
experiment_name: Название эксперимента
config: Конфигурация эксперимента
@@ -85,35 +86,35 @@ def print_experiment_info(experiment_name: str, config: dict):
def save_experiment_results(results: dict, filepath: str):
"""
Сохраняет результаты эксперимента в файл.
Args:
results: Словарь с результатами
filepath: Путь для сохранения
"""
import json
with open(filepath, 'w', encoding='utf-8') as f:
with open(filepath, "w", encoding="utf-8") as f:
json.dump(results, f, ensure_ascii=False, indent=2)
print(f"✅ Результаты эксперимента сохранены: {filepath}")
def load_experiment_results(filepath: str) -> dict:
"""
Загружает результаты эксперимента из файла.
Args:
filepath: Путь к файлу с результатами
Returns:
dict: Загруженные результаты
"""
import json
if not os.path.exists(filepath):
return {}
with open(filepath, 'r', encoding='utf-8') as f:
with open(filepath, "r", encoding="utf-8") as f:
return json.load(f)
@@ -121,42 +122,39 @@ class ExperimentLogger:
"""
Логгер для экспериментов.
"""
def __init__(self, experiment_name: str):
self.experiment_name = experiment_name
self.metrics = {}
def log_metric(self, name: str, value: float):
"""Логирует метрику."""
if name not in self.metrics:
self.metrics[name] = []
self.metrics[name].append(value)
print(f"📈 {name}: {value:.4f}")
def log_step(self, step: int, loss: float, **kwargs):
"""Логирует шаг обучения."""
print(f"📊 Step {step}: loss={loss:.4f}", end="")
for key, value in kwargs.items():
print(f", {key}={value:.4f}", end="")
print()
def log_epoch(self, epoch: int, train_loss: float, val_loss: float = None):
"""Логирует завершение эпохи."""
print(f"🎯 Epoch {epoch}: train_loss={train_loss:.4f}", end="")
if val_loss is not None:
print(f", val_loss={val_loss:.4f}", end="")
print()
def save_logs(self, filepath: str):
"""Сохраняет логи эксперимента."""
import json
logs = {
"experiment_name": self.experiment_name,
"metrics": self.metrics
}
with open(filepath, 'w', encoding='utf-8') as f:
logs = {"experiment_name": self.experiment_name, "metrics": self.metrics}
with open(filepath, "w", encoding="utf-8") as f:
json.dump(logs, f, ensure_ascii=False, indent=2)
print(f"✅ Логи эксперимента сохранены: {filepath}")

7
hf-proxy/src/hf_proxy/__init__.py
View File

@@ -27,16 +27,13 @@ __all__ = [
# Основные классы адаптера
"HFAdapter",
"HFGPTAdapter",
# Конфигурации
"HFAdapterConfig",
"HFAdapterConfig",
"HFPretrainedConfig",
# Адаптеры токенизаторов
"HFTokenizerAdapter",
"create_hf_tokenizer",
"create_hf_tokenizer",
"convert_to_hf_format",
# Утилиты
"HFUtils",
"TokenizerWrapper",

145
hf-proxy/src/hf_proxy/hf_adapter.py
View File

@@ -6,12 +6,12 @@ import torch
import torch.nn as nn
from typing import Optional, Tuple, Union, List
from transformers import (
PreTrainedModel,
PreTrainedModel,
GPT2LMHeadModel,
GPT2Config,
GenerationConfig,
LogitsProcessorList,
StoppingCriteriaList
StoppingCriteriaList,
)
from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
@@ -24,38 +24,39 @@ class HFGPTAdapter(PreTrainedModel):
Адаптер для модели GPT из библиотеки llm.
Позволяет использовать кастомные GPT модели с HuggingFace Transformers.
"""
config_class = HFPretrainedConfig
def __init__(self, config: HFPretrainedConfig, llm_model: Optional[GPT] = None):
"""
Инициализация адаптера.
Args:
config: Конфигурация HuggingFace
llm_model: Опционально, предварительно созданная модель llm
"""
super().__init__(config)
# Преобразуем HF конфигурацию в формат llm
llm_config = self._hf_to_llm_config(config)
# Создаем или используем переданную модель
if llm_model is None:
self.llm_model = GPT(llm_config)
else:
self.llm_model = llm_model
# Устанавливаем веса если они есть в конфигурации
if hasattr(config, 'state_dict') and config.state_dict is not None:
if hasattr(config, "state_dict") and config.state_dict is not None:
self.llm_model.load_state_dict(config.state_dict)
def _hf_to_llm_config(self, hf_config: HFPretrainedConfig) -> dict:
"""
Преобразует конфигурацию HF в формат llm.
Args:
hf_config: Конфигурация HuggingFace
Returns:
dict: Конфигурация для llm модели
"""
@@ -67,7 +68,7 @@ class HFGPTAdapter(PreTrainedModel):
"max_position_embeddings": hf_config.max_position_embeddings,
"dropout": hf_config.hidden_dropout_prob,
}
def forward(
self,
input_ids: Optional[torch.Tensor] = None,
@@ -78,11 +79,11 @@ class HFGPTAdapter(PreTrainedModel):
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
**kwargs
**kwargs,
) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
"""
Прямой проход модели.
Args:
input_ids: Входные токены [batch_size, seq_len]
attention_mask: Маска внимания [batch_size, seq_len]
@@ -92,38 +93,39 @@ class HFGPTAdapter(PreTrainedModel):
output_attentions: Возвращать веса внимания
output_hidden_states: Возвращать скрытые состояния
return_dict: Возвращать словарь вместо кортежа
Returns:
CausalLMOutputWithCrossAttentions или кортеж
"""
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
return_dict = (
return_dict if return_dict is not None else self.config.use_return_dict
)
# Основной forward pass
outputs = self.llm_model(input_ids)
if isinstance(outputs, tuple):
logits = outputs[0]
else:
logits = outputs
loss = None
if labels is not None:
# Сдвигаем логиты и метки для языкового моделирования
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
# Вычисляем cross-entropy loss
loss_fct = nn.CrossEntropyLoss()
loss = loss_fct(
shift_logits.view(-1, shift_logits.size(-1)),
shift_labels.view(-1)
shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)
)
if not return_dict:
output = (logits,)
if loss is not None:
output = (loss,) + output
return output
return CausalLMOutputWithCrossAttentions(
loss=loss,
logits=logits,
@@ -132,30 +134,27 @@ class HFGPTAdapter(PreTrainedModel):
attentions=None,
cross_attentions=None,
)
def prepare_inputs_for_generation(
self,
input_ids: torch.Tensor,
past_key_values: Optional[Tuple] = None,
**kwargs
self, input_ids: torch.Tensor, past_key_values: Optional[Tuple] = None, **kwargs
) -> dict:
"""
Подготавливает входные данные для генерации.
Args:
input_ids: Входные токены
past_key_values: Кешированные ключи и значения
Returns:
dict: Подготовленные входные данные
"""
# Наша простая реализация пока не поддерживает past_key_values
return {"input_ids": input_ids}
def can_generate(self) -> bool:
"""Проверяет, может ли модель генерировать текст."""
return True
def generate(
self,
input_ids: Optional[torch.Tensor] = None,
@@ -163,32 +162,32 @@ class HFGPTAdapter(PreTrainedModel):
generation_config: Optional[GenerationConfig] = None,
logits_processor: Optional[LogitsProcessorList] = None,
stopping_criteria: Optional[StoppingCriteriaList] = None,
**kwargs
**kwargs,
) -> torch.Tensor:
"""
Генерация текста с поддержкой HuggingFace интерфейса.
Args:
input_ids: Входные токены
attention_mask: Маска внимания
generation_config: Конфигурация генерации
logits_processor: Процессоры логитов
stopping_criteria: Критерии остановки
Returns:
torch.Tensor: Сгенерированные токены
"""
# Извлекаем обязательные параметры из kwargs или используем значения по умолчанию
max_new_tokens = kwargs.pop('max_new_tokens', 50)
do_sample = kwargs.pop('do_sample', True)
max_new_tokens = kwargs.pop("max_new_tokens", 50)
do_sample = kwargs.pop("do_sample", True)
# Используем встроенную генерацию llm модели
return self.llm_model.generate(
x=input_ids,
max_new_tokens=max_new_tokens,
do_sample=do_sample,
attention_mask=attention_mask,
**kwargs
**kwargs,
)
@@ -196,64 +195,66 @@ class HFAdapter:
"""
Основной класс адаптера для преобразования моделей llm в формат HuggingFace.
"""
@staticmethod
def from_llm_model(
llm_model: GPT,
hf_config: Optional[HFAdapterConfig] = None
llm_model: GPT, hf_config: Optional[HFAdapterConfig] = None
) -> HFGPTAdapter:
"""
Создает адаптер из существующей llm модели.
Args:
llm_model: Обученная модель из библиотеки llm
hf_config: Конфигурация для HuggingFace
Returns:
HFGPTAdapter: Адаптированная модель
"""
if hf_config is None:
# Создаем конфигурацию из модели llm
hf_config = HFAdapterConfig.from_llm_config(llm_model.config)
# Преобразуем в PretrainedConfig
pretrained_config = HFPretrainedConfig(**hf_config.to_dict())
return HFGPTAdapter(pretrained_config, llm_model)
@staticmethod
def from_pretrained(
model_path: str,
hf_config: Optional[HFAdapterConfig] = None
model_path: str, hf_config: Optional[HFAdapterConfig] = None
) -> HFGPTAdapter:
"""
Загружает модель из чекпоинта и создает адаптер.
Args:
model_path: Путь к сохраненной модели
hf_config: Конфигурация для HuggingFace
Returns:
HFGPTAdapter: Адаптированная модель
"""
# Загружаем состояние модели
state_dict = torch.load(model_path, map_location='cpu')
state_dict = torch.load(model_path, map_location="cpu")
# Определяем конфигурацию из состояния модели или используем переданную
if hf_config is None:
# Пытаемся определить конфигурацию из состояния модели
# Это упрощенный подход - в реальности нужно сохранять конфигурацию отдельно
vocab_size = state_dict.get('_token_embeddings._embedding.weight', torch.zeros(50257, 768)).shape[0]
embed_dim = state_dict.get('_token_embeddings._embedding.weight', torch.zeros(50257, 768)).shape[1]
vocab_size = state_dict.get(
"_token_embeddings._embedding.weight", torch.zeros(50257, 768)
).shape[0]
embed_dim = state_dict.get(
"_token_embeddings._embedding.weight", torch.zeros(50257, 768)
).shape[1]
hf_config = HFAdapterConfig(
vocab_size=vocab_size,
hidden_size=embed_dim,
# Остальные параметры можно установить по умолчанию
)
pretrained_config = HFPretrainedConfig(**hf_config.to_dict())
# Создаем модель llm и загружаем веса
llm_config = {
"vocab_size": hf_config.vocab_size,
@@ -263,21 +264,17 @@ class HFAdapter:
"max_position_embeddings": hf_config.max_position_embeddings,
"dropout": hf_config.hidden_dropout_prob,
}
llm_model = GPT(llm_config)
llm_model.load_state_dict(state_dict)
return HFGPTAdapter(pretrained_config, llm_model)
@staticmethod
def save_pretrained(
model: HFGPTAdapter,
save_directory: str,
**kwargs
):
def save_pretrained(model: HFGPTAdapter, save_directory: str, **kwargs):
"""
Сохраняет адаптированную модель в формате HuggingFace.
Args:
model: Адаптированная модель
save_directory: Директория для сохранения
@@ -285,19 +282,19 @@ class HFAdapter:
"""
import os
import json
# Создаем директорию если не существует
os.makedirs(save_directory, exist_ok=True)
# Сохраняем конфигурацию
config_path = os.path.join(save_directory, "config.json")
with open(config_path, 'w', encoding='utf-8') as f:
with open(config_path, "w", encoding="utf-8") as f:
json.dump(model.config.to_dict(), f, indent=2, ensure_ascii=False)
# Сохраняем веса модели
model_path = os.path.join(save_directory, "pytorch_model.bin")
torch.save(model.llm_model.state_dict(), model_path)
# Сохраняем токенизатор если передан
if hasattr(kwargs, 'tokenizer') and kwargs['tokenizer'] is not None:
kwargs['tokenizer'].save_pretrained(save_directory)
if hasattr(kwargs, "tokenizer") and kwargs["tokenizer"] is not None:
kwargs["tokenizer"].save_pretrained(save_directory)

38
hf-proxy/src/hf_proxy/hf_config.py
View File

@@ -6,11 +6,12 @@ from dataclasses import dataclass, field
from typing import Dict, Any, Optional
from transformers import PretrainedConfig
@dataclass
class HFAdapterConfig:
"""
Конфигурация для адаптера HuggingFace.
Параметры:
model_type: Тип модели (gpt, llama, etc.)
vocab_size: Размер словаря
@@ -28,6 +29,7 @@ class HFAdapterConfig:
eos_token_id: ID токена конца строки
bos_token_id: ID токена начала строки
"""
model_type: str = "gpt"
vocab_size: int = 50257
hidden_size: int = 768
@@ -43,49 +45,50 @@ class HFAdapterConfig:
pad_token_id: int = 50256
eos_token_id: int = 50256
bos_token_id: int = 50256
# Дополнительные параметры для совместимости
architectures: list = field(default_factory=lambda: ["GPT2LMHeadModel"])
torch_dtype: str = "float32"
transformers_version: str = "4.44.0"
def to_dict(self) -> Dict[str, Any]:
"""Преобразует конфигурацию в словарь."""
return {
k: v for k, v in self.__dict__.items()
if not k.startswith('_') and not callable(v)
k: v
for k, v in self.__dict__.items()
if not k.startswith("_") and not callable(v)
}
@classmethod
def from_llm_config(cls, llm_config: Dict[str, Any]) -> "HFAdapterConfig":
"""
Создает конфигурацию HF из конфигурации llm.
Args:
llm_config: Конфигурация модели из библиотеки llm
Returns:
HFAdapterConfig: Конфигурация для HuggingFace
"""
# Маппинг параметров из llm в HF формат
mapping = {
"embed_dim": "hidden_size",
"num_layers": "num_hidden_layers",
"num_layers": "num_hidden_layers",
"num_heads": "num_attention_heads",
"max_position_embeddings": "max_position_embeddings",
"dropout": "hidden_dropout_prob",
"vocab_size": "vocab_size"
"vocab_size": "vocab_size",
}
hf_config_dict = {}
for llm_key, hf_key in mapping.items():
if llm_key in llm_config:
hf_config_dict[hf_key] = llm_config[llm_key]
# Устанавливаем промежуточный размер (обычно 4x hidden_size)
if "hidden_size" in hf_config_dict:
hf_config_dict["intermediate_size"] = hf_config_dict["hidden_size"] * 4
return cls(**hf_config_dict)
@@ -94,8 +97,9 @@ class HFPretrainedConfig(PretrainedConfig):
Конфигурация для предобученных моделей HuggingFace.
Наследуется от PretrainedConfig для полной совместимости.
"""
model_type = "gpt"
def __init__(
self,
vocab_size=50257,
@@ -112,15 +116,15 @@ class HFPretrainedConfig(PretrainedConfig):
pad_token_id=50256,
eos_token_id=50256,
bos_token_id=50256,
**kwargs
**kwargs,
):
super().__init__(
pad_token_id=pad_token_id,
eos_token_id=eos_token_id,
bos_token_id=bos_token_id,
**kwargs
**kwargs,
)
self.vocab_size = vocab_size
self.hidden_size = hidden_size
self.num_hidden_layers = num_hidden_layers

261
hf-proxy/src/hf_proxy/hf_tokenizer.py
View File

@@ -12,84 +12,82 @@ class HFTokenizerAdapter:
Упрощенный адаптер для кастомных токенизаторов llm.
Предоставляет совместимый с HuggingFace интерфейс.
"""
def __init__(self, llm_tokenizer: BaseTokenizer):
"""
Инициализация адаптера.
Args:
llm_tokenizer: Кастомный токенизатор из llm
"""
self.llm_tokenizer = llm_tokenizer
# Получаем словарь и размер
self._vocab = llm_tokenizer.get_vocab()
self.vocab_size = llm_tokenizer.get_vocab_size()
# Устанавливаем специальные токены
self.pad_token = getattr(llm_tokenizer, 'pad_token', '<pad>')
self.unk_token = getattr(llm_tokenizer, 'unk_token', '<unk>')
self.bos_token = getattr(llm_tokenizer, 'bos_token', '<bos>')
self.eos_token = getattr(llm_tokenizer, 'eos_token', '<eos>')
self.pad_token = getattr(llm_tokenizer, "pad_token", "<pad>")
self.unk_token = getattr(llm_tokenizer, "unk_token", "<unk>")
self.bos_token = getattr(llm_tokenizer, "bos_token", "<bos>")
self.eos_token = getattr(llm_tokenizer, "eos_token", "<eos>")
# Сохраняем ID специальных токенов
self.pad_token_id = getattr(llm_tokenizer, 'pad_token_id', 0)
self.unk_token_id = getattr(llm_tokenizer, 'unk_token_id', 1)
self.bos_token_id = getattr(llm_tokenizer, 'bos_token_id', 2)
self.eos_token_id = getattr(llm_tokenizer, 'eos_token_id', 3)
self.pad_token_id = getattr(llm_tokenizer, "pad_token_id", 0)
self.unk_token_id = getattr(llm_tokenizer, "unk_token_id", 1)
self.bos_token_id = getattr(llm_tokenizer, "bos_token_id", 2)
self.eos_token_id = getattr(llm_tokenizer, "eos_token_id", 3)
def __call__(self, text: str, **kwargs):
"""
Вызов токенизатора с параметрами как у HuggingFace.
Args:
text: Входной текст
**kwargs: Параметры токенизации
Returns:
dict: Словарь с токенами
"""
return_tensors = kwargs.get('return_tensors', None)
padding = kwargs.get('padding', False)
truncation = kwargs.get('truncation', False)
max_length = kwargs.get('max_length', None)
add_special_tokens = kwargs.get('add_special_tokens', True)
return_tensors = kwargs.get("return_tensors", None)
padding = kwargs.get("padding", False)
truncation = kwargs.get("truncation", False)
max_length = kwargs.get("max_length", None)
add_special_tokens = kwargs.get("add_special_tokens", True)
# Кодируем текст
#input_ids = self.llm_tokenizer.encode(
# text,
# input_ids = self.llm_tokenizer.encode(
# text,
# add_special_tokens=add_special_tokens
#)
# )
if isinstance(text, str):
input_ids = self.llm_tokenizer.encode(
text,
add_special_tokens=add_special_tokens
text, add_special_tokens=add_special_tokens
)
input_ids = [input_ids] # <-- оборачиваем в batch
else:
# Список строк, батч-режим!
input_ids = [
self.llm_tokenizer.encode(
t,
add_special_tokens=add_special_tokens
) for t in text
self.llm_tokenizer.encode(t, add_special_tokens=add_special_tokens)
for t in text
]
# Применяем truncation
if truncation and max_length is not None and len(input_ids) > max_length:
input_ids = input_ids[:max_length]
# Применяем padding
if padding and max_length is not None and len(input_ids) < max_length:
input_ids = input_ids + [self.pad_token_id] * (max_length - len(input_ids))
# Конвертируем в тензоры если нужно
if return_tensors == "pt":
import torch
input_ids = torch.tensor([input_ids])
return {"input_ids": input_ids}
def encode(
self,
text: str,
@@ -99,11 +97,11 @@ class HFTokenizerAdapter:
truncation: bool = False,
max_length: Optional[int] = None,
return_tensors: Optional[str] = None,
**kwargs
**kwargs,
) -> Union[List[int], List[List[int]]]:
"""
Кодирует текст в последовательность токенов.
Args:
text: Входной текст
text_pair: Второй текст (для парных задач)
@@ -112,84 +110,91 @@ class HFTokenizerAdapter:
truncation: Обрезать последовательность
max_length: Максимальная длина
return_tensors: Возвращать тензоры
Returns:
Список токенов или список списков токенов
"""
# Кодируем основной текст
token_ids = self.llm_tokenizer.encode(
text,
add_special_tokens=add_special_tokens
text, add_special_tokens=add_special_tokens
)
# Обрабатываем text_pair если есть
if text_pair is not None:
pair_ids = self.llm_tokenizer.encode(
text_pair,
add_special_tokens=False
)
pair_ids = self.llm_tokenizer.encode(text_pair, add_special_tokens=False)
token_ids.extend(pair_ids)
# Применяем truncation
if truncation and max_length is not None and len(token_ids) > max_length:
token_ids = token_ids[:max_length]
# Применяем padding
if padding and max_length is not None and len(token_ids) < max_length:
token_ids = token_ids + [self.pad_token_id] * (max_length - len(token_ids))
# Конвертируем в тензоры если нужно
if return_tensors == "pt":
import torch
return torch.tensor([token_ids])
elif return_tensors == "np":
import numpy as np
return np.array([token_ids])
return token_ids
def decode(
self,
token_ids: Union[int, List[int], List[List[int]]],
skip_special_tokens: bool = True,
**kwargs
**kwargs,
) -> str:
"""
Декодирует последовательность токенов в текст.
Args:
token_ids: ID токенов
skip_special_tokens: Пропускать специальные токены
Returns:
str: Декодированный текст
"""
# Обрабатываем разные форматы входных данных
if isinstance(token_ids, int):
token_ids = [token_ids]
elif isinstance(token_ids, list) and len(token_ids) > 0 and isinstance(token_ids[0], list):
elif (
isinstance(token_ids, list)
and len(token_ids) > 0
and isinstance(token_ids[0], list)
):
# Список списков - берем первый элемент
token_ids = token_ids[0]
# Фильтруем специальные токены если нужно
if skip_special_tokens:
special_ids = {self.pad_token_id, self.unk_token_id, self.bos_token_id, self.eos_token_id}
special_ids = {
self.pad_token_id,
self.unk_token_id,
self.bos_token_id,
self.eos_token_id,
}
token_ids = [tid for tid in token_ids if tid not in special_ids]
return self.llm_tokenizer.decode(token_ids)
def tokenize(self, text: str, **kwargs) -> List[str]:
"""
Токенизирует текст в список строковых токенов.
Args:
text: Входной текст
Returns:
List[str]: Список токенов
"""
return self.llm_tokenizer.tokenize(text)
def pad(
self,
encoded_inputs,
@@ -202,7 +207,7 @@ class HFTokenizerAdapter:
):
"""
Pad a list of encoded inputs.
Args:
encoded_inputs: List of encoded inputs
padding: Padding strategy
@@ -211,7 +216,7 @@ class HFTokenizerAdapter:
return_attention_mask: Return attention mask
return_tensors: Return tensors
verbose: Verbose mode
Returns:
Padded inputs
"""
@@ -224,47 +229,62 @@ class HFTokenizerAdapter:
# Обрабатываем разные типы данных
if isinstance(input_ids, int):
seq_len = 1
elif hasattr(input_ids, 'shape'):
seq_len = input_ids.shape[-1] if len(input_ids.shape) > 1 else len(input_ids)
elif hasattr(input_ids, "shape"):
seq_len = (
input_ids.shape[-1]
if len(input_ids.shape) > 1
else len(input_ids)
)
else:
seq_len = len(input_ids)
max_len = max(max_len, seq_len)
if max_length is not None:
max_len = min(max_len, max_length)
# Применяем padding
for item in encoded_inputs:
input_ids = item["input_ids"]
# Получаем текущую длину
if isinstance(input_ids, int):
current_len = 1
elif hasattr(input_ids, 'shape'):
current_len = input_ids.shape[-1] if len(input_ids.shape) > 1 else len(input_ids)
elif hasattr(input_ids, "shape"):
current_len = (
input_ids.shape[-1]
if len(input_ids.shape) > 1
else len(input_ids)
)
else:
current_len = len(input_ids)
if current_len < max_len:
# Дополняем pad_token_id
padding_length = max_len - current_len
# Обрабатываем разные типы данных
if isinstance(input_ids, int):
item["input_ids"] = [input_ids] + [self.pad_token_id] * padding_length
elif hasattr(input_ids, 'shape'):
item["input_ids"] = [input_ids] + [
self.pad_token_id
] * padding_length
elif hasattr(input_ids, "shape"):
import torch
padding_tensor = torch.full((padding_length,), self.pad_token_id, dtype=input_ids.dtype)
padding_tensor = torch.full(
(padding_length,), self.pad_token_id, dtype=input_ids.dtype
)
item["input_ids"] = torch.cat([input_ids, padding_tensor])
else:
item["input_ids"] = input_ids + [self.pad_token_id] * padding_length
item["input_ids"] = (
input_ids + [self.pad_token_id] * padding_length
)
# Добавляем attention_mask если требуется
if "attention_mask" in item:
mask = item["attention_mask"]
if isinstance(mask, int):
item["attention_mask"] = [mask] + [0] * padding_length
elif hasattr(mask, 'shape'):
elif hasattr(mask, "shape"):
padding_mask = torch.zeros(padding_length, dtype=mask.dtype)
item["attention_mask"] = torch.cat([mask, padding_mask])
else:
@@ -272,44 +292,49 @@ class HFTokenizerAdapter:
elif return_attention_mask:
if isinstance(input_ids, int):
item["attention_mask"] = [1] + [0] * padding_length
elif hasattr(input_ids, 'shape'):
elif hasattr(input_ids, "shape"):
attention_mask = torch.ones(current_len, dtype=torch.long)
padding_mask = torch.zeros(padding_length, dtype=torch.long)
item["attention_mask"] = torch.cat([attention_mask, padding_mask])
item["attention_mask"] = torch.cat(
[attention_mask, padding_mask]
)
else:
item["attention_mask"] = [1] * current_len + [0] * padding_length
item["attention_mask"] = [1] * current_len + [
0
] * padding_length
# Конвертируем в тензоры если требуется
if return_tensors == "pt":
import torch
for key in list(encoded_inputs[0].keys()):
if isinstance(encoded_inputs[0][key], list):
for i in range(len(encoded_inputs)):
encoded_inputs[i][key] = torch.tensor(encoded_inputs[i][key])
return encoded_inputs
def get_vocab(self) -> Dict[str, int]:
"""Возвращает словарь токенизатора."""
return self._vocab
def __len__(self) -> int:
"""Возвращает размер словаря."""
return self.vocab_size
def save_pretrained(self, save_directory: str, **kwargs):
"""
Сохраняет токенизатор в формате HuggingFace.
Args:
save_directory: Директория для сохранения
**kwargs: Дополнительные параметры
"""
import os
# Создаем директорию если не существует
os.makedirs(save_directory, exist_ok=True)
# Сохраняем конфигурацию токенизатора
tokenizer_config = {
"tokenizer_class": self.__class__.__name__,
@@ -324,77 +349,81 @@ class HFTokenizerAdapter:
"bos_token_id": self.bos_token_id,
"eos_token_id": self.eos_token_id,
}
config_path = os.path.join(save_directory, "tokenizer_config.json")
with open(config_path, 'w', encoding='utf-8') as f:
with open(config_path, "w", encoding="utf-8") as f:
json.dump(tokenizer_config, f, ensure_ascii=False, indent=2)
# Сохраняем словарь
vocab_path = os.path.join(save_directory, "vocab.json")
with open(vocab_path, 'w', encoding='utf-8') as f:
with open(vocab_path, "w", encoding="utf-8") as f:
json.dump(self._vocab, f, ensure_ascii=False, indent=2)
print(f"✅ Токенизатор сохранен в {save_directory}")
@classmethod
def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
"""
Загружает адаптированный токенизатор.
Args:
pretrained_model_name_or_path: Путь к сохраненному токенизатору
**kwargs: Дополнительные параметры
Returns:
HFTokenizerAdapter: Загруженный адаптер
"""
import os
# Проверяем, является ли путь директорией с файлами токенизатора
if os.path.isdir(pretrained_model_name_or_path):
# Загружаем из директории
config_path = os.path.join(pretrained_model_name_or_path, "tokenizer_config.json")
config_path = os.path.join(
pretrained_model_name_or_path, "tokenizer_config.json"
)
vocab_path = os.path.join(pretrained_model_name_or_path, "vocab.json")
if not os.path.exists(config_path) or not os.path.exists(vocab_path):
raise FileNotFoundError(
f"Файлы токенизатора не найдены в {pretrained_model_name_or_path}"
)
# Загружаем конфигурацию
with open(config_path, 'r', encoding='utf-8') as f:
with open(config_path, "r", encoding="utf-8") as f:
config = json.load(f)
# Определяем тип токенизатора llm
llm_tokenizer_type = config.get("llm_tokenizer_type", "BPETokenizer")
if llm_tokenizer_type == "BPETokenizer":
# Создаем BPETokenizer и загружаем словарь
llm_tokenizer = BPETokenizer()
# Загружаем словарь
with open(vocab_path, 'r', encoding='utf-8') as f:
with open(vocab_path, "r", encoding="utf-8") as f:
vocab = json.load(f)
llm_tokenizer.vocab = vocab
llm_tokenizer.inverse_vocab = {v: k for k, v in vocab.items()}
llm_tokenizer.vocab_size = len(vocab)
# Устанавливаем специальные токены
llm_tokenizer.pad_token = config.get("pad_token", "<pad>")
llm_tokenizer.unk_token = config.get("unk_token", "<unk>")
llm_tokenizer.bos_token = config.get("bos_token", "<bos>")
llm_tokenizer.eos_token = config.get("eos_token", "<eos>")
llm_tokenizer.pad_token_id = config.get("pad_token_id", 0)
llm_tokenizer.unk_token_id = config.get("unk_token_id", 1)
llm_tokenizer.bos_token_id = config.get("bos_token_id", 2)
llm_tokenizer.eos_token_id = config.get("eos_token_id", 3)
return cls(llm_tokenizer, **kwargs)
else:
raise ValueError(f"Неподдерживаемый тип токенизатора: {llm_tokenizer_type}")
raise ValueError(
f"Неподдерживаемый тип токенизатора: {llm_tokenizer_type}"
)
else:
# Пытаемся загрузить как файл llm токенизатора
try:
@@ -409,10 +438,10 @@ class HFTokenizerAdapter:
def create_hf_tokenizer(llm_tokenizer: BaseTokenizer) -> HFTokenizerAdapter:
"""
Создает адаптер HuggingFace для кастомного токенизатора.
Args:
llm_tokenizer: Токенизатор из библиотеки llm
Returns:
HFTokenizerAdapter: Адаптированный токенизатор
"""
@@ -422,7 +451,7 @@ def create_hf_tokenizer(llm_tokenizer: BaseTokenizer) -> HFTokenizerAdapter:
def convert_to_hf_format(llm_tokenizer: BaseTokenizer, save_directory: str):
"""
Конвертирует кастомный токенизатор в формат HuggingFace.
Args:
llm_tokenizer: Токенизатор из llm
save_directory: Директория для сохранения

163
hf-proxy/src/hf_proxy/hf_utils.py
View File

@@ -14,55 +14,57 @@ class HFUtils:
"""
Утилиты для работы с HuggingFace адаптером.
"""
@staticmethod
def create_hf_config_from_llm(llm_config: Dict[str, Any]) -> HFPretrainedConfig:
"""
Создает конфигурацию HuggingFace из конфигурации llm.
Args:
llm_config: Конфигурация модели из библиотеки llm
Returns:
HFPretrainedConfig: Конфигурация для HuggingFace
"""
adapter_config = HFAdapterConfig.from_llm_config(llm_config)
return HFPretrainedConfig(**adapter_config.to_dict())
@staticmethod
def convert_to_hf_format(
llm_model,
tokenizer = None,
model_name: str = "custom-gpt"
llm_model, tokenizer=None, model_name: str = "custom-gpt"
) -> tuple:
"""
Конвертирует llm модель в формат HuggingFace.
Args:
llm_model: Модель из библиотеки llm
tokenizer: Токенизатор (HF или кастомный)
model_name: Имя модели для сохранения
Returns:
tuple: (адаптированная модель, токенизатор)
"""
# Создаем адаптер
hf_model = HFAdapter.from_llm_model(llm_model)
# Если токенизатор не передан, создаем стандартный
if tokenizer is None:
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("gpt2")
# Устанавливаем специальные токены
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
elif hasattr(tokenizer, '__class__') and 'BPETokenizer' in str(tokenizer.__class__):
elif hasattr(tokenizer, "__class__") and "BPETokenizer" in str(
tokenizer.__class__
):
# Если передан наш кастомный токенизатор, создаем адаптер
from .hf_tokenizer import create_hf_tokenizer
tokenizer = create_hf_tokenizer(tokenizer)
return hf_model, tokenizer
@staticmethod
def push_to_hub(
model: HFGPTAdapter,
@@ -70,11 +72,11 @@ class HFUtils:
repo_name: str,
organization: Optional[str] = None,
private: bool = False,
**kwargs
**kwargs,
):
"""
Загружает модель в HuggingFace Hub.
Args:
model: Адаптированная модель
tokenizer: Токенизатор
@@ -85,23 +87,23 @@ class HFUtils:
"""
try:
from huggingface_hub import HfApi, ModelCard, create_repo
# Создаем репозиторий
if organization:
repo_id = f"{organization}/{repo_name}"
else:
repo_id = repo_name
create_repo(repo_id, private=private, exist_ok=True)
# Сохраняем модель локально
import tempfile
import os
with tempfile.TemporaryDirectory() as tmp_dir:
# Сохраняем модель
HFAdapter.save_pretrained(model, tmp_dir, tokenizer=tokenizer)
# Создаем Model Card
card = ModelCard.from_template(
model_name=repo_name,
@@ -110,46 +112,43 @@ class HFUtils:
tags=["llm", "gpt", "custom"],
)
card.save(os.path.join(tmp_dir, "README.md"))
# Загружаем в Hub
api = HfApi()
api.upload_folder(
folder_path=tmp_dir,
repo_id=repo_id,
commit_message="Initial commit with custom GPT model"
commit_message="Initial commit with custom GPT model",
)
print(f"✅ Модель успешно загружена в HuggingFace Hub: {repo_id}")
except ImportError:
raise ImportError(
"Для загрузки в HuggingFace Hub установите huggingface_hub: "
"pip install huggingface_hub"
)
@staticmethod
def load_from_hub(
repo_id: str,
**kwargs
) -> tuple:
def load_from_hub(repo_id: str, **kwargs) -> tuple:
"""
Загружает модель из HuggingFace Hub.
Args:
repo_id: ID репозитория
**kwargs: Дополнительные параметры
Returns:
tuple: (модель, токенизатор)
"""
from transformers import AutoTokenizer
# Загружаем токенизатор
tokenizer = AutoTokenizer.from_pretrained(repo_id, **kwargs)
# Загружаем конфигурацию
config = AutoConfig.from_pretrained(repo_id, **kwargs)
# Создаем модель llm на основе конфигурации
llm_config = {
"vocab_size": config.vocab_size,
@@ -159,63 +158,56 @@ class HFUtils:
"max_position_embeddings": config.max_position_embeddings,
"dropout": config.hidden_dropout_prob,
}
# Загружаем модель через адаптер
model = HFAdapter.from_pretrained(
f"{repo_id}/pytorch_model.bin",
HFAdapterConfig.from_llm_config(llm_config)
f"{repo_id}/pytorch_model.bin", HFAdapterConfig.from_llm_config(llm_config)
)
return model, tokenizer
@staticmethod
def compare_with_hf_model(
llm_model,
hf_model_name: str = "gpt2",
test_input: str = "Hello world"
llm_model, hf_model_name: str = "gpt2", test_input: str = "Hello world"
) -> Dict[str, Any]:
"""
Сравнивает llm модель с эталонной моделью из HuggingFace.
Args:
llm_model: Модель из библиотеки llm
hf_model_name: Имя модели HuggingFace для сравнения
test_input: Тестовый вход
Returns:
Dict: Результаты сравнения
"""
from transformers import AutoModelForCausalLM, AutoTokenizer
# Загружаем эталонную модель
hf_tokenizer = AutoTokenizer.from_pretrained(hf_model_name)
hf_model = AutoModelForCausalLM.from_pretrained(hf_model_name)
# Подготавливаем входные данные
inputs = hf_tokenizer(test_input, return_tensors="pt")
# Получаем логиты от обеих моделей
with torch.no_grad():
hf_logits = hf_model(**inputs).logits
llm_logits = llm_model(inputs['input_ids'])
llm_logits = llm_model(inputs["input_ids"])
# Сравниваем результаты
hf_probs = torch.softmax(hf_logits[0, -1], dim=-1)
llm_probs = torch.softmax(llm_logits[0, -1], dim=-1)
# Вычисляем метрики
kl_divergence = torch.nn.functional.kl_div(
torch.log(llm_probs + 1e-8),
hf_probs,
reduction='batchmean'
torch.log(llm_probs + 1e-8), hf_probs, reduction="batchmean"
)
cosine_similarity = torch.nn.functional.cosine_similarity(
hf_logits.flatten(),
llm_logits.flatten(),
dim=0
hf_logits.flatten(), llm_logits.flatten(), dim=0
)
return {
"kl_divergence": kl_divergence.item(),
"cosine_similarity": cosine_similarity.item(),
@@ -228,58 +220,52 @@ class TokenizerWrapper:
"""
Обертка для токенизатора с дополнительными утилитами.
"""
def __init__(self, tokenizer):
self.tokenizer = tokenizer
def encode_batch(self, texts: List[str], **kwargs) -> Dict[str, torch.Tensor]:
"""
Кодирует батч текстов.
Args:
texts: Список текстов
**kwargs: Дополнительные параметры токенизации
Returns:
Dict: Токенизированные данные
"""
return self.tokenizer(
texts,
padding=True,
truncation=True,
return_tensors="pt",
**kwargs
texts, padding=True, truncation=True, return_tensors="pt", **kwargs
)
def decode_batch(self, token_ids: torch.Tensor, **kwargs) -> List[str]:
"""
Декодирует батч токенов.
Args:
token_ids: Тензор с токенами
**kwargs: Дополнительные параметры декодирования
Returns:
List[str]: Декодированные тексты
"""
if token_ids.dim() == 1:
token_ids = token_ids.unsqueeze(0)
texts = []
for i in range(token_ids.size(0)):
text = self.tokenizer.decode(
token_ids[i],
skip_special_tokens=True,
**kwargs
token_ids[i], skip_special_tokens=True, **kwargs
)
texts.append(text)
return texts
def get_vocab_size(self) -> int:
"""Возвращает размер словаря."""
return len(self.tokenizer)
def get_special_tokens(self) -> Dict[str, int]:
"""Возвращает специальные токены."""
return {
@@ -290,36 +276,27 @@ class TokenizerWrapper:
}
def create_hf_pipeline(
llm_model,
tokenizer=None,
device: str = "auto",
**kwargs
):
def create_hf_pipeline(llm_model, tokenizer=None, device: str = "auto", **kwargs):
"""
Создает HuggingFace pipeline из llm модели.
Args:
llm_model: Модель из библиотеки llm
tokenizer: Токенизатор
device: Устройство для вычислений
**kwargs: Дополнительные параметры pipeline
Returns:
transformers.Pipeline: Готовый pipeline
"""
from transformers import pipeline
# Конвертируем модель в HF формат
hf_model, tokenizer = HFUtils.convert_to_hf_format(llm_model, tokenizer)
# Создаем pipeline
pipe = pipeline(
"text-generation",
model=hf_model,
tokenizer=tokenizer,
device=device,
**kwargs
"text-generation", model=hf_model, tokenizer=tokenizer, device=device, **kwargs
)
return pipe

18
llm/src/llm/core/base_model.py
View File

@@ -19,23 +19,25 @@ from abc import ABC, abstractmethod
from typing import Optional, Tuple
import torch
class BaseModel(nn.Module, ABC):
"""
Абстрактный класс — стандарт для всех архитектур LLM.
Научная идея:
Реализация унифицированного входа/выхода для поддержки построения и обучения любых современных языковых моделей.
Args:
config (dict): Параметры архитектуры (размерность эмбеддингов, число слоев, heads и т.д.)
Attributes:
config (dict): Конфиг модели
"""
def __init__(self, config: dict):
"""
Инициализация модели.
Args:
config (dict): Настройки архитектуры модели (размеры слоев, типы блоков и т.д.)
"""
@@ -43,10 +45,12 @@ class BaseModel(nn.Module, ABC):
self.config = config
@abstractmethod
def forward(self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
def forward(
self, input_ids: torch.Tensor, attention_mask: Optional[torch.Tensor] = None
) -> torch.Tensor:
"""
Прямой проход — получение логитов для входных токенов.
Args:
input_ids (Tensor[int]): Индексы токенов [batch, seq_len]
attention_mask (Optional[Tensor[bool]]): Маска разрешенных позиций (если требуется) [batch, seq_len]
@@ -59,7 +63,7 @@ class BaseModel(nn.Module, ABC):
def generate(self, input_ids: torch.Tensor, max_length: int = 50) -> torch.Tensor:
"""
Генерация текста (авторегрессивно, greedy или sampling).
Args:
input_ids (Tensor[int]): Начальные токены [batch, start_len]
max_length (int): Максимальная длина последовательности

10
llm/src/llm/core/cached_decoder.py
View File

@@ -6,6 +6,7 @@ from .feed_forward import FeedForward
from .multi_head_attention import MultiHeadAttention
from .rope import RoPE
class CachedDecoder(nn.Module):
"""
Универсальный декодерный блок для современных LLM (GPT, LLaMA, др.), поддерживает кэширование key-value для эффективной генерации.
@@ -28,7 +29,7 @@ class CachedDecoder(nn.Module):
norm_layer (тип nn.Module): Normalization слой (LayerNorm или RMSNorm)
dropout (float): Dropout
rope (RoPE|None): Экземпляр RoPE (для LLaMA)
Пример (GPT2 style):
>>> decoder = CachedDecoder(
... feed_forward_layer=FeedForward(...),
@@ -36,6 +37,7 @@ class CachedDecoder(nn.Module):
... num_heads=4, emb_size=256, head_size=64, max_seq_len=128)
>>> out, cache = decoder(x, use_cache=True)
"""
def __init__(
self,
feed_forward_layer: nn.Module,
@@ -49,7 +51,7 @@ class CachedDecoder(nn.Module):
):
"""
Инициализация декодера с кэшированием.
Поведение аналогично блоку TransformerDecoderLayer,
но с гибкой возможностью подмены любых подкомпонент (активация, norm, позиции).
@@ -85,7 +87,7 @@ class CachedDecoder(nn.Module):
):
"""
Прямой проход с поддержкой кэша.
Args:
x (Tensor[float]): [batch, seq_len, emb_size] — скрытые состояния
mask (Optional[Tensor]): маска внимания (или causal mask), shape [seq_len, seq_len]
@@ -111,4 +113,4 @@ class CachedDecoder(nn.Module):
if use_cache:
return (result, kv_caches)
else:
return (result, None)
return (result, None)

21
llm/src/llm/core/decoder.py
View File

@@ -3,6 +3,7 @@ import torch
from .feed_forward import FeedForward
from .multi_head_attention import MultiHeadAttention
class Decoder(nn.Module):
"""
Базовый автогерессивный блок-декодер трансформера (без кэша KV).
@@ -24,12 +25,14 @@ class Decoder(nn.Module):
>>> out = decoder(x)
>>> print(out.shape) # torch.Size([1, 10, 512])
"""
def __init__(self,
def __init__(
self,
num_heads: int,
emb_size: int,
head_size: int,
max_seq_len: int,
dropout: float = 0.1
dropout: float = 0.1,
):
"""
Инициализация декодера.
@@ -43,11 +46,11 @@ class Decoder(nn.Module):
"""
super().__init__()
self._heads = MultiHeadAttention(
num_heads=num_heads,
emb_size=emb_size,
head_size=head_size,
max_seq_len=max_seq_len,
dropout=dropout
num_heads=num_heads,
emb_size=emb_size,
head_size=head_size,
max_seq_len=max_seq_len,
dropout=dropout,
)
self._ff = FeedForward(emb_size=emb_size, dropout=dropout)
self._norm1 = nn.LayerNorm(emb_size)
@@ -73,7 +76,7 @@ class Decoder(nn.Module):
# Self-Attention блок
attention, _ = self._heads(x, mask, use_cache=False, cache=None)
out = self._norm1(attention + x)
# FeedForward блок
ffn_out = self._ff(out)
return self._norm2(ffn_out + out)
return self._norm2(ffn_out + out)

19
llm/src/llm/core/feed_forward.py
View File

@@ -16,7 +16,7 @@ class FeedForward(nn.Module):
- После внимания каждому токену применяется одинаковая двухслойная нейросеть.
- Дает глубокую нелинейность; позволяет модели не только сопоставлять, но и моделировать сложные связи между токенами.
- Изначально предложен в «Attention is All You Need» (Vaswani et al., 2017).
Формула:
FFN(x) = Dropout(W2·act(W1·x))
где act — ReLU, GELU и др., обычно expansion x4.
@@ -33,22 +33,23 @@ class FeedForward(nn.Module):
- Добавляет нелинейность в архитектуру трансформера
- Обеспечивает взаимодействие между различными размерностями эмбеддингов
- Работает независимо для каждого токена в последовательности
Args:
emb_size (int): размерность входных эмбеддингов
dropout (float): вероятность(dropout)
activation (str): нелинейная функция (relu, gelu, gelu_exact)
Пример:
>>> ff = FeedForward(emb_size=512, dropout=0.1)
>>> x = torch.randn(32, 10, 512)
>>> output = ff(x)
>>> print(output.shape) # torch.Size([32, 10, 512])
"""
def __init__(self, emb_size: int, dropout: float = 0.1, activation: str = "relu"):
"""
Инициализация слоя Feed Forward Network.
Args:
emb_size: Размерность входных эмбеддингов
dropout: Вероятность dropout для регуляризации (по умолчанию: 0.1)
@@ -73,23 +74,23 @@ class FeedForward(nn.Module):
def forward(self, x: torch.Tensor):
"""
Прямой проход через слой Feed Forward Network.
Args:
x: Входной тензор размерности [batch_size, seq_len, emb_size]
Returns:
Тензор той же размерности, что и входной
"""
# Сохраняем dtype входных данных
input_dtype = x.dtype
# Приводим веса к нужному типу если необходимо
if input_dtype != self._layer1.weight.dtype:
self._layer1 = self._layer1.to(dtype=input_dtype)
self._layer2 = self._layer2.to(dtype=input_dtype)
# Пропустим тензор x по очереди через все созданные слои
x = self._layer1(x)
x = self._activation(x)
x = self._layer2(x)
return self._dropout(x)
return self._dropout(x)

12
llm/src/llm/core/gelu.py
View File

@@ -1,6 +1,7 @@
import torch
from torch import nn
class GELU(nn.Module):
"""
Гауссовская Эрф-активация (GELU, Gaussian Error Linear Unit).
@@ -17,11 +18,14 @@ class GELU(nn.Module):
>>> y = gelu(torch.tensor([-1.0, 0.0, 1.0]))
>>> print(y)
"""
def __init__(self):
super().__init__()
self.sqrt_2_over_pi = torch.sqrt(torch.tensor(2.0) / math.pi)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return 0.5 * x * (1 + torch.tanh(
self.sqrt_2_over_pi * (x + 0.044715 * torch.pow(x, 3))
))
return (
0.5
* x
* (1 + torch.tanh(self.sqrt_2_over_pi * (x + 0.044715 * torch.pow(x, 3))))
)

44
llm/src/llm/core/head_attention.py
View File

@@ -4,6 +4,7 @@ import torch.nn.functional as F
from math import sqrt
from .rope import RoPE
class HeadAttention(nn.Module):
"""
Одноголовый механизм внимания (scaled dot-product attention) — фундаментальный строительный блок всех современных Transformer.
@@ -11,11 +12,11 @@ class HeadAttention(nn.Module):
Научная суть:
- Attention учит модель самостоятельно "выбирать" важные связи между словами, независимо от их положения.
- Механизм causal mask гарантирует невозможность "заглядывания в будущее" при генерации (авторегрессия).
Формула:
Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) · V
(Q — запросы, K — ключи, V — значения; d_k — размерность ключа)
Поддерживает Rotary Position Encoding (RoPE) для относительного позиционного кодирования.
Args:
@@ -28,14 +29,17 @@ class HeadAttention(nn.Module):
- Использует нижнетреугольную маску для предотвращения "заглядывания в будущее"
- Автоматически адаптируется к разным версиям PyTorch
- Поддерживает batch-обработку входных данных
Пример использования:
>>> attention = HeadAttention(emb_size=64, head_size=32, max_seq_len=128)
>>> x = torch.randn(1, 10, 64)
>>> output, _ = attention(x)
>>> print(output.shape) # torch.Size([1, 10, 32])
"""
def __init__(self, emb_size: int, head_size: int, max_seq_len: int, rope: RoPE = None):
def __init__(
self, emb_size: int, head_size: int, max_seq_len: int, rope: RoPE = None
):
super().__init__()
self._emb_size = emb_size
self._head_size = head_size
@@ -49,21 +53,25 @@ class HeadAttention(nn.Module):
# Создание causal маски
mask = torch.tril(torch.ones(max_seq_len, max_seq_len))
self.register_buffer('_tril_mask', mask.bool() if hasattr(torch, 'bool') else mask.byte())
self.register_buffer(
"_tril_mask", mask.bool() if hasattr(torch, "bool") else mask.byte()
)
def forward(self, x: torch.Tensor, use_cache: bool = True, cache: tuple = None) -> tuple:
def forward(
self, x: torch.Tensor, use_cache: bool = True, cache: tuple = None
) -> tuple:
"""
Прямой проход через слой внимания.
Аргументы:
x (torch.Tensor): Входной тензор формы [batch_size, seq_len, emb_size]
Возвращает:
torch.Tensor: Выходной тензор формы [batch_size, seq_len, head_size]
Исключения:
ValueError: Если длина последовательности превышает max_seq_len
Пример внутренних преобразований:
Для входа x.shape = [2, 5, 64]:
1. Q/K/V преобразования -> [2, 5, 32]
@@ -73,7 +81,9 @@ class HeadAttention(nn.Module):
"""
seq_len = x.shape[1]
if seq_len > self._max_seq_len:
raise ValueError(f"Длина последовательности {seq_len} превышает максимум {self._max_seq_len}")
raise ValueError(
f"Длина последовательности {seq_len} превышает максимум {self._max_seq_len}"
)
k = self._k(x) # [B, T, hs]
q = self._q(x) # [B, T, hs]
@@ -88,16 +98,18 @@ class HeadAttention(nn.Module):
k_cache, v_cache = cache
k = torch.cat([k_cache, k], dim=1) # [B, cache_len + T, hs]
v = torch.cat([v_cache, v], dim=1) # [B, cache_len + T, hs]
scores = q @ k.transpose(-2, -1) / sqrt(self._head_size)
if cache is None:
scores = scores.masked_fill(~self._tril_mask[:seq_len, :seq_len], float('-inf'))
scores = scores.masked_fill(
~self._tril_mask[:seq_len, :seq_len], float("-inf")
)
weights = F.softmax(scores, dim=-1)
x_out = weights @ v # [B, T, hs]
if use_cache is True:
return (x_out, (k, v))
else:
return (x_out, None)
return (x_out, None)

57
llm/src/llm/core/multi_head_attention.py
View File

@@ -3,6 +3,7 @@ import torch
from .head_attention import HeadAttention
from .rope import RoPE
class MultiHeadAttention(nn.Module):
"""
Мультиголовый (многоголовый) механизм внимания — ключевой компонент любого Transformer.
@@ -12,7 +13,7 @@ class MultiHeadAttention(nn.Module):
чтобы видеть разные связи в последовательности (разный контекст, локально/глобально).
- Каждый attention блок работает независимо, выход конкатенируется.
- Механизм предложен в статье "Attention is All You Need" (Vaswani et al., 2017).
Формула внимания для одной головы:
Attention(Q, K, V) = softmax(QK^T/sqrt(d_k))·V
Мультиголовый:
@@ -32,7 +33,16 @@ class MultiHeadAttention(nn.Module):
>>> out, cache = mha(x)
>>> print(out.shape)
"""
def __init__(self, num_heads: int, emb_size: int, head_size: int, max_seq_len: int, rope: RoPE = None, dropout: float = 0.1):
def __init__(
self,
num_heads: int,
emb_size: int,
head_size: int,
max_seq_len: int,
rope: RoPE = None,
dropout: float = 0.1,
):
"""
Инициализация многоголового внимания.
@@ -49,18 +59,27 @@ class MultiHeadAttention(nn.Module):
- max_seq_len зависит от задачи (512 для BERT, 2048 для GPT-3)
"""
super().__init__()
self._heads = nn.ModuleList([
HeadAttention(
emb_size=emb_size,
head_size=head_size,
max_seq_len=max_seq_len,
rope=rope,
) for _ in range(num_heads)
])
self._heads = nn.ModuleList(
[
HeadAttention(
emb_size=emb_size,
head_size=head_size,
max_seq_len=max_seq_len,
rope=rope,
)
for _ in range(num_heads)
]
)
self._layer = nn.Linear(head_size * num_heads, emb_size)
self._dropout = nn.Dropout(dropout)
def forward(self, x: torch.Tensor, mask: torch.Tensor = None, use_cache: bool = True, cache: list = None):
def forward(
self,
x: torch.Tensor,
mask: torch.Tensor = None,
use_cache: bool = True,
cache: list = None,
):
"""
Прямой проход (forward):
Для каждого токена оценивает "важность" остальных токенов сразу через несколько attention-блоков.
@@ -76,7 +95,7 @@ class MultiHeadAttention(nn.Module):
4. Линейная проекция:
- Выход: [batch_size, seq_len, emb_size]
5. Применение dropout
Args:
x (Tensor[float]): [batch, seq_len, emb_size] — вход
mask (Optional[Tensor[bool]]): маска позиции [seq_len, seq_len]
@@ -94,11 +113,11 @@ class MultiHeadAttention(nn.Module):
Пример преобразований для emb_size=512, num_heads=8:
Вход: [4, 100, 512]
-> Каждая голова: [4, 100, 64]
-> После внимания: 8 x [4, 100, 64]
-> После внимания: 8 x [4, 100, 64]
-> Конкатенация: [4, 100, 512]
-> Проекция: [4, 100, 512]
-> Dropout: [4, 100, 512]
Пример:
>>> out, caches = mha(x)
>>> out.shape # [batch, seq_len, emb_size]
@@ -109,20 +128,20 @@ class MultiHeadAttention(nn.Module):
head_cache = cache[i] if cache is not None else None
result = head(x, use_cache=use_cache, cache=head_cache)
attention_results.append(result)
outputs, caches = zip(*attention_results)
attention_outputs = list(outputs)
kv_caches = list(caches)
# 2. Объединяем результаты всех голов
concatenated_attention = torch.cat(attention_outputs, dim=-1)
# 3. Проецируем в пространство эмбеддингов
projected_output = self._layer(concatenated_attention)
# 4. Применяем dropout для регуляризации
final_output = self._dropout(projected_output)
if use_cache is True:
return (final_output, kv_caches)
else:

24
llm/src/llm/core/positional_embeddings.py
View File

@@ -1,13 +1,14 @@
import torch
from torch import nn, Tensor
class PositionalEmbeddings(nn.Module):
"""
Обучаемые позиционные эмбеддинги (learnable positional embeddings).
Позиционные эмбеддинги используются в нейросетях для передачи информации
Позиционные эмбеддинги используются в нейросетях для передачи информации
о позиции элементов в последовательности (например, в Transformer).
Научная суть:
- Трансформеры не используют рекуррентность, а значит сами по себе не различают порядок слов.
- Позиционные эмбеддинги добавляются к токеновым, чтобы сеть понимала, в каком месте последовательности находится каждый токен.
@@ -16,7 +17,7 @@ class PositionalEmbeddings(nn.Module):
Args:
max_seq_len (int): максимальная длина последовательности
emb_size (int): размер вектора позиции
Пример использования:
>>> pos_encoder = PositionalEmbeddings(max_seq_len=100, emb_size=256)
>>> # Получить эмбеддинги для последовательности из 10 элементов
@@ -36,23 +37,22 @@ class PositionalEmbeddings(nn.Module):
self.max_seq_len = max_seq_len
self.emb_size = emb_size
self.embedding = nn.Embedding(
num_embeddings=max_seq_len,
embedding_dim=emb_size
num_embeddings=max_seq_len, embedding_dim=emb_size
)
def forward(self, seq_len: int, start_pos: int = 0) -> Tensor:
"""
Возвращает позиционные эмбеддинги для заданной длины последовательности.
Args:
seq_len (int): Длина последовательности (1 <= seq_len <= max_seq_len)
Returns:
Tensor: Тензор позиционных эмбеддингов формы [seq_len, emb_size]
Raises:
IndexError: Если seq_len выходит за допустимые границы
Пример:
>>> pos_encoder = PositionalEmbeddings(100, 64)
>>> emb = pos_encoder(10) # Тензор 10x64
@@ -62,5 +62,9 @@ class PositionalEmbeddings(nn.Module):
if start_pos == 0:
positions = torch.arange(seq_len, device=self.embedding.weight.device)
else:
positions = torch.arange(start=start_pos, end=start_pos + seq_len, device=self.embedding.weight.device)
positions = torch.arange(
start=start_pos,
end=start_pos + seq_len,
device=self.embedding.weight.device,
)
return self.embedding(positions)

22
llm/src/llm/core/rms_norm.py
View File

@@ -33,23 +33,23 @@ class RMSNorm(nn.Module):
- Упрощенный вариант LayerNorm без вычисления среднего, только деление на rms.
- Лучшая численная стабильность на больших моделях, меньше вычислений.
- Применяется в LLaMA, PaLM и др.
Формула:
RMSNorm(x) = (x / sqrt(mean(x²) + eps)) * w (w — обучаемый вектор)
Args:
dim (int): размер последнего измерения (обычно emb_size)
eps (float): для численной устойчивости
Пример:
>>> norm = RMSNorm(emb_size)
>>> out = norm(x)
"""
def __init__(self, dim: int, eps: float = 1e-6):
"""
Инициализация RMSNorm слоя.
Args:
dim: Размерность нормализуемого измерения
eps: Малое значение для численной стабильности (по умолчанию 1e-6)
@@ -57,27 +57,27 @@ class RMSNorm(nn.Module):
super().__init__()
self._eps = eps
self._w = nn.Parameter(torch.ones(dim))
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Прямой проход через RMSNorm слой.
Args:
x: Входной тензор формы [..., dim]
Returns:
Нормализованный тензор той же формы, что и входной
Формула:
output = w * (x / sqrt(mean(x²) + eps))
"""
# Вычисление RMS (Root Mean Square) по последнему измерению
rms = (x.pow(2).mean(-1, keepdim=True) + self._eps) ** 0.5
# Нормализация и масштабирование
norm_x = x / rms
return self._w * norm_x
def extra_repr(self) -> str:
"""Строковое представление для отладки."""
return f'dim={self._w.shape[0]}, eps={self._eps}'
return f"dim={self._w.shape[0]}, eps={self._eps}"

36
llm/src/llm/core/rope.py
View File

@@ -26,72 +26,72 @@ from typing import Optional
class RoPE(nn.Module):
"""
Rotary Positional Embeddings (RoPE) для механизма внимания.
Кодирует позиционную информацию через вращение векторов запросов и ключей
в многомерном пространстве с использованием синусов и косинусов.
Args:
head_size: Размерность головы внимания (должен быть четным)
max_seq_len: Максимальная длина последовательности
base: Базовое значение для вычисления частот (по умолчанию 10000)
Attributes:
cos_matrix: Буферизованная матрица косинусов формы [max_seq_len, head_size//2]
sin_matrix: Буферизованная матрица синусов формы [max_seq_len, head_size//2]
"""
def __init__(self, head_size: int, max_seq_len: int, base: int = 10_000):
"""
Инициализация RoPE эмбеддингов.
Args:
head_size: Размерность головы внимания (должен быть четным)
max_seq_len: Максимальная поддерживаемая длина последовательности
base: Базовое значение для вычисления частот (типично 10000)
Raises:
AssertionError: Если head_size не четный
"""
super().__init__()
assert head_size % 2 == 0, "head_size должен быть четным"
# Вычисление частот: θ_i = base^(-2i/d) для i ∈ [0, d/2-1]
freqs = 1.0 / (base ** (2 * torch.arange(head_size // 2).float() / head_size))
# Позиции от 0 до max_seq_len-1
positions = torch.arange(max_seq_len).float()
# Внешнее произведение: m * θ_i для всех позиций и частот
freq_matrix = positions.unsqueeze(1) * freqs.unsqueeze(0)
# Предвычисление матриц косинусов и синусов
self.register_buffer('cos_matrix', torch.cos(freq_matrix))
self.register_buffer('sin_matrix', torch.sin(freq_matrix))
self.register_buffer("cos_matrix", torch.cos(freq_matrix))
self.register_buffer("sin_matrix", torch.sin(freq_matrix))
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Применение ротационного позиционного кодирования к входному тензору.
Args:
x: Входной тензор формы [batch_size, seq_len, head_size]
Returns:
Тензор с примененным RoPE формы [batch_size, seq_len, head_size]
Алгоритм:
1. Разделение векторов на четные и нечетные компоненты
2. Применение вращения через синусы и косинусы
3. Объединение компонент обратно
"""
seq_len = x.size(1)
# Берем нужную часть матриц и приводим к типу x
cos = self.cos_matrix[:seq_len].to(x.dtype) # [seq_len, head_size//2]
sin = self.sin_matrix[:seq_len].to(x.dtype) # [seq_len, head_size//2]
# Разделяем на четные и нечетные компоненты
x_even = x[:, :, 0::2] # [batch_size, seq_len, head_size//2]
x_odd = x[:, :, 1::2] # [batch_size, seq_len, head_size//2]
x_odd = x[:, :, 1::2] # [batch_size, seq_len, head_size//2]
# Применяем поворот: q' = q * cos(mθ) + rotate(q) * sin(mθ)
x_rotated_even = x_even * cos - x_odd * sin
@@ -101,4 +101,4 @@ class RoPE(nn.Module):
x_rotated = torch.stack([x_rotated_even, x_rotated_odd], dim=-1)
x_rotated = x_rotated.flatten(-2) # [batch_size, seq_len, head_size]
return x_rotated
return x_rotated

8
llm/src/llm/core/silu.py
View File

@@ -1,19 +1,21 @@
import torch
from torch import nn
class SiLU(nn.Module):
"""
SiLU (Swish) — современная активационная функция для нейросетей.
Научная суть:
- Формула: $SiLU(x) = x * \sigm(x)$, где $\sigm(x)$ — сигмоида.
- Более гладкая альтернатива ReLU, улучшает поток градиентов в глубоких сетях.
- Используется во многих «state-of-the-art» архитектурах (SwiGLU, PaLM, LLaMA).
- Также известна как Swish (Ramachandran et al, 2017).
- Также известна как Swish (Ramachandran et al, 2017).
Пример:
>>> act = SiLU()
>>> x = torch.tensor([-1.0, 0.0, 1.0])
>>> print(act(x))
"""
def forward(self, x: torch.Tensor):
return torch.sigmoid(x) * x
return torch.sigmoid(x) * x

34
llm/src/llm/core/swi_glu.py
View File

@@ -27,13 +27,13 @@ class SwiGLU(nn.Module):
SwiGLU (Swish-Gated Linear Unit) — современная нелинейность для архитектур LLM (LLaMA, PaLM).
Реализация SwiGLU активационной функции.
Состоит из трех линейных слоев и активации SiLU:
1. Gate слой + SiLU активация
2. Up слой (линейное преобразование)
3. Element-wise multiplication gate и up
4. Down слой (линейная проекция)
Научная суть:
- Сохраняет преимущества GLU (раздельные гейтом и телом) + мощность Swish/SiLU активации.
- Дает надежную гладкую активацию, хорошо работает на больших масштабах.
@@ -50,11 +50,11 @@ class SwiGLU(nn.Module):
>>> ff = SwiGLU(emb_size=512, dropout=0.1)
>>> y = ff(torch.randn(2,10,512))
"""
def __init__(self, emb_size: int, dropout: float = 0.1):
"""
Инициализация SwiGLU слоя.
Args:
emb_size: Размерность входных/выходных эмбеддингов
dropout: Вероятность dropout (по умолчанию 0.1)
@@ -69,13 +69,13 @@ class SwiGLU(nn.Module):
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Прямой проход через SwiGLU слой.
Args:
x: Входной тензор формы [batch_size, seq_len, emb_size]
Returns:
Выходной тензор формы [batch_size, seq_len, emb_size]
Алгоритм:
1. gate = SiLU(linear_gate(x))
2. up = linear_up(x)
@@ -83,19 +83,19 @@ class SwiGLU(nn.Module):
4. apply dropout
"""
# Gate ветвь: линейное преобразование + активация
gate_out = self._gate(x) # [batch, seq, 4*emb]
activation_out = self._activation(gate_out) # [batch, seq, 4*emb]
gate_out = self._gate(x) # [batch, seq, 4*emb]
activation_out = self._activation(gate_out) # [batch, seq, 4*emb]
# Up ветвь: линейное преобразование
up_out = self._up(x) # [batch, seq, 4*emb]
up_out = self._up(x) # [batch, seq, 4*emb]
# Element-wise multiplication (gating mechanism)
out = up_out * activation_out # поэлементное умножение!
out = up_out * activation_out # поэлементное умножение!
# Final projection and dropout
out = self._down(out) # [batch, seq, emb]
out = self._down(out) # [batch, seq, emb]
return self._dropout(out)
def extra_repr(self) -> str:
"""Строковое представление для отладки."""
return f'emb_size={self._gate.in_features}, dropout={self._dropout.p}'
return f"emb_size={self._gate.in_features}, dropout={self._dropout.p}"

16
llm/src/llm/core/token_embeddings.py
View File

@@ -2,13 +2,14 @@ import torch
from torch import nn
from torch import Tensor
class TokenEmbeddings(nn.Module):
"""
Токеновые эмбеддинги — обучаемые векторные представления для каждого токена словаря.
Преобразует целочисленные индексы токенов в обучаемые векторные представления фиксированного размера.
Обычно используется как первый слой в нейронных сетях для задач NLP.
Научная суть:
- Первый шаг для любого NLP-модуля: вместо индекса токена подаём его dense-вектор.
- Эти вектора изучаются в процессе обучения и отражают скрытые взаимосвязи между токенами.
@@ -22,18 +23,18 @@ class TokenEmbeddings(nn.Module):
Примечание:
- Индексы должны быть в диапазоне [0, vocab_size-1]
- Эмбеддинги инициализируются случайно и обучаются в процессе тренировки модели
Пример:
>>> emb = TokenEmbeddings(vocab_size=10000, emb_size=256)
>>> tokens = torch.tensor([[1, 2, 3]])
>>> vecs = emb(tokens)
>>> vecs.shape # torch.Size([1, 3, 256])
"""
def __init__(self, vocab_size: int, emb_size: int):
super().__init__()
self._embedding = nn.Embedding(
num_embeddings=vocab_size,
embedding_dim=emb_size
num_embeddings=vocab_size, embedding_dim=emb_size
)
def forward(self, x: Tensor) -> Tensor:
@@ -55,10 +56,7 @@ if __name__ == "__main__":
embedding = TokenEmbeddings(vocab_size=100, emb_size=128)
# Создаем тензор с индексами в пределах vocab_size (0-99)
tensor = torch.tensor([
[11, 45, 76, 34],
[34, 67, 45, 54]
])
tensor = torch.tensor([[11, 45, 76, 34], [34, 67, 45, 54]])
# Проверяем индексы
if (tensor >= 100).any():
@@ -66,4 +64,4 @@ if __name__ == "__main__":
output = embedding(tensor)
print("Embeddings shape:", output.shape)
print(f"{output.shape} | {output.mean().item():.11f}") # Формат как в ТЗ
print(f"{output.shape} | {output.mean().item():.11f}") # Формат как в ТЗ

138
llm/src/llm/models/gpt/gpt.py
View File

@@ -34,10 +34,10 @@ from llm.core.positional_embeddings import PositionalEmbeddings
class GPT(BaseModel):
"""
Original GPT (Generative Pre-trained Transformer) модель.
Первая версия трансформерной архитектуры от OpenAI, предназначенная
для генеративного предобучения на текстовых данных.
Args:
config: Словарь конфигурации с параметрами:
- vocab_size: Размер словаря токенов
@@ -46,7 +46,7 @@ class GPT(BaseModel):
- num_layers: Количество декодерных слоев
- max_position_embeddings: Максимальная длина последовательности
- dropout: Вероятность dropout
Attributes:
_token_embeddings: Слой векторных представлений токенов
_position_embeddings: Слой позиционных эмбеддингов
@@ -54,30 +54,34 @@ class GPT(BaseModel):
_norm: Финальный слой нормализации
_linear: Выходной линейный слой
"""
def __init__(self, config):
super().__init__(config)
# Инициализация слоев
self._max_seq_len = config["max_position_embeddings"]
self._token_embeddings = TokenEmbeddings(
vocab_size=config["vocab_size"],
emb_size=config["embed_dim"]
vocab_size=config["vocab_size"], emb_size=config["embed_dim"]
)
self._position_embeddings = PositionalEmbeddings(
max_seq_len=config["max_position_embeddings"],
emb_size=config["embed_dim"]
max_seq_len=config["max_position_embeddings"], emb_size=config["embed_dim"]
)
self._dropout = nn.Dropout(config["dropout"])
# head_size = emb_size // num_heads
self._decoders = nn.ModuleList([Decoder(
num_heads=config["num_heads"],
emb_size=config["embed_dim"],
head_size=config["embed_dim"] // config["num_heads"],
max_seq_len=config["max_position_embeddings"],
dropout=config["dropout"]
) for _ in range(config["num_layers"])])
self._decoders = nn.ModuleList(
[
Decoder(
num_heads=config["num_heads"],
emb_size=config["embed_dim"],
head_size=config["embed_dim"] // config["num_heads"],
max_seq_len=config["max_position_embeddings"],
dropout=config["dropout"],
)
for _ in range(config["num_layers"])
]
)
self._linear = nn.Linear(config["embed_dim"], config["vocab_size"])
@property
def max_seq_len(self):
"""Возвращает максимальную длину последовательности."""
@@ -85,57 +89,60 @@ class GPT(BaseModel):
def forward(self, x: torch.Tensor, attention_mask=None) -> torch.Tensor:
"""Прямой проход через GPT
Args:
x: Входной тензор [batch_size, seq_len]
Returns:
Тензор логитов [batch_size, seq_len, vocab_size]
"""
# Проверка длины последовательности
if x.size(1) > self._max_seq_len:
raise ValueError(f"Длина последовательности {x.size(1)} превышает максимальную {self._max_seq_len}")
raise ValueError(
f"Длина последовательности {x.size(1)} превышает максимальную {self._max_seq_len}"
)
# Эмбеддинги токенов и позиций
tok_out = self._token_embeddings(x) # [batch, seq_len, emb_size]
pos_out = self._position_embeddings(x.size(1)) # [seq_len, emb_size]
# Комбинирование
out = self._dropout(tok_out + pos_out.unsqueeze(0)) # [batch, seq_len, emb_size]
out = self._dropout(
tok_out + pos_out.unsqueeze(0)
) # [batch, seq_len, emb_size]
# Стек декодеров
for decoder in self._decoders:
out = decoder(out)
return self._linear(out) # [batch, seq_len, vocab_size]
# def forward(self, input_ids, attention_mask=None):
# B, T = input_ids.size()
# pos = torch.arange(0, T, device=input_ids.device).unsqueeze(0)
#
# x = self.token_emb(input_ids) + self.pos_emb(pos)
#
# for block in self.blocks:
# x = block(x, attention_mask)
#
# x = self.ln_f(x)
# logits = self.head(x)
# return logits
# def forward(self, input_ids, attention_mask=None):
# B, T = input_ids.size()
# pos = torch.arange(0, T, device=input_ids.device).unsqueeze(0)
#
# x = self.token_emb(input_ids) + self.pos_emb(pos)
#
# for block in self.blocks:
# x = block(x, attention_mask)
#
# x = self.ln_f(x)
# logits = self.head(x)
# return logits
def generate(self,
x: torch.Tensor,
max_new_tokens: int,
def generate(
self,
x: torch.Tensor,
max_new_tokens: int,
do_sample: bool,
temperature: float = 1.0,
top_k: int = None,
top_p: float = None,
attention_mask: torch.Tensor = None, # Добавляем для совместимости с HF
**kwargs # Игнорируем остальные параметры
**kwargs, # Игнорируем остальные параметры
) -> torch.Tensor:
"""Авторегрессивная генерация текста.
Параметры:
x: Входной тензор с индексами токенов формы [batch_size, seq_len],
где batch_size - размер батча, seq_len - длина последовательности.
@@ -157,9 +164,9 @@ class GPT(BaseModel):
- Гарантируется, что хотя бы один токен остаётся (даже если его вероятность > top_p)
- None: отключено (по умолчанию)
- Должен быть в диапазоне (0, 1]
Возвращает:
torch.Tensor: Тензор с расширенной последовательностью токенов формы
torch.Tensor: Тензор с расширенной последовательностью токенов формы
[batch_size, seq_len + max_new_tokens]
Исключения:
@@ -172,7 +179,7 @@ class GPT(BaseModel):
Примеры:
>>> # Жадная генерация
>>> output = model.generate(input_ids, max_new_tokens=10, do_sample=False)
>>>
>>>
>>> # Вероятностная генерация с top-k
>>> output = model.generate(input_ids, max_new_tokens=10, do_sample=True, top_k=50)
>>>
@@ -180,11 +187,11 @@ class GPT(BaseModel):
>>> output = model.generate(input_ids, max_new_tokens=10, do_sample=True, top_p=0.9)
>>>
>>> # Комбинация температуры и top-k
>>> output = model.generate(input_ids, max_new_tokens=10, do_sample=True,
>>> output = model.generate(input_ids, max_new_tokens=10, do_sample=True,
... temperature=0.7, top_k=50)
Примечания:
1. Для детерминированных результатов в режиме сэмплирования
1. Для детерминированных результатов в режиме сэмплирования
зафиксируйте random seed (torch.manual_seed).
2. Температура влияет только на режим сэмплирования (do_sample=True).
3. Одновременное использование top_k и top_p запрещено.
@@ -204,7 +211,7 @@ class GPT(BaseModel):
Должна быть > 0 (по умолчанию: 1.0)
Returns:
torch.Tensor: Тензор с расширенной последовательностью токенов формы
torch.Tensor: Тензор с расширенной последовательностью токенов формы
[batch_size, seq_len + max_new_tokens]
Raises:
@@ -222,13 +229,13 @@ class GPT(BaseModel):
>>> output = model.generate(input_ids, max_new_tokens=10, do_sample=True, temperature=1.5)
Note:
Для детерминированных результатов в режиме сэмплирования
Для детерминированных результатов в режиме сэмплирования
зафиксируйте random seed (torch.manual_seed).
Температура влияет только на режим сэмплирования (do_sample=True).
"""
for _ in range(max_new_tokens):
# 1. Обрезаем вход, если последовательность слишком длинная
x_cond = x[:, -self._max_seq_len:]
x_cond = x[:, -self._max_seq_len :]
# 2. Передаем последовательность в метод forward класса GPT и полуаем логиты.
logits = self.forward(x_cond)
@@ -250,9 +257,14 @@ class GPT(BaseModel):
vocab_size = logits_scaled.size(-1)
# создаём маску: True, если токен НЕ в topk_indices
mask = torch.ones_like(logits_scaled, dtype=torch.bool if hasattr(torch, 'bool') else torch.uint8)
mask.scatter_(1, topk_indices, False if hasattr(torch, 'bool') else 0) # False там, где top-k индексы
masked_logits[mask] = float('-inf')
mask = torch.ones_like(
logits_scaled,
dtype=torch.bool if hasattr(torch, "bool") else torch.uint8,
)
mask.scatter_(
1, topk_indices, False if hasattr(torch, "bool") else 0
) # False там, где top-k индексы
masked_logits[mask] = float("-inf")
logits_scaled = masked_logits
@@ -260,36 +272,42 @@ class GPT(BaseModel):
# 1. Применим softmax, чтобы получить вероятности:
probs = F.softmax(logits_scaled, dim=-1) # [B, vocab_size]
# 2. Отсортируем токены по убыванию вероятностей:
sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=-1)
sorted_probs, sorted_indices = torch.sort(
probs, descending=True, dim=-1
)
# 3. Посчитаем кумулятивную сумму вероятностей:
cum_probs = torch.cumsum(sorted_probs, dim=-1) # [B, vocab_size]
# 4. Определим маску: оставить токены, пока сумма < top_p
sorted_mask = (cum_probs <= top_p) # [B, vocab_size]
sorted_mask = cum_probs <= top_p # [B, vocab_size]
# Гарантируем, что хотя бы первый токен останется
sorted_mask[:, 0] = True
# 5. Преобразуем маску обратно в оригинальный порядок:
# Создаём полную маску из False
mask = torch.zeros_like(probs, dtype=torch.bool if hasattr(torch, 'bool') else torch.uint8)
mask = torch.zeros_like(
probs, dtype=torch.bool if hasattr(torch, "bool") else torch.uint8
)
# Устанавливаем True в местах нужных токенов
mask.scatter_(dim=1, index=sorted_indices, src=sorted_mask)
# 6. Зануляем логиты токенов вне топ-p:
logits_scaled[~mask] = float('-inf')
logits_scaled[~mask] = float("-inf")
# 4. Применяем Softmax
probs = F.softmax(logits_scaled, dim=-1) # [batch_size, vocab_size]
if do_sample == True:
# 5. Если do_sample равен True, то отбираем токен случайно с помощью torch.multinomial
next_token = torch.multinomial(probs, num_samples=1) # [batch_size, 1]
else:
# 5. Если do_sample равен False, то выбираем токен с максимальной вероятностью
next_token = torch.argmax(probs, dim=-1, keepdim=True) # [batch_size, 1]
next_token = torch.argmax(
probs, dim=-1, keepdim=True
) # [batch_size, 1]
# 6. Добавляем его к последовательности
x = torch.cat([x, next_token], dim=1) # [batch_size, seq_len+1]
return x
# def generate(self, input_ids, max_length=50):
# for _ in range(max_length):
# logits = self.forward(input_ids)

94
llm/src/llm/models/gpt/gpt2.py
View File

@@ -27,6 +27,7 @@ from llm.core.positional_embeddings import PositionalEmbeddings
from llm.core.cached_decoder import CachedDecoder
from llm.core.feed_forward import FeedForward
class GPT2(BaseModel):
"""
GPT2 — автогерессивная языковая модель, архитектура Transformer, предложенная OpenAI.
@@ -35,7 +36,7 @@ class GPT2(BaseModel):
- Масштабируемый автогерессивный трансформер для предсказания токенов слева направо.
- Главное отличие от классической GPT: порядок layer normalization ПЕРЕД attention и FFN.
- Используется GELU, efficient KV-cache, несет наследие классической GPT, но делает архитектуру глубже/шире.
Args:
config (dict): параметры архитектуры (vocab_size, embed_dim, num_heads, num_layers, max_position_embeddings, dropout)
@@ -44,37 +45,43 @@ class GPT2(BaseModel):
>>> logits = model(input_ids)
>>> out = model.generate(input_ids, max_length=20)
"""
def __init__(self, config):
super().__init__(config)
# Инициализация слоев
self._max_seq_len = config["max_position_embeddings"]
self._token_embeddings = TokenEmbeddings(
vocab_size=config["vocab_size"],
emb_size=config["embed_dim"]
vocab_size=config["vocab_size"], emb_size=config["embed_dim"]
)
self._position_embeddings = PositionalEmbeddings(
max_seq_len=config["max_position_embeddings"],
emb_size=config["embed_dim"]
max_seq_len=config["max_position_embeddings"], emb_size=config["embed_dim"]
)
self._dropout = nn.Dropout(config["dropout"])
# head_size = emb_size // num_heads
self._decoders = nn.ModuleList([CachedDecoder(
num_heads=config["num_heads"],
emb_size=config["embed_dim"],
head_size=config["embed_dim"] // config["num_heads"],
feed_forward_layer=FeedForward(
emb_size=config["embed_dim"],
dropout=config["dropout"],
activation="gelu"
),
max_seq_len=config["max_position_embeddings"],
dropout=config["dropout"]
) for _ in range(config["num_layers"])])
self._decoders = nn.ModuleList(
[
CachedDecoder(
num_heads=config["num_heads"],
emb_size=config["embed_dim"],
head_size=config["embed_dim"] // config["num_heads"],
feed_forward_layer=FeedForward(
emb_size=config["embed_dim"],
dropout=config["dropout"],
activation="gelu",
),
max_seq_len=config["max_position_embeddings"],
dropout=config["dropout"],
)
for _ in range(config["num_layers"])
]
)
self._norm = nn.LayerNorm(config["embed_dim"])
self._linear = nn.Linear(config["embed_dim"], config["vocab_size"])
def forward(self, x: torch.Tensor, use_cache: bool = True, cache: list = None) -> tuple:
def forward(
self, x: torch.Tensor, use_cache: bool = True, cache: list = None
) -> tuple:
"""
Прямой проход GPT2:
- Все слои работают как autoregressive transformer (masked self-attention).
@@ -91,9 +98,10 @@ class GPT2(BaseModel):
"""
# Проверка длины последовательности (только при отсутствии кэша)
if cache is None and x.size(1) > self._max_seq_len:
raise ValueError(f"Длина последовательности {x.size(1)} превышает максимальную {self.max_seq_len}")
raise ValueError(
f"Длина последовательности {x.size(1)} превышает максимальную {self.max_seq_len}"
)
# Вычисление start_pos из кэша (если кэш передан)
if cache is not None:
# При кэше обрабатываем только один токен (последний)
@@ -111,11 +119,15 @@ class GPT2(BaseModel):
# Эмбеддинги токенов и позиций
tok_out = self._token_embeddings(x) # [batch, seq_len, emb_size]
pos_out = self._position_embeddings(seq_len, start_pos=start_pos) # [seq_len, emb_size]
pos_out = self._position_embeddings(
seq_len, start_pos=start_pos
) # [seq_len, emb_size]
# Комбинирование
out = self._dropout(tok_out + pos_out.unsqueeze(0)) # [batch, seq_len, emb_size]
out = self._dropout(
tok_out + pos_out.unsqueeze(0)
) # [batch, seq_len, emb_size]
# Стек декодеров с передачей кэша
new_cache = []
for i, decoder in enumerate(self._decoders):
@@ -131,21 +143,22 @@ class GPT2(BaseModel):
out = self._norm(out)
logits = self._linear(out)
# Возвращаем результат с учетом use_cache
if use_cache:
return (logits, new_cache)
else:
return (logits, None)
def generate(self,
x: torch.Tensor,
max_new_tokens: int,
def generate(
self,
x: torch.Tensor,
max_new_tokens: int,
do_sample: bool,
temperature: float = 1.0,
top_k: int = None,
top_p: float = None,
use_cache: bool = True
use_cache: bool = True,
) -> torch.Tensor:
"""
Генерация текста с использованием autoregressive трансформера (GPT2).
@@ -174,10 +187,10 @@ class GPT2(BaseModel):
else:
# Первая итерация или кэш отключен - передаем всю последовательность
x_input = x
# Прямой проход с кэшем
logits, new_cache = self.forward(x_input, use_cache=use_cache, cache=cache)
# Обновляем кэш для следующей итерации
if use_cache:
cache = new_cache
@@ -200,7 +213,7 @@ class GPT2(BaseModel):
# создаём маску: 1, если токен НЕ в topk_indices
mask = torch.ones_like(logits_scaled, dtype=torch.uint8)
mask.scatter_(1, topk_indices, 0) # 0 там, где top-k индексы
masked_logits[mask.byte()] = float('-inf')
masked_logits[mask.byte()] = float("-inf")
logits_scaled = masked_logits
@@ -208,7 +221,9 @@ class GPT2(BaseModel):
# 1. Применим softmax, чтобы получить вероятности:
probs = F.softmax(logits_scaled, dim=-1) # [B, vocab_size]
# 2. Отсортируем токены по убыванию вероятностей:
sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=-1)
sorted_probs, sorted_indices = torch.sort(
probs, descending=True, dim=-1
)
# 3. Посчитаем кумулятивную сумму вероятностей:
cum_probs = torch.cumsum(sorted_probs, dim=-1) # [B, vocab_size]
# 4. Определим маску: оставить токены, пока сумма < top_p
@@ -221,23 +236,24 @@ class GPT2(BaseModel):
# Устанавливаем 1 в местах нужных токенов
mask.scatter_(dim=1, index=sorted_indices, src=sorted_mask)
# 6. Зануляем логиты токенов вне топ-p:
logits_scaled[~mask] = float('-inf')
logits_scaled[~mask] = float("-inf")
# 4. Применяем Softmax
probs = F.softmax(logits_scaled, dim=-1) # [batch_size, vocab_size]
if do_sample == True:
# 5. Если do_sample равен True, то отбираем токен случайно с помощью torch.multinomial
next_token = torch.multinomial(probs, num_samples=1) # [batch_size, 1]
else:
# 5. Если do_sample равен False, то выбираем токен с максимальной вероятностью
next_token = torch.argmax(probs, dim=-1, keepdim=True) # [batch_size, 1]
next_token = torch.argmax(
probs, dim=-1, keepdim=True
) # [batch_size, 1]
# 6. Добавляем его к последовательности
x = torch.cat([x, next_token], dim=1) # [batch_size, seq_len+1]
return x
@property
def max_seq_len(self) -> int:
return self._max_seq_len
return self._max_seq_len

131
llm/src/llm/models/llama/llama.py
View File

@@ -10,7 +10,6 @@ from llm.core.rope import RoPE
from llm.core.cached_decoder import CachedDecoder
class Llama(BaseModel):
"""
LLaMA (Large Language Model Meta AI) — высокоэффективная масштабируемая языковая модель, разработанная Meta AI Research.
@@ -29,38 +28,45 @@ class Llama(BaseModel):
>>> logits, cache = model(input_ids, use_cache=True)
>>> out = model.generate(input_ids, max_new_tokens=20)
"""
def __init__(self,config):
def __init__(self, config):
super().__init__(config)
# Инициализация слоев
self._max_seq_len = config["max_position_embeddings"]
self._token_embeddings = TokenEmbeddings(
vocab_size=config["vocab_size"],
emb_size=config["embed_dim"]
vocab_size=config["vocab_size"], emb_size=config["embed_dim"]
)
self._position_embeddings = RoPE(
head_size=config["embed_dim"] // config["num_heads"],
max_seq_len=config["max_position_embeddings"]
max_seq_len=config["max_position_embeddings"],
)
self._dropout = nn.Dropout(config["dropout"])
self._decoders = nn.ModuleList([CachedDecoder(
norm_layer=RMSNorm,
num_heads=config["num_heads"],
emb_size=config["embed_dim"],
head_size=config["embed_dim"] // config["num_heads"],
feed_forward_layer=SwiGLU(
emb_size=config["embed_dim"],
dropout=config["dropout"],
),
max_seq_len=config["max_position_embeddings"],
rope=self._position_embeddings,
dropout=config["dropout"],
) for _ in range(config["num_layers"])])
self._decoders = nn.ModuleList(
[
CachedDecoder(
norm_layer=RMSNorm,
num_heads=config["num_heads"],
emb_size=config["embed_dim"],
head_size=config["embed_dim"] // config["num_heads"],
feed_forward_layer=SwiGLU(
emb_size=config["embed_dim"],
dropout=config["dropout"],
),
max_seq_len=config["max_position_embeddings"],
rope=self._position_embeddings,
dropout=config["dropout"],
)
for _ in range(config["num_layers"])
]
)
self._norm = RMSNorm(config["embed_dim"])
self._linear = nn.Linear(config["embed_dim"], config["vocab_size"])
def forward(self, x: torch.Tensor, use_cache: bool = True, cache: list = None) -> tuple:
def forward(
self, x: torch.Tensor, use_cache: bool = True, cache: list = None
) -> tuple:
"""
Прямой проход через LLaMA (inference/train): авторегрессионное предсказание токенов.
@@ -76,11 +82,12 @@ class Llama(BaseModel):
"""
# Проверка длины последовательности (только при отсутствии кэша)
if cache is None and x.size(1) > self._max_seq_len:
raise ValueError(f"Длина последовательности {x.size(1)} превышает максимальную {self.max_seq_len}")
raise ValueError(
f"Длина последовательности {x.size(1)} превышает максимальную {self.max_seq_len}"
)
# Вычисление start_pos из кэша (если кэш передан)
#if cache is not None:
# if cache is not None:
# # При кэше обрабатываем только один токен (последний)
# seq_len = 1
# # Вычисляем start_pos из самого нижнего уровня кэша
@@ -89,18 +96,18 @@ class Llama(BaseModel):
# start_pos = key_cache.size(1) # cache_len
# else:
# start_pos = 0
#else:
# else:
# # Без кэша работаем как раньше
# start_pos = 0
# seq_len = x.size(1)
# Эмбеддинги токенов и позиций
tok_out = self._token_embeddings(x) # [batch, seq_len, emb_size]
#pos_out = self._position_embeddings(x) # [batch, seq_len, emb_size]
# pos_out = self._position_embeddings(x) # [batch, seq_len, emb_size]
# Комбинирование
out = self._dropout(tok_out) # [batch, seq_len, emb_size]
# Стек декодеров с передачей кэша
new_cache = []
for i, decoder in enumerate(self._decoders):
@@ -116,42 +123,43 @@ class Llama(BaseModel):
out = self._norm(out)
logits = self._linear(out)
# Возвращаем результат с учетом use_cache
if use_cache:
return (logits, new_cache)
else:
return (logits, None)
def generate(self,
x: torch.Tensor,
max_new_tokens: int,
def generate(
self,
x: torch.Tensor,
max_new_tokens: int,
do_sample: bool,
temperature: float = 1.0,
top_k: int = None,
top_p: float = None,
use_cache: bool = True
use_cache: bool = True,
) -> torch.Tensor:
"""
Генерация текста c помощью LLaMA (autoregressive Transformer).
Поддерживается:
- greedy и вероятностное сэмплирование (top-k, top-p, temperature)
- кэш attention для ускорения генерации длинных последовательностей
Генерация текста c помощью LLaMA (autoregressive Transformer).
Поддерживается:
- greedy и вероятностное сэмплирование (top-k, top-p, temperature)
- кэш attention для ускорения генерации длинных последовательностей
Args:
x (Tensor[int]): начальная последовательность [batch, seq_len]
max_new_tokens (int): сколько новых токенов сгенерировать
do_sample (bool): использовать стохастику (True) или жадный выбор (False)
temperature (float): масштаб для softmax (важно для sampling)
top_k (int|None): ограничение на количество кандидатов (top-k sampling)
top_p (float|None): nucleus sampling
use_cache (bool): ускоряет autoregressive при длинной генерации
Returns:
output (Tensor[int]): [batch, seq_len + max_new_tokens]
Пример:
>>> prompt = tokenizer.encode('Meta AI', return_tensors="pt")
>>> generated = model.generate(prompt, max_new_tokens=30, do_sample=True)
>>> print(tokenizer.decode(generated[0]))
Args:
x (Tensor[int]): начальная последовательность [batch, seq_len]
max_new_tokens (int): сколько новых токенов сгенерировать
do_sample (bool): использовать стохастику (True) или жадный выбор (False)
temperature (float): масштаб для softmax (важно для sampling)
top_k (int|None): ограничение на количество кандидатов (top-k sampling)
top_p (float|None): nucleus sampling
use_cache (bool): ускоряет autoregressive при длинной генерации
Returns:
output (Tensor[int]): [batch, seq_len + max_new_tokens]
Пример:
>>> prompt = tokenizer.encode('Meta AI', return_tensors="pt")
>>> generated = model.generate(prompt, max_new_tokens=30, do_sample=True)
>>> print(tokenizer.decode(generated[0]))
"""
cache = None
@@ -162,10 +170,10 @@ class Llama(BaseModel):
else:
# Первая итерация или кэш отключен - передаем всю последовательность
x_input = x
# Прямой проход с кэшем
logits, new_cache = self.forward(x_input, use_cache=use_cache, cache=cache)
# Обновляем кэш для следующей итерации
if use_cache:
cache = new_cache
@@ -188,7 +196,7 @@ class Llama(BaseModel):
# создаём маску: 1, если токен НЕ в topk_indices
mask = torch.ones_like(logits_scaled, dtype=torch.uint8)
mask.scatter_(1, topk_indices, 0) # 0 там, где top-k индексы
masked_logits[mask.byte()] = float('-inf')
masked_logits[mask.byte()] = float("-inf")
logits_scaled = masked_logits
@@ -196,7 +204,9 @@ class Llama(BaseModel):
# 1. Применим softmax, чтобы получить вероятности:
probs = F.softmax(logits_scaled, dim=-1) # [B, vocab_size]
# 2. Отсортируем токены по убыванию вероятностей:
sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=-1)
sorted_probs, sorted_indices = torch.sort(
probs, descending=True, dim=-1
)
# 3. Посчитаем кумулятивную сумму вероятностей:
cum_probs = torch.cumsum(sorted_probs, dim=-1) # [B, vocab_size]
# 4. Определим маску: оставить токены, пока сумма < top_p
@@ -209,25 +219,24 @@ class Llama(BaseModel):
# Устанавливаем 1 в местах нужных токенов
mask.scatter_(dim=1, index=sorted_indices, src=sorted_mask)
# 6. Зануляем логиты токенов вне топ-p:
logits_scaled[~mask] = float('-inf')
logits_scaled[~mask] = float("-inf")
# 4. Применяем Softmax
probs = F.softmax(logits_scaled, dim=-1) # [batch_size, vocab_size]
if do_sample == True:
# 5. Если do_sample равен True, то отбираем токен случайно с помощью torch.multinomial
next_token = torch.multinomial(probs, num_samples=1) # [batch_size, 1]
else:
# 5. Если do_sample равен False, то выбираем токен с максимальной вероятностью
next_token = torch.argmax(probs, dim=-1, keepdim=True) # [batch_size, 1]
next_token = torch.argmax(
probs, dim=-1, keepdim=True
) # [batch_size, 1]
# 6. Добавляем его к последовательности
x = torch.cat([x, next_token], dim=1) # [batch_size, seq_len+1]
return x
@property
def max_seq_len(self) -> int:
return self._max_seq_len
return self._max_seq_len

98
llm/src/llm/tokenizers/base_tokenizer.py
View File

@@ -10,92 +10,94 @@ import json
class BaseTokenizer(ABC):
"""
Абстрактный базовый класс для всех токенизаторов.
Определяет общий интерфейс для токенизации текста.
"""
def __init__(self):
self.vocab: Dict[str, int] = {}
self.inverse_vocab: Dict[int, str] = {}
self.vocab_size: int = 0
# Специальные токены
self.pad_token = "<pad>"
self.unk_token = "<unk>"
self.bos_token = "<bos>"
self.eos_token = "<eos>"
self.pad_token_id: Optional[int] = None
self.unk_token_id: Optional[int] = None
self.bos_token_id: Optional[int] = None
self.eos_token_id: Optional[int] = None
@abstractmethod
def train(self, texts: List[str], vocab_size: int = 1000, **kwargs):
"""
Обучение токенизатора на текстах.
Args:
texts: Список текстов для обучения
vocab_size: Желаемый размер словаря
**kwargs: Дополнительные параметры обучения
"""
pass
@abstractmethod
def encode(self, text: str, **kwargs) -> List[int]:
"""
Кодирование текста в последовательность токенов.
Args:
text: Входной текст
**kwargs: Дополнительные параметры кодирования
Returns:
List[int]: Список идентификаторов токенов
"""
pass
@abstractmethod
def decode(self, tokens: List[int], **kwargs) -> str:
"""
Декодирование последовательности токенов в текст.
Args:
tokens: Список идентификаторов токенов
**kwargs: Дополнительные параметры декодирования
Returns:
str: Декодированный текст
"""
pass
def tokenize(self, text: str, **kwargs) -> List[str]:
"""
Токенизация текста в список строковых токенов.
Args:
text: Входной текст
**kwargs: Дополнительные параметры
Returns:
List[str]: Список токенов
"""
token_ids = self.encode(text, **kwargs)
return [self.inverse_vocab.get(token_id, self.unk_token) for token_id in token_ids]
return [
self.inverse_vocab.get(token_id, self.unk_token) for token_id in token_ids
]
def get_vocab(self) -> Dict[str, int]:
"""Возвращает словарь токенизатора."""
return self.vocab.copy()
def get_vocab_size(self) -> int:
"""Возвращает размер словаря."""
return self.vocab_size
def add_special_tokens(self, special_tokens: List[str]):
"""
Добавляет специальные токены в словарь.
Args:
special_tokens: Список специальных токенов
"""
@@ -105,70 +107,70 @@ class BaseTokenizer(ABC):
self.vocab[token] = token_id
self.inverse_vocab[token_id] = token
self.vocab_size += 1
# Обновляем ID специальных токенов
self.pad_token_id = self.vocab.get(self.pad_token)
self.unk_token_id = self.vocab.get(self.unk_token)
self.bos_token_id = self.vocab.get(self.bos_token)
self.eos_token_id = self.vocab.get(self.eos_token)
def save(self, filepath: str):
"""
Сохраняет токенизатор в файл.
Args:
filepath: Путь для сохранения
"""
config = {
'vocab': self.vocab,
'vocab_size': self.vocab_size,
'pad_token': self.pad_token,
'unk_token': self.unk_token,
'bos_token': self.bos_token,
'eos_token': self.eos_token,
'tokenizer_type': self.__class__.__name__
"vocab": self.vocab,
"vocab_size": self.vocab_size,
"pad_token": self.pad_token,
"unk_token": self.unk_token,
"bos_token": self.bos_token,
"eos_token": self.eos_token,
"tokenizer_type": self.__class__.__name__,
}
with open(filepath, 'w', encoding='utf-8') as f:
with open(filepath, "w", encoding="utf-8") as f:
json.dump(config, f, ensure_ascii=False, indent=2)
@classmethod
def load(cls, filepath: str):
"""
Загружает токенизатор из файла.
Args:
filepath: Путь к файлу
Returns:
BaseTokenizer: Загруженный токенизатор
"""
with open(filepath, 'r', encoding='utf-8') as f:
with open(filepath, "r", encoding="utf-8") as f:
config = json.load(f)
# Создаем экземпляр токенизатора
tokenizer = cls()
tokenizer.vocab = config['vocab']
tokenizer.vocab_size = config['vocab_size']
tokenizer.pad_token = config['pad_token']
tokenizer.unk_token = config['unk_token']
tokenizer.bos_token = config['bos_token']
tokenizer.eos_token = config['eos_token']
tokenizer.vocab = config["vocab"]
tokenizer.vocab_size = config["vocab_size"]
tokenizer.pad_token = config["pad_token"]
tokenizer.unk_token = config["unk_token"]
tokenizer.bos_token = config["bos_token"]
tokenizer.eos_token = config["eos_token"]
# Создаем обратный словарь
tokenizer.inverse_vocab = {v: k for k, v in tokenizer.vocab.items()}
# Обновляем ID специальных токенов
tokenizer.pad_token_id = tokenizer.vocab.get(tokenizer.pad_token)
tokenizer.unk_token_id = tokenizer.vocab.get(tokenizer.unk_token)
tokenizer.bos_token_id = tokenizer.vocab.get(tokenizer.bos_token)
tokenizer.eos_token_id = tokenizer.vocab.get(tokenizer.eos_token)
return tokenizer
def __len__(self) -> int:
"""Возвращает размер словаря."""
return self.vocab_size
def __repr__(self) -> str:
return f"{self.__class__.__name__}(vocab_size={self.vocab_size})"

272
llm/src/llm/tokenizers/bpe_tokenizer copy.py
View File

@@ -13,26 +13,26 @@ from .base_tokenizer import BaseTokenizer
class BPETokenizer(BaseTokenizer):
"""
BPE токенизатор для обработки текста.
Реализует алгоритм Byte Pair Encoding для создания субсловных токенов.
Примеры использования:
>>> tokenizer = BPETokenizer()
>>> tokenizer.train(["пример текста для обучения"], vocab_size=1000)
>>> tokens = tokenizer.encode("новый текст")
>>> text = tokenizer.decode(tokens)
"""
def __init__(self):
super().__init__()
self.merges: Dict[Tuple[str, str], int] = {}
self.pattern = r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""
self.compiled_pattern = re.compile(self.pattern, re.UNICODE)
def train(self, texts: List[str], vocab_size: int = 1000, **kwargs):
"""
Обучение BPE токенизатора на текстах.
Args:
texts: Список текстов для обучения
vocab_size: Желаемый размер словаря
@@ -42,37 +42,40 @@ class BPETokenizer(BaseTokenizer):
"""
# Инициализация базового словаря
self._initialize_vocab()
# Добавляем специальные токены если указаны
special_tokens = kwargs.get('special_tokens', [self.pad_token, self.unk_token, self.bos_token, self.eos_token])
special_tokens = kwargs.get(
"special_tokens",
[self.pad_token, self.unk_token, self.bos_token, self.eos_token],
)
self.add_special_tokens(special_tokens)
# Предобработка текстов
words = self._preprocess_texts(texts)
# Получаем начальные токены
vocab = self._get_initial_vocab(words)
# Выполняем BPE мерджи
self._perform_merges(vocab, vocab_size, kwargs.get('min_frequency', 2))
self._perform_merges(vocab, vocab_size, kwargs.get("min_frequency", 2))
# Строим финальный словарь
self._build_final_vocab()
def _initialize_vocab(self):
"""Инициализирует базовый словарь."""
self.vocab.clear()
self.inverse_vocab.clear()
self.merges.clear()
self.vocab_size = 0
def _preprocess_texts(self, texts: List[str]) -> List[List[str]]:
"""
Предобработка текстов для обучения.
Args:
texts: Список текстов
Returns:
List[List[str]]: Предобработанные слова
"""
@@ -84,14 +87,14 @@ class BPETokenizer(BaseTokenizer):
tokens = self.compiled_pattern.findall(text)
words.append(tokens)
return words
def _get_initial_vocab(self, words: List[List[str]]) -> Dict[str, int]:
"""
Создает начальный словарь из символов.
Args:
words: Список токенизированных текстов
Returns:
Dict[str, int]: Начальный словарь частот
"""
@@ -99,43 +102,45 @@ class BPETokenizer(BaseTokenizer):
for word_list in words:
for word in word_list:
# Разбиваем слово на символы и добавляем специальный символ конца слова
chars = list(word) + ['</w>']
vocab.update([''.join(chars[i:i+1]) for i in range(len(chars))])
chars = list(word) + ["</w>"]
vocab.update(["".join(chars[i : i + 1]) for i in range(len(chars))])
return vocab
def _perform_merges(self, vocab: Dict[str, int], target_vocab_size: int, min_frequency: int):
def _perform_merges(
self, vocab: Dict[str, int], target_vocab_size: int, min_frequency: int
):
"""
Выполняет BPE мерджи до достижения целевого размера словаря.
Args:
vocab: Начальный словарь
target_vocab_size: Целевой размер словаря
min_frequency: Минимальная частота для мерджа
"""
current_vocab_size = len(vocab) + len(self.vocab)
while current_vocab_size < target_vocab_size:
# Находим наиболее частую пару
pairs = self._get_stats(vocab)
if not pairs:
break
best_pair = max(pairs, key=pairs.get)
if pairs[best_pair] < min_frequency:
break
# Выполняем мердж
vocab = self._merge_vocab(vocab, best_pair)
self.merges[best_pair] = len(self.merges)
current_vocab_size += 1
def _get_stats(self, vocab: Dict[str, int]) -> Dict[Tuple[str, str], int]:
"""
Собирает статистику по парам символов.
Args:
vocab: Словарь токенов
Returns:
Dict[Tuple[str, str], int]: Частоты пар
"""
@@ -145,107 +150,113 @@ class BPETokenizer(BaseTokenizer):
for i in range(len(symbols) - 1):
pairs[symbols[i], symbols[i + 1]] += freq
return pairs
def _merge_vocab(self, vocab: Dict[str, int], pair: Tuple[str, str]) -> Dict[str, int]:
def _merge_vocab(
self, vocab: Dict[str, int], pair: Tuple[str, str]
) -> Dict[str, int]:
"""
Объединяет пару символов в словаре.
Args:
vocab: Исходный словарь
pair: Пара для объединения
Returns:
Dict[str, int]: Обновленный словарь
"""
new_vocab = {}
bigram = re.compile(r'(?<!\\S)' + re.escape(pair[0]) + r' ' + re.escape(pair[1]) + r'(?!\\S)')
bigram = re.compile(
r"(?<!\\S)" + re.escape(pair[0]) + r" " + re.escape(pair[1]) + r"(?!\\S)"
)
replacement = pair[0] + pair[1]
for word in vocab:
new_word = bigram.sub(replacement, word)
new_vocab[new_word] = vocab[word]
return new_vocab
def _build_final_vocab(self):
"""Строит финальный словарь токенизатора."""
# Собираем все уникальные токены из мерджей
all_tokens = set()
# Добавляем специальные токены
all_tokens.update([self.pad_token, self.unk_token, self.bos_token, self.eos_token])
all_tokens.update(
[self.pad_token, self.unk_token, self.bos_token, self.eos_token]
)
# Добавляем токены из мерджей
for pair in self.merges:
all_tokens.update(pair)
# Создаем словарь
for i, token in enumerate(sorted(all_tokens)):
self.vocab[token] = i
self.inverse_vocab = {v: k for k, v in self.vocab.items()}
self.vocab_size = len(self.vocab)
# Обновляем ID специальных токенов
self.pad_token_id = self.vocab.get(self.pad_token)
self.unk_token_id = self.vocab.get(self.unk_token)
self.bos_token_id = self.vocab.get(self.bos_token)
self.eos_token_id = self.vocab.get(self.eos_token)
def encode(self, text: str, **kwargs) -> List[int]:
"""
Кодирует текст в последовательность токенов.
Args:
text: Входной текст
**kwargs: Дополнительные параметры
- add_special_tokens: Добавлять специальные токены
Returns:
List[int]: Список идентификаторов токенов
"""
add_special_tokens = kwargs.get('add_special_tokens', False)
add_special_tokens = kwargs.get("add_special_tokens", False)
# Токенизация текста
tokens = self.compiled_pattern.findall(text)
# Применяем BPE к каждому токену
bpe_tokens = []
for token in tokens:
# Преобразуем токен в BPE представление
bpe_token = self._apply_bpe(token)
bpe_tokens.extend(bpe_token)
# Конвертируем в ID
token_ids = []
for token in bpe_tokens:
token_id = self.vocab.get(token, self.unk_token_id)
if token_id is not None:
token_ids.append(token_id)
# Добавляем специальные токены если нужно
if add_special_tokens:
if self.bos_token_id is not None:
token_ids.insert(0, self.bos_token_id)
if self.eos_token_id is not None:
token_ids.append(self.eos_token_id)
return token_ids
def _apply_bpe(self, token: str) -> List[str]:
"""
Применяет BPE к одному токену.
Args:
token: Входной токен
Returns:
List[str]: Список BPE токенов
"""
# Простая реализация - в реальной реализации нужно применять обученные мерджи
word = token + '</w>'
tokens = [word[i:i+1] for i in range(len(word))]
word = token + "</w>"
tokens = [word[i : i + 1] for i in range(len(word))]
# Применяем мерджи (упрощенная версия)
# В полной реализации нужно применять все обученные мерджи
for pair in self.merges:
@@ -256,109 +267,114 @@ class BPETokenizer(BaseTokenizer):
del tokens[i + 1]
else:
i += 1
return tokens
def decode(self, tokens: List[int], **kwargs) -> str:
"""
Декодирует последовательность токенов в текст.
Args:
tokens: Список идентификаторов токенов
**kwargs: Дополнительные параметры
- skip_special_tokens: Пропускать специальные токены
Returns:
str: Декодированный текст
"""
skip_special_tokens = kwargs.get('skip_special_tokens', True)
skip_special_tokens = kwargs.get("skip_special_tokens", True)
# Конвертируем ID в токены
token_strings = []
for token_id in tokens:
token = self.inverse_vocab.get(token_id, self.unk_token)
# Пропускаем специальные токены если нужно
if skip_special_tokens and token in [self.pad_token, self.unk_token, self.bos_token, self.eos_token]:
if skip_special_tokens and token in [
self.pad_token,
self.unk_token,
self.bos_token,
self.eos_token,
]:
continue
token_strings.append(token)
# Объединяем токены в текст
text = ''.join(token_strings)
text = "".join(token_strings)
# Убираем маркер конца слова
text = text.replace('</w>', ' ')
text = text.replace("</w>", " ")
return text.strip()
def save(self, filepath: str):
"""
Сохраняет BPE токенизатор в файл.
Args:
filepath: Путь для сохранения
"""
import json
config = {
'vocab': self.vocab,
'merges': {f"{k[0]} {k[1]}": v for k, v in self.merges.items()},
'vocab_size': self.vocab_size,
'pad_token': self.pad_token,
'unk_token': self.unk_token,
'bos_token': self.bos_token,
'eos_token': self.eos_token,
'pattern': self.pattern,
'tokenizer_type': self.__class__.__name__
"vocab": self.vocab,
"merges": {f"{k[0]} {k[1]}": v for k, v in self.merges.items()},
"vocab_size": self.vocab_size,
"pad_token": self.pad_token,
"unk_token": self.unk_token,
"bos_token": self.bos_token,
"eos_token": self.eos_token,
"pattern": self.pattern,
"tokenizer_type": self.__class__.__name__,
}
with open(filepath, 'w', encoding='utf-8') as f:
with open(filepath, "w", encoding="utf-8") as f:
json.dump(config, f, ensure_ascii=False, indent=2)
@classmethod
def load(cls, filepath: str):
"""
Загружает BPE токенизатор из файла.
Args:
filepath: Путь к файлу
Returns:
BPETokenizer: Загруженный токенизатор
"""
import json
with open(filepath, 'r', encoding='utf-8') as f:
with open(filepath, "r", encoding="utf-8") as f:
config = json.load(f)
tokenizer = cls()
tokenizer.vocab = config['vocab']
tokenizer.vocab_size = config['vocab_size']
tokenizer.pad_token = config['pad_token']
tokenizer.unk_token = config['unk_token']
tokenizer.bos_token = config['bos_token']
tokenizer.eos_token = config['eos_token']
tokenizer.pattern = config.get('pattern', tokenizer.pattern)
tokenizer.vocab = config["vocab"]
tokenizer.vocab_size = config["vocab_size"]
tokenizer.pad_token = config["pad_token"]
tokenizer.unk_token = config["unk_token"]
tokenizer.bos_token = config["bos_token"]
tokenizer.eos_token = config["eos_token"]
tokenizer.pattern = config.get("pattern", tokenizer.pattern)
tokenizer.compiled_pattern = re.compile(tokenizer.pattern, re.UNICODE)
# Восстанавливаем мерджи
merges = config.get('merges', {})
merges = config.get("merges", {})
tokenizer.merges = {}
for k, v in merges.items():
parts = k.split()
if len(parts) == 2:
tokenizer.merges[(parts[0], parts[1])] = v
# Создаем обратный словарь
tokenizer.inverse_vocab = {v: k for k, v in tokenizer.vocab.items()}
# Обновляем ID специальных токенов
tokenizer.pad_token_id = tokenizer.vocab.get(tokenizer.pad_token)
tokenizer.unk_token_id = tokenizer.vocab.get(tokenizer.unk_token)
tokenizer.bos_token_id = tokenizer.vocab.get(tokenizer.bos_token)
tokenizer.eos_token_id = tokenizer.vocab.get(tokenizer.eos_token)
return tokenizer
@@ -367,62 +383,72 @@ class SimpleBPETokenizer(BPETokenizer):
"""
Упрощенная версия BPE токенизатора для демонстрации.
"""
def train(self, texts: List[str], vocab_size: int = 1000, **kwargs):
"""Упрощенное обучение для демонстрации."""
# Инициализация базового словаря
self._initialize_vocab()
# Добавляем базовые токены
special_tokens = [self.pad_token, self.unk_token, self.bos_token, self.eos_token]
special_tokens = [
self.pad_token,
self.unk_token,
self.bos_token,
self.eos_token,
]
self.add_special_tokens(special_tokens)
# Простая реализация - собираем все символы
all_chars = set()
for text in texts:
all_chars.update(text)
# Добавляем символы в словарь
for char in sorted(all_chars):
if char not in self.vocab:
self.vocab[char] = len(self.vocab)
self.inverse_vocab = {v: k for k, v in self.vocab.items()}
self.vocab_size = len(self.vocab)
# Обновляем ID специальных токенов
self.pad_token_id = self.vocab.get(self.pad_token)
self.unk_token_id = self.vocab.get(self.unk_token)
self.bos_token_id = self.vocab.get(self.bos_token)
self.eos_token_id = self.vocab.get(self.eos_token)
def encode(self, text: str, **kwargs) -> List[int]:
"""Упрощенное кодирование - разбиваем на символы."""
add_special_tokens = kwargs.get('add_special_tokens', False)
add_special_tokens = kwargs.get("add_special_tokens", False)
token_ids = []
for char in text:
token_id = self.vocab.get(char, self.unk_token_id)
if token_id is not None:
token_ids.append(token_id)
if add_special_tokens:
if self.bos_token_id is not None:
token_ids.insert(0, self.bos_token_id)
if self.eos_token_id is not None:
token_ids.append(self.eos_token_id)
return token_ids
def decode(self, tokens: List[int], **kwargs) -> str:
"""Упрощенное декодирование."""
skip_special_tokens = kwargs.get('skip_special_tokens', True)
skip_special_tokens = kwargs.get("skip_special_tokens", True)
chars = []
for token_id in tokens:
char = self.inverse_vocab.get(token_id, self.unk_token)
if skip_special_tokens and char in [self.pad_token, self.unk_token, self.bos_token, self.eos_token]:
if skip_special_tokens and char in [
self.pad_token,
self.unk_token,
self.bos_token,
self.eos_token,
]:
continue
chars.append(char)
return ''.join(chars)
return "".join(chars)

156
llm/src/llm/tokenizers/bpe_tokenizer.py
View File

@@ -11,26 +11,26 @@ from .base_tokenizer import BaseTokenizer
class BPETokenizer(BaseTokenizer):
"""
BPE токенизатор для обработки текста.
Реализует алгоритм Byte Pair Encoding для создания субсловных токенов.
Использует вашу реализацию BPE.
Примеры использования:
>>> tokenizer = BPETokenizer()
>>> tokenizer.train(["пример текста для обучения"], vocab_size=1000)
>>> tokens = tokenizer.encode("новый текст")
>>> text = tokenizer.decode(tokens)
"""
def __init__(self):
super().__init__()
self.merges: Dict[Tuple[str, str], int] = {}
self.vocab_list: List[str] = []
def train(self, texts: List[str], vocab_size: int = 1000, **kwargs):
"""
Обучение BPE токенизатора на текстах.
Args:
texts: Список текстов для обучения
vocab_size: Желаемый размер словаря
@@ -39,7 +39,7 @@ class BPETokenizer(BaseTokenizer):
"""
# Объединяем все тексты в одну строку для обучения
combined_text = " ".join(texts)
# 1. Получаем уникальные токены (символы)
unique_tokens = sorted(set(combined_text))
tokens = unique_tokens.copy()
@@ -61,7 +61,10 @@ class BPETokenizer(BaseTokenizer):
break # нет пар — выходим
# Находим самую частую пару (в случае равенства — та, что встретилась первой)
most_frequent_pair = max(pair_freq.items(), key=lambda x: (x[1], -self._pair_first_index(sequence, x[0])))[0]
most_frequent_pair = max(
pair_freq.items(),
key=lambda x: (x[1], -self._pair_first_index(sequence, x[0])),
)[0]
# Создаем новый токен
new_token = most_frequent_pair[0] + most_frequent_pair[1]
@@ -71,45 +74,51 @@ class BPETokenizer(BaseTokenizer):
new_sequence = []
while i < len(sequence):
if i < len(sequence) - 1 and (sequence[i], sequence[i + 1]) == most_frequent_pair:
if (
i < len(sequence) - 1
and (sequence[i], sequence[i + 1]) == most_frequent_pair
):
new_sequence.append(new_token)
i += 2 # пропускаем два символа — заменённую пару
else:
new_sequence.append(sequence[i])
i += 1
sequence = new_sequence
# 4. Создаем словари
self.vocab_list = tokens.copy()
self.vocab = dict(zip(tokens, range(vocab_size)))
self.inverse_vocab = dict(zip(range(vocab_size), tokens))
self.vocab_size = len(self.vocab)
# Добавляем специальные токены если указаны
special_tokens = kwargs.get('special_tokens', [self.pad_token, self.unk_token, self.bos_token, self.eos_token])
special_tokens = kwargs.get(
"special_tokens",
[self.pad_token, self.unk_token, self.bos_token, self.eos_token],
)
self.add_special_tokens(special_tokens)
def _pair_first_index(self, sequence, pair):
"""Находит первый индекс пары в последовательности."""
for i in range(len(sequence) - 1):
if (sequence[i], sequence[i + 1]) == pair:
return i
return float('inf') # если пара не найдена (в теории не должно случиться)
return float("inf") # если пара не найдена (в теории не должно случиться)
def encode(self, text: str, **kwargs) -> List[int]:
"""
Кодирует текст в последовательность токенов.
Args:
text: Входной текст
**kwargs: Дополнительные параметры
- add_special_tokens: Добавлять специальные токены
Returns:
List[int]: Список идентификаторов токенов
"""
add_special_tokens = kwargs.get('add_special_tokens', False)
add_special_tokens = kwargs.get("add_special_tokens", False)
# 1. Разбиваем текст на токены-символы
sequence = list(text)
# 2. Инициализация пустого списка токенов
@@ -119,7 +128,9 @@ class BPETokenizer(BaseTokenizer):
while i < len(text):
# 3.1 Найти все токены в словаре, начинающиеся с text[i]
start_char = text[i]
result = [token for token in self.vocab_list if token.startswith(start_char)]
result = [
token for token in self.vocab_list if token.startswith(start_char)
]
# 3.2 Выбрать самый длинный подходящий токен
find_token = self._find_max_matching_token(text[i:], result)
if find_token is None:
@@ -134,19 +145,19 @@ class BPETokenizer(BaseTokenizer):
# 4. Заменить токены на их ID
token_ids = self._tokens_to_ids(tokens)
# Заменяем -1 на unk_token_id
token_ids = [tid if tid != -1 else self.unk_token_id for tid in token_ids]
# Добавляем специальные токены если нужно
if add_special_tokens:
if self.bos_token_id is not None:
token_ids.insert(0, self.bos_token_id)
if self.eos_token_id is not None:
token_ids.append(self.eos_token_id)
return token_ids
def _find_max_matching_token(self, text: str, tokens: list) -> Optional[str]:
"""Находит самый длинный токен из списка, с которого начинается текст"""
matching = [token for token in tokens if text.startswith(token)]
@@ -161,33 +172,41 @@ class BPETokenizer(BaseTokenizer):
else:
ids.append(-1) # Специальное значение
return ids
def decode(self, tokens: List[int], **kwargs) -> str:
"""
Декодирует последовательность токенов в текст.
Args:
tokens: Список идентификаторов токенов
**kwargs: Дополнительные параметры
- skip_special_tokens: Пропускать специальные токены
Returns:
str: Декодированный текст
"""
skip_special_tokens = kwargs.get('skip_special_tokens', True)
skip_special_tokens = kwargs.get("skip_special_tokens", True)
# Фильтруем специальные токены если нужно
if skip_special_tokens:
tokens = [tid for tid in tokens if tid not in [
self.pad_token_id, self.unk_token_id, self.bos_token_id, self.eos_token_id
]]
tokens = [
tid
for tid in tokens
if tid
not in [
self.pad_token_id,
self.unk_token_id,
self.bos_token_id,
self.eos_token_id,
]
]
# Конвертируем ID в токены
token_strings = self._ids_to_tokens(tokens)
# Объединяем токены в текст
return ''.join(token_strings)
return "".join(token_strings)
def _ids_to_tokens(self, ids: List[int]) -> List[str]:
"""Конвертирует список Ids в их tokens"""
tokens = []
@@ -197,76 +216,76 @@ class BPETokenizer(BaseTokenizer):
else:
tokens.append(self.unk_token) # Специальное значение
return tokens
def save(self, filepath: str):
"""
Сохраняет токенизатор в файл.
Args:
filepath: Путь для сохранения
"""
import json
# Преобразуем кортежи в строки для JSON сериализации
merges_serializable = {f"{k[0]},{k[1]}": v for k, v in self.merges.items()}
config = {
'vocab': self.vocab,
'vocab_size': self.vocab_size,
'pad_token': self.pad_token,
'unk_token': self.unk_token,
'bos_token': self.bos_token,
'eos_token': self.eos_token,
'tokenizer_type': self.__class__.__name__,
'merges': merges_serializable,
'vocab_list': self.vocab_list
"vocab": self.vocab,
"vocab_size": self.vocab_size,
"pad_token": self.pad_token,
"unk_token": self.unk_token,
"bos_token": self.bos_token,
"eos_token": self.eos_token,
"tokenizer_type": self.__class__.__name__,
"merges": merges_serializable,
"vocab_list": self.vocab_list,
}
with open(filepath, 'w', encoding='utf-8') as f:
with open(filepath, "w", encoding="utf-8") as f:
json.dump(config, f, ensure_ascii=False, indent=2)
@classmethod
def load(cls, filepath: str):
"""
Загружает токенизатор из файла.
Args:
filepath: Путь к файлу
Returns:
BPETokenizer: Загруженный токенизатор
"""
import json
with open(filepath, 'r', encoding='utf-8') as f:
with open(filepath, "r", encoding="utf-8") as f:
config = json.load(f)
# Создаем экземпляр токенизатора
tokenizer = cls()
tokenizer.vocab = config['vocab']
tokenizer.vocab_size = config['vocab_size']
tokenizer.pad_token = config['pad_token']
tokenizer.unk_token = config['unk_token']
tokenizer.bos_token = config['bos_token']
tokenizer.eos_token = config['eos_token']
tokenizer.vocab_list = config['vocab_list']
tokenizer.vocab = config["vocab"]
tokenizer.vocab_size = config["vocab_size"]
tokenizer.pad_token = config["pad_token"]
tokenizer.unk_token = config["unk_token"]
tokenizer.bos_token = config["bos_token"]
tokenizer.eos_token = config["eos_token"]
tokenizer.vocab_list = config["vocab_list"]
# Восстанавливаем кортежи из строк
tokenizer.merges = {}
for k, v in config['merges'].items():
parts = k.split(',')
for k, v in config["merges"].items():
parts = k.split(",")
if len(parts) == 2:
tokenizer.merges[(parts[0], parts[1])] = v
# Создаем обратный словарь
tokenizer.inverse_vocab = {v: k for k, v in tokenizer.vocab.items()}
# Обновляем ID специальных токенов
tokenizer.pad_token_id = tokenizer.vocab.get(tokenizer.pad_token)
tokenizer.unk_token_id = tokenizer.vocab.get(tokenizer.unk_token)
tokenizer.bos_token_id = tokenizer.vocab.get(tokenizer.bos_token)
tokenizer.eos_token_id = tokenizer.vocab.get(tokenizer.eos_token)
return tokenizer
@@ -275,4 +294,5 @@ class SimpleBPETokenizer(BPETokenizer):
Упрощенная версия BPE токенизатора для демонстрации.
Наследует вашу реализацию, но может быть упрощена при необходимости.
"""
pass

69
llm/src/llm/training/dataset.py
View File

@@ -12,7 +12,7 @@ class TextDataset(Dataset):
def __init__(self, texts: List[str], tokenizer: Any, block_size: int = 128):
"""
Инициализация датасета.
Args:
texts: Список текстов для обучения
tokenizer: Токенизатор с методами encode/decode
@@ -25,15 +25,15 @@ class TextDataset(Dataset):
for text in texts:
# Кодируем текст в токены
input_ids = tokenizer.encode(text, add_special_tokens=False)
# Обрезаем или дополняем до нужной длины
if len(input_ids) > block_size:
input_ids = input_ids[:block_size]
else:
# Дополняем pad_token_id
pad_token_id = getattr(tokenizer, 'pad_token_id', 0)
pad_token_id = getattr(tokenizer, "pad_token_id", 0)
input_ids = input_ids + [pad_token_id] * (block_size - len(input_ids))
self.examples.append(input_ids)
def __len__(self):
@@ -50,33 +50,35 @@ class StreamingTextDataset(Dataset):
Датасет для потоковой обработки больших текстов.
Токенизация происходит на лету, что экономит память.
"""
def __init__(self, texts: List[str], tokenizer: Any, block_size: int = 128):
self.texts = texts
self.tokenizer = tokenizer
self.block_size = block_size
# Получаем pad_token_id из токенизатора
self.pad_token_id = getattr(tokenizer, 'pad_token_id', 0)
self.pad_token_id = getattr(tokenizer, "pad_token_id", 0)
def __len__(self):
return len(self.texts)
def __getitem__(self, idx):
text = self.texts[idx]
# Токенизация на лету
input_ids = self.tokenizer.encode(text, add_special_tokens=False)
# Обрезаем или дополняем до нужной длины
if len(input_ids) > self.block_size:
input_ids = input_ids[:self.block_size]
input_ids = input_ids[: self.block_size]
else:
input_ids = input_ids + [self.pad_token_id] * (self.block_size - len(input_ids))
input_ids = input_ids + [self.pad_token_id] * (
self.block_size - len(input_ids)
)
input_ids = torch.tensor(input_ids, dtype=torch.long)
labels = input_ids.clone()
return {"input_ids": input_ids, "labels": labels}
@@ -84,9 +86,15 @@ class TextDatasetWithSpecialTokens(TextDataset):
"""
Расширенная версия TextDataset с поддержкой специальных токенов.
"""
def __init__(self, texts: List[str], tokenizer: Any, block_size: int = 128,
add_bos: bool = False, add_eos: bool = False):
def __init__(
self,
texts: List[str],
tokenizer: Any,
block_size: int = 128,
add_bos: bool = False,
add_eos: bool = False,
):
"""
Args:
texts: Список текстов
@@ -104,33 +112,38 @@ class TextDatasetWithSpecialTokens(TextDataset):
for text in texts:
# Кодируем с специальными токенами
input_ids = tokenizer.encode(
text,
add_special_tokens=True,
add_bos_token=add_bos,
add_eos_token=eos
text, add_special_tokens=True, add_bos_token=add_bos, add_eos_token=eos
)
# Учитываем специальные токены при обрезке/дополнении
effective_block_size = block_size
if add_bos:
effective_block_size -= 1
if add_eos:
effective_block_size -= 1
if len(input_ids) > effective_block_size:
input_ids = input_ids[:effective_block_size]
# Добавляем специальные токены если нужно
if add_bos and hasattr(tokenizer, 'bos_token_id') and tokenizer.bos_token_id is not None:
if (
add_bos
and hasattr(tokenizer, "bos_token_id")
and tokenizer.bos_token_id is not None
):
input_ids = [tokenizer.bos_token_id] + input_ids
if add_eos and hasattr(tokenizer, 'eos_token_id') and tokenizer.eos_token_id is not None:
if (
add_eos
and hasattr(tokenizer, "eos_token_id")
and tokenizer.eos_token_id is not None
):
input_ids = input_ids + [tokenizer.eos_token_id]
# Дополняем до полной длины
pad_token_id = getattr(tokenizer, 'pad_token_id', 0)
pad_token_id = getattr(tokenizer, "pad_token_id", 0)
if len(input_ids) < block_size:
input_ids = input_ids + [pad_token_id] * (block_size - len(input_ids))
self.examples.append(input_ids)
def __len__(self):

1
llm/src/llm/training/optimizer.py
View File

@@ -1,5 +1,6 @@
import torch.optim as optim
def get_optimizer(model, lr=3e-4, weight_decay=0.01, optimizer_type="adamw"):
"""
Возвращает оптимизатор для обучения модели.

7
llm/src/llm/training/scheduler.py
View File

@@ -1,5 +1,6 @@
from torch.optim.lr_scheduler import LambdaLR
def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps):
"""
Линейный планировщик обучения с warmup.
@@ -8,6 +9,10 @@ def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_st
def lr_lambda(current_step):
if current_step < num_warmup_steps:
return float(current_step) / float(max(1, num_warmup_steps))
return max(0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps)))
return max(
0.0,
float(num_training_steps - current_step)
/ float(max(1, num_training_steps - num_warmup_steps)),
)
return LambdaLR(optimizer, lr_lambda)

36
llm/src/llm/training/trainer.py
View File

@@ -5,15 +5,29 @@ from tqdm import tqdm
from llm.training.optimizer import get_optimizer
from llm.training.scheduler import get_linear_schedule_with_warmup
class Trainer:
"""
Универсальный класс обучения LLM (GPT, LLaMA, Mistral и т.д.)
"""
def __init__(self, model, train_dataset, val_dataset=None, lr=3e-4, batch_size=8, num_epochs=3, warmup_steps=100):
def __init__(
self,
model,
train_dataset,
val_dataset=None,
lr=3e-4,
batch_size=8,
num_epochs=3,
warmup_steps=100,
):
self.model = model
self.train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
self.val_loader = DataLoader(val_dataset, batch_size=batch_size) if val_dataset else None
self.train_loader = DataLoader(
train_dataset, batch_size=batch_size, shuffle=True
)
self.val_loader = (
DataLoader(val_dataset, batch_size=batch_size) if val_dataset else None
)
self.optimizer = get_optimizer(model, lr=lr)
self.scheduler = None
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -29,24 +43,28 @@ class Trainer:
# Сдвигаем логиты и метки для языкового моделирования
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
# Вычисляем cross-entropy loss
loss = F.cross_entropy(
shift_logits.view(-1, shift_logits.size(-1)),
shift_labels.view(-1),
ignore_index=-100 # Игнорируем padding tokens
ignore_index=-100, # Игнорируем padding tokens
)
return loss
def train(self):
total_steps = len(self.train_loader) * self.num_epochs
self.scheduler = get_linear_schedule_with_warmup(self.optimizer, self.warmup_steps, total_steps)
self.scheduler = get_linear_schedule_with_warmup(
self.optimizer, self.warmup_steps, total_steps
)
for epoch in range(self.num_epochs):
self.model.train()
total_loss = 0
progress_bar = tqdm(self.train_loader, desc=f"Epoch {epoch+1}/{self.num_epochs}")
progress_bar = tqdm(
self.train_loader, desc=f"Epoch {epoch+1}/{self.num_epochs}"
)
for batch in progress_bar:
self.optimizer.zero_grad()
@@ -59,7 +77,7 @@ class Trainer:
logits = outputs[0]
else:
logits = outputs
# Trainer вычисляет loss
loss = self.compute_lm_loss(logits, labels)
loss.backward()
@@ -85,7 +103,7 @@ class Trainer:
for batch in self.val_loader:
input_ids = batch["input_ids"].to(self.device)
labels = batch["labels"].to(self.device)
outputs = self.model(input_ids)
if isinstance(outputs, tuple):
logits = outputs[0]

4
llm/tests/conftest.py
View File

@@ -58,7 +58,7 @@ def gpt_config(vocab_size, embed_dim, num_heads, num_layers):
"num_heads": num_heads,
"num_layers": num_layers,
"max_position_embeddings": 1024,
"dropout": 0.1
"dropout": 0.1,
}
@@ -68,12 +68,14 @@ def random_inputs(batch_size, seq_len, vocab_size):
input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
return input_ids
@pytest.fixture
def random_float_inputs(batch_size, seq_len, embed_dim):
"""Generate random floating point input tensors for testing feed forward."""
inputs = torch.randn(batch_size, seq_len, embed_dim)
return inputs
@pytest.fixture
def random_embeddings(batch_size, seq_len, embed_dim):
"""Generate random embedding tensors for testing attention modules."""

171
llm/tests/core/test_decoder.py
View File

@@ -9,180 +9,233 @@ from llm.core.decoder import Decoder
class TestDecoder:
"""Test cases for Decoder."""
def test_initialization(self, embed_dim, num_heads):
"""Test that Decoder can be initialized."""
head_size = embed_dim // num_heads
max_seq_len = 1024
decoder = Decoder(num_heads=num_heads, emb_size=embed_dim, head_size=head_size, max_seq_len=max_seq_len)
decoder = Decoder(
num_heads=num_heads,
emb_size=embed_dim,
head_size=head_size,
max_seq_len=max_seq_len,
)
assert decoder is not None
# Check internal components
assert hasattr(decoder, '_heads')
assert hasattr(decoder, '_ff')
assert hasattr(decoder, '_norm1')
assert hasattr(decoder, '_norm2')
assert hasattr(decoder, "_heads")
assert hasattr(decoder, "_ff")
assert hasattr(decoder, "_norm1")
assert hasattr(decoder, "_norm2")
def test_forward_pass(self, embed_dim, num_heads, random_embeddings):
"""Test forward pass of Decoder."""
head_size = embed_dim // num_heads
max_seq_len = 1024
decoder = Decoder(num_heads=num_heads, emb_size=embed_dim, head_size=head_size, max_seq_len=max_seq_len)
decoder = Decoder(
num_heads=num_heads,
emb_size=embed_dim,
head_size=head_size,
max_seq_len=max_seq_len,
)
# Forward pass
output = decoder(random_embeddings)
# Check output shape
assert output.shape == random_embeddings.shape
assert isinstance(output, torch.Tensor)
def test_forward_with_causal_mask(self, embed_dim, num_heads, random_embeddings):
"""Test forward pass with causal mask."""
head_size = embed_dim // num_heads
max_seq_len = 1024
decoder = Decoder(num_heads=num_heads, emb_size=embed_dim, head_size=head_size, max_seq_len=max_seq_len)
decoder = Decoder(
num_heads=num_heads,
emb_size=embed_dim,
head_size=head_size,
max_seq_len=max_seq_len,
)
batch_size, seq_len = random_embeddings.shape[:2]
# Create causal mask
mask = torch.tril(torch.ones(seq_len, seq_len))
# Forward pass with causal mask
output = decoder(random_embeddings, mask=mask)
# Check output shape
assert output.shape == random_embeddings.shape
def test_residual_connections(self, embed_dim, num_heads, random_embeddings):
"""Test that residual connections are properly applied."""
head_size = embed_dim // num_heads
max_seq_len = 1024
decoder = Decoder(num_heads=num_heads, emb_size=embed_dim, head_size=head_size, max_seq_len=max_seq_len)
decoder = Decoder(
num_heads=num_heads,
emb_size=embed_dim,
head_size=head_size,
max_seq_len=max_seq_len,
)
output = decoder(random_embeddings)
# With residual connections and layer norm, the output shouldn't be
# too different from input (in terms of scale/distribution)
input_norm = random_embeddings.norm(dim=-1).mean()
output_norm = output.norm(dim=-1).mean()
# Norms should be of similar magnitude (not exact due to transformations)
assert 0.1 < (output_norm / input_norm) < 10.0
def test_layer_norm(self, embed_dim, num_heads, random_embeddings):
"""Test that layer normalization is applied."""
head_size = embed_dim // num_heads
max_seq_len = 1024
decoder = Decoder(num_heads=num_heads, emb_size=embed_dim, head_size=head_size, max_seq_len=max_seq_len)
decoder = Decoder(
num_heads=num_heads,
emb_size=embed_dim,
head_size=head_size,
max_seq_len=max_seq_len,
)
output = decoder(random_embeddings)
# Check that output has reasonable statistics (due to layer norm)
# Mean should be close to 0, std close to 1 for each sequence position
output_mean = output.mean(dim=-1)
output_std = output.std(dim=-1)
# These are approximate checks since the data goes through multiple transformations
assert torch.allclose(output_mean, torch.zeros_like(output_mean), atol=1.0)
assert torch.allclose(output_std, torch.ones_like(output_std), atol=2.0)
def test_gradient_flow(self, embed_dim, num_heads, random_embeddings):
"""Test that gradients flow through Decoder."""
head_size = embed_dim // num_heads
max_seq_len = 1024
decoder = Decoder(num_heads=num_heads, emb_size=embed_dim, head_size=head_size, max_seq_len=max_seq_len)
decoder = Decoder(
num_heads=num_heads,
emb_size=embed_dim,
head_size=head_size,
max_seq_len=max_seq_len,
)
# Forward pass
output = decoder(random_embeddings)
# Create a dummy loss and backward pass
loss = output.sum()
loss.backward()
# Check that gradients are computed for learnable parameters
# in attention and feed forward components
assert decoder._heads._layer.weight.grad is not None
assert decoder._ff._layer1.weight.grad is not None
assert decoder._norm1.weight.grad is not None
assert decoder._norm2.weight.grad is not None
def test_device_consistency(self, embed_dim, num_heads, random_embeddings, device):
"""Test that Decoder works on correct device."""
head_size = embed_dim // num_heads
max_seq_len = 1024
decoder = Decoder(num_heads=num_heads, emb_size=embed_dim, head_size=head_size, max_seq_len=max_seq_len).to(device)
decoder = Decoder(
num_heads=num_heads,
emb_size=embed_dim,
head_size=head_size,
max_seq_len=max_seq_len,
).to(device)
inputs = random_embeddings.to(device)
# Forward pass
output = decoder(inputs)
# Check device consistency
assert output.device == device
assert decoder._heads._layer.weight.device == device
def test_different_configurations(self):
"""Test Decoder with different configurations."""
test_cases = [
(64, 2), # embed_dim=64, num_heads=2
(64, 2), # embed_dim=64, num_heads=2
(128, 4), # embed_dim=128, num_heads=4
(256, 8), # embed_dim=256, num_heads=8
]
for embed_dim, num_heads in test_cases:
head_size = embed_dim // num_heads
max_seq_len = 1024
decoder = Decoder(num_heads=num_heads, emb_size=embed_dim, head_size=head_size, max_seq_len=max_seq_len)
decoder = Decoder(
num_heads=num_heads,
emb_size=embed_dim,
head_size=head_size,
max_seq_len=max_seq_len,
)
batch_size, seq_len = 2, 16
inputs = torch.randn(batch_size, seq_len, embed_dim)
output = decoder(inputs)
assert output.shape == inputs.shape
@pytest.mark.parametrize("batch_size,seq_len", [(1, 8), (2, 16), (4, 32)])
def test_different_input_shapes(self, embed_dim, num_heads, batch_size, seq_len):
"""Test Decoder with different input shapes."""
head_size = embed_dim // num_heads
max_seq_len = 1024
decoder = Decoder(num_heads=num_heads, emb_size=embed_dim, head_size=head_size, max_seq_len=max_seq_len)
decoder = Decoder(
num_heads=num_heads,
emb_size=embed_dim,
head_size=head_size,
max_seq_len=max_seq_len,
)
inputs = torch.randn(batch_size, seq_len, embed_dim)
output = decoder(inputs)
assert output.shape == (batch_size, seq_len, embed_dim)
def test_training_vs_evaluation(self, embed_dim, num_heads, random_embeddings):
"""Test that Decoder behaves differently in train vs eval mode."""
head_size = embed_dim // num_heads
max_seq_len = 1024
decoder = Decoder(num_heads=num_heads, emb_size=embed_dim, head_size=head_size, max_seq_len=max_seq_len, dropout=0.5)
decoder = Decoder(
num_heads=num_heads,
emb_size=embed_dim,
head_size=head_size,
max_seq_len=max_seq_len,
dropout=0.5,
)
# Training mode
decoder.train()
output_train = decoder(random_embeddings)
# Evaluation mode
decoder.eval()
output_eval = decoder(random_embeddings)
# Outputs should be different due to dropout
assert not torch.allclose(output_train, output_eval)
def test_parameter_initialization(self, embed_dim, num_heads):
"""Test that parameters are properly initialized."""
head_size = embed_dim // num_heads
max_seq_len = 1024
decoder = Decoder(num_heads=num_heads, emb_size=embed_dim, head_size=head_size, max_seq_len=max_seq_len)
decoder = Decoder(
num_heads=num_heads,
emb_size=embed_dim,
head_size=head_size,
max_seq_len=max_seq_len,
)
# Check that various components have non-zero parameters
assert not torch.allclose(
decoder._heads._layer.weight,
torch.zeros_like(decoder._heads._layer.weight)
decoder._heads._layer.weight, torch.zeros_like(decoder._heads._layer.weight)
)
assert not torch.allclose(
decoder._ff._layer1.weight,
torch.zeros_like(decoder._ff._layer1.weight)
decoder._ff._layer1.weight, torch.zeros_like(decoder._ff._layer1.weight)
)
assert not torch.allclose(
decoder._norm1.weight,
torch.zeros_like(decoder._norm1.weight)
decoder._norm1.weight, torch.zeros_like(decoder._norm1.weight)
)

114
llm/tests/core/test_feed_forward.py
View File

@@ -10,168 +10,178 @@ from llm.core.feed_forward import FeedForward
class TestFeedForward:
"""Test cases for FeedForward."""
def test_initialization(self, embed_dim):
"""Test that FeedForward can be initialized."""
ff = FeedForward(embed_dim)
assert ff is not None
# Check internal layers
assert hasattr(ff, '_layer1')
assert hasattr(ff, '_layer2')
assert hasattr(ff, '_activation')
assert hasattr(ff, '_dropout')
assert hasattr(ff, "_layer1")
assert hasattr(ff, "_layer2")
assert hasattr(ff, "_activation")
assert hasattr(ff, "_dropout")
# Check layer dimensions
expected_hidden_dim = embed_dim * 4 # Default expansion factor
assert ff._layer1.weight.shape == (expected_hidden_dim, embed_dim)
assert ff._layer2.weight.shape == (embed_dim, expected_hidden_dim)
def test_forward_pass(self, embed_dim, random_float_inputs):
"""Test forward pass of FeedForward."""
ff = FeedForward(embed_dim)
# Forward pass
output = ff(random_float_inputs)
# Check output shape
assert output.shape == random_float_inputs.shape
assert isinstance(output, torch.Tensor)
def test_custom_hidden_dim(self, embed_dim):
"""Test FeedForward with custom hidden dimension."""
# FeedForward doesn't support custom hidden_dim in current implementation
# This test is not applicable
ff = FeedForward(embed_dim)
# Check layer dimensions (fixed 4x expansion)
expected_hidden_dim = embed_dim * 4
assert ff._layer1.weight.shape == (expected_hidden_dim, embed_dim)
assert ff._layer2.weight.shape == (embed_dim, expected_hidden_dim)
def test_dropout(self, embed_dim, random_float_inputs):
"""Test that dropout is applied during training."""
ff = FeedForward(embed_dim, dropout=0.5)
ff.train() # Set to training mode
output = ff(random_float_inputs)
# In training mode with dropout, some values should be zeroed
# This is probabilistic, so we can't assert exact zeros,
# but we can check the structure is preserved
assert output.shape == random_float_inputs.shape
def test_no_dropout_in_eval(self, embed_dim, random_float_inputs):
"""Test that dropout is not applied during evaluation."""
ff = FeedForward(embed_dim, dropout=0.5)
ff.eval() # Set to evaluation mode
# Run forward pass multiple times - outputs should be identical
output1 = ff(random_float_inputs)
output2 = ff(random_float_inputs)
assert torch.allclose(output1, output2)
def test_activation_function(self, embed_dim, random_float_inputs):
"""Test that activation function is applied."""
ff = FeedForward(embed_dim)
# Manually compute expected output without dropout for deterministic comparison
hidden = ff._layer1(random_float_inputs)
activated = ff._activation(hidden)
expected_output = ff._layer2(activated)
# Compare with forward pass in eval mode (no dropout)
ff.eval()
actual_output = ff(random_float_inputs)
assert torch.allclose(actual_output, expected_output, rtol=1e-4)
def test_gradient_flow(self, embed_dim, random_float_inputs):
"""Test that gradients flow through FeedForward."""
ff = FeedForward(embed_dim)
# Forward pass
output = ff(random_float_inputs)
# Create a dummy loss and backward pass
loss = output.sum()
loss.backward()
# Check that gradients are computed for learnable parameters
assert ff._layer1.weight.grad is not None
assert ff._layer2.weight.grad is not None
assert not torch.allclose(ff._layer1.weight.grad,
torch.zeros_like(ff._layer1.weight.grad))
assert not torch.allclose(ff._layer2.weight.grad,
torch.zeros_like(ff._layer2.weight.grad))
assert not torch.allclose(
ff._layer1.weight.grad, torch.zeros_like(ff._layer1.weight.grad)
)
assert not torch.allclose(
ff._layer2.weight.grad, torch.zeros_like(ff._layer2.weight.grad)
)
def test_device_consistency(self, embed_dim, random_float_inputs, device):
"""Test that FeedForward works on correct device."""
ff = FeedForward(embed_dim).to(device)
inputs = random_float_inputs.to(device)
# Forward pass
output = ff(inputs)
# Check device consistency
assert output.device == device
assert ff._layer1.weight.device == device
assert ff._layer2.weight.device == device
def test_different_embed_dims(self):
"""Test FeedForward with different embedding dimensions."""
test_cases = [64, 128, 256, 512]
for embed_dim in test_cases:
ff = FeedForward(embed_dim)
batch_size, seq_len = 2, 16
inputs = torch.randn(batch_size, seq_len, embed_dim)
output = ff(inputs)
assert output.shape == inputs.shape
@pytest.mark.parametrize("batch_size,seq_len", [(1, 8), (2, 16), (4, 32)])
def test_different_input_shapes(self, embed_dim, batch_size, seq_len):
"""Test FeedForward with different input shapes."""
ff = FeedForward(embed_dim)
inputs = torch.randn(batch_size, seq_len, embed_dim)
output = ff(inputs)
assert output.shape == (batch_size, seq_len, embed_dim)
def test_non_linearity(self, embed_dim, random_float_inputs):
"""Test that FeedForward introduces non-linearity."""
ff = FeedForward(embed_dim)
# Create a simple linear transformation for comparison
linear_layer = nn.Linear(embed_dim, embed_dim)
# Copy weights to make comparison fair
with torch.no_grad():
linear_layer.weight.copy_(ff._layer2.weight @ ff._layer1.weight)
if linear_layer.bias is not None:
linear_layer.bias.zero_()
linear_output = linear_layer(random_float_inputs)
ff_output = ff(random_float_inputs)
# FeedForward output should be different from pure linear transformation
# due to activation function
assert not torch.allclose(ff_output, linear_output, rtol=1e-4)
def test_parameter_initialization(self, embed_dim):
"""Test that parameters are properly initialized."""
ff = FeedForward(embed_dim)
# Check that weights are not all zeros
assert not torch.allclose(ff._layer1.weight, torch.zeros_like(ff._layer1.weight))
assert not torch.allclose(ff._layer2.weight, torch.zeros_like(ff._layer2.weight))
assert not torch.allclose(
ff._layer1.weight, torch.zeros_like(ff._layer1.weight)
)
assert not torch.allclose(
ff._layer2.weight, torch.zeros_like(ff._layer2.weight)
)
# Check that biases are not all zeros (they should be initialized with some values)
if ff._layer1.bias is not None:
assert not torch.allclose(ff._layer1.bias, torch.zeros_like(ff._layer1.bias))
assert not torch.allclose(
ff._layer1.bias, torch.zeros_like(ff._layer1.bias)
)
if ff._layer2.bias is not None:
assert not torch.allclose(ff._layer2.bias, torch.zeros_like(ff._layer2.bias))
assert not torch.allclose(
ff._layer2.bias, torch.zeros_like(ff._layer2.bias)
)

128
llm/tests/core/test_multi_head_attention.py
View File

@@ -9,157 +9,181 @@ from llm.core.multi_head_attention import MultiHeadAttention
class TestMultiHeadAttention:
"""Test cases for MultiHeadAttention."""
def test_initialization(self, embed_dim, num_heads):
"""Test that MultiHeadAttention can be initialized."""
head_size = embed_dim // num_heads
attention = MultiHeadAttention(num_heads, embed_dim, head_size, max_seq_len=1024)
attention = MultiHeadAttention(
num_heads, embed_dim, head_size, max_seq_len=1024
)
assert attention is not None
# Check internal attributes
assert len(attention._heads) == num_heads
assert attention._layer.in_features == embed_dim
assert attention._layer.out_features == embed_dim
def test_forward_pass(self, embed_dim, num_heads, random_embeddings):
"""Test forward pass of MultiHeadAttention."""
head_size = embed_dim // num_heads
attention = MultiHeadAttention(num_heads, embed_dim, head_size, max_seq_len=1024)
attention = MultiHeadAttention(
num_heads, embed_dim, head_size, max_seq_len=1024
)
# Forward pass
output, _ = attention(random_embeddings)
# Check output shape
assert output.shape == random_embeddings.shape
assert isinstance(output, torch.Tensor)
def test_forward_with_mask(self, embed_dim, num_heads, random_embeddings):
"""Test forward pass with attention mask."""
head_size = embed_dim // num_heads
attention = MultiHeadAttention(num_heads, embed_dim, head_size, max_seq_len=1024)
attention = MultiHeadAttention(
num_heads, embed_dim, head_size, max_seq_len=1024
)
# Create a simple mask
seq_len = random_embeddings.shape[1]
mask = torch.tril(torch.ones(seq_len, seq_len)) # Causal mask
# Forward pass with mask
output, _ = attention(random_embeddings, mask=mask)
# Check output shape
assert output.shape == random_embeddings.shape
def test_causal_mask(self, embed_dim, num_heads, random_embeddings):
"""Test that causal mask prevents attending to future positions."""
head_size = embed_dim // num_heads
attention = MultiHeadAttention(num_heads, embed_dim, head_size, max_seq_len=1024)
attention = MultiHeadAttention(
num_heads, embed_dim, head_size, max_seq_len=1024
)
# Create causal mask
seq_len = random_embeddings.shape[1]
causal_mask = torch.tril(torch.ones(seq_len, seq_len))
# Forward pass with causal mask
output, _ = attention(random_embeddings, mask=causal_mask)
# Check output shape
assert output.shape == random_embeddings.shape
def test_attention_weights_normalization(self, embed_dim, num_heads, random_embeddings):
def test_attention_weights_normalization(
self, embed_dim, num_heads, random_embeddings
):
"""Test that attention weights are properly normalized."""
head_size = embed_dim // num_heads
attention = MultiHeadAttention(num_heads, embed_dim, head_size, max_seq_len=1024)
attention = MultiHeadAttention(
num_heads, embed_dim, head_size, max_seq_len=1024
)
# Forward pass
output, _ = attention(random_embeddings)
# Check output shape
assert output.shape == random_embeddings.shape
def test_gradient_flow(self, embed_dim, num_heads, random_embeddings):
"""Test that gradients flow through MultiHeadAttention."""
head_size = embed_dim // num_heads
attention = MultiHeadAttention(num_heads, embed_dim, head_size, max_seq_len=1024)
attention = MultiHeadAttention(
num_heads, embed_dim, head_size, max_seq_len=1024
)
# Forward pass
output, _ = attention(random_embeddings)
# Create a dummy loss and backward pass
loss = output.sum()
loss.backward()
# Check that gradients are computed for learnable parameters
assert attention._layer.weight.grad is not None
if len(attention._heads) > 0:
assert attention._heads[0]._q.weight.grad is not None
def test_device_consistency(self, embed_dim, num_heads, random_embeddings, device):
"""Test that MultiHeadAttention works on correct device."""
head_size = embed_dim // num_heads
attention = MultiHeadAttention(num_heads, embed_dim, head_size, max_seq_len=1024).to(device)
attention = MultiHeadAttention(
num_heads, embed_dim, head_size, max_seq_len=1024
).to(device)
inputs = random_embeddings.to(device)
# Forward pass
output, _ = attention(inputs)
# Check device consistency
assert output.device == device
assert attention._layer.weight.device == device
def test_different_embed_dim_and_heads(self):
"""Test MultiHeadAttention with different embed_dim and num_heads combinations."""
test_cases = [
(64, 2), # embed_dim=64, num_heads=2
(64, 2), # embed_dim=64, num_heads=2
(128, 4), # embed_dim=128, num_heads=4
(256, 8), # embed_dim=256, num_heads=8
(512, 16), # embed_dim=512, num_heads=16
(512, 16), # embed_dim=512, num_heads=16
]
for embed_dim, num_heads in test_cases:
head_size = embed_dim // num_heads
attention = MultiHeadAttention(num_heads, embed_dim, head_size, max_seq_len=1024)
attention = MultiHeadAttention(
num_heads, embed_dim, head_size, max_seq_len=1024
)
batch_size, seq_len = 2, 16
inputs = torch.randn(batch_size, seq_len, embed_dim)
output, _ = attention(inputs)
assert output.shape == inputs.shape
def test_attention_output_range(self, embed_dim, num_heads, random_embeddings):
"""Test that attention output is in reasonable range."""
head_size = embed_dim // num_heads
attention = MultiHeadAttention(num_heads, embed_dim, head_size, max_seq_len=1024)
attention = MultiHeadAttention(
num_heads, embed_dim, head_size, max_seq_len=1024
)
output, _ = attention(random_embeddings)
# Output shouldn't have extreme values
assert output.abs().max() < 100 # Reasonable upper bound
@pytest.mark.parametrize("batch_size,seq_len", [(1, 8), (2, 16), (4, 32)])
def test_different_input_shapes(self, embed_dim, num_heads, batch_size, seq_len):
"""Test MultiHeadAttention with different input shapes."""
head_size = embed_dim // num_heads
attention = MultiHeadAttention(num_heads, embed_dim, head_size, max_seq_len=1024)
attention = MultiHeadAttention(
num_heads, embed_dim, head_size, max_seq_len=1024
)
inputs = torch.randn(batch_size, seq_len, embed_dim)
output, _ = attention(inputs)
assert output.shape == (batch_size, seq_len, embed_dim)
def test_parameter_sharing(self, embed_dim, num_heads):
"""Test that parameters are properly shared across the sequence."""
head_size = embed_dim // num_heads
attention = MultiHeadAttention(num_heads, embed_dim, head_size, max_seq_len=1024, dropout=0.0) # No dropout for deterministic test
attention = MultiHeadAttention(
num_heads, embed_dim, head_size, max_seq_len=1024, dropout=0.0
) # No dropout for deterministic test
# Create two identical sequences
seq_len = 10
base_sequence = torch.randn(1, seq_len, embed_dim)
identical_sequence = base_sequence.clone()
# Set to eval mode to disable dropout
attention.eval()
with torch.no_grad():
output1, _ = attention(base_sequence)
output2, _ = attention(identical_sequence)
# With identical inputs and same parameters, outputs should be identical
assert torch.allclose(output1, output2, rtol=1e-5)

81
llm/tests/core/test_positional_embeddings.py
View File

@@ -10,127 +10,134 @@ from llm.core.positional_embeddings import PositionalEmbeddings
class TestPositionalEmbeddings:
"""Test cases for PositionalEmbeddings."""
def test_initialization(self, embed_dim):
"""Test that PositionalEmbeddings can be initialized."""
max_seq_len = 1024
embeddings = PositionalEmbeddings(max_seq_len, embed_dim)
assert embeddings is not None
# Check that positional embeddings are created
assert hasattr(embeddings, 'embedding')
assert hasattr(embeddings, "embedding")
assert embeddings.embedding.weight.shape == (max_seq_len, embed_dim)
def test_forward_pass(self, embed_dim):
"""Test forward pass of PositionalEmbeddings."""
max_seq_len = 1024
seq_len = 64
embeddings = PositionalEmbeddings(max_seq_len, embed_dim)
# Forward pass - takes sequence length, not input tensor
output = embeddings(seq_len)
# Check output shape
expected_shape = (seq_len, embed_dim)
assert output.shape == expected_shape
assert isinstance(output, torch.Tensor)
def test_positional_encoding_values(self, embed_dim):
"""Test that positional encoding values are computed correctly."""
max_seq_len = 10
embeddings = PositionalEmbeddings(max_seq_len, embed_dim)
# Get embeddings for all positions
pe = embeddings(max_seq_len) # Shape: [max_seq_len, embed_dim]
# Check that different positions have different embeddings
# (since these are learnable embeddings, not fixed sine/cosine)
for pos in range(max_seq_len):
for i in range(pos + 1, max_seq_len):
assert not torch.allclose(pe[pos], pe[i], rtol=1e-4)
def test_different_sequence_lengths(self, embed_dim):
"""Test PositionalEmbeddings with different sequence lengths."""
test_cases = [
(10, 5), # seq_len < max_seq_len
(10, 5), # seq_len < max_seq_len
(10, 10), # seq_len == max_seq_len
]
for max_seq_len, seq_len in test_cases:
embeddings = PositionalEmbeddings(max_seq_len, embed_dim)
# Get embeddings for specific sequence length
output = embeddings(seq_len)
# Output should have shape [seq_len, embed_dim]
assert output.shape == (seq_len, embed_dim)
def test_gradient_flow(self, embed_dim):
"""Test that gradients flow through PositionalEmbeddings."""
max_seq_len = 64
seq_len = 32
embeddings = PositionalEmbeddings(max_seq_len, embed_dim)
# Forward pass
output = embeddings(seq_len)
# Create a dummy loss and backward pass
loss = output.sum()
loss.backward()
# Positional embeddings should have gradients (they're learnable)
assert embeddings.embedding.weight.grad is not None
assert not torch.allclose(embeddings.embedding.weight.grad,
torch.zeros_like(embeddings.embedding.weight.grad))
assert not torch.allclose(
embeddings.embedding.weight.grad,
torch.zeros_like(embeddings.embedding.weight.grad),
)
def test_device_consistency(self, embed_dim, device):
"""Test that PositionalEmbeddings works on correct device."""
max_seq_len = 64
seq_len = 32
embeddings = PositionalEmbeddings(max_seq_len, embed_dim).to(device)
# Forward pass
output = embeddings(seq_len)
# Check device consistency
assert output.device == device
assert embeddings.embedding.weight.device == device
def test_reproducibility(self, embed_dim):
"""Test that positional embeddings are reproducible."""
max_seq_len = 100
embeddings1 = PositionalEmbeddings(max_seq_len, embed_dim)
embeddings2 = PositionalEmbeddings(max_seq_len, embed_dim)
# Different instances should have different embeddings (random initialization)
assert not torch.allclose(embeddings1.embedding.weight, embeddings2.embedding.weight)
assert not torch.allclose(
embeddings1.embedding.weight, embeddings2.embedding.weight
)
# But same instance should produce same output for same input
seq_len = 50
output1 = embeddings1(seq_len)
output2 = embeddings1(seq_len) # Same instance, same input
assert torch.allclose(output1, output2)
def test_positional_pattern(self, embed_dim):
"""Test that positional embeddings create a meaningful pattern."""
max_seq_len = 50
embeddings = PositionalEmbeddings(max_seq_len, embed_dim)
pe = embeddings(max_seq_len) # Get all positional embeddings
# Check that different positions have different embeddings
# (with high probability due to random initialization)
assert not torch.allclose(pe[0], pe[1], rtol=1e-4)
assert not torch.allclose(pe[10], pe[20], rtol=1e-4)
@pytest.mark.parametrize("max_seq_len,seq_len,embed_dim", [
(64, 10, 64),
(128, 50, 128),
(256, 100, 256),
])
@pytest.mark.parametrize(
"max_seq_len,seq_len,embed_dim",
[
(64, 10, 64),
(128, 50, 128),
(256, 100, 256),
],
)
def test_different_configurations(self, max_seq_len, seq_len, embed_dim):
"""Test PositionalEmbeddings with different configurations."""
embeddings = PositionalEmbeddings(max_seq_len, embed_dim)
output = embeddings(seq_len)
assert output.shape == (seq_len, embed_dim)

74
llm/tests/core/test_token_embeddings.py
View File

@@ -9,99 +9,103 @@ from llm.core.token_embeddings import TokenEmbeddings
class TestTokenEmbeddings:
"""Test cases for TokenEmbeddings."""
def test_initialization(self, vocab_size, embed_dim):
"""Test that TokenEmbeddings can be initialized."""
embeddings = TokenEmbeddings(vocab_size, embed_dim)
assert embeddings is not None
# Check embedding layer
assert hasattr(embeddings, '_embedding')
assert hasattr(embeddings, "_embedding")
assert embeddings._embedding.weight.shape == (vocab_size, embed_dim)
def test_forward_pass(self, vocab_size, embed_dim, random_inputs):
"""Test forward pass of TokenEmbeddings."""
embeddings = TokenEmbeddings(vocab_size, embed_dim)
# Forward pass
output = embeddings(random_inputs)
# Check output shape
assert output.shape == (random_inputs.shape[0], random_inputs.shape[1], embed_dim)
assert output.shape == (
random_inputs.shape[0],
random_inputs.shape[1],
embed_dim,
)
assert isinstance(output, torch.Tensor)
def test_embedding_weights(self, vocab_size, embed_dim):
"""Test that embedding weights are properly initialized."""
embeddings = TokenEmbeddings(vocab_size, embed_dim)
weights = embeddings._embedding.weight
assert weights.requires_grad is True
# Check that weights are not all zeros
assert not torch.allclose(weights, torch.zeros_like(weights))
def test_different_vocab_sizes(self):
"""Test TokenEmbeddings with different vocabulary sizes."""
test_cases = [
(100, 128),
(1000, 256),
(50000, 512)
]
test_cases = [(100, 128), (1000, 256), (50000, 512)]
for vocab_size, embed_dim in test_cases:
embeddings = TokenEmbeddings(vocab_size, embed_dim)
assert embeddings._embedding.weight.shape == (vocab_size, embed_dim)
def test_gradient_flow(self, vocab_size, embed_dim, random_inputs):
"""Test that gradients flow through TokenEmbeddings."""
embeddings = TokenEmbeddings(vocab_size, embed_dim)
# Forward pass
output = embeddings(random_inputs)
# Create a dummy loss and backward pass
loss = output.sum()
loss.backward()
# Check that gradients are computed
assert embeddings._embedding.weight.grad is not None
assert not torch.allclose(embeddings._embedding.weight.grad,
torch.zeros_like(embeddings._embedding.weight.grad))
assert not torch.allclose(
embeddings._embedding.weight.grad,
torch.zeros_like(embeddings._embedding.weight.grad),
)
def test_device_consistency(self, vocab_size, embed_dim, random_inputs, device):
"""Test that TokenEmbeddings works on correct device."""
embeddings = TokenEmbeddings(vocab_size, embed_dim).to(device)
inputs = random_inputs.to(device)
# Forward pass
output = embeddings(inputs)
# Check device consistency
assert output.device == device
assert embeddings._embedding.weight.device == device
def test_embedding_lookup(self, vocab_size, embed_dim):
"""Test specific embedding lookups."""
embeddings = TokenEmbeddings(vocab_size, embed_dim)
# Test lookup for specific tokens
test_tokens = torch.tensor([[0, 1, 2], [vocab_size - 1, vocab_size - 2, vocab_size - 3]])
test_tokens = torch.tensor(
[[0, 1, 2], [vocab_size - 1, vocab_size - 2, vocab_size - 3]]
)
output = embeddings(test_tokens)
# Check shape
assert output.shape == (2, 3, embed_dim)
# Check that different tokens have different embeddings
# (with high probability due to random initialization)
assert not torch.allclose(output[0, 0], output[0, 1], rtol=1e-4)
@pytest.mark.parametrize("batch_size,seq_len", [(1, 1), (2, 10), (8, 64)])
def test_different_input_shapes(self, vocab_size, embed_dim, batch_size, seq_len):
"""Test TokenEmbeddings with different input shapes."""
embeddings = TokenEmbeddings(vocab_size, embed_dim)
inputs = torch.randint(0, vocab_size, (batch_size, seq_len))
output = embeddings(inputs)
assert output.shape == (batch_size, seq_len, embed_dim)

200
llm/tests/models/test_gpt.py
View File

@@ -9,162 +9,156 @@ from llm.models.gpt import GPT
class TestGPT:
"""Test cases for GPT model."""
def test_initialization(self, gpt_config):
"""Test that GPT can be initialized."""
model = GPT(gpt_config)
assert model is not None
# Check that model has required components
assert hasattr(model, '_token_embeddings')
assert hasattr(model, '_position_embeddings')
assert hasattr(model, '_decoders')
assert hasattr(model, '_linear')
assert hasattr(model, '_dropout')
assert hasattr(model, "_token_embeddings")
assert hasattr(model, "_position_embeddings")
assert hasattr(model, "_decoders")
assert hasattr(model, "_linear")
assert hasattr(model, "_dropout")
# Check number of decoder layers
assert len(model._decoders) == gpt_config['num_layers']
assert len(model._decoders) == gpt_config["num_layers"]
def test_forward_pass(self, gpt_config, random_inputs):
"""Test forward pass of GPT."""
model = GPT(gpt_config)
# Forward pass
logits = model(random_inputs)
# Check output shape
batch_size, seq_len = random_inputs.shape
vocab_size = gpt_config['vocab_size']
vocab_size = gpt_config["vocab_size"]
assert logits.shape == (batch_size, seq_len, vocab_size)
assert isinstance(logits, torch.Tensor)
def test_forward_with_attention_mask(self, gpt_config, random_inputs, attention_mask):
def test_forward_with_attention_mask(
self, gpt_config, random_inputs, attention_mask
):
"""Test forward pass with attention mask."""
model = GPT(gpt_config)
# Forward pass with mask
logits = model(random_inputs, attention_mask=attention_mask)
# Check output shape
batch_size, seq_len = random_inputs.shape
vocab_size = gpt_config['vocab_size']
vocab_size = gpt_config["vocab_size"]
assert logits.shape == (batch_size, seq_len, vocab_size)
def test_generate_text(self, gpt_config):
"""Test text generation."""
model = GPT(gpt_config)
model.eval() # Set to evaluation mode for generation
# Create initial input
batch_size = 2
initial_seq_len = 5
input_ids = torch.randint(0, gpt_config['vocab_size'], (batch_size, initial_seq_len))
input_ids = torch.randint(
0, gpt_config["vocab_size"], (batch_size, initial_seq_len)
)
# Generate text
with torch.no_grad():
generated = model.generate(
x=input_ids,
max_new_tokens=10,
do_sample=False # Use greedy for deterministic testing
do_sample=False, # Use greedy for deterministic testing
)
# Check output shape
expected_seq_len = initial_seq_len + 10
assert generated.shape == (batch_size, expected_seq_len)
# Check that initial sequence is preserved
assert torch.allclose(generated[:, :initial_seq_len], input_ids)
def test_generate_with_temperature(self, gpt_config):
"""Test text generation with temperature sampling."""
model = GPT(gpt_config)
model.eval()
# Create initial input
input_ids = torch.randint(0, gpt_config['vocab_size'], (1, 3))
input_ids = torch.randint(0, gpt_config["vocab_size"], (1, 3))
# Generate with temperature
with torch.no_grad():
generated = model.generate(
x=input_ids,
max_new_tokens=5,
do_sample=True,
temperature=0.8
x=input_ids, max_new_tokens=5, do_sample=True, temperature=0.8
)
assert generated.shape == (1, 8) # 3 initial + 5 new tokens
def test_generate_with_top_k(self, gpt_config):
"""Test text generation with top-k sampling."""
model = GPT(gpt_config)
model.eval()
# Create initial input
input_ids = torch.randint(0, gpt_config['vocab_size'], (1, 3))
input_ids = torch.randint(0, gpt_config["vocab_size"], (1, 3))
# Generate with top-k
with torch.no_grad():
generated = model.generate(
x=input_ids,
max_new_tokens=5,
do_sample=True,
top_k=10
x=input_ids, max_new_tokens=5, do_sample=True, top_k=10
)
assert generated.shape == (1, 8)
def test_generate_with_top_p(self, gpt_config):
"""Test text generation with top-p (nucleus) sampling."""
model = GPT(gpt_config)
model.eval()
# Create initial input
input_ids = torch.randint(0, gpt_config['vocab_size'], (1, 3))
input_ids = torch.randint(0, gpt_config["vocab_size"], (1, 3))
# Generate with top-p
with torch.no_grad():
generated = model.generate(
x=input_ids,
max_new_tokens=5,
do_sample=True,
top_p=0.9
x=input_ids, max_new_tokens=5, do_sample=True, top_p=0.9
)
assert generated.shape == (1, 8)
def test_gradient_flow(self, gpt_config, random_inputs):
"""Test that gradients flow through GPT."""
model = GPT(gpt_config)
# Forward pass
logits = model(random_inputs)
# Create a dummy loss and backward pass
targets = torch.randint(0, gpt_config['vocab_size'], random_inputs.shape)
targets = torch.randint(0, gpt_config["vocab_size"], random_inputs.shape)
loss = torch.nn.functional.cross_entropy(
logits.view(-1, logits.size(-1)),
targets.view(-1)
logits.view(-1, logits.size(-1)), targets.view(-1)
)
loss.backward()
# Check that gradients are computed for various components
assert model._token_embeddings._embedding.weight.grad is not None
assert model._linear.weight.grad is not None
if len(model._decoders) > 0:
assert model._decoders[0]._heads._heads[0]._q.weight.grad is not None
def test_device_consistency(self, gpt_config, random_inputs, device):
"""Test that GPT works on correct device."""
model = GPT(gpt_config).to(device)
inputs = random_inputs.to(device)
# Forward pass
logits = model(inputs)
# Check device consistency
assert logits.device == device
assert model._token_embeddings._embedding.weight.device == device
def test_different_configurations(self):
"""Test GPT with different configurations."""
test_configs = [
@@ -174,7 +168,7 @@ class TestGPT:
"num_heads": 2,
"num_layers": 2,
"max_position_embeddings": 256,
"dropout": 0.1
"dropout": 0.1,
},
{
"vocab_size": 5000,
@@ -182,7 +176,7 @@ class TestGPT:
"num_heads": 4,
"num_layers": 4,
"max_position_embeddings": 512,
"dropout": 0.1
"dropout": 0.1,
},
{
"vocab_size": 10000,
@@ -190,98 +184,94 @@ class TestGPT:
"num_heads": 8,
"num_layers": 6,
"max_position_embeddings": 1024,
"dropout": 0.1
}
"dropout": 0.1,
},
]
for config in test_configs:
model = GPT(config)
batch_size, seq_len = 2, 16
inputs = torch.randint(0, config['vocab_size'], (batch_size, seq_len))
inputs = torch.randint(0, config["vocab_size"], (batch_size, seq_len))
logits = model(inputs)
expected_shape = (batch_size, seq_len, config['vocab_size'])
expected_shape = (batch_size, seq_len, config["vocab_size"])
assert logits.shape == expected_shape
@pytest.mark.parametrize("batch_size,seq_len", [(1, 8), (2, 16), (4, 32)])
def test_different_input_shapes(self, gpt_config, batch_size, seq_len):
"""Test GPT with different input shapes."""
model = GPT(gpt_config)
inputs = torch.randint(0, gpt_config['vocab_size'], (batch_size, seq_len))
inputs = torch.randint(0, gpt_config["vocab_size"], (batch_size, seq_len))
logits = model(inputs)
expected_shape = (batch_size, seq_len, gpt_config['vocab_size'])
expected_shape = (batch_size, seq_len, gpt_config["vocab_size"])
assert logits.shape == expected_shape
def test_training_vs_evaluation(self, gpt_config, random_inputs):
"""Test that GPT behaves differently in train vs eval mode."""
model = GPT(gpt_config)
# Training mode
model.train()
output_train = model(random_inputs)
# Evaluation mode
model.eval()
output_eval = model(random_inputs)
# Outputs should be different due to dropout
assert not torch.allclose(output_train, output_eval)
def test_parameter_count(self, gpt_config):
"""Test that GPT has reasonable number of parameters."""
model = GPT(gpt_config)
total_params = sum(p.numel() for p in model.parameters())
# For a small GPT model, parameters should be in reasonable range
vocab_size = gpt_config['vocab_size']
embed_dim = gpt_config['embed_dim']
num_layers = gpt_config['num_layers']
num_heads = gpt_config['num_heads']
vocab_size = gpt_config["vocab_size"]
embed_dim = gpt_config["embed_dim"]
num_layers = gpt_config["num_layers"]
num_heads = gpt_config["num_heads"]
# Rough estimate: token_embeddings + output_layer + (attention + ff) * layers
expected_min = vocab_size * embed_dim * 2 # embeddings and output
expected_max = expected_min * 10 # Allow for decoder parameters
assert expected_min < total_params < expected_max
def test_causal_attention(self, gpt_config):
"""Test that GPT uses causal attention during generation."""
model = GPT(gpt_config)
model.eval()
# Create input with known pattern
input_ids = torch.tensor([[1, 2, 3]]).long()
with torch.no_grad():
# Get logits for next token prediction
logits = model(input_ids)
# The model should only attend to previous tokens (causal)
# We can't directly test attention masks in the public API,
# but we can verify the generation works correctly
generated = model.generate(
x=input_ids,
max_new_tokens=3,
do_sample=False
)
generated = model.generate(x=input_ids, max_new_tokens=3, do_sample=False)
# Generated sequence should be longer than input
assert generated.shape[1] == input_ids.shape[1] + 3
def test_output_distribution(self, gpt_config, random_inputs):
"""Test that GPT output has proper distribution."""
model = GPT(gpt_config)
logits = model(random_inputs)
# Logits should not have extreme values
assert logits.abs().max() < 100
# Softmax should produce valid probabilities
probs = torch.softmax(logits, dim=-1)
assert torch.allclose(probs.sum(dim=-1), torch.ones_like(probs.sum(dim=-1)))

130
llm/tests/test_basic.py
View File

@@ -11,25 +11,25 @@ import os
def test_gpt_model_creation():
"""Test that GPT model can be created and forward pass works."""
from llm.models.gpt import GPT
config = {
"vocab_size": 1000,
"embed_dim": 128,
"num_heads": 4,
"num_layers": 2,
"max_position_embeddings": 256,
"dropout": 0.1
"dropout": 0.1,
}
model = GPT(config)
# Test forward pass
batch_size, seq_len = 2, 16
input_ids = torch.randint(0, config["vocab_size"], (batch_size, seq_len))
with torch.no_grad():
logits = model(input_ids)
assert logits.shape == (batch_size, seq_len, config["vocab_size"])
print("✅ GPT model creation and forward pass test passed")
@@ -37,27 +37,21 @@ def test_gpt_model_creation():
def test_bpe_tokenizer_basic():
"""Test basic BPE tokenizer functionality."""
from llm.tokenizers import BPETokenizer
tokenizer = BPETokenizer()
# Train on simple texts
texts = [
"hello world",
"test tokenization",
"simple example"
]
texts = ["hello world", "test tokenization", "simple example"]
tokenizer.train(
texts=texts,
vocab_size=50,
special_tokens=["<pad>", "<unk>", "<bos>", "<eos>"]
texts=texts, vocab_size=50, special_tokens=["<pad>", "<unk>", "<bos>", "<eos>"]
)
# Test encoding/decoding
text = "hello world"
tokens = tokenizer.encode(text)
decoded = tokenizer.decode(tokens)
assert isinstance(tokens, list)
assert isinstance(decoded, str)
assert len(tokens) > 0
@@ -67,18 +61,18 @@ def test_bpe_tokenizer_basic():
def test_token_embeddings():
"""Test token embeddings."""
from llm.core.token_embeddings import TokenEmbeddings
vocab_size = 1000
embed_dim = 128
embeddings = TokenEmbeddings(vocab_size, embed_dim)
# Test forward pass
batch_size, seq_len = 2, 16
input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
output = embeddings(input_ids)
assert output.shape == (batch_size, seq_len, embed_dim)
print("✅ Token embeddings test passed")
@@ -86,20 +80,20 @@ def test_token_embeddings():
def test_multi_head_attention():
"""Test multi-head attention."""
from llm.core.multi_head_attention import MultiHeadAttention
num_heads = 4
emb_size = 128
head_size = emb_size // num_heads
max_seq_len = 256
attention = MultiHeadAttention(num_heads, emb_size, head_size, max_seq_len)
# Test forward pass
batch_size, seq_len = 2, 16
inputs = torch.randn(batch_size, seq_len, emb_size)
output, _ = attention(inputs)
assert output.shape == inputs.shape
print("✅ Multi-head attention test passed")
@@ -107,17 +101,17 @@ def test_multi_head_attention():
def test_feed_forward():
"""Test feed forward network."""
from llm.core.feed_forward import FeedForward
embed_dim = 128
ff = FeedForward(embed_dim)
# Test forward pass
batch_size, seq_len = 2, 16
inputs = torch.randn(batch_size, seq_len, embed_dim)
output = ff(inputs)
assert output.shape == inputs.shape
print("✅ Feed forward test passed")
@@ -125,29 +119,25 @@ def test_feed_forward():
def test_gpt_generation():
"""Test GPT text generation."""
from llm.models.gpt import GPT
config = {
"vocab_size": 1000,
"embed_dim": 128,
"num_heads": 4,
"num_layers": 2,
"max_position_embeddings": 256,
"dropout": 0.1
"dropout": 0.1,
}
model = GPT(config)
model.eval()
# Test greedy generation
input_ids = torch.randint(0, config["vocab_size"], (1, 5))
with torch.no_grad():
generated = model.generate(
x=input_ids,
max_new_tokens=3,
do_sample=False
)
generated = model.generate(x=input_ids, max_new_tokens=3, do_sample=False)
assert generated.shape == (1, 8) # 5 initial + 3 new tokens
print("✅ GPT generation test passed")
@@ -155,50 +145,48 @@ def test_gpt_generation():
def test_bpe_tokenizer_save_load():
"""Test BPE tokenizer save/load functionality."""
from llm.tokenizers import BPETokenizer
tokenizer = BPETokenizer()
# Train on simple texts
texts = ["hello world", "test save load"]
tokenizer.train(
texts=texts,
vocab_size=30,
special_tokens=["<pad>", "<unk>", "<bos>", "<eos>"]
texts=texts, vocab_size=30, special_tokens=["<pad>", "<unk>", "<bos>", "<eos>"]
)
with tempfile.TemporaryDirectory() as temp_dir:
save_path = os.path.join(temp_dir, "test_tokenizer.json")
# Save tokenizer
tokenizer.save(save_path)
assert os.path.exists(save_path)
# Load tokenizer
loaded_tokenizer = BPETokenizer.load(save_path)
# Test that vocab size is the same
assert tokenizer.get_vocab_size() == loaded_tokenizer.get_vocab_size()
# Test that vocabularies are the same
assert tokenizer.get_vocab() == loaded_tokenizer.get_vocab()
# Test that both can encode/decode (even if tokens differ due to BPE state)
text = "hello world"
original_tokens = tokenizer.encode(text)
loaded_tokens = loaded_tokenizer.encode(text)
# Both should produce valid token lists
assert isinstance(original_tokens, list)
assert isinstance(loaded_tokens, list)
assert len(original_tokens) > 0
assert len(loaded_tokens) > 0
# Both should be able to decode
original_decoded = tokenizer.decode(original_tokens)
loaded_decoded = loaded_tokenizer.decode(loaded_tokens)
assert isinstance(original_decoded, str)
assert isinstance(loaded_decoded, str)
print("✅ BPE tokenizer save/load test passed")
@@ -206,18 +194,16 @@ def test_gpt_with_tokenizer():
"""Test GPT model with tokenizer integration."""
from llm.models.gpt import GPT
from llm.tokenizers import BPETokenizer
# Create and train tokenizer
tokenizer = BPETokenizer()
texts = ["hello world", "test integration"]
tokenizer.train(
texts=texts,
vocab_size=50,
special_tokens=["<pad>", "<unk>", "<bos>", "<eos>"]
texts=texts, vocab_size=50, special_tokens=["<pad>", "<unk>", "<bos>", "<eos>"]
)
vocab_size = tokenizer.get_vocab_size()
# Create GPT model with tokenizer's vocab size
config = {
"vocab_size": vocab_size,
@@ -225,19 +211,19 @@ def test_gpt_with_tokenizer():
"num_heads": 4,
"num_layers": 2,
"max_position_embeddings": 256,
"dropout": 0.1
"dropout": 0.1,
}
model = GPT(config)
# Test with tokenized input
text = "hello world"
tokens = tokenizer.encode(text, add_special_tokens=False)
input_ids = torch.tensor([tokens])
with torch.no_grad():
logits = model(input_ids)
assert logits.shape == (1, len(tokens), vocab_size)
print("✅ GPT with tokenizer integration test passed")
@@ -245,7 +231,7 @@ def test_gpt_with_tokenizer():
def run_all_tests():
"""Run all basic tests."""
print("🧪 Running basic tests for llm library...")
test_gpt_model_creation()
test_bpe_tokenizer_basic()
test_token_embeddings()
@@ -254,7 +240,7 @@ def run_all_tests():
test_gpt_generation()
test_bpe_tokenizer_save_load()
test_gpt_with_tokenizer()
print("🎉 All basic tests passed!")

16
llm/tests/tokenizers/test_base_tokenizer.py
View File

@@ -8,15 +8,15 @@ from llm.tokenizers import BaseTokenizer
class ConcreteTokenizer(BaseTokenizer):
"""Concrete implementation for testing BaseTokenizer."""
def train(self, texts: list, vocab_size: int = 1000, **kwargs):
"""Dummy implementation for testing."""
pass
def encode(self, text: str, **kwargs) -> list:
"""Dummy implementation for testing."""
return [1, 2, 3]
def decode(self, tokens: list, **kwargs) -> str:
"""Dummy implementation for testing."""
return "decoded text"
@@ -24,33 +24,33 @@ class ConcreteTokenizer(BaseTokenizer):
class TestBaseTokenizer:
"""Test cases for BaseTokenizer."""
def test_initialization(self):
"""Test that BaseTokenizer can be initialized through concrete class."""
tokenizer = ConcreteTokenizer()
assert tokenizer is not None
assert tokenizer.vocab == {}
assert tokenizer.vocab_size == 0
def test_encode_implemented(self):
"""Test that encode method works in concrete implementation."""
tokenizer = ConcreteTokenizer()
result = tokenizer.encode("test text")
assert result == [1, 2, 3]
def test_decode_implemented(self):
"""Test that decode method works in concrete implementation."""
tokenizer = ConcreteTokenizer()
result = tokenizer.decode([1, 2, 3])
assert result == "decoded text"
def test_get_vocab_size(self):
"""Test that get_vocab_size method works."""
tokenizer = ConcreteTokenizer()
tokenizer.vocab = {"a": 0, "b": 1, "c": 2}
tokenizer.vocab_size = 3
assert tokenizer.get_vocab_size() == 3
def test_get_vocab(self):
"""Test that get_vocab method works."""
tokenizer = ConcreteTokenizer()

72
llm/tests/tokenizers/test_bpe_tokenizer.py
View File

@@ -10,18 +10,18 @@ from llm.tokenizers import BPETokenizer
class TestBPETokenizer:
"""Test cases for BPETokenizer."""
@pytest.fixture
def sample_texts(self):
"""Sample texts for training tokenizer."""
return [
"Искусственный интеллект",
"Нейронные сети",
"Нейронные сети",
"Машинное обучение",
"Глубокое обучение",
"Трансформеры"
"Трансформеры",
]
@pytest.fixture
def trained_tokenizer(self, sample_texts):
"""Create and train a BPE tokenizer."""
@@ -29,128 +29,130 @@ class TestBPETokenizer:
tokenizer.train(
texts=sample_texts,
vocab_size=100,
special_tokens=["<pad>", "<unk>", "<bos>", "<eos>"]
special_tokens=["<pad>", "<unk>", "<bos>", "<eos>"],
)
return tokenizer
def test_initialization(self):
"""Test that BPETokenizer can be initialized."""
tokenizer = BPETokenizer()
assert tokenizer is not None
def test_train_tokenizer(self, sample_texts):
"""Test that tokenizer can be trained."""
tokenizer = BPETokenizer()
tokenizer.train(
texts=sample_texts,
vocab_size=50,
special_tokens=["<pad>", "<unk>", "<bos>", "<eos>"]
special_tokens=["<pad>", "<unk>", "<bos>", "<eos>"],
)
assert tokenizer.get_vocab_size() > 0
assert len(tokenizer.get_vocab()) == tokenizer.get_vocab_size()
def test_encode_decode(self, trained_tokenizer):
"""Test encoding and decoding text."""
text = "Искусственный интеллект"
# Encode text
tokens = trained_tokenizer.encode(text)
assert isinstance(tokens, list)
assert len(tokens) > 0
assert all(isinstance(token, int) for token in tokens)
# Decode tokens
decoded_text = trained_tokenizer.decode(tokens)
assert isinstance(decoded_text, str)
# Decoded text should be similar to original (may have special tokens)
assert len(decoded_text) > 0
def test_encode_with_special_tokens(self, trained_tokenizer):
"""Test encoding with special tokens."""
text = "Нейронные сети"
# Without special tokens
tokens_no_special = trained_tokenizer.encode(text, add_special_tokens=False)
# With special tokens
tokens_with_special = trained_tokenizer.encode(text, add_special_tokens=True)
# Should have more tokens when special tokens are added
assert len(tokens_with_special) >= len(tokens_no_special)
def test_vocab_size(self, trained_tokenizer):
"""Test vocabulary size."""
vocab_size = trained_tokenizer.get_vocab_size()
assert isinstance(vocab_size, int)
assert vocab_size > 0
vocab = trained_tokenizer.get_vocab()
assert isinstance(vocab, dict)
assert len(vocab) == vocab_size
def test_special_tokens(self, trained_tokenizer):
"""Test that special tokens are in vocabulary."""
vocab = trained_tokenizer.get_vocab()
# Check that special tokens are in vocabulary
special_tokens = ["<pad>", "<unk>", "<bos>", "<eos>"]
for token in special_tokens:
assert token in vocab
assert isinstance(vocab[token], int)
def test_save_load(self, trained_tokenizer, sample_texts):
"""Test saving and loading tokenizer."""
with tempfile.TemporaryDirectory() as temp_dir:
save_path = os.path.join(temp_dir, "test_tokenizer.json")
# Save tokenizer
trained_tokenizer.save(save_path)
assert os.path.exists(save_path)
# Load tokenizer
loaded_tokenizer = BPETokenizer.load(save_path)
assert loaded_tokenizer is not None
# Check that loaded tokenizer works the same
original_vocab = trained_tokenizer.get_vocab()
loaded_vocab = loaded_tokenizer.get_vocab()
assert original_vocab == loaded_vocab
assert trained_tokenizer.get_vocab_size() == loaded_tokenizer.get_vocab_size()
assert (
trained_tokenizer.get_vocab_size() == loaded_tokenizer.get_vocab_size()
)
# Test encoding consistency
text = sample_texts[0]
original_tokens = trained_tokenizer.encode(text)
loaded_tokens = loaded_tokenizer.encode(text)
assert original_tokens == loaded_tokens
def test_unknown_tokens(self, trained_tokenizer):
"""Test handling of unknown tokens."""
# Use text that likely contains unknown subwords
text = "xyzabc123" # Random text that shouldn't be in training data
tokens = trained_tokenizer.encode(text)
assert len(tokens) > 0
# Should be able to decode back (even if it's mostly unk tokens)
decoded = trained_tokenizer.decode(tokens)
assert isinstance(decoded, str)
def test_empty_text(self, trained_tokenizer):
"""Test encoding and decoding empty text."""
tokens = trained_tokenizer.encode("")
assert isinstance(tokens, list)
decoded = trained_tokenizer.decode([])
assert decoded == ""
def test_tokenize_method(self, trained_tokenizer):
"""Test the tokenize method."""
text = "Искусственный интеллект"
tokens = trained_tokenizer.tokenize(text)
assert isinstance(tokens, list)
assert len(tokens) > 0
assert all(isinstance(token, str) for token in tokens)