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 hf_proxy import HFAdapter, HFTokenizerAdapter, create_hf_pipeline
from shared.configs import ( from shared.configs import TEST_PROMPTS, GENERATION_CONFIG, PATHS
TEST_PROMPTS, GENERATION_CONFIG, PATHS from shared.data import print_experiment_info, ensure_directories, ExperimentLogger
)
from shared.data import (
print_experiment_info, ensure_directories, ExperimentLogger
)
def load_hf_model_and_tokenizer() -> tuple: def load_hf_model_and_tokenizer() -> tuple:
""" """
Загружает модель и токенизатор в формате HuggingFace. Загружает модель и токенизатор в формате HuggingFace.
Returns: Returns:
tuple: (hf_model, hf_tokenizer, model_config) tuple: (hf_model, hf_tokenizer, model_config)
""" """
# Используем упрощенную версию модели # Используем упрощенную версию модели
model_path = "checkpoints/hf_simple_trained" model_path = "checkpoints/hf_simple_trained"
tokenizer_path = "checkpoints/hf_simple_tokenizer" tokenizer_path = "checkpoints/hf_simple_tokenizer"
# Проверяем существование файлов # Проверяем существование файлов
if not os.path.exists(model_path): if not os.path.exists(model_path):
raise FileNotFoundError( raise FileNotFoundError(
f"Модель не найдена: {model_path}\n" f"Модель не найдена: {model_path}\n"
f"Сначала обучите модель: uv run python experiments/hf_integration/simple_hf_training.py" f"Сначала обучите модель: uv run python experiments/hf_integration/simple_hf_training.py"
) )
if not os.path.exists(tokenizer_path): if not os.path.exists(tokenizer_path):
raise FileNotFoundError( raise FileNotFoundError(f"Токенизатор не найден: {tokenizer_path}")
f"Токенизатор не найден: {tokenizer_path}"
)
# Загружаем адаптированный токенизатор # Загружаем адаптированный токенизатор
print("🔧 Загрузка адаптированного токенизатора...") print("🔧 Загрузка адаптированного токенизатора...")
hf_tokenizer = HFTokenizerAdapter.from_pretrained(tokenizer_path) hf_tokenizer = HFTokenizerAdapter.from_pretrained(tokenizer_path)
print(f"✅ Токенизатор загружен (vocab_size={hf_tokenizer.vocab_size})") print(f"✅ Токенизатор загружен (vocab_size={hf_tokenizer.vocab_size})")
# Загружаем конфигурацию модели # Загружаем конфигурацию модели
import json import json
config_path = os.path.join(model_path, "config.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) model_config = json.load(f)
# Загружаем модель через HFAdapter с правильной конфигурацией # Загружаем модель через HFAdapter с правильной конфигурацией
print("🔧 Загрузка адаптированной модели...") print("🔧 Загрузка адаптированной модели...")
model_bin_path = os.path.join(model_path, "pytorch_model.bin") model_bin_path = os.path.join(model_path, "pytorch_model.bin")
# Создаем конфигурацию из сохраненного config.json # Создаем конфигурацию из сохраненного config.json
from hf_proxy import HFAdapterConfig from hf_proxy import HFAdapterConfig
hf_config = HFAdapterConfig( hf_config = HFAdapterConfig(
vocab_size=model_config["vocab_size"], vocab_size=model_config["vocab_size"],
hidden_size=model_config["hidden_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"], num_attention_heads=model_config["num_attention_heads"],
max_position_embeddings=model_config["max_position_embeddings"], max_position_embeddings=model_config["max_position_embeddings"],
hidden_dropout_prob=model_config.get("hidden_dropout_prob", 0.1), 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 = HFAdapter.from_pretrained(model_bin_path, hf_config=hf_config)
hf_model.eval() hf_model.eval()
print("✅ Модель загружена") print("✅ Модель загружена")
return hf_model, hf_tokenizer, model_config return hf_model, hf_tokenizer, model_config
def test_hf_pipeline(hf_model, hf_tokenizer): def test_hf_pipeline(hf_model, hf_tokenizer):
""" """
Тестирует создание HuggingFace pipeline. Тестирует создание HuggingFace pipeline.
Args: Args:
hf_model: Адаптированная модель hf_model: Адаптированная модель
hf_tokenizer: Адаптированный токенизатор hf_tokenizer: Адаптированный токенизатор
""" """
print("\n🧪 Тестирование HuggingFace pipeline...") print("\n🧪 Тестирование HuggingFace pipeline...")
try: try:
# Создаем pipeline # Создаем pipeline
pipe = create_hf_pipeline( pipe = create_hf_pipeline(
@@ -97,23 +95,23 @@ def test_hf_pipeline(hf_model, hf_tokenizer):
device="cpu", device="cpu",
max_length=50, max_length=50,
do_sample=True, do_sample=True,
temperature=0.7 temperature=0.7,
) )
print("✅ HuggingFace pipeline создан") print("✅ HuggingFace pipeline создан")
# Тестируем pipeline # Тестируем pipeline
test_prompts = TEST_PROMPTS[:3] test_prompts = TEST_PROMPTS[:3]
for prompt in test_prompts: for prompt in test_prompts:
print(f"\n🔤 Промпт: '{prompt}'") print(f"\n🔤 Промпт: '{prompt}'")
try: try:
result = pipe(prompt, max_new_tokens=20) result = pipe(prompt, max_new_tokens=20)
print(f"🎯 Результат: {result[0]['generated_text']}") print(f"🎯 Результат: {result[0]['generated_text']}")
except Exception as e: except Exception as e:
print(f"❌ Ошибка в pipeline: {e}") print(f"❌ Ошибка в pipeline: {e}")
except Exception as e: except Exception as e:
print(f"❌ Ошибка создания pipeline: {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: def generate_with_hf_model(hf_model, hf_tokenizer, prompt: str, config: dict) -> str:
""" """
Генерирует текст через адаптированную модель HF. Генерирует текст через адаптированную модель HF.
Args: Args:
hf_model: Адаптированная модель hf_model: Адаптированная модель
hf_tokenizer: Адаптированный токенизатор hf_tokenizer: Адаптированный токенизатор
prompt: Входной текст prompt: Входной текст
config: Конфигурация генерации config: Конфигурация генерации
Returns: Returns:
str: Сгенерированный текст str: Сгенерированный текст
""" """
print(f"🔤 Промпт: '{prompt}'") print(f"🔤 Промпт: '{prompt}'")
print(f"📊 Параметры: max_tokens={config['max_new_tokens']}, " print(
f"temp={config['temperature']}, sample={config['do_sample']}") f"📊 Параметры: max_tokens={config['max_new_tokens']}, "
f"temp={config['temperature']}, sample={config['do_sample']}"
)
# Кодируем через адаптированный токенизатор # Кодируем через адаптированный токенизатор
inputs = hf_tokenizer(prompt, return_tensors="pt") inputs = hf_tokenizer(prompt, return_tensors="pt")
print(f"🎯 Токены промпта: {inputs['input_ids'].tolist()[0]}") print(f"🎯 Токены промпта: {inputs['input_ids'].tolist()[0]}")
print("🔄 Генерация через HF адаптер...") print("🔄 Генерация через HF адаптер...")
# Генерируем через адаптированную модель # Генерируем через адаптированную модель
with torch.no_grad(): with torch.no_grad():
generated_ids = hf_model.generate( generated_ids = hf_model.generate(
input_ids=inputs['input_ids'], input_ids=inputs["input_ids"],
max_new_tokens=config["max_new_tokens"], max_new_tokens=config["max_new_tokens"],
do_sample=config["do_sample"], do_sample=config["do_sample"],
temperature=config["temperature"], temperature=config["temperature"],
top_k=config["top_k"], 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) generated_text = hf_tokenizer.decode(generated_ids[0], skip_special_tokens=True)
return generated_text return generated_text
def test_different_hf_strategies(hf_model, hf_tokenizer, prompt: str): def test_different_hf_strategies(hf_model, hf_tokenizer, prompt: str):
""" """
Тестирует разные стратегии генерации через HF интерфейс. Тестирует разные стратегии генерации через HF интерфейс.
Args: Args:
hf_model: Адаптированная модель hf_model: Адаптированная модель
hf_tokenizer: Адаптированный токенизатор hf_tokenizer: Адаптированный токенизатор
@@ -169,32 +169,38 @@ def test_different_hf_strategies(hf_model, hf_tokenizer, prompt: str):
""" """
print(f"\n🎭 Сравнение стратегий генерации через HF для промпта: '{prompt}'") print(f"\n🎭 Сравнение стратегий генерации через HF для промпта: '{prompt}'")
print("=" * 70) print("=" * 70)
strategies = [ strategies = [
{"name": "🎯 Жадный поиск", "do_sample": False, "temperature": 1.0}, {"name": "🎯 Жадный поиск", "do_sample": False, "temperature": 1.0},
{"name": "🎲 Вероятностная (temp=0.7)", "do_sample": True, "temperature": 0.7}, {"name": "🎲 Вероятностная (temp=0.7)", "do_sample": True, "temperature": 0.7},
{"name": "🔥 Случайная (temp=1.2)", "do_sample": True, "temperature": 1.2}, {"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: for strategy in strategies:
print(f"\n{strategy['name']}:") print(f"\n{strategy['name']}:")
try: try:
config = GENERATION_CONFIG.copy() config = GENERATION_CONFIG.copy()
config.update({ config.update(
"do_sample": strategy["do_sample"], {
"temperature": strategy["temperature"], "do_sample": strategy["do_sample"],
"max_new_tokens": 20 "temperature": strategy["temperature"],
}) "max_new_tokens": 20,
}
)
generated = generate_with_hf_model(hf_model, hf_tokenizer, prompt, config) 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" 📤 Промпт: '{prompt}'")
print(f" 🎯 Сгенерировано: '{generated_part}'") print(f" 🎯 Сгенерировано: '{generated_part}'")
print(f" 📄 Полный текст: '{generated}'") print(f" 📄 Полный текст: '{generated}'")
except Exception as e: except Exception as e:
print(f" ❌ Ошибка: {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): def analyze_hf_tokenization(hf_tokenizer, texts: list):
""" """
Анализирует токенизацию через адаптированный токенизатор. Анализирует токенизацию через адаптированный токенизатор.
Args: Args:
hf_tokenizer: Адаптированный токенизатор hf_tokenizer: Адаптированный токенизатор
texts: Список текстов для анализа texts: Список текстов для анализа
""" """
print(f"\n🔍 Анализ токенизации через HF адаптер:") print(f"\n🔍 Анализ токенизации через HF адаптер:")
print("=" * 60) print("=" * 60)
for i, text in enumerate(texts): for i, text in enumerate(texts):
print(f"\nТекст {i+1}: '{text}'") print(f"\nТекст {i+1}: '{text}'")
# Токенизация через адаптер # Токенизация через адаптер
inputs = hf_tokenizer(text, return_tensors="pt") 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) token_strings = hf_tokenizer.tokenize(text)
print(f" Токены (ID): {tokens}") print(f" Токены (ID): {tokens}")
print(f" Токены (текст): {token_strings}") print(f" Токены (текст): {token_strings}")
print(f" Количество токенов: {len(tokens)}") print(f" Количество токенов: {len(tokens)}")
# Декодирование обратно # Декодирование обратно
decoded = hf_tokenizer.decode(tokens) decoded = hf_tokenizer.decode(tokens)
print(f" Декодированный: '{decoded}'") print(f" Декодированный: '{decoded}'")
if text == decoded: if text == decoded:
print(f" ✅ Декодирование корректно") print(f" ✅ Декодирование корректно")
else: else:
@@ -235,51 +241,55 @@ def analyze_hf_tokenization(hf_tokenizer, texts: list):
def interactive_hf_generation(hf_model, hf_tokenizer): def interactive_hf_generation(hf_model, hf_tokenizer):
""" """
Режим интерактивной генерации через HF интерфейс. Режим интерактивной генерации через HF интерфейс.
Args: Args:
hf_model: Адаптированная модель hf_model: Адаптированная модель
hf_tokenizer: Адаптированный токенизатор hf_tokenizer: Адаптированный токенизатор
""" """
print(f"\n💬 Интерактивная генерация через HF (для выхода введите 'exit')") print(f"\n💬 Интерактивная генерация через HF (для выхода введите 'exit')")
print("-" * 60) print("-" * 60)
while True: while True:
try: try:
user_input = input("\n🔤 Введите промпт: ").strip() user_input = input("\n🔤 Введите промпт: ").strip()
if user_input.lower() in ['exit', 'quit', 'выход']: if user_input.lower() in ["exit", "quit", "выход"]:
break break
if not user_input: if not user_input:
continue continue
# Запрашиваем параметры # Запрашиваем параметры
try: try:
max_tokens = int(input("📏 Макс. токенов [50]: ") or "50") max_tokens = int(input("📏 Макс. токенов [50]: ") or "50")
temperature = float(input("🌡️ Температура [0.7]: ") or "0.7") temperature = float(input("🌡️ Температура [0.7]: ") or "0.7")
do_sample_input = input("🎲 Сэмплирование (y/n) [y]: ").lower() do_sample_input = input("🎲 Сэмплирование (y/n) [y]: ").lower()
do_sample = do_sample_input != 'n' do_sample = do_sample_input != "n"
except: except:
max_tokens = 50 max_tokens = 50
temperature = 0.7 temperature = 0.7
do_sample = True do_sample = True
print("⚠️ Использую параметры по умолчанию") print("⚠️ Использую параметры по умолчанию")
config = GENERATION_CONFIG.copy() config = GENERATION_CONFIG.copy()
config.update({ config.update(
"max_new_tokens": max_tokens, {
"temperature": temperature, "max_new_tokens": max_tokens,
"do_sample": do_sample "temperature": temperature,
}) "do_sample": do_sample,
}
generated = generate_with_hf_model(hf_model, hf_tokenizer, user_input, config) )
generated_part = generated[len(user_input):] generated = generate_with_hf_model(
hf_model, hf_tokenizer, user_input, config
)
generated_part = generated[len(user_input) :]
print(f"\n🎯 Результат:") print(f"\n🎯 Результат:")
print(f" 📤 Промпт: '{user_input}'") print(f" 📤 Промпт: '{user_input}'")
print(f" 🎯 Сгенерировано: '{generated_part}'") print(f" 🎯 Сгенерировано: '{generated_part}'")
print(f" 📄 Полный текст: '{generated}'") print(f" 📄 Полный текст: '{generated}'")
except KeyboardInterrupt: except KeyboardInterrupt:
print("\n👋 Завершение работы...") print("\n👋 Завершение работы...")
break break
@@ -295,76 +305,79 @@ def main():
"model": "GPT через HFAdapter", "model": "GPT через HFAdapter",
"tokenizer": "BPE через HFTokenizerAdapter", "tokenizer": "BPE через HFTokenizerAdapter",
"инструменты": "HuggingFace pipeline & генерация", "инструменты": "HuggingFace pipeline & генерация",
"стратегия": "интеграция с HF экосистемой" "стратегия": "интеграция с HF экосистемой",
} }
print_experiment_info(experiment_name, experiment_config) print_experiment_info(experiment_name, experiment_config)
ensure_directories() ensure_directories()
logger = ExperimentLogger(experiment_name) logger = ExperimentLogger(experiment_name)
try: try:
# Загружаем модель и токенизатор в HF формате # Загружаем модель и токенизатор в HF формате
hf_model, hf_tokenizer, model_config = load_hf_model_and_tokenizer() hf_model, hf_tokenizer, model_config = load_hf_model_and_tokenizer()
# === Анализ токенизации === # === Анализ токенизации ===
analysis_texts = [ analysis_texts = [
"Искусственный интеллект", "Искусственный интеллект",
"Нейронные сети", "Нейронные сети",
"Машинное обучение" "Машинное обучение",
] ]
analyze_hf_tokenization(hf_tokenizer, analysis_texts) analyze_hf_tokenization(hf_tokenizer, analysis_texts)
# === Тестирование HF pipeline === # === Тестирование HF pipeline ===
test_hf_pipeline(hf_model, hf_tokenizer) test_hf_pipeline(hf_model, hf_tokenizer)
# === Генерация с разными промптами === # === Генерация с разными промптами ===
print(f"\n🎯 Генерация текста через HF адаптер") print(f"\n🎯 Генерация текста через HF адаптер")
print("=" * 60) print("=" * 60)
for i, prompt in enumerate(TEST_PROMPTS): for i, prompt in enumerate(TEST_PROMPTS):
print(f"\n📝 Пример {i+1}/{len(TEST_PROMPTS)}") print(f"\n📝 Пример {i+1}/{len(TEST_PROMPTS)}")
print("-" * 40) print("-" * 40)
try: 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"📤 Промпт: '{prompt}'")
print(f"🎯 Сгенерировано: '{generated_part}'") print(f"🎯 Сгенерировано: '{generated_part}'")
print(f"📄 Полный текст: '{generated}'") print(f"📄 Полный текст: '{generated}'")
print(f"📏 Длина: {len(generated)} символов") print(f"📏 Длина: {len(generated)} символов")
# Логируем успешную генерацию # Логируем успешную генерацию
logger.log_metric(f"hf_generation_length_{i}", len(generated)) logger.log_metric(f"hf_generation_length_{i}", len(generated))
except Exception as e: except Exception as e:
print(f"❌ Ошибка при генерации: {e}") print(f"❌ Ошибка при генерации: {e}")
continue continue
# === Сравнение стратегий генерации === # === Сравнение стратегий генерации ===
test_prompt = "Искусственный" test_prompt = "Искусственный"
test_different_hf_strategies(hf_model, hf_tokenizer, test_prompt) test_different_hf_strategies(hf_model, hf_tokenizer, test_prompt)
# === Интерактивная генерация === # === Интерактивная генерация ===
interactive_hf_generation(hf_model, hf_tokenizer) interactive_hf_generation(hf_model, hf_tokenizer)
# === Сохранение результатов === # === Сохранение результатов ===
logger.save_logs("checkpoints/hf_integration_generation_logs.json") logger.save_logs("checkpoints/hf_integration_generation_logs.json")
print(f"\n🎉 Эксперимент с HF интеграцией завершен успешно!") print(f"\n🎉 Эксперимент с HF интеграцией завершен успешно!")
print(f"\n📚 Достигнутая интеграция:") print(f"\n📚 Достигнутая интеграция:")
print(f" ✅ Загрузка модели и токенизатора в HF формате") print(f" ✅ Загрузка модели и токенизатора в HF формате")
print(f" ✅ Использование HF pipeline") print(f" ✅ Использование HF pipeline")
print(f" ✅ Генерация через стандартные HF интерфейсы") print(f" ✅ Генерация через стандартные HF интерфейсы")
print(f" ✅ Совместимость с HF экосистемой") print(f" ✅ Совместимость с HF экосистемой")
except FileNotFoundError as e: except FileNotFoundError as e:
print(f"❌ {e}") print(f"❌ {e}")
except Exception as e: except Exception as e:
print(f"❌ Ошибка в эксперименте: {e}") print(f"❌ Ошибка в эксперименте: {e}")
import traceback import traceback
traceback.print_exc() 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 hf_proxy import HFAdapter, HFTokenizerAdapter
from shared.configs import ( from shared.configs import (
TRAIN_TEXTS, BASE_GPT_CONFIG, BPE_CONFIG, TRAIN_TEXTS,
TRAINING_CONFIG, PATHS, TEST_PROMPTS BASE_GPT_CONFIG,
BPE_CONFIG,
TRAINING_CONFIG,
PATHS,
TEST_PROMPTS,
) )
def create_dataset(hf_tokenizer, texts, max_length=128): def create_dataset(hf_tokenizer, texts, max_length=128):
""" """
Создает простой датасет для обучения. Создает простой датасет для обучения.
Args: Args:
hf_tokenizer: Адаптированный токенизатор hf_tokenizer: Адаптированный токенизатор
texts: Список текстов texts: Список текстов
max_length: Максимальная длина последовательности max_length: Максимальная длина последовательности
Returns: Returns:
list: Список тензоров input_ids list: Список тензоров input_ids
""" """
dataset = [] dataset = []
for text in texts: for text in texts:
# Токенизируем текст # Токенизируем текст
inputs = hf_tokenizer( inputs = hf_tokenizer(
text, text,
max_length=max_length, max_length=max_length,
truncation=True, truncation=True,
padding=False, padding=False,
return_tensors="pt" return_tensors="pt",
) )
input_ids = inputs['input_ids'][0] input_ids = inputs["input_ids"][0]
# Создаем метки для языкового моделирования # Создаем метки для языкового моделирования
labels = input_ids.clone() labels = input_ids.clone()
dataset.append({ dataset.append({"input_ids": input_ids, "labels": labels})
'input_ids': input_ids,
'labels': labels
})
return dataset return dataset
def manual_training_loop(hf_model, hf_tokenizer, train_texts, val_texts, config): def manual_training_loop(hf_model, hf_tokenizer, train_texts, val_texts, config):
""" """
Ручной цикл обучения без использования Trainer. Ручной цикл обучения без использования Trainer.
Args: Args:
hf_model: Адаптированная модель hf_model: Адаптированная модель
hf_tokenizer: Адаптированный токенизатор hf_tokenizer: Адаптированный токенизатор
train_texts: Тексты для обучения train_texts: Тексты для обучения
val_texts: Тексты для валидации val_texts: Тексты для валидации
config: Конфигурация обучения config: Конфигурация обучения
Returns: Returns:
dict: Результаты обучения dict: Результаты обучения
""" """
print("🎯 Запуск ручного обучения...") print("🎯 Запуск ручного обучения...")
# Создаем датасеты # Создаем датасеты
train_dataset = create_dataset(hf_tokenizer, train_texts) train_dataset = create_dataset(hf_tokenizer, train_texts)
val_dataset = create_dataset(hf_tokenizer, val_texts) val_dataset = create_dataset(hf_tokenizer, val_texts)
print(f"📊 Данные: {len(train_dataset)} train, {len(val_dataset)} validation") print(f"📊 Данные: {len(train_dataset)} train, {len(val_dataset)} validation")
# Оптимизатор # Оптимизатор
optimizer = torch.optim.AdamW( optimizer = torch.optim.AdamW(hf_model.parameters(), lr=config["learning_rate"])
hf_model.parameters(),
lr=config["learning_rate"]
)
# Функция потерь # Функция потерь
loss_fn = nn.CrossEntropyLoss() loss_fn = nn.CrossEntropyLoss()
# Обучение # Обучение
hf_model.train() hf_model.train()
train_losses = [] train_losses = []
val_losses = [] val_losses = []
for epoch in range(config["num_epochs"]): for epoch in range(config["num_epochs"]):
print(f"\n📅 Эпоха {epoch + 1}/{config['num_epochs']}") print(f"\n📅 Эпоха {epoch + 1}/{config['num_epochs']}")
# Обучение # Обучение
epoch_train_loss = 0 epoch_train_loss = 0
for i, batch in enumerate(train_dataset): for i, batch in enumerate(train_dataset):
optimizer.zero_grad() optimizer.zero_grad()
input_ids = batch['input_ids'].unsqueeze(0) # [1, seq_len] input_ids = batch["input_ids"].unsqueeze(0) # [1, seq_len]
labels = batch['labels'].unsqueeze(0) # [1, seq_len] labels = batch["labels"].unsqueeze(0) # [1, seq_len]
# Forward pass # Forward pass
outputs = hf_model(input_ids=input_ids, labels=labels) outputs = hf_model(input_ids=input_ids, labels=labels)
loss = outputs.loss loss = outputs.loss
# Backward pass # Backward pass
loss.backward() loss.backward()
optimizer.step() optimizer.step()
epoch_train_loss += loss.item() epoch_train_loss += loss.item()
if i % 5 == 0: if i % 5 == 0:
print(f" Batch {i}/{len(train_dataset)}: loss = {loss.item():.4f}") print(f" Batch {i}/{len(train_dataset)}: loss = {loss.item():.4f}")
avg_train_loss = epoch_train_loss / len(train_dataset) avg_train_loss = epoch_train_loss / len(train_dataset)
train_losses.append(avg_train_loss) train_losses.append(avg_train_loss)
print(f" 📊 Средняя train loss: {avg_train_loss:.4f}") print(f" 📊 Средняя train loss: {avg_train_loss:.4f}")
# Валидация # Валидация
hf_model.eval() hf_model.eval()
epoch_val_loss = 0 epoch_val_loss = 0
with torch.no_grad(): with torch.no_grad():
for batch in val_dataset: for batch in val_dataset:
input_ids = batch['input_ids'].unsqueeze(0) input_ids = batch["input_ids"].unsqueeze(0)
labels = batch['labels'].unsqueeze(0) labels = batch["labels"].unsqueeze(0)
outputs = hf_model(input_ids=input_ids, labels=labels) outputs = hf_model(input_ids=input_ids, labels=labels)
epoch_val_loss += outputs.loss.item() epoch_val_loss += outputs.loss.item()
avg_val_loss = epoch_val_loss / len(val_dataset) avg_val_loss = epoch_val_loss / len(val_dataset)
val_losses.append(avg_val_loss) val_losses.append(avg_val_loss)
print(f" 📊 Средняя val loss: {avg_val_loss:.4f}") print(f" 📊 Средняя val loss: {avg_val_loss:.4f}")
hf_model.train() hf_model.train()
return { return {
'train_losses': train_losses, "train_losses": train_losses,
'val_losses': val_losses, "val_losses": val_losses,
'final_train_loss': train_losses[-1], "final_train_loss": train_losses[-1],
'final_val_loss': val_losses[-1] "final_val_loss": val_losses[-1],
} }
def test_generation_after_training(hf_model, hf_tokenizer, test_prompts): def test_generation_after_training(hf_model, hf_tokenizer, test_prompts):
""" """
Тестирует генерацию после обучения. Тестирует генерацию после обучения.
Args: Args:
hf_model: Обученная модель hf_model: Обученная модель
hf_tokenizer: Токенизатор hf_tokenizer: Токенизатор
@@ -161,24 +159,24 @@ def test_generation_after_training(hf_model, hf_tokenizer, test_prompts):
""" """
print("\n🧪 Тестирование генерации после обучения...") print("\n🧪 Тестирование генерации после обучения...")
hf_model.eval() hf_model.eval()
for prompt in test_prompts[:3]: for prompt in test_prompts[:3]:
print(f"\n🔤 Промпт: '{prompt}'") print(f"\n🔤 Промпт: '{prompt}'")
try: try:
inputs = hf_tokenizer(prompt, return_tensors="pt") inputs = hf_tokenizer(prompt, return_tensors="pt")
with torch.no_grad(): with torch.no_grad():
generated = hf_model.generate( generated = hf_model.generate(
input_ids=inputs['input_ids'], input_ids=inputs["input_ids"],
max_new_tokens=20, max_new_tokens=20,
do_sample=True, 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}'") print(f"🎯 Результат: '{generated_text}'")
except Exception as e: except Exception as e:
print(f"❌ Ошибка генерации: {e}") print(f"❌ Ошибка генерации: {e}")
@@ -188,96 +186,102 @@ def main():
print("=" * 60) print("=" * 60)
print("🚀 УПРОЩЕННОЕ ОБУЧЕНИЕ GPT С HF-PROXY") print("🚀 УПРОЩЕННОЕ ОБУЧЕНИЕ GPT С HF-PROXY")
print("=" * 60) print("=" * 60)
try: try:
# === Подготовка данных === # === Подготовка данных ===
print("🔧 Подготовка данных...") print("🔧 Подготовка данных...")
train_texts = TRAIN_TEXTS[:10] # Используем меньше данных для быстрого тестирования train_texts = TRAIN_TEXTS[
:10
] # Используем меньше данных для быстрого тестирования
val_texts = TRAIN_TEXTS[10:12] val_texts = TRAIN_TEXTS[10:12]
print(f"📊 Данные: {len(train_texts)} train, {len(val_texts)} validation") print(f"📊 Данные: {len(train_texts)} train, {len(val_texts)} validation")
# === Подготовка токенизатора === # === Подготовка токенизатора ===
print("🔧 Подготовка токенизатора...") print("🔧 Подготовка токенизатора...")
llm_tokenizer = BPETokenizer() llm_tokenizer = BPETokenizer()
llm_tokenizer.train( llm_tokenizer.train(
texts=train_texts, texts=train_texts,
vocab_size=BPE_CONFIG["vocab_size"], vocab_size=BPE_CONFIG["vocab_size"],
special_tokens=BPE_CONFIG["special_tokens"] special_tokens=BPE_CONFIG["special_tokens"],
) )
hf_tokenizer = HFTokenizerAdapter(llm_tokenizer) hf_tokenizer = HFTokenizerAdapter(llm_tokenizer)
print(f"✅ Токенизатор создан (vocab_size={hf_tokenizer.vocab_size})") print(f"✅ Токенизатор создан (vocab_size={hf_tokenizer.vocab_size})")
# === Подготовка модели === # === Подготовка модели ===
print("🔧 Подготовка модели...") print("🔧 Подготовка модели...")
model_config = BASE_GPT_CONFIG.copy() model_config = BASE_GPT_CONFIG.copy()
model_config["vocab_size"] = hf_tokenizer.vocab_size model_config["vocab_size"] = hf_tokenizer.vocab_size
llm_model = GPT(model_config) llm_model = GPT(model_config)
hf_model = HFAdapter.from_llm_model(llm_model) hf_model = HFAdapter.from_llm_model(llm_model)
print(f"✅ Модель создана") print(f"✅ Модель создана")
# === Тестирование до обучения === # === Тестирование до обучения ===
print("\n🧪 Тестирование до обучения...") print("\n🧪 Тестирование до обучения...")
test_generation_after_training(hf_model, hf_tokenizer, TEST_PROMPTS) test_generation_after_training(hf_model, hf_tokenizer, TEST_PROMPTS)
# === Обучение === # === Обучение ===
print(f"\n🎯 Обучение модели...") print(f"\n🎯 Обучение модели...")
training_config = { training_config = {
"learning_rate": TRAINING_CONFIG["learning_rate"], "learning_rate": TRAINING_CONFIG["learning_rate"],
"num_epochs": 2, # Меньше эпох для быстрого тестирования "num_epochs": 2, # Меньше эпох для быстрого тестирования
"batch_size": TRAINING_CONFIG["batch_size"] "batch_size": TRAINING_CONFIG["batch_size"],
} }
results = manual_training_loop( results = manual_training_loop(
hf_model, hf_tokenizer, train_texts, val_texts, training_config hf_model, hf_tokenizer, train_texts, val_texts, training_config
) )
print(f"\n📊 Результаты обучения:") print(f"\n📊 Результаты обучения:")
print(f" Final train loss: {results['final_train_loss']:.4f}") print(f" Final train loss: {results['final_train_loss']:.4f}")
print(f" Final val loss: {results['final_val_loss']:.4f}") print(f" Final val loss: {results['final_val_loss']:.4f}")
# === Тестирование после обучения === # === Тестирование после обучения ===
print("\n🧪 Тестирование после обучения...") print("\n🧪 Тестирование после обучения...")
test_generation_after_training(hf_model, hf_tokenizer, TEST_PROMPTS) test_generation_after_training(hf_model, hf_tokenizer, TEST_PROMPTS)
# === Сохранение модели === # === Сохранение модели ===
print(f"\n💾 Сохранение модели...") print(f"\n💾 Сохранение модели...")
# Создаем директории # Создаем директории
os.makedirs("checkpoints/hf_simple_trained", exist_ok=True) os.makedirs("checkpoints/hf_simple_trained", exist_ok=True)
os.makedirs("checkpoints/hf_simple_tokenizer", exist_ok=True) os.makedirs("checkpoints/hf_simple_tokenizer", exist_ok=True)
# Сохраняем токенизатор # Сохраняем токенизатор
hf_tokenizer.save_pretrained("checkpoints/hf_simple_tokenizer") hf_tokenizer.save_pretrained("checkpoints/hf_simple_tokenizer")
print("✅ Токенизатор сохранен") print("✅ Токенизатор сохранен")
# Сохраняем модель # Сохраняем модель
HFAdapter.save_pretrained( HFAdapter.save_pretrained(
hf_model, hf_model, "checkpoints/hf_simple_trained", tokenizer=hf_tokenizer
"checkpoints/hf_simple_trained",
tokenizer=hf_tokenizer
) )
print("✅ Модель сохранена") print("✅ Модель сохранена")
# Сохраняем результаты # Сохраняем результаты
results_path = "checkpoints/simple_training_results.json" results_path = "checkpoints/simple_training_results.json"
with open(results_path, 'w', encoding='utf-8') as f: with open(results_path, "w", encoding="utf-8") as f:
json.dump({ json.dump(
'training_config': training_config, {
'model_config': model_config, "training_config": training_config,
'results': results "model_config": model_config,
}, f, indent=2, ensure_ascii=False) "results": results,
},
f,
indent=2,
ensure_ascii=False,
)
print(f"✅ Результаты сохранены в {results_path}") print(f"✅ Результаты сохранены в {results_path}")
print(f"\n🎉 Упрощенное обучение завершено успешно!") print(f"\n🎉 Упрощенное обучение завершено успешно!")
print(f"\n💡 Для использования обученной модели:") print(f"\n💡 Для использования обученной модели:")
print(f" uv run python experiments/hf_integration/generate_with_hf_tools.py") print(f" uv run python experiments/hf_integration/generate_with_hf_tools.py")
except Exception as e: except Exception as e:
print(f"❌ Ошибка в эксперименте: {e}") print(f"❌ Ошибка в эксперименте: {e}")
import traceback import traceback
traceback.print_exc() 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 hf_proxy import HFAdapter, HFTokenizerAdapter
from shared.configs import ( from shared.configs import (
TRAIN_TEXTS, BASE_GPT_CONFIG, BPE_CONFIG, TRAIN_TEXTS,
TEST_PROMPTS, GENERATION_CONFIG BASE_GPT_CONFIG,
BPE_CONFIG,
TEST_PROMPTS,
GENERATION_CONFIG,
) )
def test_basic_hf_integration(): def test_basic_hf_integration():
"""Тестирует базовую интеграцию hf-proxy.""" """Тестирует базовую интеграцию hf-proxy."""
print("🧪 Тестирование базовой интеграции hf-proxy...") print("🧪 Тестирование базовой интеграции hf-proxy...")
# === Подготовка токенизатора === # === Подготовка токенизатора ===
print("1. Подготовка токенизатора...") print("1. Подготовка токенизатора...")
llm_tokenizer = BPETokenizer() llm_tokenizer = BPETokenizer()
llm_tokenizer.train( llm_tokenizer.train(
texts=TRAIN_TEXTS, texts=TRAIN_TEXTS,
vocab_size=BPE_CONFIG["vocab_size"], vocab_size=BPE_CONFIG["vocab_size"],
special_tokens=BPE_CONFIG["special_tokens"] special_tokens=BPE_CONFIG["special_tokens"],
) )
hf_tokenizer = HFTokenizerAdapter(llm_tokenizer) hf_tokenizer = HFTokenizerAdapter(llm_tokenizer)
print(f" ✅ Токенизатор создан (vocab_size={hf_tokenizer.vocab_size})") print(f" ✅ Токенизатор создан (vocab_size={hf_tokenizer.vocab_size})")
# === Подготовка модели === # === Подготовка модели ===
print("2. Подготовка модели...") print("2. Подготовка модели...")
model_config = BASE_GPT_CONFIG.copy() model_config = BASE_GPT_CONFIG.copy()
model_config["vocab_size"] = hf_tokenizer.vocab_size model_config["vocab_size"] = hf_tokenizer.vocab_size
llm_model = GPT(model_config) llm_model = GPT(model_config)
hf_model = HFAdapter.from_llm_model(llm_model) hf_model = HFAdapter.from_llm_model(llm_model)
print(f" ✅ Модель создана") print(f" ✅ Модель создана")
# === Тестирование токенизации === # === Тестирование токенизации ===
print("3. Тестирование токенизации...") print("3. Тестирование токенизации...")
test_texts = ["Искусственный интеллект", "Нейронные сети"] test_texts = ["Искусственный интеллект", "Нейронные сети"]
for text in test_texts: for text in test_texts:
print(f" 📝 Текст: '{text}'") print(f" 📝 Текст: '{text}'")
# Оригинальный токенизатор # Оригинальный токенизатор
original_tokens = llm_tokenizer.encode(text) original_tokens = llm_tokenizer.encode(text)
print(f" Оригинальный: {len(original_tokens)} токенов") print(f" Оригинальный: {len(original_tokens)} токенов")
# HF адаптер # HF адаптер
hf_inputs = hf_tokenizer(text, return_tensors="pt") hf_inputs = hf_tokenizer(text, return_tensors="pt")
print(f" HF адаптер: {hf_inputs['input_ids'].shape}") 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}'") print(f" Декодированный: '{decoded}'")
# === Тестирование forward pass === # === Тестирование forward pass ===
print("4. Тестирование forward pass...") print("4. Тестирование forward pass...")
for text in test_texts: for text in test_texts:
hf_inputs = hf_tokenizer(text, return_tensors="pt") hf_inputs = hf_tokenizer(text, return_tensors="pt")
with torch.no_grad(): with torch.no_grad():
outputs = hf_model(**hf_inputs) outputs = hf_model(**hf_inputs)
print(f" 📝 '{text}' -> logits: {outputs.logits.shape}") print(f" 📝 '{text}' -> logits: {outputs.logits.shape}")
# === Тестирование генерации === # === Тестирование генерации ===
print("5. Тестирование генерации...") print("5. Тестирование генерации...")
hf_model.eval() hf_model.eval()
for prompt in TEST_PROMPTS[:3]: for prompt in TEST_PROMPTS[:3]:
print(f" 🔤 Промпт: '{prompt}'") print(f" 🔤 Промпт: '{prompt}'")
try: try:
inputs = hf_tokenizer(prompt, return_tensors="pt") inputs = hf_tokenizer(prompt, return_tensors="pt")
with torch.no_grad(): with torch.no_grad():
generated = hf_model.generate( generated = hf_model.generate(
input_ids=inputs['input_ids'], input_ids=inputs["input_ids"],
max_new_tokens=10, max_new_tokens=10,
do_sample=True, 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}'") print(f" 🎯 Результат: '{generated_text}'")
except Exception as e: except Exception as e:
print(f" ❌ Ошибка: {e}") print(f" ❌ Ошибка: {e}")
# === Тестирование сохранения/загрузки === # === Тестирование сохранения/загрузки ===
print("6. Тестирование сохранения/загрузки...") print("6. Тестирование сохранения/загрузки...")
try: try:
# Сохраняем токенизатор # Сохраняем токенизатор
hf_tokenizer.save_pretrained("test_save/tokenizer") hf_tokenizer.save_pretrained("test_save/tokenizer")
print(" ✅ Токенизатор сохранен") print(" ✅ Токенизатор сохранен")
# Сохраняем модель # Сохраняем модель
HFAdapter.save_pretrained(hf_model, "test_save/model", tokenizer=hf_tokenizer) HFAdapter.save_pretrained(hf_model, "test_save/model", tokenizer=hf_tokenizer)
print(" ✅ Модель сохранена") print(" ✅ Модель сохранена")
# Загружаем токенизатор # Загружаем токенизатор
loaded_tokenizer = HFTokenizerAdapter.from_pretrained("test_save/tokenizer") loaded_tokenizer = HFTokenizerAdapter.from_pretrained("test_save/tokenizer")
print(f" ✅ Токенизатор загружен (vocab_size={loaded_tokenizer.vocab_size})") print(f" ✅ Токенизатор загружен (vocab_size={loaded_tokenizer.vocab_size})")
# Загружаем модель # Загружаем модель
model_path = os.path.join("test_save/model", "pytorch_model.bin") model_path = os.path.join("test_save/model", "pytorch_model.bin")
loaded_model = HFAdapter.from_pretrained(model_path) loaded_model = HFAdapter.from_pretrained(model_path)
print(" ✅ Модель загружена") print(" ✅ Модель загружена")
# Проверяем работоспособность загруженной модели # Проверяем работоспособность загруженной модели
test_input = hf_tokenizer("Тест", return_tensors="pt") test_input = hf_tokenizer("Тест", return_tensors="pt")
with torch.no_grad(): with torch.no_grad():
loaded_outputs = loaded_model(**test_input) loaded_outputs = loaded_model(**test_input)
print(f" ✅ Загруженная модель работает (logits: {loaded_outputs.logits.shape})") print(
f" ✅ Загруженная модель работает (logits: {loaded_outputs.logits.shape})"
)
except Exception as e: except Exception as e:
print(f" ❌ Ошибка сохранения/загрузки: {e}") print(f" ❌ Ошибка сохранения/загрузки: {e}")
print("\n🎉 Базовое тестирование hf-proxy завершено!") print("\n🎉 Базовое тестирование hf-proxy завершено!")
def test_hf_tokenizer_methods(): def test_hf_tokenizer_methods():
"""Тестирует различные методы HF токенизатора.""" """Тестирует различные методы HF токенизатора."""
print("\n🧪 Тестирование методов HF токенизатора...") print("\n🧪 Тестирование методов HF токенизатора...")
# Создаем токенизатор # Создаем токенизатор
llm_tokenizer = BPETokenizer() llm_tokenizer = BPETokenizer()
llm_tokenizer.train( llm_tokenizer.train(
texts=TRAIN_TEXTS[:5], texts=TRAIN_TEXTS[:5],
vocab_size=500, vocab_size=500,
special_tokens=BPE_CONFIG["special_tokens"] special_tokens=BPE_CONFIG["special_tokens"],
) )
hf_tokenizer = HFTokenizerAdapter(llm_tokenizer) hf_tokenizer = HFTokenizerAdapter(llm_tokenizer)
test_text = "Искусственный интеллект и машинное обучение" test_text = "Искусственный интеллект и машинное обучение"
# Тестируем разные методы # Тестируем разные методы
print("1. Метод __call__:") print("1. Метод __call__:")
result = hf_tokenizer(test_text, return_tensors="pt") result = hf_tokenizer(test_text, return_tensors="pt")
print(f" Результат: {result}") print(f" Результат: {result}")
print("2. Метод encode:") print("2. Метод encode:")
encoded = hf_tokenizer.encode(test_text) encoded = hf_tokenizer.encode(test_text)
print(f" Закодировано: {encoded}") print(f" Закодировано: {encoded}")
print("3. Метод decode:") print("3. Метод decode:")
decoded = hf_tokenizer.decode(encoded) decoded = hf_tokenizer.decode(encoded)
print(f" Декодировано: '{decoded}'") print(f" Декодировано: '{decoded}'")
print("4. Метод tokenize:") print("4. Метод tokenize:")
tokens = hf_tokenizer.tokenize(test_text) tokens = hf_tokenizer.tokenize(test_text)
print(f" Токены: {tokens}") print(f" Токены: {tokens}")
print("5. Метод get_vocab:") print("5. Метод get_vocab:")
vocab = hf_tokenizer.get_vocab() vocab = hf_tokenizer.get_vocab()
print(f" Размер словаря: {len(vocab)}") print(f" Размер словаря: {len(vocab)}")
print("✅ Все методы токенизатора работают!") print("✅ Все методы токенизатора работают!")
@@ -176,14 +181,14 @@ def main():
print("=" * 60) print("=" * 60)
print("🧪 ТЕСТИРОВАНИЕ HF-PROXY") print("🧪 ТЕСТИРОВАНИЕ HF-PROXY")
print("=" * 60) print("=" * 60)
try: try:
# Тестируем базовую интеграцию # Тестируем базовую интеграцию
test_basic_hf_integration() test_basic_hf_integration()
# Тестируем методы токенизатора # Тестируем методы токенизатора
test_hf_tokenizer_methods() test_hf_tokenizer_methods()
print("\n" + "=" * 60) print("\n" + "=" * 60)
print("🎉 ВСЕ ТЕСТЫ ПРОЙДЕНЫ УСПЕШНО!") print("🎉 ВСЕ ТЕСТЫ ПРОЙДЕНЫ УСПЕШНО!")
print("=" * 60) print("=" * 60)
@@ -195,10 +200,11 @@ def main():
print(" ✅ Генерация текста") print(" ✅ Генерация текста")
print(" ✅ Сохранение и загрузка моделей") print(" ✅ Сохранение и загрузка моделей")
print(" ✅ Все методы HF токенизатора") print(" ✅ Все методы HF токенизатора")
except Exception as e: except Exception as e:
print(f"\n❌ Ошибка в тестировании: {e}") print(f"\n❌ Ошибка в тестировании: {e}")
import traceback import traceback
traceback.print_exc() 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 hf_proxy import HFAdapter, HFTokenizerAdapter
from shared.configs import ( from shared.configs import (
TRAIN_TEXTS, BASE_GPT_CONFIG, BPE_CONFIG, TRAIN_TEXTS,
TRAINING_CONFIG, PATHS, TEST_PROMPTS BASE_GPT_CONFIG,
BPE_CONFIG,
TRAINING_CONFIG,
PATHS,
TEST_PROMPTS,
) )
from shared.data import ( from shared.data import (
load_training_data, ensure_directories, load_training_data,
print_experiment_info, ExperimentLogger ensure_directories,
print_experiment_info,
ExperimentLogger,
) )
def setup_hf_training(): def setup_hf_training():
""" """
Настраивает окружение для обучения через HuggingFace Trainer. Настраивает окружение для обучения через HuggingFace Trainer.
Returns: Returns:
tuple: (hf_model, hf_tokenizer, llm_tokenizer, model_config) tuple: (hf_model, hf_tokenizer, llm_tokenizer, model_config)
""" """
print("🔧 Настройка HuggingFace обучения...") print("🔧 Настройка HuggingFace обучения...")
# === Подготовка данных === # === Подготовка данных ===
train_texts, val_texts = load_training_data() train_texts, val_texts = load_training_data()
print(f"📊 Данные: {len(train_texts)} train, {len(val_texts)} validation") print(f"📊 Данные: {len(train_texts)} train, {len(val_texts)} validation")
# === Обучение/загрузка токенизатора === # === Обучение/загрузка токенизатора ===
if os.path.exists(PATHS["bpe_tokenizer"]): if os.path.exists(PATHS["bpe_tokenizer"]):
print("📝 Загрузка BPE токенизатора...") print("📝 Загрузка BPE токенизатора...")
@@ -50,55 +56,55 @@ def setup_hf_training():
llm_tokenizer.train( llm_tokenizer.train(
texts=TRAIN_TEXTS, texts=TRAIN_TEXTS,
vocab_size=BPE_CONFIG["vocab_size"], vocab_size=BPE_CONFIG["vocab_size"],
special_tokens=BPE_CONFIG["special_tokens"] special_tokens=BPE_CONFIG["special_tokens"],
) )
llm_tokenizer.save(PATHS["bpe_tokenizer"]) llm_tokenizer.save(PATHS["bpe_tokenizer"])
print(f"✅ Токенизатор обучен и сохранен") print(f"✅ Токенизатор обучен и сохранен")
# === Создание адаптера токенизатора === # === Создание адаптера токенизатора ===
print("🔧 Создание адаптера HuggingFace для токенизатора...") print("🔧 Создание адаптера HuggingFace для токенизатора...")
hf_tokenizer = HFTokenizerAdapter(llm_tokenizer) hf_tokenizer = HFTokenizerAdapter(llm_tokenizer)
print(f"✅ Адаптер токенизатора создан") print(f"✅ Адаптер токенизатора создан")
# === Инициализация модели === # === Инициализация модели ===
model_config = BASE_GPT_CONFIG.copy() model_config = BASE_GPT_CONFIG.copy()
model_config["vocab_size"] = llm_tokenizer.get_vocab_size() model_config["vocab_size"] = llm_tokenizer.get_vocab_size()
print("🔧 Создание GPT модели...") print("🔧 Создание GPT модели...")
llm_model = GPT(model_config) llm_model = GPT(model_config)
# === Создание адаптера модели === # === Создание адаптера модели ===
print("🔧 Создание адаптера HuggingFace для модели...") print("🔧 Создание адаптера HuggingFace для модели...")
hf_model = HFAdapter.from_llm_model(llm_model) hf_model = HFAdapter.from_llm_model(llm_model)
print(f"✅ Адаптер модели создан") print(f"✅ Адаптер модели создан")
return hf_model, hf_tokenizer, llm_tokenizer, model_config, train_texts, val_texts return hf_model, hf_tokenizer, llm_tokenizer, model_config, train_texts, val_texts
def test_hf_integration(hf_model, hf_tokenizer, llm_tokenizer): def test_hf_integration(hf_model, hf_tokenizer, llm_tokenizer):
""" """
Тестирует интеграцию с HuggingFace инструментами. Тестирует интеграцию с HuggingFace инструментами.
Args: Args:
hf_model: Адаптированная модель hf_model: Адаптированная модель
hf_tokenizer: Адаптированный токенизатор hf_tokenizer: Адаптированный токенизатор
llm_tokenizer: Оригинальный токенизатор llm_tokenizer: Оригинальный токенизатор
""" """
print("\n🧪 Тестирование интеграции с HuggingFace...") print("\n🧪 Тестирование интеграции с HuggingFace...")
test_texts = ["Искусственный интеллект", "Нейронные сети"] test_texts = ["Искусственный интеллект", "Нейронные сети"]
for text in test_texts: for text in test_texts:
print(f"\n🔤 Текст: '{text}'") print(f"\n🔤 Текст: '{text}'")
# Тестируем адаптированный токенизатор # Тестируем адаптированный токенизатор
hf_inputs = hf_tokenizer(text, return_tensors="pt") hf_inputs = hf_tokenizer(text, return_tensors="pt")
print(f" HF токенизатор: {hf_inputs['input_ids'].shape}") print(f" HF токенизатор: {hf_inputs['input_ids'].shape}")
# Тестируем оригинальный токенизатор для сравнения # Тестируем оригинальный токенизатор для сравнения
original_tokens = llm_tokenizer.encode(text) original_tokens = llm_tokenizer.encode(text)
print(f" Оригинальный токенизатор: {len(original_tokens)} токенов") print(f" Оригинальный токенизатор: {len(original_tokens)} токенов")
# Тестируем forward pass через адаптированную модель # Тестируем forward pass через адаптированную модель
try: try:
with torch.no_grad(): with torch.no_grad():
@@ -114,28 +120,35 @@ def main():
experiment_name = "Обучение GPT через HF Trainer (с hf-proxy)" experiment_name = "Обучение GPT через HF Trainer (с hf-proxy)"
experiment_config = { experiment_config = {
"model": "GPT через HFAdapter", "model": "GPT через HFAdapter",
"tokenizer": "BPE через HFTokenizerAdapter", "tokenizer": "BPE через HFTokenizerAdapter",
"trainer": "HuggingFace Trainer", "trainer": "HuggingFace Trainer",
"vocab_size": BPE_CONFIG["vocab_size"], "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) print_experiment_info(experiment_name, experiment_config)
ensure_directories() ensure_directories()
logger = ExperimentLogger(experiment_name) logger = ExperimentLogger(experiment_name)
try: 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) test_hf_integration(hf_model, hf_tokenizer, llm_tokenizer)
# === Подготовка датасетов HuggingFace === # === Подготовка датасетов HuggingFace ===
print(f"\n📊 Подготовка датасетов HuggingFace...") print(f"\n📊 Подготовка датасетов HuggingFace...")
from datasets import Dataset from datasets import Dataset
def tokenize_function(examples): def tokenize_function(examples):
"""Функция токенизации для HF datasets.""" """Функция токенизации для HF datasets."""
# Используем адаптированный токенизатор # Используем адаптированный токенизатор
@@ -147,11 +160,11 @@ def main():
) )
tokenized["labels"] = tokenized["input_ids"].copy() tokenized["labels"] = tokenized["input_ids"].copy()
return tokenized return tokenized
# Создаем датасеты # Создаем датасеты
train_dataset = Dataset.from_dict({"text": train_texts}) train_dataset = Dataset.from_dict({"text": train_texts})
val_dataset = Dataset.from_dict({"text": val_texts}) val_dataset = Dataset.from_dict({"text": val_texts})
# Токенизируем # Токенизируем
train_dataset = train_dataset.map( train_dataset = train_dataset.map(
tokenize_function, tokenize_function,
@@ -163,26 +176,26 @@ def main():
batched=True, batched=True,
remove_columns=val_dataset.column_names, remove_columns=val_dataset.column_names,
) )
print(f" Train датасет: {len(train_dataset)} примеров") print(f" Train датасет: {len(train_dataset)} примеров")
print(f" Validation датасет: {len(val_dataset)} примеров") print(f" Validation датасет: {len(val_dataset)} примеров")
# === Настройка HuggingFace Trainer === # === Настройка HuggingFace Trainer ===
print(f"\n🔧 Настройка HuggingFace Trainer...") print(f"\n🔧 Настройка HuggingFace Trainer...")
from transformers import ( from transformers import (
Trainer, Trainer,
TrainingArguments, TrainingArguments,
DataCollatorForLanguageModeling DataCollatorForLanguageModeling,
) )
# Data collator для языкового моделирования # Data collator для языкового моделирования
data_collator = DataCollatorForLanguageModeling( data_collator = DataCollatorForLanguageModeling(
tokenizer=hf_tokenizer, tokenizer=hf_tokenizer,
mlm=False, mlm=False,
pad_to_multiple_of=8, pad_to_multiple_of=8,
) )
# Аргументы обучения # Аргументы обучения
training_args = TrainingArguments( training_args = TrainingArguments(
output_dir=PATHS["hf_model"], output_dir=PATHS["hf_model"],
@@ -204,7 +217,7 @@ def main():
dataloader_pin_memory=False, dataloader_pin_memory=False,
report_to=None, report_to=None,
) )
# Создаем Trainer # Создаем Trainer
trainer = Trainer( trainer = Trainer(
model=hf_model, model=hf_model,
@@ -213,84 +226,87 @@ def main():
eval_dataset=val_dataset, eval_dataset=val_dataset,
data_collator=data_collator, data_collator=data_collator,
) )
print("✅ HuggingFace Trainer настроен") print("✅ HuggingFace Trainer настроен")
# === Запуск обучения === # === Запуск обучения ===
print(f"\n🎯 Запуск обучения через HuggingFace Trainer...") print(f"\n🎯 Запуск обучения через HuggingFace Trainer...")
train_result = trainer.train() train_result = trainer.train()
# Сохраняем лучшую модель # Сохраняем лучшую модель
trainer.save_model() trainer.save_model()
hf_tokenizer.save_pretrained(PATHS["hf_model"]) hf_tokenizer.save_pretrained(PATHS["hf_model"])
print("✅ Обучение завершено успешно!") print("✅ Обучение завершено успешно!")
print(f"📊 Final train loss: {train_result.metrics['train_loss']:.4f}") print(f"📊 Final train loss: {train_result.metrics['train_loss']:.4f}")
if "eval_loss" in train_result.metrics: if "eval_loss" in train_result.metrics:
print(f"📊 Final eval loss: {train_result.metrics['eval_loss']:.4f}") print(f"📊 Final eval loss: {train_result.metrics['eval_loss']:.4f}")
# === Сохранение через hf-proxy === # === Сохранение через hf-proxy ===
print(f"\n💾 Сохранение через hf-proxy...") print(f"\n💾 Сохранение через hf-proxy...")
from hf_proxy import convert_to_hf_format from hf_proxy import convert_to_hf_format
# Сохраняем токенизатор в HF формате # Сохраняем токенизатор в HF формате
hf_tokenizer_dir = PATHS["hf_tokenizer"] hf_tokenizer_dir = PATHS["hf_tokenizer"]
hf_tokenizer.save_pretrained(hf_tokenizer_dir) hf_tokenizer.save_pretrained(hf_tokenizer_dir)
# Сохраняем модель через hf-proxy # Сохраняем модель через hf-proxy
hf_proxy_dir = PATHS["hf_proxy_model"] hf_proxy_dir = PATHS["hf_proxy_model"]
HFAdapter.save_pretrained(hf_model, hf_proxy_dir, tokenizer=hf_tokenizer) HFAdapter.save_pretrained(hf_model, hf_proxy_dir, tokenizer=hf_tokenizer)
print(f"✅ Модель сохранена в HF формате:") print(f"✅ Модель сохранена в HF формате:")
print(f" - {PATHS['hf_model']}: стандартный HF формат") print(f" - {PATHS['hf_model']}: стандартный HF формат")
print(f" - {hf_proxy_dir}: через hf-proxy") print(f" - {hf_proxy_dir}: через hf-proxy")
print(f" - {hf_tokenizer_dir}: токенизатор в HF формате") print(f" - {hf_tokenizer_dir}: токенизатор в HF формате")
# === Тестирование генерации === # === Тестирование генерации ===
print(f"\n🧪 Тестирование генерации после обучения...") print(f"\n🧪 Тестирование генерации после обучения...")
hf_model.eval() hf_model.eval()
for prompt in TEST_PROMPTS[:3]: for prompt in TEST_PROMPTS[:3]:
print(f"\n🔤 Промпт: '{prompt}'") print(f"\n🔤 Промпт: '{prompt}'")
try: try:
inputs = hf_tokenizer(prompt, return_tensors="pt") inputs = hf_tokenizer(prompt, return_tensors="pt")
with torch.no_grad(): with torch.no_grad():
generated = hf_model.generate( generated = hf_model.generate(
input_ids=inputs['input_ids'], input_ids=inputs["input_ids"],
max_new_tokens=20, max_new_tokens=20,
do_sample=True, 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}'") print(f"🎯 Результат: '{generated_text}'")
except Exception as e: except Exception as e:
print(f"❌ Ошибка генерации: {e}") print(f"❌ Ошибка генерации: {e}")
# === Сохранение результатов === # === Сохранение результатов ===
results = { results = {
"experiment": experiment_name, "experiment": experiment_name,
"model_config": model_config, "model_config": model_config,
"training_config": TRAINING_CONFIG, "training_config": TRAINING_CONFIG,
"final_loss": train_result.metrics.get('train_loss', 'N/A'), "final_loss": train_result.metrics.get("train_loss", "N/A"),
"eval_loss": train_result.metrics.get('eval_loss', 'N/A') "eval_loss": train_result.metrics.get("eval_loss", "N/A"),
} }
logger.save_logs("checkpoints/hf_integration_training_logs.json") logger.save_logs("checkpoints/hf_integration_training_logs.json")
print(f"\n🎉 Эксперимент с HF интеграцией завершен успешно!") print(f"\n🎉 Эксперимент с HF интеграцией завершен успешно!")
print(f"\n💡 Для использования обученной модели:") print(f"\n💡 Для использования обученной модели:")
print(f" uv run python experiments/hf_integration/generate_with_hf_tools.py") print(f" uv run python experiments/hf_integration/generate_with_hf_tools.py")
except Exception as e: except Exception as e:
print(f"❌ Ошибка в эксперименте: {e}") print(f"❌ Ошибка в эксперименте: {e}")
import traceback import traceback
traceback.print_exc() 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.models.gpt import GPT2
from llm.tokenizers import BPETokenizer from llm.tokenizers import BPETokenizer
from shared.configs import ( from shared.configs import BASE_GPT_CONFIG, TEST_PROMPTS, GENERATION_CONFIG, PATHS
BASE_GPT_CONFIG, TEST_PROMPTS, GENERATION_CONFIG, PATHS from shared.data import print_experiment_info, ensure_directories, ExperimentLogger
)
from shared.data import (
print_experiment_info, ensure_directories, ExperimentLogger
)
def load_model_and_tokenizer() -> tuple: def load_model_and_tokenizer() -> tuple:
""" """
Загружает обученную модель и токенизатор. Загружает обученную модель и токенизатор.
Returns: Returns:
tuple: (модель, токенизатор, конфигурация) tuple: (модель, токенизатор, конфигурация)
""" """
@@ -36,62 +32,60 @@ def load_model_and_tokenizer() -> tuple:
f"Модель не найдена: {PATHS['gpt_bpe_model']}\n" f"Модель не найдена: {PATHS['gpt_bpe_model']}\n"
f"Сначала обучите модель: uv run python experiments/llm_only/train_gpt_bpe.py" f"Сначала обучите модель: uv run python experiments/llm_only/train_gpt_bpe.py"
) )
if not os.path.exists(PATHS["bpe_tokenizer"]): if not os.path.exists(PATHS["bpe_tokenizer"]):
raise FileNotFoundError( raise FileNotFoundError(f"Токенизатор не найден: {PATHS['bpe_tokenizer']}")
f"Токенизатор не найден: {PATHS['bpe_tokenizer']}"
)
# Загружаем конфигурацию модели # Загружаем конфигурацию модели
import json 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) model_config = json.load(f)
# Загружаем токенизатор # Загружаем токенизатор
print("🔧 Загрузка BPE токенизатора...") print("🔧 Загрузка BPE токенизатора...")
tokenizer = BPETokenizer.load(PATHS["bpe_tokenizer"]) tokenizer = BPETokenizer.load(PATHS["bpe_tokenizer"])
print(f"✅ Токенизатор загружен (vocab_size={tokenizer.get_vocab_size()})") print(f"✅ Токенизатор загружен (vocab_size={tokenizer.get_vocab_size()})")
# Загружаем модель # Загружаем модель
print("🔧 Загрузка GPT2 модели...") print("🔧 Загрузка GPT2 модели...")
model = GPT2(model_config) 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() model.eval()
print("✅ Модель загружена") print("✅ Модель загружена")
return model, tokenizer, model_config return model, tokenizer, model_config
def generate_text( def generate_text(
model: GPT2, model: GPT2, tokenizer: BPETokenizer, prompt: str, config: dict
tokenizer: BPETokenizer,
prompt: str,
config: dict
) -> str: ) -> str:
""" """
Генерирует текст на основе промпта. Генерирует текст на основе промпта.
Args: Args:
model: Обученная GPT модель model: Обученная GPT модель
tokenizer: BPE токенизатор tokenizer: BPE токенизатор
prompt: Входной текст prompt: Входной текст
config: Конфигурация генерации config: Конфигурация генерации
Returns: Returns:
str: Сгенерированный текст str: Сгенерированный текст
""" """
print(f"🔤 Промпт: '{prompt}'") print(f"🔤 Промпт: '{prompt}'")
print(f"📊 Параметры: max_tokens={config['max_new_tokens']}, " print(
f"temp={config['temperature']}, sample={config['do_sample']}") 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_ids = tokenizer.encode(prompt, add_special_tokens=False)
input_tensor = torch.tensor([input_ids], dtype=torch.long) input_tensor = torch.tensor([input_ids], dtype=torch.long)
print(f"🎯 Токены промпта: {input_ids}") print(f"🎯 Токены промпта: {input_ids}")
print(f"🎯 Токены (текст): {tokenizer.tokenize(prompt)}") print(f"🎯 Токены (текст): {tokenizer.tokenize(prompt)}")
print("🔄 Генерация...") print("🔄 Генерация...")
# Генерируем текст # Генерируем текст
with torch.no_grad(): with torch.no_grad():
generated_ids = model.generate( generated_ids = model.generate(
@@ -100,19 +94,19 @@ def generate_text(
do_sample=config["do_sample"], do_sample=config["do_sample"],
temperature=config["temperature"], temperature=config["temperature"],
top_k=config["top_k"], top_k=config["top_k"],
top_p=config["top_p"] top_p=config["top_p"],
) )
# Декодируем результат # Декодируем результат
generated_text = tokenizer.decode(generated_ids[0].tolist()) generated_text = tokenizer.decode(generated_ids[0].tolist())
return generated_text return generated_text
def test_different_strategies(model: GPT2, tokenizer: BPETokenizer, prompt: str): def test_different_strategies(model: GPT2, tokenizer: BPETokenizer, prompt: str):
""" """
Тестирует разные стратегии генерации на одном промпте. Тестирует разные стратегии генерации на одном промпте.
Args: Args:
model: Обученная модель model: Обученная модель
tokenizer: BPE токенизатор tokenizer: BPE токенизатор
@@ -120,32 +114,38 @@ def test_different_strategies(model: GPT2, tokenizer: BPETokenizer, prompt: str)
""" """
print(f"\n🎭 Сравнение стратегий генерации для промпта: '{prompt}'") print(f"\n🎭 Сравнение стратегий генерации для промпта: '{prompt}'")
print("=" * 60) print("=" * 60)
strategies = [ strategies = [
{"name": "🎯 Жадный поиск", "do_sample": False, "temperature": 1.0}, {"name": "🎯 Жадный поиск", "do_sample": False, "temperature": 1.0},
{"name": "🎲 Вероятностная (temp=0.7)", "do_sample": True, "temperature": 0.7}, {"name": "🎲 Вероятностная (temp=0.7)", "do_sample": True, "temperature": 0.7},
{"name": "🔥 Случайная (temp=1.2)", "do_sample": True, "temperature": 1.2}, {"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: for strategy in strategies:
print(f"\n{strategy['name']}:") print(f"\n{strategy['name']}:")
try: try:
config = GENERATION_CONFIG.copy() config = GENERATION_CONFIG.copy()
config.update({ config.update(
"do_sample": strategy["do_sample"], {
"temperature": strategy["temperature"], "do_sample": strategy["do_sample"],
"max_new_tokens": 20 "temperature": strategy["temperature"],
}) "max_new_tokens": 20,
}
)
generated = generate_text(model, tokenizer, prompt, config) generated = generate_text(model, tokenizer, prompt, config)
# Выделяем сгенерированную часть # Выделяем сгенерированную часть
generated_part = generated[len(prompt):] generated_part = generated[len(prompt) :]
print(f" 📤 Промпт: '{prompt}'") print(f" 📤 Промпт: '{prompt}'")
print(f" 🎯 Сгенерировано: '{generated_part}'") print(f" 🎯 Сгенерировано: '{generated_part}'")
print(f" 📄 Полный текст: '{generated}'") print(f" 📄 Полный текст: '{generated}'")
except Exception as e: except Exception as e:
print(f" ❌ Ошибка: {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): def analyze_tokenization(tokenizer: BPETokenizer, texts: list):
""" """
Анализирует токенизацию различных текстов. Анализирует токенизацию различных текстов.
Args: Args:
tokenizer: BPE токенизатор tokenizer: BPE токенизатор
texts: Список текстов для анализа texts: Список текстов для анализа
""" """
print(f"\n🔍 Анализ токенизации BPE:") print(f"\n🔍 Анализ токенизации BPE:")
print("=" * 50) print("=" * 50)
for i, text in enumerate(texts): for i, text in enumerate(texts):
print(f"\nТекст {i+1}: '{text}'") print(f"\nТекст {i+1}: '{text}'")
# Токенизация # Токенизация
tokens = tokenizer.encode(text, add_special_tokens=False) tokens = tokenizer.encode(text, add_special_tokens=False)
token_strings = tokenizer.tokenize(text) token_strings = tokenizer.tokenize(text)
print(f" Токены (ID): {tokens}") print(f" Токены (ID): {tokens}")
print(f" Токены (текст): {token_strings}") print(f" Токены (текст): {token_strings}")
print(f" Количество токенов: {len(tokens)}") print(f" Количество токенов: {len(tokens)}")
print(f" Эффективность: {len(text)} символов → {len(tokens)} токенов") print(f" Эффективность: {len(text)} символов → {len(tokens)} токенов")
# Декодирование обратно # Декодирование обратно
decoded = tokenizer.decode(tokens) decoded = tokenizer.decode(tokens)
if text == decoded: if text == decoded:
@@ -184,51 +184,53 @@ def analyze_tokenization(tokenizer: BPETokenizer, texts: list):
def interactive_generation(model: GPT2, tokenizer: BPETokenizer): def interactive_generation(model: GPT2, tokenizer: BPETokenizer):
""" """
Режим интерактивной генерации. Режим интерактивной генерации.
Args: Args:
model: Обученная модель model: Обученная модель
tokenizer: BPE токенизатор tokenizer: BPE токенизатор
""" """
print(f"\n💬 Интерактивная генерация (для выхода введите 'exit')") print(f"\n💬 Интерактивная генерация (для выхода введите 'exit')")
print("-" * 50) print("-" * 50)
while True: while True:
try: try:
user_input = input("\n🔤 Введите промпт: ").strip() user_input = input("\n🔤 Введите промпт: ").strip()
if user_input.lower() in ['exit', 'quit', 'выход']: if user_input.lower() in ["exit", "quit", "выход"]:
break break
if not user_input: if not user_input:
continue continue
# Запрашиваем параметры # Запрашиваем параметры
try: try:
max_tokens = int(input("📏 Макс. токенов [50]: ") or "50") max_tokens = int(input("📏 Макс. токенов [50]: ") or "50")
temperature = float(input("🌡️ Температура [0.7]: ") or "0.7") temperature = float(input("🌡️ Температура [0.7]: ") or "0.7")
do_sample_input = input("🎲 Сэмплирование (y/n) [y]: ").lower() do_sample_input = input("🎲 Сэмплирование (y/n) [y]: ").lower()
do_sample = do_sample_input != 'n' do_sample = do_sample_input != "n"
except: except:
max_tokens = 50 max_tokens = 50
temperature = 0.7 temperature = 0.7
do_sample = True do_sample = True
print("⚠️ Использую параметры по умолчанию") print("⚠️ Использую параметры по умолчанию")
config = GENERATION_CONFIG.copy() config = GENERATION_CONFIG.copy()
config.update({ config.update(
"max_new_tokens": max_tokens, {
"temperature": temperature, "max_new_tokens": max_tokens,
"do_sample": do_sample "temperature": temperature,
}) "do_sample": do_sample,
}
)
generated = generate_text(model, tokenizer, user_input, config) 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"\n🎯 Результат:")
print(f" 📤 Промпт: '{user_input}'") print(f" 📤 Промпт: '{user_input}'")
print(f" 🎯 Сгенерировано: '{generated_part}'") print(f" 🎯 Сгенерировано: '{generated_part}'")
print(f" 📄 Полный текст: '{generated}'") print(f" 📄 Полный текст: '{generated}'")
except KeyboardInterrupt: except KeyboardInterrupt:
print("\n👋 Завершение работы...") print("\n👋 Завершение работы...")
break break
@@ -244,68 +246,69 @@ def main():
"model": "GPT2 с BPE токенизатором", "model": "GPT2 с BPE токенизатором",
"стратегия": "автономная генерация", "стратегия": "автономная генерация",
"вход": "промпты", "вход": "промпты",
"выход": "сгенерированный текст" "выход": "сгенерированный текст",
} }
print_experiment_info(experiment_name, experiment_config) print_experiment_info(experiment_name, experiment_config)
ensure_directories() ensure_directories()
logger = ExperimentLogger(experiment_name) logger = ExperimentLogger(experiment_name)
try: try:
# Загружаем модель и токенизатор # Загружаем модель и токенизатор
model, tokenizer, model_config = load_model_and_tokenizer() model, tokenizer, model_config = load_model_and_tokenizer()
# === Анализ токенизации === # === Анализ токенизации ===
analysis_texts = [ analysis_texts = [
"Искусственный интеллект", "Искусственный интеллект",
"Нейронные сети", "Нейронные сети",
"Машинное обучение", "Машинное обучение",
] ]
analyze_tokenization(tokenizer, analysis_texts) analyze_tokenization(tokenizer, analysis_texts)
# === Генерация с разными промптами === # === Генерация с разными промптами ===
print(f"\n🎯 Генерация текста с разными промптами") print(f"\n🎯 Генерация текста с разными промптами")
print("=" * 60) print("=" * 60)
for i, prompt in enumerate(TEST_PROMPTS): for i, prompt in enumerate(TEST_PROMPTS):
print(f"\n📝 Пример {i+1}/{len(TEST_PROMPTS)}") print(f"\n📝 Пример {i+1}/{len(TEST_PROMPTS)}")
print("-" * 40) print("-" * 40)
try: try:
generated = generate_text(model, tokenizer, prompt, GENERATION_CONFIG) generated = generate_text(model, tokenizer, prompt, GENERATION_CONFIG)
# Выделяем сгенерированную часть # Выделяем сгенерированную часть
generated_part = generated[len(prompt):] generated_part = generated[len(prompt) :]
print(f"📤 Промпт: '{prompt}'") print(f"📤 Промпт: '{prompt}'")
print(f"🎯 Сгенерировано: '{generated_part}'") print(f"🎯 Сгенерировано: '{generated_part}'")
print(f"📄 Полный текст: '{generated}'") print(f"📄 Полный текст: '{generated}'")
print(f"📏 Длина: {len(generated)} символов") print(f"📏 Длина: {len(generated)} символов")
# Логируем успешную генерацию # Логируем успешную генерацию
logger.log_metric(f"generation_length_{i}", len(generated)) logger.log_metric(f"generation_length_{i}", len(generated))
except Exception as e: except Exception as e:
print(f"❌ Ошибка при генерации: {e}") print(f"❌ Ошибка при генерации: {e}")
continue continue
# === Сравнение стратегий генерации === # === Сравнение стратегий генерации ===
test_prompt = "Искусственный" test_prompt = "Искусственный"
test_different_strategies(model, tokenizer, test_prompt) test_different_strategies(model, tokenizer, test_prompt)
# === Интерактивная генерация === # === Интерактивная генерация ===
interactive_generation(model, tokenizer) interactive_generation(model, tokenizer)
# === Сохранение результатов === # === Сохранение результатов ===
logger.save_logs("checkpoints/llm_only_generation_logs.json") logger.save_logs("checkpoints/llm_only_generation_logs.json")
print(f"\n🎉 Эксперимент генерации завершен успешно!") print(f"\n🎉 Эксперимент генерации завершен успешно!")
except FileNotFoundError as e: except FileNotFoundError as e:
print(f"❌ {e}") print(f"❌ {e}")
except Exception as e: except Exception as e:
print(f"❌ Ошибка в эксперименте: {e}") print(f"❌ Ошибка в эксперименте: {e}")
import traceback import traceback
traceback.print_exc() 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.models.gpt import GPT
from llm.tokenizers import BPETokenizer from llm.tokenizers import BPETokenizer
from shared.configs import ( from shared.configs import BASE_GPT_CONFIG, TEST_PROMPTS, GENERATION_CONFIG, PATHS
BASE_GPT_CONFIG, TEST_PROMPTS, GENERATION_CONFIG, PATHS from shared.data import print_experiment_info, ensure_directories, ExperimentLogger
)
from shared.data import (
print_experiment_info, ensure_directories, ExperimentLogger
)
def load_model_and_tokenizer() -> tuple: def load_model_and_tokenizer() -> tuple:
""" """
Загружает обученную модель и токенизатор. Загружает обученную модель и токенизатор.
Returns: Returns:
tuple: (модель, токенизатор, конфигурация) tuple: (модель, токенизатор, конфигурация)
""" """
@@ -36,62 +32,60 @@ def load_model_and_tokenizer() -> tuple:
f"Модель не найдена: {PATHS['gpt_bpe_model']}\n" f"Модель не найдена: {PATHS['gpt_bpe_model']}\n"
f"Сначала обучите модель: uv run python experiments/llm_only/train_gpt_bpe.py" f"Сначала обучите модель: uv run python experiments/llm_only/train_gpt_bpe.py"
) )
if not os.path.exists(PATHS["bpe_tokenizer"]): if not os.path.exists(PATHS["bpe_tokenizer"]):
raise FileNotFoundError( raise FileNotFoundError(f"Токенизатор не найден: {PATHS['bpe_tokenizer']}")
f"Токенизатор не найден: {PATHS['bpe_tokenizer']}"
)
# Загружаем конфигурацию модели # Загружаем конфигурацию модели
import json 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) model_config = json.load(f)
# Загружаем токенизатор # Загружаем токенизатор
print("🔧 Загрузка BPE токенизатора...") print("🔧 Загрузка BPE токенизатора...")
tokenizer = BPETokenizer.load(PATHS["bpe_tokenizer"]) tokenizer = BPETokenizer.load(PATHS["bpe_tokenizer"])
print(f"✅ Токенизатор загружен (vocab_size={tokenizer.get_vocab_size()})") print(f"✅ Токенизатор загружен (vocab_size={tokenizer.get_vocab_size()})")
# Загружаем модель # Загружаем модель
print("🔧 Загрузка GPT модели...") print("🔧 Загрузка GPT модели...")
model = GPT(model_config) 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() model.eval()
print("✅ Модель загружена") print("✅ Модель загружена")
return model, tokenizer, model_config return model, tokenizer, model_config
def generate_text( def generate_text(
model: GPT, model: GPT, tokenizer: BPETokenizer, prompt: str, config: dict
tokenizer: BPETokenizer,
prompt: str,
config: dict
) -> str: ) -> str:
""" """
Генерирует текст на основе промпта. Генерирует текст на основе промпта.
Args: Args:
model: Обученная GPT модель model: Обученная GPT модель
tokenizer: BPE токенизатор tokenizer: BPE токенизатор
prompt: Входной текст prompt: Входной текст
config: Конфигурация генерации config: Конфигурация генерации
Returns: Returns:
str: Сгенерированный текст str: Сгенерированный текст
""" """
print(f"🔤 Промпт: '{prompt}'") print(f"🔤 Промпт: '{prompt}'")
print(f"📊 Параметры: max_tokens={config['max_new_tokens']}, " print(
f"temp={config['temperature']}, sample={config['do_sample']}") 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_ids = tokenizer.encode(prompt, add_special_tokens=False)
input_tensor = torch.tensor([input_ids], dtype=torch.long) input_tensor = torch.tensor([input_ids], dtype=torch.long)
print(f"🎯 Токены промпта: {input_ids}") print(f"🎯 Токены промпта: {input_ids}")
print(f"🎯 Токены (текст): {tokenizer.tokenize(prompt)}") print(f"🎯 Токены (текст): {tokenizer.tokenize(prompt)}")
print("🔄 Генерация...") print("🔄 Генерация...")
# Генерируем текст # Генерируем текст
with torch.no_grad(): with torch.no_grad():
generated_ids = model.generate( generated_ids = model.generate(
@@ -100,19 +94,19 @@ def generate_text(
do_sample=config["do_sample"], do_sample=config["do_sample"],
temperature=config["temperature"], temperature=config["temperature"],
top_k=config["top_k"], top_k=config["top_k"],
top_p=config["top_p"] top_p=config["top_p"],
) )
# Декодируем результат # Декодируем результат
generated_text = tokenizer.decode(generated_ids[0].tolist()) generated_text = tokenizer.decode(generated_ids[0].tolist())
return generated_text return generated_text
def test_different_strategies(model: GPT, tokenizer: BPETokenizer, prompt: str): def test_different_strategies(model: GPT, tokenizer: BPETokenizer, prompt: str):
""" """
Тестирует разные стратегии генерации на одном промпте. Тестирует разные стратегии генерации на одном промпте.
Args: Args:
model: Обученная модель model: Обученная модель
tokenizer: BPE токенизатор tokenizer: BPE токенизатор
@@ -120,32 +114,38 @@ def test_different_strategies(model: GPT, tokenizer: BPETokenizer, prompt: str):
""" """
print(f"\n🎭 Сравнение стратегий генерации для промпта: '{prompt}'") print(f"\n🎭 Сравнение стратегий генерации для промпта: '{prompt}'")
print("=" * 60) print("=" * 60)
strategies = [ strategies = [
{"name": "🎯 Жадный поиск", "do_sample": False, "temperature": 1.0}, {"name": "🎯 Жадный поиск", "do_sample": False, "temperature": 1.0},
{"name": "🎲 Вероятностная (temp=0.7)", "do_sample": True, "temperature": 0.7}, {"name": "🎲 Вероятностная (temp=0.7)", "do_sample": True, "temperature": 0.7},
{"name": "🔥 Случайная (temp=1.2)", "do_sample": True, "temperature": 1.2}, {"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: for strategy in strategies:
print(f"\n{strategy['name']}:") print(f"\n{strategy['name']}:")
try: try:
config = GENERATION_CONFIG.copy() config = GENERATION_CONFIG.copy()
config.update({ config.update(
"do_sample": strategy["do_sample"], {
"temperature": strategy["temperature"], "do_sample": strategy["do_sample"],
"max_new_tokens": 20 "temperature": strategy["temperature"],
}) "max_new_tokens": 20,
}
)
generated = generate_text(model, tokenizer, prompt, config) generated = generate_text(model, tokenizer, prompt, config)
# Выделяем сгенерированную часть # Выделяем сгенерированную часть
generated_part = generated[len(prompt):] generated_part = generated[len(prompt) :]
print(f" 📤 Промпт: '{prompt}'") print(f" 📤 Промпт: '{prompt}'")
print(f" 🎯 Сгенерировано: '{generated_part}'") print(f" 🎯 Сгенерировано: '{generated_part}'")
print(f" 📄 Полный текст: '{generated}'") print(f" 📄 Полный текст: '{generated}'")
except Exception as e: except Exception as e:
print(f" ❌ Ошибка: {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): def analyze_tokenization(tokenizer: BPETokenizer, texts: list):
""" """
Анализирует токенизацию различных текстов. Анализирует токенизацию различных текстов.
Args: Args:
tokenizer: BPE токенизатор tokenizer: BPE токенизатор
texts: Список текстов для анализа texts: Список текстов для анализа
""" """
print(f"\n🔍 Анализ токенизации BPE:") print(f"\n🔍 Анализ токенизации BPE:")
print("=" * 50) print("=" * 50)
for i, text in enumerate(texts): for i, text in enumerate(texts):
print(f"\nТекст {i+1}: '{text}'") print(f"\nТекст {i+1}: '{text}'")
# Токенизация # Токенизация
tokens = tokenizer.encode(text, add_special_tokens=False) tokens = tokenizer.encode(text, add_special_tokens=False)
token_strings = tokenizer.tokenize(text) token_strings = tokenizer.tokenize(text)
print(f" Токены (ID): {tokens}") print(f" Токены (ID): {tokens}")
print(f" Токены (текст): {token_strings}") print(f" Токены (текст): {token_strings}")
print(f" Количество токенов: {len(tokens)}") print(f" Количество токенов: {len(tokens)}")
print(f" Эффективность: {len(text)} символов → {len(tokens)} токенов") print(f" Эффективность: {len(text)} символов → {len(tokens)} токенов")
# Декодирование обратно # Декодирование обратно
decoded = tokenizer.decode(tokens) decoded = tokenizer.decode(tokens)
if text == decoded: if text == decoded:
@@ -184,51 +184,53 @@ def analyze_tokenization(tokenizer: BPETokenizer, texts: list):
def interactive_generation(model: GPT, tokenizer: BPETokenizer): def interactive_generation(model: GPT, tokenizer: BPETokenizer):
""" """
Режим интерактивной генерации. Режим интерактивной генерации.
Args: Args:
model: Обученная модель model: Обученная модель
tokenizer: BPE токенизатор tokenizer: BPE токенизатор
""" """
print(f"\n💬 Интерактивная генерация (для выхода введите 'exit')") print(f"\n💬 Интерактивная генерация (для выхода введите 'exit')")
print("-" * 50) print("-" * 50)
while True: while True:
try: try:
user_input = input("\n🔤 Введите промпт: ").strip() user_input = input("\n🔤 Введите промпт: ").strip()
if user_input.lower() in ['exit', 'quit', 'выход']: if user_input.lower() in ["exit", "quit", "выход"]:
break break
if not user_input: if not user_input:
continue continue
# Запрашиваем параметры # Запрашиваем параметры
try: try:
max_tokens = int(input("📏 Макс. токенов [50]: ") or "50") max_tokens = int(input("📏 Макс. токенов [50]: ") or "50")
temperature = float(input("🌡️ Температура [0.7]: ") or "0.7") temperature = float(input("🌡️ Температура [0.7]: ") or "0.7")
do_sample_input = input("🎲 Сэмплирование (y/n) [y]: ").lower() do_sample_input = input("🎲 Сэмплирование (y/n) [y]: ").lower()
do_sample = do_sample_input != 'n' do_sample = do_sample_input != "n"
except: except:
max_tokens = 50 max_tokens = 50
temperature = 0.7 temperature = 0.7
do_sample = True do_sample = True
print("⚠️ Использую параметры по умолчанию") print("⚠️ Использую параметры по умолчанию")
config = GENERATION_CONFIG.copy() config = GENERATION_CONFIG.copy()
config.update({ config.update(
"max_new_tokens": max_tokens, {
"temperature": temperature, "max_new_tokens": max_tokens,
"do_sample": do_sample "temperature": temperature,
}) "do_sample": do_sample,
}
)
generated = generate_text(model, tokenizer, user_input, config) 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"\n🎯 Результат:")
print(f" 📤 Промпт: '{user_input}'") print(f" 📤 Промпт: '{user_input}'")
print(f" 🎯 Сгенерировано: '{generated_part}'") print(f" 🎯 Сгенерировано: '{generated_part}'")
print(f" 📄 Полный текст: '{generated}'") print(f" 📄 Полный текст: '{generated}'")
except KeyboardInterrupt: except KeyboardInterrupt:
print("\n👋 Завершение работы...") print("\n👋 Завершение работы...")
break break
@@ -244,68 +246,69 @@ def main():
"model": "GPT с BPE токенизатором", "model": "GPT с BPE токенизатором",
"стратегия": "автономная генерация", "стратегия": "автономная генерация",
"вход": "промпты", "вход": "промпты",
"выход": "сгенерированный текст" "выход": "сгенерированный текст",
} }
print_experiment_info(experiment_name, experiment_config) print_experiment_info(experiment_name, experiment_config)
ensure_directories() ensure_directories()
logger = ExperimentLogger(experiment_name) logger = ExperimentLogger(experiment_name)
try: try:
# Загружаем модель и токенизатор # Загружаем модель и токенизатор
model, tokenizer, model_config = load_model_and_tokenizer() model, tokenizer, model_config = load_model_and_tokenizer()
# === Анализ токенизации === # === Анализ токенизации ===
analysis_texts = [ analysis_texts = [
"Искусственный интеллект", "Искусственный интеллект",
"Нейронные сети", "Нейронные сети",
"Машинное обучение", "Машинное обучение",
] ]
analyze_tokenization(tokenizer, analysis_texts) analyze_tokenization(tokenizer, analysis_texts)
# === Генерация с разными промптами === # === Генерация с разными промптами ===
print(f"\n🎯 Генерация текста с разными промптами") print(f"\n🎯 Генерация текста с разными промптами")
print("=" * 60) print("=" * 60)
for i, prompt in enumerate(TEST_PROMPTS): for i, prompt in enumerate(TEST_PROMPTS):
print(f"\n📝 Пример {i+1}/{len(TEST_PROMPTS)}") print(f"\n📝 Пример {i+1}/{len(TEST_PROMPTS)}")
print("-" * 40) print("-" * 40)
try: try:
generated = generate_text(model, tokenizer, prompt, GENERATION_CONFIG) generated = generate_text(model, tokenizer, prompt, GENERATION_CONFIG)
# Выделяем сгенерированную часть # Выделяем сгенерированную часть
generated_part = generated[len(prompt):] generated_part = generated[len(prompt) :]
print(f"📤 Промпт: '{prompt}'") print(f"📤 Промпт: '{prompt}'")
print(f"🎯 Сгенерировано: '{generated_part}'") print(f"🎯 Сгенерировано: '{generated_part}'")
print(f"📄 Полный текст: '{generated}'") print(f"📄 Полный текст: '{generated}'")
print(f"📏 Длина: {len(generated)} символов") print(f"📏 Длина: {len(generated)} символов")
# Логируем успешную генерацию # Логируем успешную генерацию
logger.log_metric(f"generation_length_{i}", len(generated)) logger.log_metric(f"generation_length_{i}", len(generated))
except Exception as e: except Exception as e:
print(f"❌ Ошибка при генерации: {e}") print(f"❌ Ошибка при генерации: {e}")
continue continue
# === Сравнение стратегий генерации === # === Сравнение стратегий генерации ===
test_prompt = "Искусственный" test_prompt = "Искусственный"
test_different_strategies(model, tokenizer, test_prompt) test_different_strategies(model, tokenizer, test_prompt)
# === Интерактивная генерация === # === Интерактивная генерация ===
interactive_generation(model, tokenizer) interactive_generation(model, tokenizer)
# === Сохранение результатов === # === Сохранение результатов ===
logger.save_logs("checkpoints/llm_only_generation_logs.json") logger.save_logs("checkpoints/llm_only_generation_logs.json")
print(f"\n🎉 Эксперимент генерации завершен успешно!") print(f"\n🎉 Эксперимент генерации завершен успешно!")
except FileNotFoundError as e: except FileNotFoundError as e:
print(f"❌ {e}") print(f"❌ {e}")
except Exception as e: except Exception as e:
print(f"❌ Ошибка в эксперименте: {e}") print(f"❌ Ошибка в эксперименте: {e}")
import traceback import traceback
traceback.print_exc() 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 llm.training.trainer import Trainer
from shared.configs import ( from shared.configs import (
TRAIN_TEXTS, BASE_GPT_CONFIG, BPE_CONFIG, TRAIN_TEXTS,
TRAINING_CONFIG, PATHS, TEST_PROMPTS BASE_GPT_CONFIG,
BPE_CONFIG,
TRAINING_CONFIG,
PATHS,
TEST_PROMPTS,
) )
from shared.data import ( from shared.data import (
load_training_data, ensure_directories, load_training_data,
print_experiment_info, ExperimentLogger ensure_directories,
print_experiment_info,
ExperimentLogger,
) )
def train_bpe_tokenizer(texts: list, config: dict) -> BPETokenizer: def train_bpe_tokenizer(texts: list, config: dict) -> BPETokenizer:
""" """
Обучает BPE токенизатор на текстах. Обучает BPE токенизатор на текстах.
Args: Args:
texts: Список текстов для обучения texts: Список текстов для обучения
config: Конфигурация токенизатора config: Конфигурация токенизатора
Returns: Returns:
BPETokenizer: Обученный токенизатор BPETokenizer: Обученный токенизатор
""" """
print("🔧 Обучение BPE токенизатора...") print("🔧 Обучение BPE токенизатора...")
tokenizer = BPETokenizer() tokenizer = BPETokenizer()
tokenizer.train( tokenizer.train(
texts=texts, texts=texts,
vocab_size=config["vocab_size"], 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) os.makedirs(os.path.dirname(PATHS["bpe_tokenizer"]), exist_ok=True)
tokenizer.save(PATHS["bpe_tokenizer"]) tokenizer.save(PATHS["bpe_tokenizer"])
print(f"✅ BPE токенизатор обучен и сохранен: {PATHS['bpe_tokenizer']}") print(f"✅ BPE токенизатор обучен и сохранен: {PATHS['bpe_tokenizer']}")
print(f"📊 Размер словаря: {tokenizer.get_vocab_size()}") print(f"📊 Размер словаря: {tokenizer.get_vocab_size()}")
return tokenizer return tokenizer
def test_tokenizer(tokenizer: BPETokenizer, texts: list): def test_tokenizer(tokenizer: BPETokenizer, texts: list):
""" """
Тестирует токенизатор на примерах. Тестирует токенизатор на примерах.
Args: Args:
tokenizer: Обученный токенизатор tokenizer: Обученный токенизатор
texts: Список тестовых текстов texts: Список тестовых текстов
""" """
print("\n🧪 Тестирование токенизатора:") print("\n🧪 Тестирование токенизатора:")
for i, text in enumerate(texts[:3]): for i, text in enumerate(texts[:3]):
print(f"\nПример {i+1}:") print(f"\nПример {i+1}:")
print(f" Исходный текст: '{text}'") print(f" Исходный текст: '{text}'")
# Кодирование # Кодирование
tokens = tokenizer.encode(text) tokens = tokenizer.encode(text)
token_strings = tokenizer.tokenize(text) token_strings = tokenizer.tokenize(text)
print(f" Токены (ID): {tokens}") print(f" Токены (ID): {tokens}")
print(f" Токены (текст): {token_strings}") print(f" Токены (текст): {token_strings}")
print(f" Количество токенов: {len(tokens)}") print(f" Количество токенов: {len(tokens)}")
# Декодирование # Декодирование
decoded = tokenizer.decode(tokens) decoded = tokenizer.decode(tokens)
print(f" Декодированный: '{decoded}'") print(f" Декодированный: '{decoded}'")
if text == decoded: if text == decoded:
print(" ✅ Кодирование/декодирование корректно") print(" ✅ Кодирование/декодирование корректно")
else: else:
@@ -95,22 +101,22 @@ def main():
experiment_name = "Обучение GPT2 с BPE токенизатором (только llm)" experiment_name = "Обучение GPT2 с BPE токенизатором (только llm)"
experiment_config = { experiment_config = {
"model": "GPT2", "model": "GPT2",
"tokenizer": "BPE", "tokenizer": "BPE",
"vocab_size": BPE_CONFIG["vocab_size"], "vocab_size": BPE_CONFIG["vocab_size"],
"training_epochs": TRAINING_CONFIG["num_epochs"], "training_epochs": TRAINING_CONFIG["num_epochs"],
"batch_size": TRAINING_CONFIG["batch_size"], "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) print_experiment_info(experiment_name, experiment_config)
ensure_directories() ensure_directories()
logger = ExperimentLogger(experiment_name) logger = ExperimentLogger(experiment_name)
try: try:
# === Подготовка данных === # === Подготовка данных ===
train_texts, val_texts = load_training_data() train_texts, val_texts = load_training_data()
print(f"📊 Данные: {len(train_texts)} train, {len(val_texts)} validation") print(f"📊 Данные: {len(train_texts)} train, {len(val_texts)} validation")
# === Обучение токенизатора === # === Обучение токенизатора ===
if os.path.exists(PATHS["bpe_tokenizer"]): if os.path.exists(PATHS["bpe_tokenizer"]):
print("📝 Загрузка предварительно обученного токенизатора...") print("📝 Загрузка предварительно обученного токенизатора...")
@@ -118,112 +124,112 @@ def main():
print(f"✅ Токенизатор загружен (vocab_size={tokenizer.get_vocab_size()})") print(f"✅ Токенизатор загружен (vocab_size={tokenizer.get_vocab_size()})")
else: else:
tokenizer = train_bpe_tokenizer(TRAIN_TEXTS, BPE_CONFIG) tokenizer = train_bpe_tokenizer(TRAIN_TEXTS, BPE_CONFIG)
# Тестируем токенизатор # Тестируем токенизатор
test_tokenizer(tokenizer, TEST_PROMPTS[:3]) test_tokenizer(tokenizer, TEST_PROMPTS[:3])
# === Инициализация модели === # === Инициализация модели ===
model_config = BASE_GPT_CONFIG.copy() model_config = BASE_GPT_CONFIG.copy()
model_config["vocab_size"] = tokenizer.get_vocab_size() model_config["vocab_size"] = tokenizer.get_vocab_size()
print(f"\n🔧 Инициализация GPT2 модели...") print(f"\n🔧 Инициализация GPT2 модели...")
print(f" Размер словаря: {model_config['vocab_size']}") print(f" Размер словаря: {model_config['vocab_size']}")
print(f" Размер эмбеддингов: {model_config['embed_dim']}") print(f" Размер эмбеддингов: {model_config['embed_dim']}")
print(f" Количество слоев: {model_config['num_layers']}") print(f" Количество слоев: {model_config['num_layers']}")
print(f" Количество голов внимания: {model_config['num_heads']}") print(f" Количество голов внимания: {model_config['num_heads']}")
model = GPT2(model_config) model = GPT2(model_config)
# === Подготовка датасета === # === Подготовка датасета ===
print(f"\n📊 Подготовка датасета...") print(f"\n📊 Подготовка датасета...")
train_dataset = TextDataset( train_dataset = TextDataset(
train_texts, train_texts, tokenizer, block_size=model_config["max_position_embeddings"]
tokenizer,
block_size=model_config["max_position_embeddings"]
) )
print(f" Размер train датасета: {len(train_dataset)} примеров") print(f" Размер train датасета: {len(train_dataset)} примеров")
# === Обучение модели === # === Обучение модели ===
print(f"\n🎯 Начало обучения GPT2 модели...") print(f"\n🎯 Начало обучения GPT2 модели...")
trainer = Trainer( trainer = Trainer(
model=model, model=model,
train_dataset=train_dataset, train_dataset=train_dataset,
lr=TRAINING_CONFIG["learning_rate"], lr=TRAINING_CONFIG["learning_rate"],
batch_size=TRAINING_CONFIG["batch_size"], batch_size=TRAINING_CONFIG["batch_size"],
num_epochs=TRAINING_CONFIG["num_epochs"], num_epochs=TRAINING_CONFIG["num_epochs"],
warmup_steps=TRAINING_CONFIG["warmup_steps"] warmup_steps=TRAINING_CONFIG["warmup_steps"],
) )
# Запускаем обучение # Запускаем обучение
trainer.train() trainer.train()
# === Сохранение модели === # === Сохранение модели ===
print(f"\n💾 Сохранение модели...") print(f"\n💾 Сохранение модели...")
os.makedirs(os.path.dirname(PATHS["gpt_bpe_model"]), exist_ok=True) os.makedirs(os.path.dirname(PATHS["gpt_bpe_model"]), exist_ok=True)
# Сохраняем модель # Сохраняем модель
torch.save(model.state_dict(), PATHS["gpt_bpe_model"]) torch.save(model.state_dict(), PATHS["gpt_bpe_model"])
# Сохраняем конфигурацию # Сохраняем конфигурацию
import json 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) json.dump(model_config, f, indent=2, ensure_ascii=False)
print(f"✅ Модель сохранена:") print(f"✅ Модель сохранена:")
print(f" - {PATHS['gpt_bpe_model']}: веса модели") print(f" - {PATHS['gpt_bpe_model']}: веса модели")
print(f" - {PATHS['gpt_bpe_config']}: конфигурация модели") print(f" - {PATHS['gpt_bpe_config']}: конфигурация модели")
print(f" - {PATHS['bpe_tokenizer']}: токенизатор") print(f" - {PATHS['bpe_tokenizer']}: токенизатор")
# === Тестирование генерации === # === Тестирование генерации ===
print(f"\n🧪 Тестирование генерации текста...") print(f"\n🧪 Тестирование генерации текста...")
model.eval() model.eval()
for prompt in TEST_PROMPTS[:3]: for prompt in TEST_PROMPTS[:3]:
print(f"\n🔤 Промпт: '{prompt}'") print(f"\n🔤 Промпт: '{prompt}'")
try: try:
# Кодируем промпт # Кодируем промпт
input_ids = tokenizer.encode(prompt, add_special_tokens=False) input_ids = tokenizer.encode(prompt, add_special_tokens=False)
input_tensor = torch.tensor([input_ids], dtype=torch.long) input_tensor = torch.tensor([input_ids], dtype=torch.long)
# Генерируем текст # Генерируем текст
with torch.no_grad(): with torch.no_grad():
generated_ids = model.generate( generated_ids = model.generate(
x=input_tensor, x=input_tensor,
max_new_tokens=20, max_new_tokens=20,
do_sample=True, do_sample=True,
temperature=0.8 temperature=0.8,
) )
# Декодируем результат # Декодируем результат
generated_text = tokenizer.decode(generated_ids[0].tolist()) 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_part}'")
print(f"📄 Полный текст: '{generated_text}'") print(f"📄 Полный текст: '{generated_text}'")
except Exception as e: except Exception as e:
print(f"❌ Ошибка генерации: {e}") print(f"❌ Ошибка генерации: {e}")
# === Сохранение результатов === # === Сохранение результатов ===
results = { results = {
"experiment": experiment_name, "experiment": experiment_name,
"model_config": model_config, "model_config": model_config,
"training_config": TRAINING_CONFIG, "training_config": TRAINING_CONFIG,
"tokenizer_vocab_size": tokenizer.get_vocab_size(), "tokenizer_vocab_size": tokenizer.get_vocab_size(),
"final_loss": "см. логи обучения" # В реальном эксперименте можно сохранить final loss "final_loss": "см. логи обучения", # В реальном эксперименте можно сохранить final loss
} }
logger.save_logs("checkpoints/llm_only_training_logs.json") logger.save_logs("checkpoints/llm_only_training_logs.json")
print(f"\n🎉 Эксперимент завершен успешно!") print(f"\n🎉 Эксперимент завершен успешно!")
print(f"\n💡 Для использования обученной модели:") print(f"\n💡 Для использования обученной модели:")
print(f" uv run python experiments/llm_only/generate_gpt_bpe.py") print(f" uv run python experiments/llm_only/generate_gpt_bpe.py")
except Exception as e: except Exception as e:
print(f"❌ Ошибка в эксперименте: {e}") print(f"❌ Ошибка в эксперименте: {e}")
import traceback import traceback
traceback.print_exc() 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 llm.training.trainer import Trainer
from shared.configs import ( from shared.configs import (
TRAIN_TEXTS, BASE_GPT_CONFIG, BPE_CONFIG, TRAIN_TEXTS,
TRAINING_CONFIG, PATHS, TEST_PROMPTS BASE_GPT_CONFIG,
BPE_CONFIG,
TRAINING_CONFIG,
PATHS,
TEST_PROMPTS,
) )
from shared.data import ( from shared.data import (
load_training_data, ensure_directories, load_training_data,
print_experiment_info, ExperimentLogger ensure_directories,
print_experiment_info,
ExperimentLogger,
) )
def train_bpe_tokenizer(texts: list, config: dict) -> BPETokenizer: def train_bpe_tokenizer(texts: list, config: dict) -> BPETokenizer:
""" """
Обучает BPE токенизатор на текстах. Обучает BPE токенизатор на текстах.
Args: Args:
texts: Список текстов для обучения texts: Список текстов для обучения
config: Конфигурация токенизатора config: Конфигурация токенизатора
Returns: Returns:
BPETokenizer: Обученный токенизатор BPETokenizer: Обученный токенизатор
""" """
print("🔧 Обучение BPE токенизатора...") print("🔧 Обучение BPE токенизатора...")
tokenizer = BPETokenizer() tokenizer = BPETokenizer()
tokenizer.train( tokenizer.train(
texts=texts, texts=texts,
vocab_size=config["vocab_size"], 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) os.makedirs(os.path.dirname(PATHS["bpe_tokenizer"]), exist_ok=True)
tokenizer.save(PATHS["bpe_tokenizer"]) tokenizer.save(PATHS["bpe_tokenizer"])
print(f"✅ BPE токенизатор обучен и сохранен: {PATHS['bpe_tokenizer']}") print(f"✅ BPE токенизатор обучен и сохранен: {PATHS['bpe_tokenizer']}")
print(f"📊 Размер словаря: {tokenizer.get_vocab_size()}") print(f"📊 Размер словаря: {tokenizer.get_vocab_size()}")
return tokenizer return tokenizer
def test_tokenizer(tokenizer: BPETokenizer, texts: list): def test_tokenizer(tokenizer: BPETokenizer, texts: list):
""" """
Тестирует токенизатор на примерах. Тестирует токенизатор на примерах.
Args: Args:
tokenizer: Обученный токенизатор tokenizer: Обученный токенизатор
texts: Список тестовых текстов texts: Список тестовых текстов
""" """
print("\n🧪 Тестирование токенизатора:") print("\n🧪 Тестирование токенизатора:")
for i, text in enumerate(texts[:3]): for i, text in enumerate(texts[:3]):
print(f"\nПример {i+1}:") print(f"\nПример {i+1}:")
print(f" Исходный текст: '{text}'") print(f" Исходный текст: '{text}'")
# Кодирование # Кодирование
tokens = tokenizer.encode(text) tokens = tokenizer.encode(text)
token_strings = tokenizer.tokenize(text) token_strings = tokenizer.tokenize(text)
print(f" Токены (ID): {tokens}") print(f" Токены (ID): {tokens}")
print(f" Токены (текст): {token_strings}") print(f" Токены (текст): {token_strings}")
print(f" Количество токенов: {len(tokens)}") print(f" Количество токенов: {len(tokens)}")
# Декодирование # Декодирование
decoded = tokenizer.decode(tokens) decoded = tokenizer.decode(tokens)
print(f" Декодированный: '{decoded}'") print(f" Декодированный: '{decoded}'")
if text == decoded: if text == decoded:
print(" ✅ Кодирование/декодирование корректно") print(" ✅ Кодирование/декодирование корректно")
else: else:
@@ -95,22 +101,22 @@ def main():
experiment_name = "Обучение GPT с BPE токенизатором (только llm)" experiment_name = "Обучение GPT с BPE токенизатором (только llm)"
experiment_config = { experiment_config = {
"model": "GPT", "model": "GPT",
"tokenizer": "BPE", "tokenizer": "BPE",
"vocab_size": BPE_CONFIG["vocab_size"], "vocab_size": BPE_CONFIG["vocab_size"],
"training_epochs": TRAINING_CONFIG["num_epochs"], "training_epochs": TRAINING_CONFIG["num_epochs"],
"batch_size": TRAINING_CONFIG["batch_size"], "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) print_experiment_info(experiment_name, experiment_config)
ensure_directories() ensure_directories()
logger = ExperimentLogger(experiment_name) logger = ExperimentLogger(experiment_name)
try: try:
# === Подготовка данных === # === Подготовка данных ===
train_texts, val_texts = load_training_data() train_texts, val_texts = load_training_data()
print(f"📊 Данные: {len(train_texts)} train, {len(val_texts)} validation") print(f"📊 Данные: {len(train_texts)} train, {len(val_texts)} validation")
# === Обучение токенизатора === # === Обучение токенизатора ===
if os.path.exists(PATHS["bpe_tokenizer"]): if os.path.exists(PATHS["bpe_tokenizer"]):
print("📝 Загрузка предварительно обученного токенизатора...") print("📝 Загрузка предварительно обученного токенизатора...")
@@ -118,112 +124,112 @@ def main():
print(f"✅ Токенизатор загружен (vocab_size={tokenizer.get_vocab_size()})") print(f"✅ Токенизатор загружен (vocab_size={tokenizer.get_vocab_size()})")
else: else:
tokenizer = train_bpe_tokenizer(TRAIN_TEXTS, BPE_CONFIG) tokenizer = train_bpe_tokenizer(TRAIN_TEXTS, BPE_CONFIG)
# Тестируем токенизатор # Тестируем токенизатор
test_tokenizer(tokenizer, TEST_PROMPTS[:3]) test_tokenizer(tokenizer, TEST_PROMPTS[:3])
# === Инициализация модели === # === Инициализация модели ===
model_config = BASE_GPT_CONFIG.copy() model_config = BASE_GPT_CONFIG.copy()
model_config["vocab_size"] = tokenizer.get_vocab_size() model_config["vocab_size"] = tokenizer.get_vocab_size()
print(f"\n🔧 Инициализация GPT модели...") print(f"\n🔧 Инициализация GPT модели...")
print(f" Размер словаря: {model_config['vocab_size']}") print(f" Размер словаря: {model_config['vocab_size']}")
print(f" Размер эмбеддингов: {model_config['embed_dim']}") print(f" Размер эмбеддингов: {model_config['embed_dim']}")
print(f" Количество слоев: {model_config['num_layers']}") print(f" Количество слоев: {model_config['num_layers']}")
print(f" Количество голов внимания: {model_config['num_heads']}") print(f" Количество голов внимания: {model_config['num_heads']}")
model = GPT(model_config) model = GPT(model_config)
# === Подготовка датасета === # === Подготовка датасета ===
print(f"\n📊 Подготовка датасета...") print(f"\n📊 Подготовка датасета...")
train_dataset = TextDataset( train_dataset = TextDataset(
train_texts, train_texts, tokenizer, block_size=model_config["max_position_embeddings"]
tokenizer,
block_size=model_config["max_position_embeddings"]
) )
print(f" Размер train датасета: {len(train_dataset)} примеров") print(f" Размер train датасета: {len(train_dataset)} примеров")
# === Обучение модели === # === Обучение модели ===
print(f"\n🎯 Начало обучения GPT модели...") print(f"\n🎯 Начало обучения GPT модели...")
trainer = Trainer( trainer = Trainer(
model=model, model=model,
train_dataset=train_dataset, train_dataset=train_dataset,
lr=TRAINING_CONFIG["learning_rate"], lr=TRAINING_CONFIG["learning_rate"],
batch_size=TRAINING_CONFIG["batch_size"], batch_size=TRAINING_CONFIG["batch_size"],
num_epochs=TRAINING_CONFIG["num_epochs"], num_epochs=TRAINING_CONFIG["num_epochs"],
warmup_steps=TRAINING_CONFIG["warmup_steps"] warmup_steps=TRAINING_CONFIG["warmup_steps"],
) )
# Запускаем обучение # Запускаем обучение
trainer.train() trainer.train()
# === Сохранение модели === # === Сохранение модели ===
print(f"\n💾 Сохранение модели...") print(f"\n💾 Сохранение модели...")
os.makedirs(os.path.dirname(PATHS["gpt_bpe_model"]), exist_ok=True) os.makedirs(os.path.dirname(PATHS["gpt_bpe_model"]), exist_ok=True)
# Сохраняем модель # Сохраняем модель
torch.save(model.state_dict(), PATHS["gpt_bpe_model"]) torch.save(model.state_dict(), PATHS["gpt_bpe_model"])
# Сохраняем конфигурацию # Сохраняем конфигурацию
import json 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) json.dump(model_config, f, indent=2, ensure_ascii=False)
print(f"✅ Модель сохранена:") print(f"✅ Модель сохранена:")
print(f" - {PATHS['gpt_bpe_model']}: веса модели") print(f" - {PATHS['gpt_bpe_model']}: веса модели")
print(f" - {PATHS['gpt_bpe_config']}: конфигурация модели") print(f" - {PATHS['gpt_bpe_config']}: конфигурация модели")
print(f" - {PATHS['bpe_tokenizer']}: токенизатор") print(f" - {PATHS['bpe_tokenizer']}: токенизатор")
# === Тестирование генерации === # === Тестирование генерации ===
print(f"\n🧪 Тестирование генерации текста...") print(f"\n🧪 Тестирование генерации текста...")
model.eval() model.eval()
for prompt in TEST_PROMPTS[:3]: for prompt in TEST_PROMPTS[:3]:
print(f"\n🔤 Промпт: '{prompt}'") print(f"\n🔤 Промпт: '{prompt}'")
try: try:
# Кодируем промпт # Кодируем промпт
input_ids = tokenizer.encode(prompt, add_special_tokens=False) input_ids = tokenizer.encode(prompt, add_special_tokens=False)
input_tensor = torch.tensor([input_ids], dtype=torch.long) input_tensor = torch.tensor([input_ids], dtype=torch.long)
# Генерируем текст # Генерируем текст
with torch.no_grad(): with torch.no_grad():
generated_ids = model.generate( generated_ids = model.generate(
x=input_tensor, x=input_tensor,
max_new_tokens=20, max_new_tokens=20,
do_sample=True, do_sample=True,
temperature=0.8 temperature=0.8,
) )
# Декодируем результат # Декодируем результат
generated_text = tokenizer.decode(generated_ids[0].tolist()) 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_part}'")
print(f"📄 Полный текст: '{generated_text}'") print(f"📄 Полный текст: '{generated_text}'")
except Exception as e: except Exception as e:
print(f"❌ Ошибка генерации: {e}") print(f"❌ Ошибка генерации: {e}")
# === Сохранение результатов === # === Сохранение результатов ===
results = { results = {
"experiment": experiment_name, "experiment": experiment_name,
"model_config": model_config, "model_config": model_config,
"training_config": TRAINING_CONFIG, "training_config": TRAINING_CONFIG,
"tokenizer_vocab_size": tokenizer.get_vocab_size(), "tokenizer_vocab_size": tokenizer.get_vocab_size(),
"final_loss": "см. логи обучения" # В реальном эксперименте можно сохранить final loss "final_loss": "см. логи обучения", # В реальном эксперименте можно сохранить final loss
} }
logger.save_logs("checkpoints/llm_only_training_logs.json") logger.save_logs("checkpoints/llm_only_training_logs.json")
print(f"\n🎉 Эксперимент завершен успешно!") print(f"\n🎉 Эксперимент завершен успешно!")
print(f"\n💡 Для использования обученной модели:") print(f"\n💡 Для использования обученной модели:")
print(f" uv run python experiments/llm_only/generate_gpt_bpe.py") print(f" uv run python experiments/llm_only/generate_gpt_bpe.py")
except Exception as e: except Exception as e:
print(f"❌ Ошибка в эксперименте: {e}") print(f"❌ Ошибка в эксперименте: {e}")
import traceback import traceback
traceback.print_exc() 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 llm.training.trainer import Trainer
from shared.configs import ( from shared.configs import (
TRAIN_TEXTS, BASE_GPT_CONFIG, BPE_CONFIG, TRAIN_TEXTS,
TRAINING_CONFIG, PATHS, TEST_PROMPTS BASE_GPT_CONFIG,
BPE_CONFIG,
TRAINING_CONFIG,
PATHS,
TEST_PROMPTS,
) )
from shared.data import ( from shared.data import (
load_training_data, ensure_directories, load_training_data,
print_experiment_info, ExperimentLogger ensure_directories,
print_experiment_info,
ExperimentLogger,
) )
def train_bpe_tokenizer(texts: list, config: dict) -> BPETokenizer: def train_bpe_tokenizer(texts: list, config: dict) -> BPETokenizer:
""" """
Обучает BPE токенизатор на текстах. Обучает BPE токенизатор на текстах.
Args: Args:
texts: Список текстов для обучения texts: Список текстов для обучения
config: Конфигурация токенизатора config: Конфигурация токенизатора
Returns: Returns:
BPETokenizer: Обученный токенизатор BPETokenizer: Обученный токенизатор
""" """
print("🔧 Обучение BPE токенизатора...") print("🔧 Обучение BPE токенизатора...")
tokenizer = BPETokenizer() tokenizer = BPETokenizer()
tokenizer.train( tokenizer.train(
texts=texts, texts=texts,
vocab_size=config["vocab_size"], 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) os.makedirs(os.path.dirname(PATHS["bpe_tokenizer"]), exist_ok=True)
tokenizer.save(PATHS["bpe_tokenizer"]) tokenizer.save(PATHS["bpe_tokenizer"])
print(f"✅ BPE токенизатор обучен и сохранен: {PATHS['bpe_tokenizer']}") print(f"✅ BPE токенизатор обучен и сохранен: {PATHS['bpe_tokenizer']}")
print(f"📊 Размер словаря: {tokenizer.get_vocab_size()}") print(f"📊 Размер словаря: {tokenizer.get_vocab_size()}")
return tokenizer return tokenizer
def test_tokenizer(tokenizer: BPETokenizer, texts: list): def test_tokenizer(tokenizer: BPETokenizer, texts: list):
""" """
Тестирует токенизатор на примерах. Тестирует токенизатор на примерах.
Args: Args:
tokenizer: Обученный токенизатор tokenizer: Обученный токенизатор
texts: Список тестовых текстов texts: Список тестовых текстов
""" """
print("\n🧪 Тестирование токенизатора:") print("\n🧪 Тестирование токенизатора:")
for i, text in enumerate(texts[:3]): for i, text in enumerate(texts[:3]):
print(f"\nПример {i+1}:") print(f"\nПример {i+1}:")
print(f" Исходный текст: '{text}'") print(f" Исходный текст: '{text}'")
# Кодирование # Кодирование
tokens = tokenizer.encode(text) tokens = tokenizer.encode(text)
token_strings = tokenizer.tokenize(text) token_strings = tokenizer.tokenize(text)
print(f" Токены (ID): {tokens}") print(f" Токены (ID): {tokens}")
print(f" Токены (текст): {token_strings}") print(f" Токены (текст): {token_strings}")
print(f" Количество токенов: {len(tokens)}") print(f" Количество токенов: {len(tokens)}")
# Декодирование # Декодирование
decoded = tokenizer.decode(tokens) decoded = tokenizer.decode(tokens)
print(f" Декодированный: '{decoded}'") print(f" Декодированный: '{decoded}'")
if text == decoded: if text == decoded:
print(" ✅ Кодирование/декодирование корректно") print(" ✅ Кодирование/декодирование корректно")
else: else:
@@ -95,22 +101,22 @@ def main():
experiment_name = "Обучение Llama с BPE токенизатором (только llm)" experiment_name = "Обучение Llama с BPE токенизатором (только llm)"
experiment_config = { experiment_config = {
"model": "Llama", "model": "Llama",
"tokenizer": "BPE", "tokenizer": "BPE",
"vocab_size": BPE_CONFIG["vocab_size"], "vocab_size": BPE_CONFIG["vocab_size"],
"training_epochs": TRAINING_CONFIG["num_epochs"], "training_epochs": TRAINING_CONFIG["num_epochs"],
"batch_size": TRAINING_CONFIG["batch_size"], "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) print_experiment_info(experiment_name, experiment_config)
ensure_directories() ensure_directories()
logger = ExperimentLogger(experiment_name) logger = ExperimentLogger(experiment_name)
try: try:
# === Подготовка данных === # === Подготовка данных ===
train_texts, val_texts = load_training_data() train_texts, val_texts = load_training_data()
print(f"📊 Данные: {len(train_texts)} train, {len(val_texts)} validation") print(f"📊 Данные: {len(train_texts)} train, {len(val_texts)} validation")
# === Обучение токенизатора === # === Обучение токенизатора ===
if os.path.exists(PATHS["bpe_tokenizer"]): if os.path.exists(PATHS["bpe_tokenizer"]):
print("📝 Загрузка предварительно обученного токенизатора...") print("📝 Загрузка предварительно обученного токенизатора...")
@@ -118,112 +124,112 @@ def main():
print(f"✅ Токенизатор загружен (vocab_size={tokenizer.get_vocab_size()})") print(f"✅ Токенизатор загружен (vocab_size={tokenizer.get_vocab_size()})")
else: else:
tokenizer = train_bpe_tokenizer(TRAIN_TEXTS, BPE_CONFIG) tokenizer = train_bpe_tokenizer(TRAIN_TEXTS, BPE_CONFIG)
# Тестируем токенизатор # Тестируем токенизатор
test_tokenizer(tokenizer, TEST_PROMPTS[:3]) test_tokenizer(tokenizer, TEST_PROMPTS[:3])
# === Инициализация модели === # === Инициализация модели ===
model_config = BASE_GPT_CONFIG.copy() model_config = BASE_GPT_CONFIG.copy()
model_config["vocab_size"] = tokenizer.get_vocab_size() model_config["vocab_size"] = tokenizer.get_vocab_size()
print(f"\n🔧 Инициализация Llama модели...") print(f"\n🔧 Инициализация Llama модели...")
print(f" Размер словаря: {model_config['vocab_size']}") print(f" Размер словаря: {model_config['vocab_size']}")
print(f" Размер эмбеддингов: {model_config['embed_dim']}") print(f" Размер эмбеддингов: {model_config['embed_dim']}")
print(f" Количество слоев: {model_config['num_layers']}") print(f" Количество слоев: {model_config['num_layers']}")
print(f" Количество голов внимания: {model_config['num_heads']}") print(f" Количество голов внимания: {model_config['num_heads']}")
model = Llama(model_config) model = Llama(model_config)
# === Подготовка датасета === # === Подготовка датасета ===
print(f"\n📊 Подготовка датасета...") print(f"\n📊 Подготовка датасета...")
train_dataset = TextDataset( train_dataset = TextDataset(
train_texts, train_texts, tokenizer, block_size=model_config["max_position_embeddings"]
tokenizer,
block_size=model_config["max_position_embeddings"]
) )
print(f" Размер train датасета: {len(train_dataset)} примеров") print(f" Размер train датасета: {len(train_dataset)} примеров")
# === Обучение модели === # === Обучение модели ===
print(f"\n🎯 Начало обучения Llama модели...") print(f"\n🎯 Начало обучения Llama модели...")
trainer = Trainer( trainer = Trainer(
model=model, model=model,
train_dataset=train_dataset, train_dataset=train_dataset,
lr=TRAINING_CONFIG["learning_rate"], lr=TRAINING_CONFIG["learning_rate"],
batch_size=TRAINING_CONFIG["batch_size"], batch_size=TRAINING_CONFIG["batch_size"],
num_epochs=TRAINING_CONFIG["num_epochs"], num_epochs=TRAINING_CONFIG["num_epochs"],
warmup_steps=TRAINING_CONFIG["warmup_steps"] warmup_steps=TRAINING_CONFIG["warmup_steps"],
) )
# Запускаем обучение # Запускаем обучение
trainer.train() trainer.train()
# === Сохранение модели === # === Сохранение модели ===
print(f"\n💾 Сохранение модели...") print(f"\n💾 Сохранение модели...")
os.makedirs(os.path.dirname(PATHS["gpt_bpe_model"]), exist_ok=True) os.makedirs(os.path.dirname(PATHS["gpt_bpe_model"]), exist_ok=True)
# Сохраняем модель # Сохраняем модель
torch.save(model.state_dict(), PATHS["gpt_bpe_model"]) torch.save(model.state_dict(), PATHS["gpt_bpe_model"])
# Сохраняем конфигурацию # Сохраняем конфигурацию
import json 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) json.dump(model_config, f, indent=2, ensure_ascii=False)
print(f"✅ Модель сохранена:") print(f"✅ Модель сохранена:")
print(f" - {PATHS['gpt_bpe_model']}: веса модели") print(f" - {PATHS['gpt_bpe_model']}: веса модели")
print(f" - {PATHS['gpt_bpe_config']}: конфигурация модели") print(f" - {PATHS['gpt_bpe_config']}: конфигурация модели")
print(f" - {PATHS['bpe_tokenizer']}: токенизатор") print(f" - {PATHS['bpe_tokenizer']}: токенизатор")
# === Тестирование генерации === # === Тестирование генерации ===
print(f"\n🧪 Тестирование генерации текста...") print(f"\n🧪 Тестирование генерации текста...")
model.eval() model.eval()
for prompt in TEST_PROMPTS[:3]: for prompt in TEST_PROMPTS[:3]:
print(f"\n🔤 Промпт: '{prompt}'") print(f"\n🔤 Промпт: '{prompt}'")
try: try:
# Кодируем промпт # Кодируем промпт
input_ids = tokenizer.encode(prompt, add_special_tokens=False) input_ids = tokenizer.encode(prompt, add_special_tokens=False)
input_tensor = torch.tensor([input_ids], dtype=torch.long) input_tensor = torch.tensor([input_ids], dtype=torch.long)
# Генерируем текст # Генерируем текст
with torch.no_grad(): with torch.no_grad():
generated_ids = model.generate( generated_ids = model.generate(
x=input_tensor, x=input_tensor,
max_new_tokens=20, max_new_tokens=20,
do_sample=True, do_sample=True,
temperature=0.8 temperature=0.8,
) )
# Декодируем результат # Декодируем результат
generated_text = tokenizer.decode(generated_ids[0].tolist()) 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_part}'")
print(f"📄 Полный текст: '{generated_text}'") print(f"📄 Полный текст: '{generated_text}'")
except Exception as e: except Exception as e:
print(f"❌ Ошибка генерации: {e}") print(f"❌ Ошибка генерации: {e}")
# === Сохранение результатов === # === Сохранение результатов ===
results = { results = {
"experiment": experiment_name, "experiment": experiment_name,
"model_config": model_config, "model_config": model_config,
"training_config": TRAINING_CONFIG, "training_config": TRAINING_CONFIG,
"tokenizer_vocab_size": tokenizer.get_vocab_size(), "tokenizer_vocab_size": tokenizer.get_vocab_size(),
"final_loss": "см. логи обучения" # В реальном эксперименте можно сохранить final loss "final_loss": "см. логи обучения", # В реальном эксперименте можно сохранить final loss
} }
logger.save_logs("checkpoints/llm_only_training_logs.json") logger.save_logs("checkpoints/llm_only_training_logs.json")
print(f"\n🎉 Эксперимент завершен успешно!") print(f"\n🎉 Эксперимент завершен успешно!")
print(f"\n💡 Для использования обученной модели:") print(f"\n💡 Для использования обученной модели:")
print(f" uv run python experiments/llm_only/generate_gpt_bpe.py") print(f" uv run python experiments/llm_only/generate_gpt_bpe.py")
except Exception as e: except Exception as e:
print(f"❌ Ошибка в эксперименте: {e}") print(f"❌ Ошибка в эксперименте: {e}")
import traceback import traceback
traceback.print_exc() traceback.print_exc()

14
experiments/shared/configs.py
View File

@@ -30,7 +30,7 @@ BASE_GPT_CONFIG = {
"num_heads": 4, "num_heads": 4,
"num_layers": 4, "num_layers": 4,
"max_position_embeddings": 128, "max_position_embeddings": 128,
"dropout": 0.1 "dropout": 0.1,
} }
# Конфигурация для маленькой модели (быстрое тестирование) # Конфигурация для маленькой модели (быстрое тестирование)
@@ -40,7 +40,7 @@ SMALL_GPT_CONFIG = {
"num_heads": 2, "num_heads": 2,
"num_layers": 2, "num_layers": 2,
"max_position_embeddings": 64, "max_position_embeddings": 64,
"dropout": 0.1 "dropout": 0.1,
} }
# Конфигурация для большой модели (качественное обучение) # Конфигурация для большой модели (качественное обучение)
@@ -50,13 +50,13 @@ LARGE_GPT_CONFIG = {
"num_heads": 8, "num_heads": 8,
"num_layers": 6, "num_layers": 6,
"max_position_embeddings": 256, "max_position_embeddings": 256,
"dropout": 0.1 "dropout": 0.1,
} }
# === Конфигурации токенизатора === # === Конфигурации токенизатора ===
BPE_CONFIG = { BPE_CONFIG = {
"vocab_size": 1000, "vocab_size": 1000,
"special_tokens": ["<pad>", "<unk>", "<bos>", "<eos>"] "special_tokens": ["<pad>", "<unk>", "<bos>", "<eos>"],
} }
# === Конфигурации обучения === # === Конфигурации обучения ===
@@ -65,7 +65,7 @@ TRAINING_CONFIG = {
"batch_size": 2, "batch_size": 2,
"num_epochs": 3, "num_epochs": 3,
"warmup_steps": 50, "warmup_steps": 50,
"gradient_clip": 1.0 "gradient_clip": 1.0,
} }
# === Конфигурации генерации === # === Конфигурации генерации ===
@@ -74,7 +74,7 @@ GENERATION_CONFIG = {
"temperature": 0.7, "temperature": 0.7,
"do_sample": True, "do_sample": True,
"top_k": None, "top_k": None,
"top_p": None "top_p": None,
} }
# === Пути для сохранения === # === Пути для сохранения ===
@@ -84,7 +84,7 @@ PATHS = {
"gpt_bpe_config": "checkpoints/gpt-bpe/config.json", "gpt_bpe_config": "checkpoints/gpt-bpe/config.json",
"hf_tokenizer": "checkpoints/hf-bpe-tokenizer", "hf_tokenizer": "checkpoints/hf-bpe-tokenizer",
"hf_model": "checkpoints/hf-trained", "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]]: def load_training_data(split_ratio: float = 0.8) -> Tuple[List[str], List[str]]:
""" """
Загружает данные для обучения и разделяет на train/validation. Загружает данные для обучения и разделяет на train/validation.
Args: Args:
split_ratio: Доля данных для обучения split_ratio: Доля данных для обучения
Returns: Returns:
Tuple: (train_texts, val_texts) Tuple: (train_texts, val_texts)
""" """
train_size = int(len(TRAIN_TEXTS) * split_ratio) train_size = int(len(TRAIN_TEXTS) * split_ratio)
train_data = TRAIN_TEXTS[:train_size] train_data = TRAIN_TEXTS[:train_size]
val_data = TRAIN_TEXTS[train_size:] val_data = TRAIN_TEXTS[train_size:]
return train_data, val_data return train_data, val_data
@@ -28,13 +28,13 @@ def ensure_directories():
"""Создает необходимые директории если они не существуют.""" """Создает необходимые директории если они не существуют."""
directories = [ directories = [
"checkpoints", "checkpoints",
"checkpoints/gpt-bpe", "checkpoints/gpt-bpe",
"checkpoints/hf-bpe-tokenizer", "checkpoints/hf-bpe-tokenizer",
"checkpoints/hf-trained", "checkpoints/hf-trained",
"checkpoints/hf-trained-proxy", "checkpoints/hf-trained-proxy",
"logs" "logs",
] ]
for directory in directories: for directory in directories:
os.makedirs(directory, exist_ok=True) os.makedirs(directory, exist_ok=True)
@@ -42,33 +42,34 @@ def ensure_directories():
def get_model_paths(experiment_type: str = "llm_only") -> dict: def get_model_paths(experiment_type: str = "llm_only") -> dict:
""" """
Возвращает пути для конкретного типа эксперимента. Возвращает пути для конкретного типа эксперимента.
Args: Args:
experiment_type: Тип эксперимента ('llm_only' или 'hf_integration') experiment_type: Тип эксперимента ('llm_only' или 'hf_integration')
Returns: Returns:
dict: Словарь с путями dict: Словарь с путями
""" """
base_paths = PATHS.copy() base_paths = PATHS.copy()
if experiment_type == "hf_integration": if experiment_type == "hf_integration":
base_paths.update({ base_paths.update(
"model": base_paths["hf_model"], {"model": base_paths["hf_model"], "tokenizer": base_paths["hf_tokenizer"]}
"tokenizer": base_paths["hf_tokenizer"] )
})
else: # llm_only else: # llm_only
base_paths.update({ base_paths.update(
"model": base_paths["gpt_bpe_model"], {
"tokenizer": base_paths["bpe_tokenizer"] "model": base_paths["gpt_bpe_model"],
}) "tokenizer": base_paths["bpe_tokenizer"],
}
)
return base_paths return base_paths
def print_experiment_info(experiment_name: str, config: dict): def print_experiment_info(experiment_name: str, config: dict):
""" """
Выводит информацию о запускаемом эксперименте. Выводит информацию о запускаемом эксперименте.
Args: Args:
experiment_name: Название эксперимента experiment_name: Название эксперимента
config: Конфигурация эксперимента config: Конфигурация эксперимента
@@ -85,35 +86,35 @@ def print_experiment_info(experiment_name: str, config: dict):
def save_experiment_results(results: dict, filepath: str): def save_experiment_results(results: dict, filepath: str):
""" """
Сохраняет результаты эксперимента в файл. Сохраняет результаты эксперимента в файл.
Args: Args:
results: Словарь с результатами results: Словарь с результатами
filepath: Путь для сохранения filepath: Путь для сохранения
""" """
import json 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) json.dump(results, f, ensure_ascii=False, indent=2)
print(f"✅ Результаты эксперимента сохранены: {filepath}") print(f"✅ Результаты эксперимента сохранены: {filepath}")
def load_experiment_results(filepath: str) -> dict: def load_experiment_results(filepath: str) -> dict:
""" """
Загружает результаты эксперимента из файла. Загружает результаты эксперимента из файла.
Args: Args:
filepath: Путь к файлу с результатами filepath: Путь к файлу с результатами
Returns: Returns:
dict: Загруженные результаты dict: Загруженные результаты
""" """
import json import json
if not os.path.exists(filepath): if not os.path.exists(filepath):
return {} return {}
with open(filepath, 'r', encoding='utf-8') as f: with open(filepath, "r", encoding="utf-8") as f:
return json.load(f) return json.load(f)
@@ -121,42 +122,39 @@ class ExperimentLogger:
""" """
Логгер для экспериментов. Логгер для экспериментов.
""" """
def __init__(self, experiment_name: str): def __init__(self, experiment_name: str):
self.experiment_name = experiment_name self.experiment_name = experiment_name
self.metrics = {} self.metrics = {}
def log_metric(self, name: str, value: float): def log_metric(self, name: str, value: float):
"""Логирует метрику.""" """Логирует метрику."""
if name not in self.metrics: if name not in self.metrics:
self.metrics[name] = [] self.metrics[name] = []
self.metrics[name].append(value) self.metrics[name].append(value)
print(f"📈 {name}: {value:.4f}") print(f"📈 {name}: {value:.4f}")
def log_step(self, step: int, loss: float, **kwargs): def log_step(self, step: int, loss: float, **kwargs):
"""Логирует шаг обучения.""" """Логирует шаг обучения."""
print(f"📊 Step {step}: loss={loss:.4f}", end="") print(f"📊 Step {step}: loss={loss:.4f}", end="")
for key, value in kwargs.items(): for key, value in kwargs.items():
print(f", {key}={value:.4f}", end="") print(f", {key}={value:.4f}", end="")
print() print()
def log_epoch(self, epoch: int, train_loss: float, val_loss: float = None): def log_epoch(self, epoch: int, train_loss: float, val_loss: float = None):
"""Логирует завершение эпохи.""" """Логирует завершение эпохи."""
print(f"🎯 Epoch {epoch}: train_loss={train_loss:.4f}", end="") print(f"🎯 Epoch {epoch}: train_loss={train_loss:.4f}", end="")
if val_loss is not None: if val_loss is not None:
print(f", val_loss={val_loss:.4f}", end="") print(f", val_loss={val_loss:.4f}", end="")
print() print()
def save_logs(self, filepath: str): def save_logs(self, filepath: str):
"""Сохраняет логи эксперимента.""" """Сохраняет логи эксперимента."""
import json import json
logs = { logs = {"experiment_name": self.experiment_name, "metrics": self.metrics}
"experiment_name": self.experiment_name,
"metrics": self.metrics with open(filepath, "w", encoding="utf-8") as f:
}
with open(filepath, 'w', encoding='utf-8') as f:
json.dump(logs, f, ensure_ascii=False, indent=2) json.dump(logs, f, ensure_ascii=False, indent=2)
print(f"✅ Логи эксперимента сохранены: {filepath}") print(f"✅ Логи эксперимента сохранены: {filepath}")

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

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

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

@@ -6,12 +6,12 @@ import torch
import torch.nn as nn import torch.nn as nn
from typing import Optional, Tuple, Union, List from typing import Optional, Tuple, Union, List
from transformers import ( from transformers import (
PreTrainedModel, PreTrainedModel,
GPT2LMHeadModel, GPT2LMHeadModel,
GPT2Config, GPT2Config,
GenerationConfig, GenerationConfig,
LogitsProcessorList, LogitsProcessorList,
StoppingCriteriaList StoppingCriteriaList,
) )
from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
@@ -24,38 +24,39 @@ class HFGPTAdapter(PreTrainedModel):
Адаптер для модели GPT из библиотеки llm. Адаптер для модели GPT из библиотеки llm.
Позволяет использовать кастомные GPT модели с HuggingFace Transformers. Позволяет использовать кастомные GPT модели с HuggingFace Transformers.
""" """
config_class = HFPretrainedConfig config_class = HFPretrainedConfig
def __init__(self, config: HFPretrainedConfig, llm_model: Optional[GPT] = None): def __init__(self, config: HFPretrainedConfig, llm_model: Optional[GPT] = None):
""" """
Инициализация адаптера. Инициализация адаптера.
Args: Args:
config: Конфигурация HuggingFace config: Конфигурация HuggingFace
llm_model: Опционально, предварительно созданная модель llm llm_model: Опционально, предварительно созданная модель llm
""" """
super().__init__(config) super().__init__(config)
# Преобразуем HF конфигурацию в формат llm # Преобразуем HF конфигурацию в формат llm
llm_config = self._hf_to_llm_config(config) llm_config = self._hf_to_llm_config(config)
# Создаем или используем переданную модель # Создаем или используем переданную модель
if llm_model is None: if llm_model is None:
self.llm_model = GPT(llm_config) self.llm_model = GPT(llm_config)
else: else:
self.llm_model = llm_model 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) self.llm_model.load_state_dict(config.state_dict)
def _hf_to_llm_config(self, hf_config: HFPretrainedConfig) -> dict: def _hf_to_llm_config(self, hf_config: HFPretrainedConfig) -> dict:
""" """
Преобразует конфигурацию HF в формат llm. Преобразует конфигурацию HF в формат llm.
Args: Args:
hf_config: Конфигурация HuggingFace hf_config: Конфигурация HuggingFace
Returns: Returns:
dict: Конфигурация для llm модели dict: Конфигурация для llm модели
""" """
@@ -67,7 +68,7 @@ class HFGPTAdapter(PreTrainedModel):
"max_position_embeddings": hf_config.max_position_embeddings, "max_position_embeddings": hf_config.max_position_embeddings,
"dropout": hf_config.hidden_dropout_prob, "dropout": hf_config.hidden_dropout_prob,
} }
def forward( def forward(
self, self,
input_ids: Optional[torch.Tensor] = None, input_ids: Optional[torch.Tensor] = None,
@@ -78,11 +79,11 @@ class HFGPTAdapter(PreTrainedModel):
output_attentions: Optional[bool] = None, output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None, output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None, return_dict: Optional[bool] = None,
**kwargs **kwargs,
) -> Union[Tuple, CausalLMOutputWithCrossAttentions]: ) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
""" """
Прямой проход модели. Прямой проход модели.
Args: Args:
input_ids: Входные токены [batch_size, seq_len] input_ids: Входные токены [batch_size, seq_len]
attention_mask: Маска внимания [batch_size, seq_len] attention_mask: Маска внимания [batch_size, seq_len]
@@ -92,38 +93,39 @@ class HFGPTAdapter(PreTrainedModel):
output_attentions: Возвращать веса внимания output_attentions: Возвращать веса внимания
output_hidden_states: Возвращать скрытые состояния output_hidden_states: Возвращать скрытые состояния
return_dict: Возвращать словарь вместо кортежа return_dict: Возвращать словарь вместо кортежа
Returns: Returns:
CausalLMOutputWithCrossAttentions или кортеж 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 # Основной forward pass
outputs = self.llm_model(input_ids) outputs = self.llm_model(input_ids)
if isinstance(outputs, tuple): if isinstance(outputs, tuple):
logits = outputs[0] logits = outputs[0]
else: else:
logits = outputs logits = outputs
loss = None loss = None
if labels is not None: if labels is not None:
# Сдвигаем логиты и метки для языкового моделирования # Сдвигаем логиты и метки для языкового моделирования
shift_logits = logits[..., :-1, :].contiguous() shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous() shift_labels = labels[..., 1:].contiguous()
# Вычисляем cross-entropy loss # Вычисляем cross-entropy loss
loss_fct = nn.CrossEntropyLoss() loss_fct = nn.CrossEntropyLoss()
loss = loss_fct( loss = loss_fct(
shift_logits.view(-1, shift_logits.size(-1)), shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)
shift_labels.view(-1)
) )
if not return_dict: if not return_dict:
output = (logits,) output = (logits,)
if loss is not None: if loss is not None:
output = (loss,) + output output = (loss,) + output
return output return output
return CausalLMOutputWithCrossAttentions( return CausalLMOutputWithCrossAttentions(
loss=loss, loss=loss,
logits=logits, logits=logits,
@@ -132,30 +134,27 @@ class HFGPTAdapter(PreTrainedModel):
attentions=None, attentions=None,
cross_attentions=None, cross_attentions=None,
) )
def prepare_inputs_for_generation( def prepare_inputs_for_generation(
self, self, input_ids: torch.Tensor, past_key_values: Optional[Tuple] = None, **kwargs
input_ids: torch.Tensor,
past_key_values: Optional[Tuple] = None,
**kwargs
) -> dict: ) -> dict:
""" """
Подготавливает входные данные для генерации. Подготавливает входные данные для генерации.
Args: Args:
input_ids: Входные токены input_ids: Входные токены
past_key_values: Кешированные ключи и значения past_key_values: Кешированные ключи и значения
Returns: Returns:
dict: Подготовленные входные данные dict: Подготовленные входные данные
""" """
# Наша простая реализация пока не поддерживает past_key_values # Наша простая реализация пока не поддерживает past_key_values
return {"input_ids": input_ids} return {"input_ids": input_ids}
def can_generate(self) -> bool: def can_generate(self) -> bool:
"""Проверяет, может ли модель генерировать текст.""" """Проверяет, может ли модель генерировать текст."""
return True return True
def generate( def generate(
self, self,
input_ids: Optional[torch.Tensor] = None, input_ids: Optional[torch.Tensor] = None,
@@ -163,32 +162,32 @@ class HFGPTAdapter(PreTrainedModel):
generation_config: Optional[GenerationConfig] = None, generation_config: Optional[GenerationConfig] = None,
logits_processor: Optional[LogitsProcessorList] = None, logits_processor: Optional[LogitsProcessorList] = None,
stopping_criteria: Optional[StoppingCriteriaList] = None, stopping_criteria: Optional[StoppingCriteriaList] = None,
**kwargs **kwargs,
) -> torch.Tensor: ) -> torch.Tensor:
""" """
Генерация текста с поддержкой HuggingFace интерфейса. Генерация текста с поддержкой HuggingFace интерфейса.
Args: Args:
input_ids: Входные токены input_ids: Входные токены
attention_mask: Маска внимания attention_mask: Маска внимания
generation_config: Конфигурация генерации generation_config: Конфигурация генерации
logits_processor: Процессоры логитов logits_processor: Процессоры логитов
stopping_criteria: Критерии остановки stopping_criteria: Критерии остановки
Returns: Returns:
torch.Tensor: Сгенерированные токены torch.Tensor: Сгенерированные токены
""" """
# Извлекаем обязательные параметры из kwargs или используем значения по умолчанию # Извлекаем обязательные параметры из kwargs или используем значения по умолчанию
max_new_tokens = kwargs.pop('max_new_tokens', 50) max_new_tokens = kwargs.pop("max_new_tokens", 50)
do_sample = kwargs.pop('do_sample', True) do_sample = kwargs.pop("do_sample", True)
# Используем встроенную генерацию llm модели # Используем встроенную генерацию llm модели
return self.llm_model.generate( return self.llm_model.generate(
x=input_ids, x=input_ids,
max_new_tokens=max_new_tokens, max_new_tokens=max_new_tokens,
do_sample=do_sample, do_sample=do_sample,
attention_mask=attention_mask, attention_mask=attention_mask,
**kwargs **kwargs,
) )
@@ -196,64 +195,66 @@ class HFAdapter:
""" """
Основной класс адаптера для преобразования моделей llm в формат HuggingFace. Основной класс адаптера для преобразования моделей llm в формат HuggingFace.
""" """
@staticmethod @staticmethod
def from_llm_model( def from_llm_model(
llm_model: GPT, llm_model: GPT, hf_config: Optional[HFAdapterConfig] = None
hf_config: Optional[HFAdapterConfig] = None
) -> HFGPTAdapter: ) -> HFGPTAdapter:
""" """
Создает адаптер из существующей llm модели. Создает адаптер из существующей llm модели.
Args: Args:
llm_model: Обученная модель из библиотеки llm llm_model: Обученная модель из библиотеки llm
hf_config: Конфигурация для HuggingFace hf_config: Конфигурация для HuggingFace
Returns: Returns:
HFGPTAdapter: Адаптированная модель HFGPTAdapter: Адаптированная модель
""" """
if hf_config is None: if hf_config is None:
# Создаем конфигурацию из модели llm # Создаем конфигурацию из модели llm
hf_config = HFAdapterConfig.from_llm_config(llm_model.config) hf_config = HFAdapterConfig.from_llm_config(llm_model.config)
# Преобразуем в PretrainedConfig # Преобразуем в PretrainedConfig
pretrained_config = HFPretrainedConfig(**hf_config.to_dict()) pretrained_config = HFPretrainedConfig(**hf_config.to_dict())
return HFGPTAdapter(pretrained_config, llm_model) return HFGPTAdapter(pretrained_config, llm_model)
@staticmethod @staticmethod
def from_pretrained( def from_pretrained(
model_path: str, model_path: str, hf_config: Optional[HFAdapterConfig] = None
hf_config: Optional[HFAdapterConfig] = None
) -> HFGPTAdapter: ) -> HFGPTAdapter:
""" """
Загружает модель из чекпоинта и создает адаптер. Загружает модель из чекпоинта и создает адаптер.
Args: Args:
model_path: Путь к сохраненной модели model_path: Путь к сохраненной модели
hf_config: Конфигурация для HuggingFace hf_config: Конфигурация для HuggingFace
Returns: Returns:
HFGPTAdapter: Адаптированная модель HFGPTAdapter: Адаптированная модель
""" """
# Загружаем состояние модели # Загружаем состояние модели
state_dict = torch.load(model_path, map_location='cpu') state_dict = torch.load(model_path, map_location="cpu")
# Определяем конфигурацию из состояния модели или используем переданную # Определяем конфигурацию из состояния модели или используем переданную
if hf_config is None: if hf_config is None:
# Пытаемся определить конфигурацию из состояния модели # Пытаемся определить конфигурацию из состояния модели
# Это упрощенный подход - в реальности нужно сохранять конфигурацию отдельно # Это упрощенный подход - в реальности нужно сохранять конфигурацию отдельно
vocab_size = state_dict.get('_token_embeddings._embedding.weight', torch.zeros(50257, 768)).shape[0] vocab_size = state_dict.get(
embed_dim = state_dict.get('_token_embeddings._embedding.weight', torch.zeros(50257, 768)).shape[1] "_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( hf_config = HFAdapterConfig(
vocab_size=vocab_size, vocab_size=vocab_size,
hidden_size=embed_dim, hidden_size=embed_dim,
# Остальные параметры можно установить по умолчанию # Остальные параметры можно установить по умолчанию
) )
pretrained_config = HFPretrainedConfig(**hf_config.to_dict()) pretrained_config = HFPretrainedConfig(**hf_config.to_dict())
# Создаем модель llm и загружаем веса # Создаем модель llm и загружаем веса
llm_config = { llm_config = {
"vocab_size": hf_config.vocab_size, "vocab_size": hf_config.vocab_size,
@@ -263,21 +264,17 @@ class HFAdapter:
"max_position_embeddings": hf_config.max_position_embeddings, "max_position_embeddings": hf_config.max_position_embeddings,
"dropout": hf_config.hidden_dropout_prob, "dropout": hf_config.hidden_dropout_prob,
} }
llm_model = GPT(llm_config) llm_model = GPT(llm_config)
llm_model.load_state_dict(state_dict) llm_model.load_state_dict(state_dict)
return HFGPTAdapter(pretrained_config, llm_model) return HFGPTAdapter(pretrained_config, llm_model)
@staticmethod @staticmethod
def save_pretrained( def save_pretrained(model: HFGPTAdapter, save_directory: str, **kwargs):
model: HFGPTAdapter,
save_directory: str,
**kwargs
):
""" """
Сохраняет адаптированную модель в формате HuggingFace. Сохраняет адаптированную модель в формате HuggingFace.
Args: Args:
model: Адаптированная модель model: Адаптированная модель
save_directory: Директория для сохранения save_directory: Директория для сохранения
@@ -285,19 +282,19 @@ class HFAdapter:
""" """
import os import os
import json import json
# Создаем директорию если не существует # Создаем директорию если не существует
os.makedirs(save_directory, exist_ok=True) os.makedirs(save_directory, exist_ok=True)
# Сохраняем конфигурацию # Сохраняем конфигурацию
config_path = os.path.join(save_directory, "config.json") 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) json.dump(model.config.to_dict(), f, indent=2, ensure_ascii=False)
# Сохраняем веса модели # Сохраняем веса модели
model_path = os.path.join(save_directory, "pytorch_model.bin") model_path = os.path.join(save_directory, "pytorch_model.bin")
torch.save(model.llm_model.state_dict(), model_path) torch.save(model.llm_model.state_dict(), model_path)
# Сохраняем токенизатор если передан # Сохраняем токенизатор если передан
if hasattr(kwargs, 'tokenizer') and kwargs['tokenizer'] is not None: if hasattr(kwargs, "tokenizer") and kwargs["tokenizer"] is not None:
kwargs['tokenizer'].save_pretrained(save_directory) 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 typing import Dict, Any, Optional
from transformers import PretrainedConfig from transformers import PretrainedConfig
@dataclass @dataclass
class HFAdapterConfig: class HFAdapterConfig:
""" """
Конфигурация для адаптера HuggingFace. Конфигурация для адаптера HuggingFace.
Параметры: Параметры:
model_type: Тип модели (gpt, llama, etc.) model_type: Тип модели (gpt, llama, etc.)
vocab_size: Размер словаря vocab_size: Размер словаря
@@ -28,6 +29,7 @@ class HFAdapterConfig:
eos_token_id: ID токена конца строки eos_token_id: ID токена конца строки
bos_token_id: ID токена начала строки bos_token_id: ID токена начала строки
""" """
model_type: str = "gpt" model_type: str = "gpt"
vocab_size: int = 50257 vocab_size: int = 50257
hidden_size: int = 768 hidden_size: int = 768
@@ -43,49 +45,50 @@ class HFAdapterConfig:
pad_token_id: int = 50256 pad_token_id: int = 50256
eos_token_id: int = 50256 eos_token_id: int = 50256
bos_token_id: int = 50256 bos_token_id: int = 50256
# Дополнительные параметры для совместимости # Дополнительные параметры для совместимости
architectures: list = field(default_factory=lambda: ["GPT2LMHeadModel"]) architectures: list = field(default_factory=lambda: ["GPT2LMHeadModel"])
torch_dtype: str = "float32" torch_dtype: str = "float32"
transformers_version: str = "4.44.0" transformers_version: str = "4.44.0"
def to_dict(self) -> Dict[str, Any]: def to_dict(self) -> Dict[str, Any]:
"""Преобразует конфигурацию в словарь.""" """Преобразует конфигурацию в словарь."""
return { return {
k: v for k, v in self.__dict__.items() k: v
if not k.startswith('_') and not callable(v) for k, v in self.__dict__.items()
if not k.startswith("_") and not callable(v)
} }
@classmethod @classmethod
def from_llm_config(cls, llm_config: Dict[str, Any]) -> "HFAdapterConfig": def from_llm_config(cls, llm_config: Dict[str, Any]) -> "HFAdapterConfig":
""" """
Создает конфигурацию HF из конфигурации llm. Создает конфигурацию HF из конфигурации llm.
Args: Args:
llm_config: Конфигурация модели из библиотеки llm llm_config: Конфигурация модели из библиотеки llm
Returns: Returns:
HFAdapterConfig: Конфигурация для HuggingFace HFAdapterConfig: Конфигурация для HuggingFace
""" """
# Маппинг параметров из llm в HF формат # Маппинг параметров из llm в HF формат
mapping = { mapping = {
"embed_dim": "hidden_size", "embed_dim": "hidden_size",
"num_layers": "num_hidden_layers", "num_layers": "num_hidden_layers",
"num_heads": "num_attention_heads", "num_heads": "num_attention_heads",
"max_position_embeddings": "max_position_embeddings", "max_position_embeddings": "max_position_embeddings",
"dropout": "hidden_dropout_prob", "dropout": "hidden_dropout_prob",
"vocab_size": "vocab_size" "vocab_size": "vocab_size",
} }
hf_config_dict = {} hf_config_dict = {}
for llm_key, hf_key in mapping.items(): for llm_key, hf_key in mapping.items():
if llm_key in llm_config: if llm_key in llm_config:
hf_config_dict[hf_key] = llm_config[llm_key] hf_config_dict[hf_key] = llm_config[llm_key]
# Устанавливаем промежуточный размер (обычно 4x hidden_size) # Устанавливаем промежуточный размер (обычно 4x hidden_size)
if "hidden_size" in hf_config_dict: if "hidden_size" in hf_config_dict:
hf_config_dict["intermediate_size"] = hf_config_dict["hidden_size"] * 4 hf_config_dict["intermediate_size"] = hf_config_dict["hidden_size"] * 4
return cls(**hf_config_dict) return cls(**hf_config_dict)
@@ -94,8 +97,9 @@ class HFPretrainedConfig(PretrainedConfig):
Конфигурация для предобученных моделей HuggingFace. Конфигурация для предобученных моделей HuggingFace.
Наследуется от PretrainedConfig для полной совместимости. Наследуется от PretrainedConfig для полной совместимости.
""" """
model_type = "gpt" model_type = "gpt"
def __init__( def __init__(
self, self,
vocab_size=50257, vocab_size=50257,
@@ -112,15 +116,15 @@ class HFPretrainedConfig(PretrainedConfig):
pad_token_id=50256, pad_token_id=50256,
eos_token_id=50256, eos_token_id=50256,
bos_token_id=50256, bos_token_id=50256,
**kwargs **kwargs,
): ):
super().__init__( super().__init__(
pad_token_id=pad_token_id, pad_token_id=pad_token_id,
eos_token_id=eos_token_id, eos_token_id=eos_token_id,
bos_token_id=bos_token_id, bos_token_id=bos_token_id,
**kwargs **kwargs,
) )
self.vocab_size = vocab_size self.vocab_size = vocab_size
self.hidden_size = hidden_size self.hidden_size = hidden_size
self.num_hidden_layers = num_hidden_layers 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. Упрощенный адаптер для кастомных токенизаторов llm.
Предоставляет совместимый с HuggingFace интерфейс. Предоставляет совместимый с HuggingFace интерфейс.
""" """
def __init__(self, llm_tokenizer: BaseTokenizer): def __init__(self, llm_tokenizer: BaseTokenizer):
""" """
Инициализация адаптера. Инициализация адаптера.
Args: Args:
llm_tokenizer: Кастомный токенизатор из llm llm_tokenizer: Кастомный токенизатор из llm
""" """
self.llm_tokenizer = llm_tokenizer self.llm_tokenizer = llm_tokenizer
# Получаем словарь и размер # Получаем словарь и размер
self._vocab = llm_tokenizer.get_vocab() self._vocab = llm_tokenizer.get_vocab()
self.vocab_size = llm_tokenizer.get_vocab_size() self.vocab_size = llm_tokenizer.get_vocab_size()
# Устанавливаем специальные токены # Устанавливаем специальные токены
self.pad_token = getattr(llm_tokenizer, 'pad_token', '<pad>') self.pad_token = getattr(llm_tokenizer, "pad_token", "<pad>")
self.unk_token = getattr(llm_tokenizer, 'unk_token', '<unk>') self.unk_token = getattr(llm_tokenizer, "unk_token", "<unk>")
self.bos_token = getattr(llm_tokenizer, 'bos_token', '<bos>') self.bos_token = getattr(llm_tokenizer, "bos_token", "<bos>")
self.eos_token = getattr(llm_tokenizer, 'eos_token', '<eos>') self.eos_token = getattr(llm_tokenizer, "eos_token", "<eos>")
# Сохраняем ID специальных токенов # Сохраняем ID специальных токенов
self.pad_token_id = getattr(llm_tokenizer, 'pad_token_id', 0) self.pad_token_id = getattr(llm_tokenizer, "pad_token_id", 0)
self.unk_token_id = getattr(llm_tokenizer, 'unk_token_id', 1) self.unk_token_id = getattr(llm_tokenizer, "unk_token_id", 1)
self.bos_token_id = getattr(llm_tokenizer, 'bos_token_id', 2) self.bos_token_id = getattr(llm_tokenizer, "bos_token_id", 2)
self.eos_token_id = getattr(llm_tokenizer, 'eos_token_id', 3) self.eos_token_id = getattr(llm_tokenizer, "eos_token_id", 3)
def __call__(self, text: str, **kwargs): def __call__(self, text: str, **kwargs):
""" """
Вызов токенизатора с параметрами как у HuggingFace. Вызов токенизатора с параметрами как у HuggingFace.
Args: Args:
text: Входной текст text: Входной текст
**kwargs: Параметры токенизации **kwargs: Параметры токенизации
Returns: Returns:
dict: Словарь с токенами dict: Словарь с токенами
""" """
return_tensors = kwargs.get('return_tensors', None) return_tensors = kwargs.get("return_tensors", None)
padding = kwargs.get('padding', False) padding = kwargs.get("padding", False)
truncation = kwargs.get('truncation', False) truncation = kwargs.get("truncation", False)
max_length = kwargs.get('max_length', None) max_length = kwargs.get("max_length", None)
add_special_tokens = kwargs.get('add_special_tokens', True) add_special_tokens = kwargs.get("add_special_tokens", True)
# Кодируем текст # Кодируем текст
#input_ids = self.llm_tokenizer.encode( # input_ids = self.llm_tokenizer.encode(
# text, # text,
# add_special_tokens=add_special_tokens # add_special_tokens=add_special_tokens
#) # )
if isinstance(text, str): if isinstance(text, str):
input_ids = self.llm_tokenizer.encode( input_ids = self.llm_tokenizer.encode(
text, text, add_special_tokens=add_special_tokens
add_special_tokens=add_special_tokens
) )
input_ids = [input_ids] # <-- оборачиваем в batch input_ids = [input_ids] # <-- оборачиваем в batch
else: else:
# Список строк, батч-режим! # Список строк, батч-режим!
input_ids = [ input_ids = [
self.llm_tokenizer.encode( self.llm_tokenizer.encode(t, add_special_tokens=add_special_tokens)
t, for t in text
add_special_tokens=add_special_tokens
) for t in text
] ]
# Применяем truncation # Применяем truncation
if truncation and max_length is not None and len(input_ids) > max_length: if truncation and max_length is not None and len(input_ids) > max_length:
input_ids = input_ids[:max_length] input_ids = input_ids[:max_length]
# Применяем padding # Применяем padding
if padding and max_length is not None and len(input_ids) < max_length: 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)) input_ids = input_ids + [self.pad_token_id] * (max_length - len(input_ids))
# Конвертируем в тензоры если нужно # Конвертируем в тензоры если нужно
if return_tensors == "pt": if return_tensors == "pt":
import torch import torch
input_ids = torch.tensor([input_ids]) input_ids = torch.tensor([input_ids])
return {"input_ids": input_ids} return {"input_ids": input_ids}
def encode( def encode(
self, self,
text: str, text: str,
@@ -99,11 +97,11 @@ class HFTokenizerAdapter:
truncation: bool = False, truncation: bool = False,
max_length: Optional[int] = None, max_length: Optional[int] = None,
return_tensors: Optional[str] = None, return_tensors: Optional[str] = None,
**kwargs **kwargs,
) -> Union[List[int], List[List[int]]]: ) -> Union[List[int], List[List[int]]]:
""" """
Кодирует текст в последовательность токенов. Кодирует текст в последовательность токенов.
Args: Args:
text: Входной текст text: Входной текст
text_pair: Второй текст (для парных задач) text_pair: Второй текст (для парных задач)
@@ -112,84 +110,91 @@ class HFTokenizerAdapter:
truncation: Обрезать последовательность truncation: Обрезать последовательность
max_length: Максимальная длина max_length: Максимальная длина
return_tensors: Возвращать тензоры return_tensors: Возвращать тензоры
Returns: Returns:
Список токенов или список списков токенов Список токенов или список списков токенов
""" """
# Кодируем основной текст # Кодируем основной текст
token_ids = self.llm_tokenizer.encode( token_ids = self.llm_tokenizer.encode(
text, text, add_special_tokens=add_special_tokens
add_special_tokens=add_special_tokens
) )
# Обрабатываем text_pair если есть # Обрабатываем text_pair если есть
if text_pair is not None: if text_pair is not None:
pair_ids = self.llm_tokenizer.encode( pair_ids = self.llm_tokenizer.encode(text_pair, add_special_tokens=False)
text_pair,
add_special_tokens=False
)
token_ids.extend(pair_ids) token_ids.extend(pair_ids)
# Применяем truncation # Применяем truncation
if truncation and max_length is not None and len(token_ids) > max_length: if truncation and max_length is not None and len(token_ids) > max_length:
token_ids = token_ids[:max_length] token_ids = token_ids[:max_length]
# Применяем padding # Применяем padding
if padding and max_length is not None and len(token_ids) < max_length: 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)) token_ids = token_ids + [self.pad_token_id] * (max_length - len(token_ids))
# Конвертируем в тензоры если нужно # Конвертируем в тензоры если нужно
if return_tensors == "pt": if return_tensors == "pt":
import torch import torch
return torch.tensor([token_ids]) return torch.tensor([token_ids])
elif return_tensors == "np": elif return_tensors == "np":
import numpy as np import numpy as np
return np.array([token_ids]) return np.array([token_ids])
return token_ids return token_ids
def decode( def decode(
self, self,
token_ids: Union[int, List[int], List[List[int]]], token_ids: Union[int, List[int], List[List[int]]],
skip_special_tokens: bool = True, skip_special_tokens: bool = True,
**kwargs **kwargs,
) -> str: ) -> str:
""" """
Декодирует последовательность токенов в текст. Декодирует последовательность токенов в текст.
Args: Args:
token_ids: ID токенов token_ids: ID токенов
skip_special_tokens: Пропускать специальные токены skip_special_tokens: Пропускать специальные токены
Returns: Returns:
str: Декодированный текст str: Декодированный текст
""" """
# Обрабатываем разные форматы входных данных # Обрабатываем разные форматы входных данных
if isinstance(token_ids, int): if isinstance(token_ids, int):
token_ids = [token_ids] 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] token_ids = token_ids[0]
# Фильтруем специальные токены если нужно # Фильтруем специальные токены если нужно
if skip_special_tokens: 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] token_ids = [tid for tid in token_ids if tid not in special_ids]
return self.llm_tokenizer.decode(token_ids) return self.llm_tokenizer.decode(token_ids)
def tokenize(self, text: str, **kwargs) -> List[str]: def tokenize(self, text: str, **kwargs) -> List[str]:
""" """
Токенизирует текст в список строковых токенов. Токенизирует текст в список строковых токенов.
Args: Args:
text: Входной текст text: Входной текст
Returns: Returns:
List[str]: Список токенов List[str]: Список токенов
""" """
return self.llm_tokenizer.tokenize(text) return self.llm_tokenizer.tokenize(text)
def pad( def pad(
self, self,
encoded_inputs, encoded_inputs,
@@ -202,7 +207,7 @@ class HFTokenizerAdapter:
): ):
""" """
Pad a list of encoded inputs. Pad a list of encoded inputs.
Args: Args:
encoded_inputs: List of encoded inputs encoded_inputs: List of encoded inputs
padding: Padding strategy padding: Padding strategy
@@ -211,7 +216,7 @@ class HFTokenizerAdapter:
return_attention_mask: Return attention mask return_attention_mask: Return attention mask
return_tensors: Return tensors return_tensors: Return tensors
verbose: Verbose mode verbose: Verbose mode
Returns: Returns:
Padded inputs Padded inputs
""" """
@@ -224,47 +229,62 @@ class HFTokenizerAdapter:
# Обрабатываем разные типы данных # Обрабатываем разные типы данных
if isinstance(input_ids, int): if isinstance(input_ids, int):
seq_len = 1 seq_len = 1
elif hasattr(input_ids, 'shape'): elif hasattr(input_ids, "shape"):
seq_len = input_ids.shape[-1] if len(input_ids.shape) > 1 else len(input_ids) seq_len = (
input_ids.shape[-1]
if len(input_ids.shape) > 1
else len(input_ids)
)
else: else:
seq_len = len(input_ids) seq_len = len(input_ids)
max_len = max(max_len, seq_len) max_len = max(max_len, seq_len)
if max_length is not None: if max_length is not None:
max_len = min(max_len, max_length) max_len = min(max_len, max_length)
# Применяем padding # Применяем padding
for item in encoded_inputs: for item in encoded_inputs:
input_ids = item["input_ids"] input_ids = item["input_ids"]
# Получаем текущую длину # Получаем текущую длину
if isinstance(input_ids, int): if isinstance(input_ids, int):
current_len = 1 current_len = 1
elif hasattr(input_ids, 'shape'): elif hasattr(input_ids, "shape"):
current_len = input_ids.shape[-1] if len(input_ids.shape) > 1 else len(input_ids) current_len = (
input_ids.shape[-1]
if len(input_ids.shape) > 1
else len(input_ids)
)
else: else:
current_len = len(input_ids) current_len = len(input_ids)
if current_len < max_len: if current_len < max_len:
# Дополняем pad_token_id # Дополняем pad_token_id
padding_length = max_len - current_len padding_length = max_len - current_len
# Обрабатываем разные типы данных # Обрабатываем разные типы данных
if isinstance(input_ids, int): if isinstance(input_ids, int):
item["input_ids"] = [input_ids] + [self.pad_token_id] * padding_length item["input_ids"] = [input_ids] + [
elif hasattr(input_ids, 'shape'): self.pad_token_id
] * padding_length
elif hasattr(input_ids, "shape"):
import torch 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]) item["input_ids"] = torch.cat([input_ids, padding_tensor])
else: 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 если требуется # Добавляем attention_mask если требуется
if "attention_mask" in item: if "attention_mask" in item:
mask = item["attention_mask"] mask = item["attention_mask"]
if isinstance(mask, int): if isinstance(mask, int):
item["attention_mask"] = [mask] + [0] * padding_length item["attention_mask"] = [mask] + [0] * padding_length
elif hasattr(mask, 'shape'): elif hasattr(mask, "shape"):
padding_mask = torch.zeros(padding_length, dtype=mask.dtype) padding_mask = torch.zeros(padding_length, dtype=mask.dtype)
item["attention_mask"] = torch.cat([mask, padding_mask]) item["attention_mask"] = torch.cat([mask, padding_mask])
else: else:
@@ -272,44 +292,49 @@ class HFTokenizerAdapter:
elif return_attention_mask: elif return_attention_mask:
if isinstance(input_ids, int): if isinstance(input_ids, int):
item["attention_mask"] = [1] + [0] * padding_length 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) attention_mask = torch.ones(current_len, dtype=torch.long)
padding_mask = torch.zeros(padding_length, 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: else:
item["attention_mask"] = [1] * current_len + [0] * padding_length item["attention_mask"] = [1] * current_len + [
0
] * padding_length
# Конвертируем в тензоры если требуется # Конвертируем в тензоры если требуется
if return_tensors == "pt": if return_tensors == "pt":
import torch import torch
for key in list(encoded_inputs[0].keys()): for key in list(encoded_inputs[0].keys()):
if isinstance(encoded_inputs[0][key], list): if isinstance(encoded_inputs[0][key], list):
for i in range(len(encoded_inputs)): for i in range(len(encoded_inputs)):
encoded_inputs[i][key] = torch.tensor(encoded_inputs[i][key]) encoded_inputs[i][key] = torch.tensor(encoded_inputs[i][key])
return encoded_inputs return encoded_inputs
def get_vocab(self) -> Dict[str, int]: def get_vocab(self) -> Dict[str, int]:
"""Возвращает словарь токенизатора.""" """Возвращает словарь токенизатора."""
return self._vocab return self._vocab
def __len__(self) -> int: def __len__(self) -> int:
"""Возвращает размер словаря.""" """Возвращает размер словаря."""
return self.vocab_size return self.vocab_size
def save_pretrained(self, save_directory: str, **kwargs): def save_pretrained(self, save_directory: str, **kwargs):
""" """
Сохраняет токенизатор в формате HuggingFace. Сохраняет токенизатор в формате HuggingFace.
Args: Args:
save_directory: Директория для сохранения save_directory: Директория для сохранения
**kwargs: Дополнительные параметры **kwargs: Дополнительные параметры
""" """
import os import os
# Создаем директорию если не существует # Создаем директорию если не существует
os.makedirs(save_directory, exist_ok=True) os.makedirs(save_directory, exist_ok=True)
# Сохраняем конфигурацию токенизатора # Сохраняем конфигурацию токенизатора
tokenizer_config = { tokenizer_config = {
"tokenizer_class": self.__class__.__name__, "tokenizer_class": self.__class__.__name__,
@@ -324,77 +349,81 @@ class HFTokenizerAdapter:
"bos_token_id": self.bos_token_id, "bos_token_id": self.bos_token_id,
"eos_token_id": self.eos_token_id, "eos_token_id": self.eos_token_id,
} }
config_path = os.path.join(save_directory, "tokenizer_config.json") 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) json.dump(tokenizer_config, f, ensure_ascii=False, indent=2)
# Сохраняем словарь # Сохраняем словарь
vocab_path = os.path.join(save_directory, "vocab.json") 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) json.dump(self._vocab, f, ensure_ascii=False, indent=2)
print(f"✅ Токенизатор сохранен в {save_directory}") print(f"✅ Токенизатор сохранен в {save_directory}")
@classmethod @classmethod
def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs): def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
""" """
Загружает адаптированный токенизатор. Загружает адаптированный токенизатор.
Args: Args:
pretrained_model_name_or_path: Путь к сохраненному токенизатору pretrained_model_name_or_path: Путь к сохраненному токенизатору
**kwargs: Дополнительные параметры **kwargs: Дополнительные параметры
Returns: Returns:
HFTokenizerAdapter: Загруженный адаптер HFTokenizerAdapter: Загруженный адаптер
""" """
import os import os
# Проверяем, является ли путь директорией с файлами токенизатора # Проверяем, является ли путь директорией с файлами токенизатора
if os.path.isdir(pretrained_model_name_or_path): 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") 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): if not os.path.exists(config_path) or not os.path.exists(vocab_path):
raise FileNotFoundError( raise FileNotFoundError(
f"Файлы токенизатора не найдены в {pretrained_model_name_or_path}" 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) config = json.load(f)
# Определяем тип токенизатора llm # Определяем тип токенизатора llm
llm_tokenizer_type = config.get("llm_tokenizer_type", "BPETokenizer") llm_tokenizer_type = config.get("llm_tokenizer_type", "BPETokenizer")
if llm_tokenizer_type == "BPETokenizer": if llm_tokenizer_type == "BPETokenizer":
# Создаем BPETokenizer и загружаем словарь # Создаем BPETokenizer и загружаем словарь
llm_tokenizer = 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) vocab = json.load(f)
llm_tokenizer.vocab = vocab llm_tokenizer.vocab = vocab
llm_tokenizer.inverse_vocab = {v: k for k, v in vocab.items()} llm_tokenizer.inverse_vocab = {v: k for k, v in vocab.items()}
llm_tokenizer.vocab_size = len(vocab) llm_tokenizer.vocab_size = len(vocab)
# Устанавливаем специальные токены # Устанавливаем специальные токены
llm_tokenizer.pad_token = config.get("pad_token", "<pad>") llm_tokenizer.pad_token = config.get("pad_token", "<pad>")
llm_tokenizer.unk_token = config.get("unk_token", "<unk>") llm_tokenizer.unk_token = config.get("unk_token", "<unk>")
llm_tokenizer.bos_token = config.get("bos_token", "<bos>") llm_tokenizer.bos_token = config.get("bos_token", "<bos>")
llm_tokenizer.eos_token = config.get("eos_token", "<eos>") llm_tokenizer.eos_token = config.get("eos_token", "<eos>")
llm_tokenizer.pad_token_id = config.get("pad_token_id", 0) llm_tokenizer.pad_token_id = config.get("pad_token_id", 0)
llm_tokenizer.unk_token_id = config.get("unk_token_id", 1) llm_tokenizer.unk_token_id = config.get("unk_token_id", 1)
llm_tokenizer.bos_token_id = config.get("bos_token_id", 2) llm_tokenizer.bos_token_id = config.get("bos_token_id", 2)
llm_tokenizer.eos_token_id = config.get("eos_token_id", 3) llm_tokenizer.eos_token_id = config.get("eos_token_id", 3)
return cls(llm_tokenizer, **kwargs) return cls(llm_tokenizer, **kwargs)
else: else:
raise ValueError(f"Неподдерживаемый тип токенизатора: {llm_tokenizer_type}") raise ValueError(
f"Неподдерживаемый тип токенизатора: {llm_tokenizer_type}"
)
else: else:
# Пытаемся загрузить как файл llm токенизатора # Пытаемся загрузить как файл llm токенизатора
try: try:
@@ -409,10 +438,10 @@ class HFTokenizerAdapter:
def create_hf_tokenizer(llm_tokenizer: BaseTokenizer) -> HFTokenizerAdapter: def create_hf_tokenizer(llm_tokenizer: BaseTokenizer) -> HFTokenizerAdapter:
""" """
Создает адаптер HuggingFace для кастомного токенизатора. Создает адаптер HuggingFace для кастомного токенизатора.
Args: Args:
llm_tokenizer: Токенизатор из библиотеки llm llm_tokenizer: Токенизатор из библиотеки llm
Returns: Returns:
HFTokenizerAdapter: Адаптированный токенизатор HFTokenizerAdapter: Адаптированный токенизатор
""" """
@@ -422,7 +451,7 @@ def create_hf_tokenizer(llm_tokenizer: BaseTokenizer) -> HFTokenizerAdapter:
def convert_to_hf_format(llm_tokenizer: BaseTokenizer, save_directory: str): def convert_to_hf_format(llm_tokenizer: BaseTokenizer, save_directory: str):
""" """
Конвертирует кастомный токенизатор в формат HuggingFace. Конвертирует кастомный токенизатор в формат HuggingFace.
Args: Args:
llm_tokenizer: Токенизатор из llm llm_tokenizer: Токенизатор из llm
save_directory: Директория для сохранения save_directory: Директория для сохранения

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

@@ -14,55 +14,57 @@ class HFUtils:
""" """
Утилиты для работы с HuggingFace адаптером. Утилиты для работы с HuggingFace адаптером.
""" """
@staticmethod @staticmethod
def create_hf_config_from_llm(llm_config: Dict[str, Any]) -> HFPretrainedConfig: def create_hf_config_from_llm(llm_config: Dict[str, Any]) -> HFPretrainedConfig:
""" """
Создает конфигурацию HuggingFace из конфигурации llm. Создает конфигурацию HuggingFace из конфигурации llm.
Args: Args:
llm_config: Конфигурация модели из библиотеки llm llm_config: Конфигурация модели из библиотеки llm
Returns: Returns:
HFPretrainedConfig: Конфигурация для HuggingFace HFPretrainedConfig: Конфигурация для HuggingFace
""" """
adapter_config = HFAdapterConfig.from_llm_config(llm_config) adapter_config = HFAdapterConfig.from_llm_config(llm_config)
return HFPretrainedConfig(**adapter_config.to_dict()) return HFPretrainedConfig(**adapter_config.to_dict())
@staticmethod @staticmethod
def convert_to_hf_format( def convert_to_hf_format(
llm_model, llm_model, tokenizer=None, model_name: str = "custom-gpt"
tokenizer = None,
model_name: str = "custom-gpt"
) -> tuple: ) -> tuple:
""" """
Конвертирует llm модель в формат HuggingFace. Конвертирует llm модель в формат HuggingFace.
Args: Args:
llm_model: Модель из библиотеки llm llm_model: Модель из библиотеки llm
tokenizer: Токенизатор (HF или кастомный) tokenizer: Токенизатор (HF или кастомный)
model_name: Имя модели для сохранения model_name: Имя модели для сохранения
Returns: Returns:
tuple: (адаптированная модель, токенизатор) tuple: (адаптированная модель, токенизатор)
""" """
# Создаем адаптер # Создаем адаптер
hf_model = HFAdapter.from_llm_model(llm_model) hf_model = HFAdapter.from_llm_model(llm_model)
# Если токенизатор не передан, создаем стандартный # Если токенизатор не передан, создаем стандартный
if tokenizer is None: if tokenizer is None:
from transformers import AutoTokenizer from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("gpt2") tokenizer = AutoTokenizer.from_pretrained("gpt2")
# Устанавливаем специальные токены # Устанавливаем специальные токены
if tokenizer.pad_token is None: if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token 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 from .hf_tokenizer import create_hf_tokenizer
tokenizer = create_hf_tokenizer(tokenizer) tokenizer = create_hf_tokenizer(tokenizer)
return hf_model, tokenizer return hf_model, tokenizer
@staticmethod @staticmethod
def push_to_hub( def push_to_hub(
model: HFGPTAdapter, model: HFGPTAdapter,
@@ -70,11 +72,11 @@ class HFUtils:
repo_name: str, repo_name: str,
organization: Optional[str] = None, organization: Optional[str] = None,
private: bool = False, private: bool = False,
**kwargs **kwargs,
): ):
""" """
Загружает модель в HuggingFace Hub. Загружает модель в HuggingFace Hub.
Args: Args:
model: Адаптированная модель model: Адаптированная модель
tokenizer: Токенизатор tokenizer: Токенизатор
@@ -85,23 +87,23 @@ class HFUtils:
""" """
try: try:
from huggingface_hub import HfApi, ModelCard, create_repo from huggingface_hub import HfApi, ModelCard, create_repo
# Создаем репозиторий # Создаем репозиторий
if organization: if organization:
repo_id = f"{organization}/{repo_name}" repo_id = f"{organization}/{repo_name}"
else: else:
repo_id = repo_name repo_id = repo_name
create_repo(repo_id, private=private, exist_ok=True) create_repo(repo_id, private=private, exist_ok=True)
# Сохраняем модель локально # Сохраняем модель локально
import tempfile import tempfile
import os import os
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
# Сохраняем модель # Сохраняем модель
HFAdapter.save_pretrained(model, tmp_dir, tokenizer=tokenizer) HFAdapter.save_pretrained(model, tmp_dir, tokenizer=tokenizer)
# Создаем Model Card # Создаем Model Card
card = ModelCard.from_template( card = ModelCard.from_template(
model_name=repo_name, model_name=repo_name,
@@ -110,46 +112,43 @@ class HFUtils:
tags=["llm", "gpt", "custom"], tags=["llm", "gpt", "custom"],
) )
card.save(os.path.join(tmp_dir, "README.md")) card.save(os.path.join(tmp_dir, "README.md"))
# Загружаем в Hub # Загружаем в Hub
api = HfApi() api = HfApi()
api.upload_folder( api.upload_folder(
folder_path=tmp_dir, folder_path=tmp_dir,
repo_id=repo_id, 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}") print(f"✅ Модель успешно загружена в HuggingFace Hub: {repo_id}")
except ImportError: except ImportError:
raise ImportError( raise ImportError(
"Для загрузки в HuggingFace Hub установите huggingface_hub: " "Для загрузки в HuggingFace Hub установите huggingface_hub: "
"pip install huggingface_hub" "pip install huggingface_hub"
) )
@staticmethod @staticmethod
def load_from_hub( def load_from_hub(repo_id: str, **kwargs) -> tuple:
repo_id: str,
**kwargs
) -> tuple:
""" """
Загружает модель из HuggingFace Hub. Загружает модель из HuggingFace Hub.
Args: Args:
repo_id: ID репозитория repo_id: ID репозитория
**kwargs: Дополнительные параметры **kwargs: Дополнительные параметры
Returns: Returns:
tuple: (модель, токенизатор) tuple: (модель, токенизатор)
""" """
from transformers import AutoTokenizer from transformers import AutoTokenizer
# Загружаем токенизатор # Загружаем токенизатор
tokenizer = AutoTokenizer.from_pretrained(repo_id, **kwargs) tokenizer = AutoTokenizer.from_pretrained(repo_id, **kwargs)
# Загружаем конфигурацию # Загружаем конфигурацию
config = AutoConfig.from_pretrained(repo_id, **kwargs) config = AutoConfig.from_pretrained(repo_id, **kwargs)
# Создаем модель llm на основе конфигурации # Создаем модель llm на основе конфигурации
llm_config = { llm_config = {
"vocab_size": config.vocab_size, "vocab_size": config.vocab_size,
@@ -159,63 +158,56 @@ class HFUtils:
"max_position_embeddings": config.max_position_embeddings, "max_position_embeddings": config.max_position_embeddings,
"dropout": config.hidden_dropout_prob, "dropout": config.hidden_dropout_prob,
} }
# Загружаем модель через адаптер # Загружаем модель через адаптер
model = HFAdapter.from_pretrained( model = HFAdapter.from_pretrained(
f"{repo_id}/pytorch_model.bin", f"{repo_id}/pytorch_model.bin", HFAdapterConfig.from_llm_config(llm_config)
HFAdapterConfig.from_llm_config(llm_config)
) )
return model, tokenizer return model, tokenizer
@staticmethod @staticmethod
def compare_with_hf_model( def compare_with_hf_model(
llm_model, llm_model, hf_model_name: str = "gpt2", test_input: str = "Hello world"
hf_model_name: str = "gpt2",
test_input: str = "Hello world"
) -> Dict[str, Any]: ) -> Dict[str, Any]:
""" """
Сравнивает llm модель с эталонной моделью из HuggingFace. Сравнивает llm модель с эталонной моделью из HuggingFace.
Args: Args:
llm_model: Модель из библиотеки llm llm_model: Модель из библиотеки llm
hf_model_name: Имя модели HuggingFace для сравнения hf_model_name: Имя модели HuggingFace для сравнения
test_input: Тестовый вход test_input: Тестовый вход
Returns: Returns:
Dict: Результаты сравнения Dict: Результаты сравнения
""" """
from transformers import AutoModelForCausalLM, AutoTokenizer from transformers import AutoModelForCausalLM, AutoTokenizer
# Загружаем эталонную модель # Загружаем эталонную модель
hf_tokenizer = AutoTokenizer.from_pretrained(hf_model_name) hf_tokenizer = AutoTokenizer.from_pretrained(hf_model_name)
hf_model = AutoModelForCausalLM.from_pretrained(hf_model_name) hf_model = AutoModelForCausalLM.from_pretrained(hf_model_name)
# Подготавливаем входные данные # Подготавливаем входные данные
inputs = hf_tokenizer(test_input, return_tensors="pt") inputs = hf_tokenizer(test_input, return_tensors="pt")
# Получаем логиты от обеих моделей # Получаем логиты от обеих моделей
with torch.no_grad(): with torch.no_grad():
hf_logits = hf_model(**inputs).logits 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) hf_probs = torch.softmax(hf_logits[0, -1], dim=-1)
llm_probs = torch.softmax(llm_logits[0, -1], dim=-1) llm_probs = torch.softmax(llm_logits[0, -1], dim=-1)
# Вычисляем метрики # Вычисляем метрики
kl_divergence = torch.nn.functional.kl_div( kl_divergence = torch.nn.functional.kl_div(
torch.log(llm_probs + 1e-8), torch.log(llm_probs + 1e-8), hf_probs, reduction="batchmean"
hf_probs,
reduction='batchmean'
) )
cosine_similarity = torch.nn.functional.cosine_similarity( cosine_similarity = torch.nn.functional.cosine_similarity(
hf_logits.flatten(), hf_logits.flatten(), llm_logits.flatten(), dim=0
llm_logits.flatten(),
dim=0
) )
return { return {
"kl_divergence": kl_divergence.item(), "kl_divergence": kl_divergence.item(),
"cosine_similarity": cosine_similarity.item(), "cosine_similarity": cosine_similarity.item(),
@@ -228,58 +220,52 @@ class TokenizerWrapper:
""" """
Обертка для токенизатора с дополнительными утилитами. Обертка для токенизатора с дополнительными утилитами.
""" """
def __init__(self, tokenizer): def __init__(self, tokenizer):
self.tokenizer = tokenizer self.tokenizer = tokenizer
def encode_batch(self, texts: List[str], **kwargs) -> Dict[str, torch.Tensor]: def encode_batch(self, texts: List[str], **kwargs) -> Dict[str, torch.Tensor]:
""" """
Кодирует батч текстов. Кодирует батч текстов.
Args: Args:
texts: Список текстов texts: Список текстов
**kwargs: Дополнительные параметры токенизации **kwargs: Дополнительные параметры токенизации
Returns: Returns:
Dict: Токенизированные данные Dict: Токенизированные данные
""" """
return self.tokenizer( return self.tokenizer(
texts, texts, padding=True, truncation=True, return_tensors="pt", **kwargs
padding=True,
truncation=True,
return_tensors="pt",
**kwargs
) )
def decode_batch(self, token_ids: torch.Tensor, **kwargs) -> List[str]: def decode_batch(self, token_ids: torch.Tensor, **kwargs) -> List[str]:
""" """
Декодирует батч токенов. Декодирует батч токенов.
Args: Args:
token_ids: Тензор с токенами token_ids: Тензор с токенами
**kwargs: Дополнительные параметры декодирования **kwargs: Дополнительные параметры декодирования
Returns: Returns:
List[str]: Декодированные тексты List[str]: Декодированные тексты
""" """
if token_ids.dim() == 1: if token_ids.dim() == 1:
token_ids = token_ids.unsqueeze(0) token_ids = token_ids.unsqueeze(0)
texts = [] texts = []
for i in range(token_ids.size(0)): for i in range(token_ids.size(0)):
text = self.tokenizer.decode( text = self.tokenizer.decode(
token_ids[i], token_ids[i], skip_special_tokens=True, **kwargs
skip_special_tokens=True,
**kwargs
) )
texts.append(text) texts.append(text)
return texts return texts
def get_vocab_size(self) -> int: def get_vocab_size(self) -> int:
"""Возвращает размер словаря.""" """Возвращает размер словаря."""
return len(self.tokenizer) return len(self.tokenizer)
def get_special_tokens(self) -> Dict[str, int]: def get_special_tokens(self) -> Dict[str, int]:
"""Возвращает специальные токены.""" """Возвращает специальные токены."""
return { return {
@@ -290,36 +276,27 @@ class TokenizerWrapper:
} }
def create_hf_pipeline( def create_hf_pipeline(llm_model, tokenizer=None, device: str = "auto", **kwargs):
llm_model,
tokenizer=None,
device: str = "auto",
**kwargs
):
""" """
Создает HuggingFace pipeline из llm модели. Создает HuggingFace pipeline из llm модели.
Args: Args:
llm_model: Модель из библиотеки llm llm_model: Модель из библиотеки llm
tokenizer: Токенизатор tokenizer: Токенизатор
device: Устройство для вычислений device: Устройство для вычислений
**kwargs: Дополнительные параметры pipeline **kwargs: Дополнительные параметры pipeline
Returns: Returns:
transformers.Pipeline: Готовый pipeline transformers.Pipeline: Готовый pipeline
""" """
from transformers import pipeline from transformers import pipeline
# Конвертируем модель в HF формат # Конвертируем модель в HF формат
hf_model, tokenizer = HFUtils.convert_to_hf_format(llm_model, tokenizer) hf_model, tokenizer = HFUtils.convert_to_hf_format(llm_model, tokenizer)
# Создаем pipeline # Создаем pipeline
pipe = pipeline( pipe = pipeline(
"text-generation", "text-generation", model=hf_model, tokenizer=tokenizer, device=device, **kwargs
model=hf_model,
tokenizer=tokenizer,
device=device,
**kwargs
) )
return pipe 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 from typing import Optional, Tuple
import torch import torch
class BaseModel(nn.Module, ABC): class BaseModel(nn.Module, ABC):
""" """
Абстрактный класс — стандарт для всех архитектур LLM. Абстрактный класс — стандарт для всех архитектур LLM.
Научная идея: Научная идея:
Реализация унифицированного входа/выхода для поддержки построения и обучения любых современных языковых моделей. Реализация унифицированного входа/выхода для поддержки построения и обучения любых современных языковых моделей.
Args: Args:
config (dict): Параметры архитектуры (размерность эмбеддингов, число слоев, heads и т.д.) config (dict): Параметры архитектуры (размерность эмбеддингов, число слоев, heads и т.д.)
Attributes: Attributes:
config (dict): Конфиг модели config (dict): Конфиг модели
""" """
def __init__(self, config: dict): def __init__(self, config: dict):
""" """
Инициализация модели. Инициализация модели.
Args: Args:
config (dict): Настройки архитектуры модели (размеры слоев, типы блоков и т.д.) config (dict): Настройки архитектуры модели (размеры слоев, типы блоков и т.д.)
""" """
@@ -43,10 +45,12 @@ class BaseModel(nn.Module, ABC):
self.config = config self.config = config
@abstractmethod @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: Args:
input_ids (Tensor[int]): Индексы токенов [batch, seq_len] input_ids (Tensor[int]): Индексы токенов [batch, seq_len]
attention_mask (Optional[Tensor[bool]]): Маска разрешенных позиций (если требуется) [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: def generate(self, input_ids: torch.Tensor, max_length: int = 50) -> torch.Tensor:
""" """
Генерация текста (авторегрессивно, greedy или sampling). Генерация текста (авторегрессивно, greedy или sampling).
Args: Args:
input_ids (Tensor[int]): Начальные токены [batch, start_len] input_ids (Tensor[int]): Начальные токены [batch, start_len]
max_length (int): Максимальная длина последовательности 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 .multi_head_attention import MultiHeadAttention
from .rope import RoPE from .rope import RoPE
class CachedDecoder(nn.Module): class CachedDecoder(nn.Module):
""" """
Универсальный декодерный блок для современных LLM (GPT, LLaMA, др.), поддерживает кэширование key-value для эффективной генерации. Универсальный декодерный блок для современных LLM (GPT, LLaMA, др.), поддерживает кэширование key-value для эффективной генерации.
@@ -28,7 +29,7 @@ class CachedDecoder(nn.Module):
norm_layer (тип nn.Module): Normalization слой (LayerNorm или RMSNorm) norm_layer (тип nn.Module): Normalization слой (LayerNorm или RMSNorm)
dropout (float): Dropout dropout (float): Dropout
rope (RoPE|None): Экземпляр RoPE (для LLaMA) rope (RoPE|None): Экземпляр RoPE (для LLaMA)
Пример (GPT2 style): Пример (GPT2 style):
>>> decoder = CachedDecoder( >>> decoder = CachedDecoder(
... feed_forward_layer=FeedForward(...), ... 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) ... num_heads=4, emb_size=256, head_size=64, max_seq_len=128)
>>> out, cache = decoder(x, use_cache=True) >>> out, cache = decoder(x, use_cache=True)
""" """
def __init__( def __init__(
self, self,
feed_forward_layer: nn.Module, feed_forward_layer: nn.Module,
@@ -49,7 +51,7 @@ class CachedDecoder(nn.Module):
): ):
""" """
Инициализация декодера с кэшированием. Инициализация декодера с кэшированием.
Поведение аналогично блоку TransformerDecoderLayer, Поведение аналогично блоку TransformerDecoderLayer,
но с гибкой возможностью подмены любых подкомпонент (активация, norm, позиции). но с гибкой возможностью подмены любых подкомпонент (активация, norm, позиции).
@@ -85,7 +87,7 @@ class CachedDecoder(nn.Module):
): ):
""" """
Прямой проход с поддержкой кэша. Прямой проход с поддержкой кэша.
Args: Args:
x (Tensor[float]): [batch, seq_len, emb_size] — скрытые состояния x (Tensor[float]): [batch, seq_len, emb_size] — скрытые состояния
mask (Optional[Tensor]): маска внимания (или causal mask), shape [seq_len, seq_len] mask (Optional[Tensor]): маска внимания (или causal mask), shape [seq_len, seq_len]
@@ -111,4 +113,4 @@ class CachedDecoder(nn.Module):
if use_cache: if use_cache:
return (result, kv_caches) return (result, kv_caches)
else: 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 .feed_forward import FeedForward
from .multi_head_attention import MultiHeadAttention from .multi_head_attention import MultiHeadAttention
class Decoder(nn.Module): class Decoder(nn.Module):
""" """
Базовый автогерессивный блок-декодер трансформера (без кэша KV). Базовый автогерессивный блок-декодер трансформера (без кэша KV).
@@ -24,12 +25,14 @@ class Decoder(nn.Module):
>>> out = decoder(x) >>> out = decoder(x)
>>> print(out.shape) # torch.Size([1, 10, 512]) >>> print(out.shape) # torch.Size([1, 10, 512])
""" """
def __init__(self,
def __init__(
self,
num_heads: int, num_heads: int,
emb_size: int, emb_size: int,
head_size: int, head_size: int,
max_seq_len: int, max_seq_len: int,
dropout: float = 0.1 dropout: float = 0.1,
): ):
""" """
Инициализация декодера. Инициализация декодера.
@@ -43,11 +46,11 @@ class Decoder(nn.Module):
""" """
super().__init__() super().__init__()
self._heads = MultiHeadAttention( self._heads = MultiHeadAttention(
num_heads=num_heads, num_heads=num_heads,
emb_size=emb_size, emb_size=emb_size,
head_size=head_size, head_size=head_size,
max_seq_len=max_seq_len, max_seq_len=max_seq_len,
dropout=dropout dropout=dropout,
) )
self._ff = FeedForward(emb_size=emb_size, dropout=dropout) self._ff = FeedForward(emb_size=emb_size, dropout=dropout)
self._norm1 = nn.LayerNorm(emb_size) self._norm1 = nn.LayerNorm(emb_size)
@@ -73,7 +76,7 @@ class Decoder(nn.Module):
# Self-Attention блок # Self-Attention блок
attention, _ = self._heads(x, mask, use_cache=False, cache=None) attention, _ = self._heads(x, mask, use_cache=False, cache=None)
out = self._norm1(attention + x) out = self._norm1(attention + x)
# FeedForward блок # FeedForward блок
ffn_out = self._ff(out) 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). - Изначально предложен в «Attention is All You Need» (Vaswani et al., 2017).
Формула: Формула:
FFN(x) = Dropout(W2·act(W1·x)) FFN(x) = Dropout(W2·act(W1·x))
где act — ReLU, GELU и др., обычно expansion x4. где act — ReLU, GELU и др., обычно expansion x4.
@@ -33,22 +33,23 @@ class FeedForward(nn.Module):
- Добавляет нелинейность в архитектуру трансформера - Добавляет нелинейность в архитектуру трансформера
- Обеспечивает взаимодействие между различными размерностями эмбеддингов - Обеспечивает взаимодействие между различными размерностями эмбеддингов
- Работает независимо для каждого токена в последовательности - Работает независимо для каждого токена в последовательности
Args: Args:
emb_size (int): размерность входных эмбеддингов emb_size (int): размерность входных эмбеддингов
dropout (float): вероятность(dropout) dropout (float): вероятность(dropout)
activation (str): нелинейная функция (relu, gelu, gelu_exact) activation (str): нелинейная функция (relu, gelu, gelu_exact)
Пример: Пример:
>>> ff = FeedForward(emb_size=512, dropout=0.1) >>> ff = FeedForward(emb_size=512, dropout=0.1)
>>> x = torch.randn(32, 10, 512) >>> x = torch.randn(32, 10, 512)
>>> output = ff(x) >>> output = ff(x)
>>> print(output.shape) # torch.Size([32, 10, 512]) >>> print(output.shape) # torch.Size([32, 10, 512])
""" """
def __init__(self, emb_size: int, dropout: float = 0.1, activation: str = "relu"): def __init__(self, emb_size: int, dropout: float = 0.1, activation: str = "relu"):
""" """
Инициализация слоя Feed Forward Network. Инициализация слоя Feed Forward Network.
Args: Args:
emb_size: Размерность входных эмбеддингов emb_size: Размерность входных эмбеддингов
dropout: Вероятность dropout для регуляризации (по умолчанию: 0.1) dropout: Вероятность dropout для регуляризации (по умолчанию: 0.1)
@@ -73,23 +74,23 @@ class FeedForward(nn.Module):
def forward(self, x: torch.Tensor): def forward(self, x: torch.Tensor):
""" """
Прямой проход через слой Feed Forward Network. Прямой проход через слой Feed Forward Network.
Args: Args:
x: Входной тензор размерности [batch_size, seq_len, emb_size] x: Входной тензор размерности [batch_size, seq_len, emb_size]
Returns: Returns:
Тензор той же размерности, что и входной Тензор той же размерности, что и входной
""" """
# Сохраняем dtype входных данных # Сохраняем dtype входных данных
input_dtype = x.dtype input_dtype = x.dtype
# Приводим веса к нужному типу если необходимо # Приводим веса к нужному типу если необходимо
if input_dtype != self._layer1.weight.dtype: if input_dtype != self._layer1.weight.dtype:
self._layer1 = self._layer1.to(dtype=input_dtype) self._layer1 = self._layer1.to(dtype=input_dtype)
self._layer2 = self._layer2.to(dtype=input_dtype) self._layer2 = self._layer2.to(dtype=input_dtype)
# Пропустим тензор x по очереди через все созданные слои # Пропустим тензор x по очереди через все созданные слои
x = self._layer1(x) x = self._layer1(x)
x = self._activation(x) x = self._activation(x)
x = self._layer2(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 import torch
from torch import nn from torch import nn
class GELU(nn.Module): class GELU(nn.Module):
""" """
Гауссовская Эрф-активация (GELU, Gaussian Error Linear Unit). Гауссовская Эрф-активация (GELU, Gaussian Error Linear Unit).
@@ -17,11 +18,14 @@ class GELU(nn.Module):
>>> y = gelu(torch.tensor([-1.0, 0.0, 1.0])) >>> y = gelu(torch.tensor([-1.0, 0.0, 1.0]))
>>> print(y) >>> print(y)
""" """
def __init__(self): def __init__(self):
super().__init__() super().__init__()
self.sqrt_2_over_pi = torch.sqrt(torch.tensor(2.0) / math.pi) self.sqrt_2_over_pi = torch.sqrt(torch.tensor(2.0) / math.pi)
def forward(self, x: torch.Tensor) -> torch.Tensor: def forward(self, x: torch.Tensor) -> torch.Tensor:
return 0.5 * x * (1 + torch.tanh( return (
self.sqrt_2_over_pi * (x + 0.044715 * torch.pow(x, 3)) 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 math import sqrt
from .rope import RoPE from .rope import RoPE
class HeadAttention(nn.Module): class HeadAttention(nn.Module):
""" """
Одноголовый механизм внимания (scaled dot-product attention) — фундаментальный строительный блок всех современных Transformer. Одноголовый механизм внимания (scaled dot-product attention) — фундаментальный строительный блок всех современных Transformer.
@@ -11,11 +12,11 @@ class HeadAttention(nn.Module):
Научная суть: Научная суть:
- Attention учит модель самостоятельно "выбирать" важные связи между словами, независимо от их положения. - Attention учит модель самостоятельно "выбирать" важные связи между словами, независимо от их положения.
- Механизм causal mask гарантирует невозможность "заглядывания в будущее" при генерации (авторегрессия). - Механизм causal mask гарантирует невозможность "заглядывания в будущее" при генерации (авторегрессия).
Формула: Формула:
Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) · V Attention(Q, K, V) = softmax(QK^T / sqrt(d_k)) · V
(Q — запросы, K — ключи, V — значения; d_k — размерность ключа) (Q — запросы, K — ключи, V — значения; d_k — размерность ключа)
Поддерживает Rotary Position Encoding (RoPE) для относительного позиционного кодирования. Поддерживает Rotary Position Encoding (RoPE) для относительного позиционного кодирования.
Args: Args:
@@ -28,14 +29,17 @@ class HeadAttention(nn.Module):
- Использует нижнетреугольную маску для предотвращения "заглядывания в будущее" - Использует нижнетреугольную маску для предотвращения "заглядывания в будущее"
- Автоматически адаптируется к разным версиям PyTorch - Автоматически адаптируется к разным версиям PyTorch
- Поддерживает batch-обработку входных данных - Поддерживает batch-обработку входных данных
Пример использования: Пример использования:
>>> attention = HeadAttention(emb_size=64, head_size=32, max_seq_len=128) >>> attention = HeadAttention(emb_size=64, head_size=32, max_seq_len=128)
>>> x = torch.randn(1, 10, 64) >>> x = torch.randn(1, 10, 64)
>>> output, _ = attention(x) >>> output, _ = attention(x)
>>> print(output.shape) # torch.Size([1, 10, 32]) >>> 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__() super().__init__()
self._emb_size = emb_size self._emb_size = emb_size
self._head_size = head_size self._head_size = head_size
@@ -49,21 +53,25 @@ class HeadAttention(nn.Module):
# Создание causal маски # Создание causal маски
mask = torch.tril(torch.ones(max_seq_len, max_seq_len)) 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] x (torch.Tensor): Входной тензор формы [batch_size, seq_len, emb_size]
Возвращает: Возвращает:
torch.Tensor: Выходной тензор формы [batch_size, seq_len, head_size] torch.Tensor: Выходной тензор формы [batch_size, seq_len, head_size]
Исключения: Исключения:
ValueError: Если длина последовательности превышает max_seq_len ValueError: Если длина последовательности превышает max_seq_len
Пример внутренних преобразований: Пример внутренних преобразований:
Для входа x.shape = [2, 5, 64]: Для входа x.shape = [2, 5, 64]:
1. Q/K/V преобразования -> [2, 5, 32] 1. Q/K/V преобразования -> [2, 5, 32]
@@ -73,7 +81,9 @@ class HeadAttention(nn.Module):
""" """
seq_len = x.shape[1] seq_len = x.shape[1]
if seq_len > self._max_seq_len: 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] k = self._k(x) # [B, T, hs]
q = self._q(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_cache, v_cache = cache
k = torch.cat([k_cache, k], dim=1) # [B, cache_len + T, hs] 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] v = torch.cat([v_cache, v], dim=1) # [B, cache_len + T, hs]
scores = q @ k.transpose(-2, -1) / sqrt(self._head_size) scores = q @ k.transpose(-2, -1) / sqrt(self._head_size)
if cache is None: 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) weights = F.softmax(scores, dim=-1)
x_out = weights @ v # [B, T, hs] x_out = weights @ v # [B, T, hs]
if use_cache is True: if use_cache is True:
return (x_out, (k, v)) return (x_out, (k, v))
else: 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 .head_attention import HeadAttention
from .rope import RoPE from .rope import RoPE
class MultiHeadAttention(nn.Module): class MultiHeadAttention(nn.Module):
""" """
Мультиголовый (многоголовый) механизм внимания — ключевой компонент любого Transformer. Мультиголовый (многоголовый) механизм внимания — ключевой компонент любого Transformer.
@@ -12,7 +13,7 @@ class MultiHeadAttention(nn.Module):
чтобы видеть разные связи в последовательности (разный контекст, локально/глобально). чтобы видеть разные связи в последовательности (разный контекст, локально/глобально).
- Каждый attention блок работает независимо, выход конкатенируется. - Каждый attention блок работает независимо, выход конкатенируется.
- Механизм предложен в статье "Attention is All You Need" (Vaswani et al., 2017). - Механизм предложен в статье "Attention is All You Need" (Vaswani et al., 2017).
Формула внимания для одной головы: Формула внимания для одной головы:
Attention(Q, K, V) = softmax(QK^T/sqrt(d_k))·V Attention(Q, K, V) = softmax(QK^T/sqrt(d_k))·V
Мультиголовый: Мультиголовый:
@@ -32,7 +33,16 @@ class MultiHeadAttention(nn.Module):
>>> out, cache = mha(x) >>> out, cache = mha(x)
>>> print(out.shape) >>> 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) - max_seq_len зависит от задачи (512 для BERT, 2048 для GPT-3)
""" """
super().__init__() super().__init__()
self._heads = nn.ModuleList([ self._heads = nn.ModuleList(
HeadAttention( [
emb_size=emb_size, HeadAttention(
head_size=head_size, emb_size=emb_size,
max_seq_len=max_seq_len, head_size=head_size,
rope=rope, max_seq_len=max_seq_len,
) for _ in range(num_heads) rope=rope,
]) )
for _ in range(num_heads)
]
)
self._layer = nn.Linear(head_size * num_heads, emb_size) self._layer = nn.Linear(head_size * num_heads, emb_size)
self._dropout = nn.Dropout(dropout) 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): Прямой проход (forward):
Для каждого токена оценивает "важность" остальных токенов сразу через несколько attention-блоков. Для каждого токена оценивает "важность" остальных токенов сразу через несколько attention-блоков.
@@ -76,7 +95,7 @@ class MultiHeadAttention(nn.Module):
4. Линейная проекция: 4. Линейная проекция:
- Выход: [batch_size, seq_len, emb_size] - Выход: [batch_size, seq_len, emb_size]
5. Применение dropout 5. Применение dropout
Args: Args:
x (Tensor[float]): [batch, seq_len, emb_size] — вход x (Tensor[float]): [batch, seq_len, emb_size] — вход
mask (Optional[Tensor[bool]]): маска позиции [seq_len, seq_len] mask (Optional[Tensor[bool]]): маска позиции [seq_len, seq_len]
@@ -94,11 +113,11 @@ class MultiHeadAttention(nn.Module):
Пример преобразований для emb_size=512, num_heads=8: Пример преобразований для emb_size=512, num_heads=8:
Вход: [4, 100, 512] Вход: [4, 100, 512]
-> Каждая голова: [4, 100, 64] -> Каждая голова: [4, 100, 64]
-> После внимания: 8 x [4, 100, 64] -> После внимания: 8 x [4, 100, 64]
-> Конкатенация: [4, 100, 512] -> Конкатенация: [4, 100, 512]
-> Проекция: [4, 100, 512] -> Проекция: [4, 100, 512]
-> Dropout: [4, 100, 512] -> Dropout: [4, 100, 512]
Пример: Пример:
>>> out, caches = mha(x) >>> out, caches = mha(x)
>>> out.shape # [batch, seq_len, emb_size] >>> 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 head_cache = cache[i] if cache is not None else None
result = head(x, use_cache=use_cache, cache=head_cache) result = head(x, use_cache=use_cache, cache=head_cache)
attention_results.append(result) attention_results.append(result)
outputs, caches = zip(*attention_results) outputs, caches = zip(*attention_results)
attention_outputs = list(outputs) attention_outputs = list(outputs)
kv_caches = list(caches) kv_caches = list(caches)
# 2. Объединяем результаты всех голов # 2. Объединяем результаты всех голов
concatenated_attention = torch.cat(attention_outputs, dim=-1) concatenated_attention = torch.cat(attention_outputs, dim=-1)
# 3. Проецируем в пространство эмбеддингов # 3. Проецируем в пространство эмбеддингов
projected_output = self._layer(concatenated_attention) projected_output = self._layer(concatenated_attention)
# 4. Применяем dropout для регуляризации # 4. Применяем dropout для регуляризации
final_output = self._dropout(projected_output) final_output = self._dropout(projected_output)
if use_cache is True: if use_cache is True:
return (final_output, kv_caches) return (final_output, kv_caches)
else: else:

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

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

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

@@ -33,23 +33,23 @@ class RMSNorm(nn.Module):
- Упрощенный вариант LayerNorm без вычисления среднего, только деление на rms. - Упрощенный вариант LayerNorm без вычисления среднего, только деление на rms.
- Лучшая численная стабильность на больших моделях, меньше вычислений. - Лучшая численная стабильность на больших моделях, меньше вычислений.
- Применяется в LLaMA, PaLM и др. - Применяется в LLaMA, PaLM и др.
Формула: Формула:
RMSNorm(x) = (x / sqrt(mean(x²) + eps)) * w (w — обучаемый вектор) RMSNorm(x) = (x / sqrt(mean(x²) + eps)) * w (w — обучаемый вектор)
Args: Args:
dim (int): размер последнего измерения (обычно emb_size) dim (int): размер последнего измерения (обычно emb_size)
eps (float): для численной устойчивости eps (float): для численной устойчивости
Пример: Пример:
>>> norm = RMSNorm(emb_size) >>> norm = RMSNorm(emb_size)
>>> out = norm(x) >>> out = norm(x)
""" """
def __init__(self, dim: int, eps: float = 1e-6): def __init__(self, dim: int, eps: float = 1e-6):
""" """
Инициализация RMSNorm слоя. Инициализация RMSNorm слоя.
Args: Args:
dim: Размерность нормализуемого измерения dim: Размерность нормализуемого измерения
eps: Малое значение для численной стабильности (по умолчанию 1e-6) eps: Малое значение для численной стабильности (по умолчанию 1e-6)
@@ -57,27 +57,27 @@ class RMSNorm(nn.Module):
super().__init__() super().__init__()
self._eps = eps self._eps = eps
self._w = nn.Parameter(torch.ones(dim)) self._w = nn.Parameter(torch.ones(dim))
def forward(self, x: torch.Tensor) -> torch.Tensor: def forward(self, x: torch.Tensor) -> torch.Tensor:
""" """
Прямой проход через RMSNorm слой. Прямой проход через RMSNorm слой.
Args: Args:
x: Входной тензор формы [..., dim] x: Входной тензор формы [..., dim]
Returns: Returns:
Нормализованный тензор той же формы, что и входной Нормализованный тензор той же формы, что и входной
Формула: Формула:
output = w * (x / sqrt(mean(x²) + eps)) output = w * (x / sqrt(mean(x²) + eps))
""" """
# Вычисление RMS (Root Mean Square) по последнему измерению # Вычисление RMS (Root Mean Square) по последнему измерению
rms = (x.pow(2).mean(-1, keepdim=True) + self._eps) ** 0.5 rms = (x.pow(2).mean(-1, keepdim=True) + self._eps) ** 0.5
# Нормализация и масштабирование # Нормализация и масштабирование
norm_x = x / rms norm_x = x / rms
return self._w * norm_x return self._w * norm_x
def extra_repr(self) -> str: 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): class RoPE(nn.Module):
""" """
Rotary Positional Embeddings (RoPE) для механизма внимания. Rotary Positional Embeddings (RoPE) для механизма внимания.
Кодирует позиционную информацию через вращение векторов запросов и ключей Кодирует позиционную информацию через вращение векторов запросов и ключей
в многомерном пространстве с использованием синусов и косинусов. в многомерном пространстве с использованием синусов и косинусов.
Args: Args:
head_size: Размерность головы внимания (должен быть четным) head_size: Размерность головы внимания (должен быть четным)
max_seq_len: Максимальная длина последовательности max_seq_len: Максимальная длина последовательности
base: Базовое значение для вычисления частот (по умолчанию 10000) base: Базовое значение для вычисления частот (по умолчанию 10000)
Attributes: Attributes:
cos_matrix: Буферизованная матрица косинусов формы [max_seq_len, head_size//2] cos_matrix: Буферизованная матрица косинусов формы [max_seq_len, head_size//2]
sin_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): def __init__(self, head_size: int, max_seq_len: int, base: int = 10_000):
""" """
Инициализация RoPE эмбеддингов. Инициализация RoPE эмбеддингов.
Args: Args:
head_size: Размерность головы внимания (должен быть четным) head_size: Размерность головы внимания (должен быть четным)
max_seq_len: Максимальная поддерживаемая длина последовательности max_seq_len: Максимальная поддерживаемая длина последовательности
base: Базовое значение для вычисления частот (типично 10000) base: Базовое значение для вычисления частот (типично 10000)
Raises: Raises:
AssertionError: Если head_size не четный AssertionError: Если head_size не четный
""" """
super().__init__() super().__init__()
assert head_size % 2 == 0, "head_size должен быть четным" assert head_size % 2 == 0, "head_size должен быть четным"
# Вычисление частот: θ_i = base^(-2i/d) для i ∈ [0, d/2-1] # Вычисление частот: θ_i = base^(-2i/d) для i ∈ [0, d/2-1]
freqs = 1.0 / (base ** (2 * torch.arange(head_size // 2).float() / head_size)) freqs = 1.0 / (base ** (2 * torch.arange(head_size // 2).float() / head_size))
# Позиции от 0 до max_seq_len-1 # Позиции от 0 до max_seq_len-1
positions = torch.arange(max_seq_len).float() positions = torch.arange(max_seq_len).float()
# Внешнее произведение: m * θ_i для всех позиций и частот # Внешнее произведение: m * θ_i для всех позиций и частот
freq_matrix = positions.unsqueeze(1) * freqs.unsqueeze(0) freq_matrix = positions.unsqueeze(1) * freqs.unsqueeze(0)
# Предвычисление матриц косинусов и синусов # Предвычисление матриц косинусов и синусов
self.register_buffer('cos_matrix', torch.cos(freq_matrix)) self.register_buffer("cos_matrix", torch.cos(freq_matrix))
self.register_buffer('sin_matrix', torch.sin(freq_matrix)) self.register_buffer("sin_matrix", torch.sin(freq_matrix))
def forward(self, x: torch.Tensor) -> torch.Tensor: def forward(self, x: torch.Tensor) -> torch.Tensor:
""" """
Применение ротационного позиционного кодирования к входному тензору. Применение ротационного позиционного кодирования к входному тензору.
Args: Args:
x: Входной тензор формы [batch_size, seq_len, head_size] x: Входной тензор формы [batch_size, seq_len, head_size]
Returns: Returns:
Тензор с примененным RoPE формы [batch_size, seq_len, head_size] Тензор с примененным RoPE формы [batch_size, seq_len, head_size]
Алгоритм: Алгоритм:
1. Разделение векторов на четные и нечетные компоненты 1. Разделение векторов на четные и нечетные компоненты
2. Применение вращения через синусы и косинусы 2. Применение вращения через синусы и косинусы
3. Объединение компонент обратно 3. Объединение компонент обратно
""" """
seq_len = x.size(1) seq_len = x.size(1)
# Берем нужную часть матриц и приводим к типу x # Берем нужную часть матриц и приводим к типу x
cos = self.cos_matrix[:seq_len].to(x.dtype) # [seq_len, head_size//2] 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] 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_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θ) # Применяем поворот: q' = q * cos(mθ) + rotate(q) * sin(mθ)
x_rotated_even = x_even * cos - x_odd * sin 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 = torch.stack([x_rotated_even, x_rotated_odd], dim=-1)
x_rotated = x_rotated.flatten(-2) # [batch_size, seq_len, head_size] 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 import torch
from torch import nn from torch import nn
class SiLU(nn.Module): class SiLU(nn.Module):
""" """
SiLU (Swish) — современная активационная функция для нейросетей. SiLU (Swish) — современная активационная функция для нейросетей.
Научная суть: Научная суть:
- Формула: $SiLU(x) = x * \sigm(x)$, где $\sigm(x)$ — сигмоида. - Формула: $SiLU(x) = x * \sigm(x)$, где $\sigm(x)$ — сигмоида.
- Более гладкая альтернатива ReLU, улучшает поток градиентов в глубоких сетях. - Более гладкая альтернатива ReLU, улучшает поток градиентов в глубоких сетях.
- Используется во многих «state-of-the-art» архитектурах (SwiGLU, PaLM, LLaMA). - Используется во многих «state-of-the-art» архитектурах (SwiGLU, PaLM, LLaMA).
- Также известна как Swish (Ramachandran et al, 2017). - Также известна как Swish (Ramachandran et al, 2017).
Пример: Пример:
>>> act = SiLU() >>> act = SiLU()
>>> x = torch.tensor([-1.0, 0.0, 1.0]) >>> x = torch.tensor([-1.0, 0.0, 1.0])
>>> print(act(x)) >>> print(act(x))
""" """
def forward(self, x: torch.Tensor): 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 (Swish-Gated Linear Unit) — современная нелинейность для архитектур LLM (LLaMA, PaLM).
Реализация SwiGLU активационной функции. Реализация SwiGLU активационной функции.
Состоит из трех линейных слоев и активации SiLU: Состоит из трех линейных слоев и активации SiLU:
1. Gate слой + SiLU активация 1. Gate слой + SiLU активация
2. Up слой (линейное преобразование) 2. Up слой (линейное преобразование)
3. Element-wise multiplication gate и up 3. Element-wise multiplication gate и up
4. Down слой (линейная проекция) 4. Down слой (линейная проекция)
Научная суть: Научная суть:
- Сохраняет преимущества GLU (раздельные гейтом и телом) + мощность Swish/SiLU активации. - Сохраняет преимущества GLU (раздельные гейтом и телом) + мощность Swish/SiLU активации.
- Дает надежную гладкую активацию, хорошо работает на больших масштабах. - Дает надежную гладкую активацию, хорошо работает на больших масштабах.
@@ -50,11 +50,11 @@ class SwiGLU(nn.Module):
>>> ff = SwiGLU(emb_size=512, dropout=0.1) >>> ff = SwiGLU(emb_size=512, dropout=0.1)
>>> y = ff(torch.randn(2,10,512)) >>> y = ff(torch.randn(2,10,512))
""" """
def __init__(self, emb_size: int, dropout: float = 0.1): def __init__(self, emb_size: int, dropout: float = 0.1):
""" """
Инициализация SwiGLU слоя. Инициализация SwiGLU слоя.
Args: Args:
emb_size: Размерность входных/выходных эмбеддингов emb_size: Размерность входных/выходных эмбеддингов
dropout: Вероятность dropout (по умолчанию 0.1) dropout: Вероятность dropout (по умолчанию 0.1)
@@ -69,13 +69,13 @@ class SwiGLU(nn.Module):
def forward(self, x: torch.Tensor) -> torch.Tensor: def forward(self, x: torch.Tensor) -> torch.Tensor:
""" """
Прямой проход через SwiGLU слой. Прямой проход через SwiGLU слой.
Args: Args:
x: Входной тензор формы [batch_size, seq_len, emb_size] x: Входной тензор формы [batch_size, seq_len, emb_size]
Returns: Returns:
Выходной тензор формы [batch_size, seq_len, emb_size] Выходной тензор формы [batch_size, seq_len, emb_size]
Алгоритм: Алгоритм:
1. gate = SiLU(linear_gate(x)) 1. gate = SiLU(linear_gate(x))
2. up = linear_up(x) 2. up = linear_up(x)
@@ -83,19 +83,19 @@ class SwiGLU(nn.Module):
4. apply dropout 4. apply dropout
""" """
# Gate ветвь: линейное преобразование + активация # Gate ветвь: линейное преобразование + активация
gate_out = self._gate(x) # [batch, seq, 4*emb] gate_out = self._gate(x) # [batch, seq, 4*emb]
activation_out = self._activation(gate_out) # [batch, seq, 4*emb] activation_out = self._activation(gate_out) # [batch, seq, 4*emb]
# Up ветвь: линейное преобразование # Up ветвь: линейное преобразование
up_out = self._up(x) # [batch, seq, 4*emb] up_out = self._up(x) # [batch, seq, 4*emb]
# Element-wise multiplication (gating mechanism) # Element-wise multiplication (gating mechanism)
out = up_out * activation_out # поэлементное умножение! out = up_out * activation_out # поэлементное умножение!
# Final projection and dropout # Final projection and dropout
out = self._down(out) # [batch, seq, emb] out = self._down(out) # [batch, seq, emb]
return self._dropout(out) return self._dropout(out)
def extra_repr(self) -> str: 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 nn
from torch import Tensor from torch import Tensor
class TokenEmbeddings(nn.Module): class TokenEmbeddings(nn.Module):
""" """
Токеновые эмбеддинги — обучаемые векторные представления для каждого токена словаря. Токеновые эмбеддинги — обучаемые векторные представления для каждого токена словаря.
Преобразует целочисленные индексы токенов в обучаемые векторные представления фиксированного размера. Преобразует целочисленные индексы токенов в обучаемые векторные представления фиксированного размера.
Обычно используется как первый слой в нейронных сетях для задач NLP. Обычно используется как первый слой в нейронных сетях для задач NLP.
Научная суть: Научная суть:
- Первый шаг для любого NLP-модуля: вместо индекса токена подаём его dense-вектор. - Первый шаг для любого NLP-модуля: вместо индекса токена подаём его dense-вектор.
- Эти вектора изучаются в процессе обучения и отражают скрытые взаимосвязи между токенами. - Эти вектора изучаются в процессе обучения и отражают скрытые взаимосвязи между токенами.
@@ -22,18 +23,18 @@ class TokenEmbeddings(nn.Module):
Примечание: Примечание:
- Индексы должны быть в диапазоне [0, vocab_size-1] - Индексы должны быть в диапазоне [0, vocab_size-1]
- Эмбеддинги инициализируются случайно и обучаются в процессе тренировки модели - Эмбеддинги инициализируются случайно и обучаются в процессе тренировки модели
Пример: Пример:
>>> emb = TokenEmbeddings(vocab_size=10000, emb_size=256) >>> emb = TokenEmbeddings(vocab_size=10000, emb_size=256)
>>> tokens = torch.tensor([[1, 2, 3]]) >>> tokens = torch.tensor([[1, 2, 3]])
>>> vecs = emb(tokens) >>> vecs = emb(tokens)
>>> vecs.shape # torch.Size([1, 3, 256]) >>> vecs.shape # torch.Size([1, 3, 256])
""" """
def __init__(self, vocab_size: int, emb_size: int): def __init__(self, vocab_size: int, emb_size: int):
super().__init__() super().__init__()
self._embedding = nn.Embedding( self._embedding = nn.Embedding(
num_embeddings=vocab_size, num_embeddings=vocab_size, embedding_dim=emb_size
embedding_dim=emb_size
) )
def forward(self, x: Tensor) -> Tensor: def forward(self, x: Tensor) -> Tensor:
@@ -55,10 +56,7 @@ if __name__ == "__main__":
embedding = TokenEmbeddings(vocab_size=100, emb_size=128) embedding = TokenEmbeddings(vocab_size=100, emb_size=128)
# Создаем тензор с индексами в пределах vocab_size (0-99) # Создаем тензор с индексами в пределах vocab_size (0-99)
tensor = torch.tensor([ tensor = torch.tensor([[11, 45, 76, 34], [34, 67, 45, 54]])
[11, 45, 76, 34],
[34, 67, 45, 54]
])
# Проверяем индексы # Проверяем индексы
if (tensor >= 100).any(): if (tensor >= 100).any():
@@ -66,4 +64,4 @@ if __name__ == "__main__":
output = embedding(tensor) output = embedding(tensor)
print("Embeddings shape:", output.shape) 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): class GPT(BaseModel):
""" """
Original GPT (Generative Pre-trained Transformer) модель. Original GPT (Generative Pre-trained Transformer) модель.
Первая версия трансформерной архитектуры от OpenAI, предназначенная Первая версия трансформерной архитектуры от OpenAI, предназначенная
для генеративного предобучения на текстовых данных. для генеративного предобучения на текстовых данных.
Args: Args:
config: Словарь конфигурации с параметрами: config: Словарь конфигурации с параметрами:
- vocab_size: Размер словаря токенов - vocab_size: Размер словаря токенов
@@ -46,7 +46,7 @@ class GPT(BaseModel):
- num_layers: Количество декодерных слоев - num_layers: Количество декодерных слоев
- max_position_embeddings: Максимальная длина последовательности - max_position_embeddings: Максимальная длина последовательности
- dropout: Вероятность dropout - dropout: Вероятность dropout
Attributes: Attributes:
_token_embeddings: Слой векторных представлений токенов _token_embeddings: Слой векторных представлений токенов
_position_embeddings: Слой позиционных эмбеддингов _position_embeddings: Слой позиционных эмбеддингов
@@ -54,30 +54,34 @@ class GPT(BaseModel):
_norm: Финальный слой нормализации _norm: Финальный слой нормализации
_linear: Выходной линейный слой _linear: Выходной линейный слой
""" """
def __init__(self, config): def __init__(self, config):
super().__init__(config) super().__init__(config)
# Инициализация слоев # Инициализация слоев
self._max_seq_len = config["max_position_embeddings"] self._max_seq_len = config["max_position_embeddings"]
self._token_embeddings = TokenEmbeddings( self._token_embeddings = TokenEmbeddings(
vocab_size=config["vocab_size"], vocab_size=config["vocab_size"], emb_size=config["embed_dim"]
emb_size=config["embed_dim"]
) )
self._position_embeddings = PositionalEmbeddings( self._position_embeddings = PositionalEmbeddings(
max_seq_len=config["max_position_embeddings"], max_seq_len=config["max_position_embeddings"], emb_size=config["embed_dim"]
emb_size=config["embed_dim"]
) )
self._dropout = nn.Dropout(config["dropout"]) self._dropout = nn.Dropout(config["dropout"])
# head_size = emb_size // num_heads # head_size = emb_size // num_heads
self._decoders = nn.ModuleList([Decoder( self._decoders = nn.ModuleList(
num_heads=config["num_heads"], [
emb_size=config["embed_dim"], Decoder(
head_size=config["embed_dim"] // config["num_heads"], num_heads=config["num_heads"],
max_seq_len=config["max_position_embeddings"], emb_size=config["embed_dim"],
dropout=config["dropout"] head_size=config["embed_dim"] // config["num_heads"],
) for _ in range(config["num_layers"])]) 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"]) self._linear = nn.Linear(config["embed_dim"], config["vocab_size"])
@property @property
def max_seq_len(self): def max_seq_len(self):
"""Возвращает максимальную длину последовательности.""" """Возвращает максимальную длину последовательности."""
@@ -85,57 +89,60 @@ class GPT(BaseModel):
def forward(self, x: torch.Tensor, attention_mask=None) -> torch.Tensor: def forward(self, x: torch.Tensor, attention_mask=None) -> torch.Tensor:
"""Прямой проход через GPT """Прямой проход через GPT
Args: Args:
x: Входной тензор [batch_size, seq_len] x: Входной тензор [batch_size, seq_len]
Returns: Returns:
Тензор логитов [batch_size, seq_len, vocab_size] Тензор логитов [batch_size, seq_len, vocab_size]
""" """
# Проверка длины последовательности # Проверка длины последовательности
if x.size(1) > self._max_seq_len: 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] tok_out = self._token_embeddings(x) # [batch, seq_len, emb_size]
pos_out = self._position_embeddings(x.size(1)) # [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: for decoder in self._decoders:
out = decoder(out) out = decoder(out)
return self._linear(out) # [batch, seq_len, vocab_size] 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): def generate(
# B, T = input_ids.size() self,
# pos = torch.arange(0, T, device=input_ids.device).unsqueeze(0) x: torch.Tensor,
# max_new_tokens: int,
# 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,
do_sample: bool, do_sample: bool,
temperature: float = 1.0, temperature: float = 1.0,
top_k: int = None, top_k: int = None,
top_p: float = None, top_p: float = None,
attention_mask: torch.Tensor = None, # Добавляем для совместимости с HF attention_mask: torch.Tensor = None, # Добавляем для совместимости с HF
**kwargs # Игнорируем остальные параметры **kwargs, # Игнорируем остальные параметры
) -> torch.Tensor: ) -> torch.Tensor:
"""Авторегрессивная генерация текста. """Авторегрессивная генерация текста.
Параметры: Параметры:
x: Входной тензор с индексами токенов формы [batch_size, seq_len], x: Входной тензор с индексами токенов формы [batch_size, seq_len],
где batch_size - размер батча, seq_len - длина последовательности. где batch_size - размер батча, seq_len - длина последовательности.
@@ -157,9 +164,9 @@ class GPT(BaseModel):
- Гарантируется, что хотя бы один токен остаётся (даже если его вероятность > top_p) - Гарантируется, что хотя бы один токен остаётся (даже если его вероятность > top_p)
- None: отключено (по умолчанию) - None: отключено (по умолчанию)
- Должен быть в диапазоне (0, 1] - Должен быть в диапазоне (0, 1]
Возвращает: Возвращает:
torch.Tensor: Тензор с расширенной последовательностью токенов формы torch.Tensor: Тензор с расширенной последовательностью токенов формы
[batch_size, seq_len + max_new_tokens] [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) >>> output = model.generate(input_ids, max_new_tokens=10, do_sample=False)
>>> >>>
>>> # Вероятностная генерация с top-k >>> # Вероятностная генерация с top-k
>>> output = model.generate(input_ids, max_new_tokens=10, do_sample=True, top_k=50) >>> 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) >>> output = model.generate(input_ids, max_new_tokens=10, do_sample=True, top_p=0.9)
>>> >>>
>>> # Комбинация температуры и top-k >>> # Комбинация температуры и 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) ... temperature=0.7, top_k=50)
Примечания: Примечания:
1. Для детерминированных результатов в режиме сэмплирования 1. Для детерминированных результатов в режиме сэмплирования
зафиксируйте random seed (torch.manual_seed). зафиксируйте random seed (torch.manual_seed).
2. Температура влияет только на режим сэмплирования (do_sample=True). 2. Температура влияет только на режим сэмплирования (do_sample=True).
3. Одновременное использование top_k и top_p запрещено. 3. Одновременное использование top_k и top_p запрещено.
@@ -204,7 +211,7 @@ class GPT(BaseModel):
Должна быть > 0 (по умолчанию: 1.0) Должна быть > 0 (по умолчанию: 1.0)
Returns: Returns:
torch.Tensor: Тензор с расширенной последовательностью токенов формы torch.Tensor: Тензор с расширенной последовательностью токенов формы
[batch_size, seq_len + max_new_tokens] [batch_size, seq_len + max_new_tokens]
Raises: Raises:
@@ -222,13 +229,13 @@ class GPT(BaseModel):
>>> output = model.generate(input_ids, max_new_tokens=10, do_sample=True, temperature=1.5) >>> output = model.generate(input_ids, max_new_tokens=10, do_sample=True, temperature=1.5)
Note: Note:
Для детерминированных результатов в режиме сэмплирования Для детерминированных результатов в режиме сэмплирования
зафиксируйте random seed (torch.manual_seed). зафиксируйте random seed (torch.manual_seed).
Температура влияет только на режим сэмплирования (do_sample=True). Температура влияет только на режим сэмплирования (do_sample=True).
""" """
for _ in range(max_new_tokens): for _ in range(max_new_tokens):
# 1. Обрезаем вход, если последовательность слишком длинная # 1. Обрезаем вход, если последовательность слишком длинная
x_cond = x[:, -self._max_seq_len:] x_cond = x[:, -self._max_seq_len :]
# 2. Передаем последовательность в метод forward класса GPT и полуаем логиты. # 2. Передаем последовательность в метод forward класса GPT и полуаем логиты.
logits = self.forward(x_cond) logits = self.forward(x_cond)
@@ -250,9 +257,14 @@ class GPT(BaseModel):
vocab_size = logits_scaled.size(-1) vocab_size = logits_scaled.size(-1)
# создаём маску: True, если токен НЕ в topk_indices # создаём маску: True, если токен НЕ в topk_indices
mask = torch.ones_like(logits_scaled, dtype=torch.bool if hasattr(torch, 'bool') else torch.uint8) mask = torch.ones_like(
mask.scatter_(1, topk_indices, False if hasattr(torch, 'bool') else 0) # False там, где top-k индексы logits_scaled,
masked_logits[mask] = float('-inf') 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 logits_scaled = masked_logits
@@ -260,36 +272,42 @@ class GPT(BaseModel):
# 1. Применим softmax, чтобы получить вероятности: # 1. Применим softmax, чтобы получить вероятности:
probs = F.softmax(logits_scaled, dim=-1) # [B, vocab_size] probs = F.softmax(logits_scaled, dim=-1) # [B, vocab_size]
# 2. Отсортируем токены по убыванию вероятностей: # 2. Отсортируем токены по убыванию вероятностей:
sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=-1) sorted_probs, sorted_indices = torch.sort(
probs, descending=True, dim=-1
)
# 3. Посчитаем кумулятивную сумму вероятностей: # 3. Посчитаем кумулятивную сумму вероятностей:
cum_probs = torch.cumsum(sorted_probs, dim=-1) # [B, vocab_size] cum_probs = torch.cumsum(sorted_probs, dim=-1) # [B, vocab_size]
# 4. Определим маску: оставить токены, пока сумма < top_p # 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 sorted_mask[:, 0] = True
# 5. Преобразуем маску обратно в оригинальный порядок: # 5. Преобразуем маску обратно в оригинальный порядок:
# Создаём полную маску из False # Создаём полную маску из 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 в местах нужных токенов # Устанавливаем True в местах нужных токенов
mask.scatter_(dim=1, index=sorted_indices, src=sorted_mask) mask.scatter_(dim=1, index=sorted_indices, src=sorted_mask)
# 6. Зануляем логиты токенов вне топ-p: # 6. Зануляем логиты токенов вне топ-p:
logits_scaled[~mask] = float('-inf') logits_scaled[~mask] = float("-inf")
# 4. Применяем Softmax # 4. Применяем Softmax
probs = F.softmax(logits_scaled, dim=-1) # [batch_size, vocab_size] probs = F.softmax(logits_scaled, dim=-1) # [batch_size, vocab_size]
if do_sample == True: if do_sample == True:
# 5. Если do_sample равен True, то отбираем токен случайно с помощью torch.multinomial # 5. Если do_sample равен True, то отбираем токен случайно с помощью torch.multinomial
next_token = torch.multinomial(probs, num_samples=1) # [batch_size, 1] next_token = torch.multinomial(probs, num_samples=1) # [batch_size, 1]
else: else:
# 5. Если do_sample равен False, то выбираем токен с максимальной вероятностью # 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. Добавляем его к последовательности # 6. Добавляем его к последовательности
x = torch.cat([x, next_token], dim=1) # [batch_size, seq_len+1] x = torch.cat([x, next_token], dim=1) # [batch_size, seq_len+1]
return x return x
# def generate(self, input_ids, max_length=50): # def generate(self, input_ids, max_length=50):
# for _ in range(max_length): # for _ in range(max_length):
# logits = self.forward(input_ids) # 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.cached_decoder import CachedDecoder
from llm.core.feed_forward import FeedForward from llm.core.feed_forward import FeedForward
class GPT2(BaseModel): class GPT2(BaseModel):
""" """
GPT2 — автогерессивная языковая модель, архитектура Transformer, предложенная OpenAI. GPT2 — автогерессивная языковая модель, архитектура Transformer, предложенная OpenAI.
@@ -35,7 +36,7 @@ class GPT2(BaseModel):
- Масштабируемый автогерессивный трансформер для предсказания токенов слева направо. - Масштабируемый автогерессивный трансформер для предсказания токенов слева направо.
- Главное отличие от классической GPT: порядок layer normalization ПЕРЕД attention и FFN. - Главное отличие от классической GPT: порядок layer normalization ПЕРЕД attention и FFN.
- Используется GELU, efficient KV-cache, несет наследие классической GPT, но делает архитектуру глубже/шире. - Используется GELU, efficient KV-cache, несет наследие классической GPT, но делает архитектуру глубже/шире.
Args: Args:
config (dict): параметры архитектуры (vocab_size, embed_dim, num_heads, num_layers, max_position_embeddings, dropout) 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) >>> logits = model(input_ids)
>>> out = model.generate(input_ids, max_length=20) >>> out = model.generate(input_ids, max_length=20)
""" """
def __init__(self, config): def __init__(self, config):
super().__init__(config) super().__init__(config)
# Инициализация слоев # Инициализация слоев
self._max_seq_len = config["max_position_embeddings"] self._max_seq_len = config["max_position_embeddings"]
self._token_embeddings = TokenEmbeddings( self._token_embeddings = TokenEmbeddings(
vocab_size=config["vocab_size"], vocab_size=config["vocab_size"], emb_size=config["embed_dim"]
emb_size=config["embed_dim"]
) )
self._position_embeddings = PositionalEmbeddings( self._position_embeddings = PositionalEmbeddings(
max_seq_len=config["max_position_embeddings"], max_seq_len=config["max_position_embeddings"], emb_size=config["embed_dim"]
emb_size=config["embed_dim"]
) )
self._dropout = nn.Dropout(config["dropout"]) self._dropout = nn.Dropout(config["dropout"])
# head_size = emb_size // num_heads # head_size = emb_size // num_heads
self._decoders = nn.ModuleList([CachedDecoder( self._decoders = nn.ModuleList(
num_heads=config["num_heads"], [
emb_size=config["embed_dim"], CachedDecoder(
head_size=config["embed_dim"] // config["num_heads"], num_heads=config["num_heads"],
feed_forward_layer=FeedForward( emb_size=config["embed_dim"],
emb_size=config["embed_dim"], head_size=config["embed_dim"] // config["num_heads"],
dropout=config["dropout"], feed_forward_layer=FeedForward(
activation="gelu" emb_size=config["embed_dim"],
), dropout=config["dropout"],
max_seq_len=config["max_position_embeddings"], activation="gelu",
dropout=config["dropout"] ),
) for _ in range(config["num_layers"])]) max_seq_len=config["max_position_embeddings"],
dropout=config["dropout"],
)
for _ in range(config["num_layers"])
]
)
self._norm = nn.LayerNorm(config["embed_dim"]) self._norm = nn.LayerNorm(config["embed_dim"])
self._linear = nn.Linear(config["embed_dim"], config["vocab_size"]) 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: Прямой проход GPT2:
- Все слои работают как autoregressive transformer (masked self-attention). - Все слои работают как autoregressive transformer (masked self-attention).
@@ -91,9 +98,10 @@ class GPT2(BaseModel):
""" """
# Проверка длины последовательности (только при отсутствии кэша) # Проверка длины последовательности (только при отсутствии кэша)
if cache is None and x.size(1) > self._max_seq_len: 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 из кэша (если кэш передан) # Вычисление start_pos из кэша (если кэш передан)
if cache is not None: if cache is not None:
# При кэше обрабатываем только один токен (последний) # При кэше обрабатываем только один токен (последний)
@@ -111,11 +119,15 @@ class GPT2(BaseModel):
# Эмбеддинги токенов и позиций # Эмбеддинги токенов и позиций
tok_out = self._token_embeddings(x) # [batch, seq_len, emb_size] 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 = [] new_cache = []
for i, decoder in enumerate(self._decoders): for i, decoder in enumerate(self._decoders):
@@ -131,21 +143,22 @@ class GPT2(BaseModel):
out = self._norm(out) out = self._norm(out)
logits = self._linear(out) logits = self._linear(out)
# Возвращаем результат с учетом use_cache # Возвращаем результат с учетом use_cache
if use_cache: if use_cache:
return (logits, new_cache) return (logits, new_cache)
else: else:
return (logits, None) return (logits, None)
def generate(self, def generate(
x: torch.Tensor, self,
max_new_tokens: int, x: torch.Tensor,
max_new_tokens: int,
do_sample: bool, do_sample: bool,
temperature: float = 1.0, temperature: float = 1.0,
top_k: int = None, top_k: int = None,
top_p: float = None, top_p: float = None,
use_cache: bool = True use_cache: bool = True,
) -> torch.Tensor: ) -> torch.Tensor:
""" """
Генерация текста с использованием autoregressive трансформера (GPT2). Генерация текста с использованием autoregressive трансформера (GPT2).
@@ -174,10 +187,10 @@ class GPT2(BaseModel):
else: else:
# Первая итерация или кэш отключен - передаем всю последовательность # Первая итерация или кэш отключен - передаем всю последовательность
x_input = x x_input = x
# Прямой проход с кэшем # Прямой проход с кэшем
logits, new_cache = self.forward(x_input, use_cache=use_cache, cache=cache) logits, new_cache = self.forward(x_input, use_cache=use_cache, cache=cache)
# Обновляем кэш для следующей итерации # Обновляем кэш для следующей итерации
if use_cache: if use_cache:
cache = new_cache cache = new_cache
@@ -200,7 +213,7 @@ class GPT2(BaseModel):
# создаём маску: 1, если токен НЕ в topk_indices # создаём маску: 1, если токен НЕ в topk_indices
mask = torch.ones_like(logits_scaled, dtype=torch.uint8) mask = torch.ones_like(logits_scaled, dtype=torch.uint8)
mask.scatter_(1, topk_indices, 0) # 0 там, где top-k индексы 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 logits_scaled = masked_logits
@@ -208,7 +221,9 @@ class GPT2(BaseModel):
# 1. Применим softmax, чтобы получить вероятности: # 1. Применим softmax, чтобы получить вероятности:
probs = F.softmax(logits_scaled, dim=-1) # [B, vocab_size] probs = F.softmax(logits_scaled, dim=-1) # [B, vocab_size]
# 2. Отсортируем токены по убыванию вероятностей: # 2. Отсортируем токены по убыванию вероятностей:
sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=-1) sorted_probs, sorted_indices = torch.sort(
probs, descending=True, dim=-1
)
# 3. Посчитаем кумулятивную сумму вероятностей: # 3. Посчитаем кумулятивную сумму вероятностей:
cum_probs = torch.cumsum(sorted_probs, dim=-1) # [B, vocab_size] cum_probs = torch.cumsum(sorted_probs, dim=-1) # [B, vocab_size]
# 4. Определим маску: оставить токены, пока сумма < top_p # 4. Определим маску: оставить токены, пока сумма < top_p
@@ -221,23 +236,24 @@ class GPT2(BaseModel):
# Устанавливаем 1 в местах нужных токенов # Устанавливаем 1 в местах нужных токенов
mask.scatter_(dim=1, index=sorted_indices, src=sorted_mask) mask.scatter_(dim=1, index=sorted_indices, src=sorted_mask)
# 6. Зануляем логиты токенов вне топ-p: # 6. Зануляем логиты токенов вне топ-p:
logits_scaled[~mask] = float('-inf') logits_scaled[~mask] = float("-inf")
# 4. Применяем Softmax # 4. Применяем Softmax
probs = F.softmax(logits_scaled, dim=-1) # [batch_size, vocab_size] probs = F.softmax(logits_scaled, dim=-1) # [batch_size, vocab_size]
if do_sample == True: if do_sample == True:
# 5. Если do_sample равен True, то отбираем токен случайно с помощью torch.multinomial # 5. Если do_sample равен True, то отбираем токен случайно с помощью torch.multinomial
next_token = torch.multinomial(probs, num_samples=1) # [batch_size, 1] next_token = torch.multinomial(probs, num_samples=1) # [batch_size, 1]
else: else:
# 5. Если do_sample равен False, то выбираем токен с максимальной вероятностью # 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. Добавляем его к последовательности # 6. Добавляем его к последовательности
x = torch.cat([x, next_token], dim=1) # [batch_size, seq_len+1] x = torch.cat([x, next_token], dim=1) # [batch_size, seq_len+1]
return x return x
@property @property
def max_seq_len(self) -> int: 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 from llm.core.cached_decoder import CachedDecoder
class Llama(BaseModel): class Llama(BaseModel):
""" """
LLaMA (Large Language Model Meta AI) — высокоэффективная масштабируемая языковая модель, разработанная Meta AI Research. LLaMA (Large Language Model Meta AI) — высокоэффективная масштабируемая языковая модель, разработанная Meta AI Research.
@@ -29,38 +28,45 @@ class Llama(BaseModel):
>>> logits, cache = model(input_ids, use_cache=True) >>> logits, cache = model(input_ids, use_cache=True)
>>> out = model.generate(input_ids, max_new_tokens=20) >>> out = model.generate(input_ids, max_new_tokens=20)
""" """
def __init__(self,config):
def __init__(self, config):
super().__init__(config) super().__init__(config)
# Инициализация слоев # Инициализация слоев
self._max_seq_len = config["max_position_embeddings"] self._max_seq_len = config["max_position_embeddings"]
self._token_embeddings = TokenEmbeddings( self._token_embeddings = TokenEmbeddings(
vocab_size=config["vocab_size"], vocab_size=config["vocab_size"], emb_size=config["embed_dim"]
emb_size=config["embed_dim"]
) )
self._position_embeddings = RoPE( self._position_embeddings = RoPE(
head_size=config["embed_dim"] // config["num_heads"], 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._dropout = nn.Dropout(config["dropout"])
self._decoders = nn.ModuleList([CachedDecoder( self._decoders = nn.ModuleList(
norm_layer=RMSNorm, [
num_heads=config["num_heads"], CachedDecoder(
emb_size=config["embed_dim"], norm_layer=RMSNorm,
head_size=config["embed_dim"] // config["num_heads"], num_heads=config["num_heads"],
feed_forward_layer=SwiGLU( emb_size=config["embed_dim"],
emb_size=config["embed_dim"], head_size=config["embed_dim"] // config["num_heads"],
dropout=config["dropout"], feed_forward_layer=SwiGLU(
), emb_size=config["embed_dim"],
max_seq_len=config["max_position_embeddings"], dropout=config["dropout"],
rope=self._position_embeddings, ),
dropout=config["dropout"], max_seq_len=config["max_position_embeddings"],
) for _ in range(config["num_layers"])]) rope=self._position_embeddings,
dropout=config["dropout"],
)
for _ in range(config["num_layers"])
]
)
self._norm = RMSNorm(config["embed_dim"]) self._norm = RMSNorm(config["embed_dim"])
self._linear = nn.Linear(config["embed_dim"], config["vocab_size"]) 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): авторегрессионное предсказание токенов. Прямой проход через LLaMA (inference/train): авторегрессионное предсказание токенов.
@@ -76,11 +82,12 @@ class Llama(BaseModel):
""" """
# Проверка длины последовательности (только при отсутствии кэша) # Проверка длины последовательности (только при отсутствии кэша)
if cache is None and x.size(1) > self._max_seq_len: 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 из кэша (если кэш передан) # Вычисление start_pos из кэша (если кэш передан)
#if cache is not None: # if cache is not None:
# # При кэше обрабатываем только один токен (последний) # # При кэше обрабатываем только один токен (последний)
# seq_len = 1 # seq_len = 1
# # Вычисляем start_pos из самого нижнего уровня кэша # # Вычисляем start_pos из самого нижнего уровня кэша
@@ -89,18 +96,18 @@ class Llama(BaseModel):
# start_pos = key_cache.size(1) # cache_len # start_pos = key_cache.size(1) # cache_len
# else: # else:
# start_pos = 0 # start_pos = 0
#else: # else:
# # Без кэша работаем как раньше # # Без кэша работаем как раньше
# start_pos = 0 # start_pos = 0
# seq_len = x.size(1) # seq_len = x.size(1)
# Эмбеддинги токенов и позиций # Эмбеддинги токенов и позиций
tok_out = self._token_embeddings(x) # [batch, seq_len, emb_size] 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] out = self._dropout(tok_out) # [batch, seq_len, emb_size]
# Стек декодеров с передачей кэша # Стек декодеров с передачей кэша
new_cache = [] new_cache = []
for i, decoder in enumerate(self._decoders): for i, decoder in enumerate(self._decoders):
@@ -116,42 +123,43 @@ class Llama(BaseModel):
out = self._norm(out) out = self._norm(out)
logits = self._linear(out) logits = self._linear(out)
# Возвращаем результат с учетом use_cache # Возвращаем результат с учетом use_cache
if use_cache: if use_cache:
return (logits, new_cache) return (logits, new_cache)
else: else:
return (logits, None) return (logits, None)
def generate(self, def generate(
x: torch.Tensor, self,
max_new_tokens: int, x: torch.Tensor,
max_new_tokens: int,
do_sample: bool, do_sample: bool,
temperature: float = 1.0, temperature: float = 1.0,
top_k: int = None, top_k: int = None,
top_p: float = None, top_p: float = None,
use_cache: bool = True use_cache: bool = True,
) -> torch.Tensor: ) -> torch.Tensor:
""" """
Генерация текста c помощью LLaMA (autoregressive Transformer). Генерация текста c помощью LLaMA (autoregressive Transformer).
Поддерживается: Поддерживается:
- greedy и вероятностное сэмплирование (top-k, top-p, temperature) - greedy и вероятностное сэмплирование (top-k, top-p, temperature)
- кэш attention для ускорения генерации длинных последовательностей - кэш attention для ускорения генерации длинных последовательностей
Args: Args:
x (Tensor[int]): начальная последовательность [batch, seq_len] x (Tensor[int]): начальная последовательность [batch, seq_len]
max_new_tokens (int): сколько новых токенов сгенерировать max_new_tokens (int): сколько новых токенов сгенерировать
do_sample (bool): использовать стохастику (True) или жадный выбор (False) do_sample (bool): использовать стохастику (True) или жадный выбор (False)
temperature (float): масштаб для softmax (важно для sampling) temperature (float): масштаб для softmax (важно для sampling)
top_k (int|None): ограничение на количество кандидатов (top-k sampling) top_k (int|None): ограничение на количество кандидатов (top-k sampling)
top_p (float|None): nucleus sampling top_p (float|None): nucleus sampling
use_cache (bool): ускоряет autoregressive при длинной генерации use_cache (bool): ускоряет autoregressive при длинной генерации
Returns: Returns:
output (Tensor[int]): [batch, seq_len + max_new_tokens] output (Tensor[int]): [batch, seq_len + max_new_tokens]
Пример: Пример:
>>> prompt = tokenizer.encode('Meta AI', return_tensors="pt") >>> prompt = tokenizer.encode('Meta AI', return_tensors="pt")
>>> generated = model.generate(prompt, max_new_tokens=30, do_sample=True) >>> generated = model.generate(prompt, max_new_tokens=30, do_sample=True)
>>> print(tokenizer.decode(generated[0])) >>> print(tokenizer.decode(generated[0]))
""" """
cache = None cache = None
@@ -162,10 +170,10 @@ class Llama(BaseModel):
else: else:
# Первая итерация или кэш отключен - передаем всю последовательность # Первая итерация или кэш отключен - передаем всю последовательность
x_input = x x_input = x
# Прямой проход с кэшем # Прямой проход с кэшем
logits, new_cache = self.forward(x_input, use_cache=use_cache, cache=cache) logits, new_cache = self.forward(x_input, use_cache=use_cache, cache=cache)
# Обновляем кэш для следующей итерации # Обновляем кэш для следующей итерации
if use_cache: if use_cache:
cache = new_cache cache = new_cache
@@ -188,7 +196,7 @@ class Llama(BaseModel):
# создаём маску: 1, если токен НЕ в topk_indices # создаём маску: 1, если токен НЕ в topk_indices
mask = torch.ones_like(logits_scaled, dtype=torch.uint8) mask = torch.ones_like(logits_scaled, dtype=torch.uint8)
mask.scatter_(1, topk_indices, 0) # 0 там, где top-k индексы 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 logits_scaled = masked_logits
@@ -196,7 +204,9 @@ class Llama(BaseModel):
# 1. Применим softmax, чтобы получить вероятности: # 1. Применим softmax, чтобы получить вероятности:
probs = F.softmax(logits_scaled, dim=-1) # [B, vocab_size] probs = F.softmax(logits_scaled, dim=-1) # [B, vocab_size]
# 2. Отсортируем токены по убыванию вероятностей: # 2. Отсортируем токены по убыванию вероятностей:
sorted_probs, sorted_indices = torch.sort(probs, descending=True, dim=-1) sorted_probs, sorted_indices = torch.sort(
probs, descending=True, dim=-1
)
# 3. Посчитаем кумулятивную сумму вероятностей: # 3. Посчитаем кумулятивную сумму вероятностей:
cum_probs = torch.cumsum(sorted_probs, dim=-1) # [B, vocab_size] cum_probs = torch.cumsum(sorted_probs, dim=-1) # [B, vocab_size]
# 4. Определим маску: оставить токены, пока сумма < top_p # 4. Определим маску: оставить токены, пока сумма < top_p
@@ -209,25 +219,24 @@ class Llama(BaseModel):
# Устанавливаем 1 в местах нужных токенов # Устанавливаем 1 в местах нужных токенов
mask.scatter_(dim=1, index=sorted_indices, src=sorted_mask) mask.scatter_(dim=1, index=sorted_indices, src=sorted_mask)
# 6. Зануляем логиты токенов вне топ-p: # 6. Зануляем логиты токенов вне топ-p:
logits_scaled[~mask] = float('-inf') logits_scaled[~mask] = float("-inf")
# 4. Применяем Softmax # 4. Применяем Softmax
probs = F.softmax(logits_scaled, dim=-1) # [batch_size, vocab_size] probs = F.softmax(logits_scaled, dim=-1) # [batch_size, vocab_size]
if do_sample == True: if do_sample == True:
# 5. Если do_sample равен True, то отбираем токен случайно с помощью torch.multinomial # 5. Если do_sample равен True, то отбираем токен случайно с помощью torch.multinomial
next_token = torch.multinomial(probs, num_samples=1) # [batch_size, 1] next_token = torch.multinomial(probs, num_samples=1) # [batch_size, 1]
else: else:
# 5. Если do_sample равен False, то выбираем токен с максимальной вероятностью # 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. Добавляем его к последовательности # 6. Добавляем его к последовательности
x = torch.cat([x, next_token], dim=1) # [batch_size, seq_len+1] x = torch.cat([x, next_token], dim=1) # [batch_size, seq_len+1]
return x return x
@property @property
def max_seq_len(self) -> int: 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): class BaseTokenizer(ABC):
""" """
Абстрактный базовый класс для всех токенизаторов. Абстрактный базовый класс для всех токенизаторов.
Определяет общий интерфейс для токенизации текста. Определяет общий интерфейс для токенизации текста.
""" """
def __init__(self): def __init__(self):
self.vocab: Dict[str, int] = {} self.vocab: Dict[str, int] = {}
self.inverse_vocab: Dict[int, str] = {} self.inverse_vocab: Dict[int, str] = {}
self.vocab_size: int = 0 self.vocab_size: int = 0
# Специальные токены # Специальные токены
self.pad_token = "<pad>" self.pad_token = "<pad>"
self.unk_token = "<unk>" self.unk_token = "<unk>"
self.bos_token = "<bos>" self.bos_token = "<bos>"
self.eos_token = "<eos>" self.eos_token = "<eos>"
self.pad_token_id: Optional[int] = None self.pad_token_id: Optional[int] = None
self.unk_token_id: Optional[int] = None self.unk_token_id: Optional[int] = None
self.bos_token_id: Optional[int] = None self.bos_token_id: Optional[int] = None
self.eos_token_id: Optional[int] = None self.eos_token_id: Optional[int] = None
@abstractmethod @abstractmethod
def train(self, texts: List[str], vocab_size: int = 1000, **kwargs): def train(self, texts: List[str], vocab_size: int = 1000, **kwargs):
""" """
Обучение токенизатора на текстах. Обучение токенизатора на текстах.
Args: Args:
texts: Список текстов для обучения texts: Список текстов для обучения
vocab_size: Желаемый размер словаря vocab_size: Желаемый размер словаря
**kwargs: Дополнительные параметры обучения **kwargs: Дополнительные параметры обучения
""" """
pass pass
@abstractmethod @abstractmethod
def encode(self, text: str, **kwargs) -> List[int]: def encode(self, text: str, **kwargs) -> List[int]:
""" """
Кодирование текста в последовательность токенов. Кодирование текста в последовательность токенов.
Args: Args:
text: Входной текст text: Входной текст
**kwargs: Дополнительные параметры кодирования **kwargs: Дополнительные параметры кодирования
Returns: Returns:
List[int]: Список идентификаторов токенов List[int]: Список идентификаторов токенов
""" """
pass pass
@abstractmethod @abstractmethod
def decode(self, tokens: List[int], **kwargs) -> str: def decode(self, tokens: List[int], **kwargs) -> str:
""" """
Декодирование последовательности токенов в текст. Декодирование последовательности токенов в текст.
Args: Args:
tokens: Список идентификаторов токенов tokens: Список идентификаторов токенов
**kwargs: Дополнительные параметры декодирования **kwargs: Дополнительные параметры декодирования
Returns: Returns:
str: Декодированный текст str: Декодированный текст
""" """
pass pass
def tokenize(self, text: str, **kwargs) -> List[str]: def tokenize(self, text: str, **kwargs) -> List[str]:
""" """
Токенизация текста в список строковых токенов. Токенизация текста в список строковых токенов.
Args: Args:
text: Входной текст text: Входной текст
**kwargs: Дополнительные параметры **kwargs: Дополнительные параметры
Returns: Returns:
List[str]: Список токенов List[str]: Список токенов
""" """
token_ids = self.encode(text, **kwargs) 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]: def get_vocab(self) -> Dict[str, int]:
"""Возвращает словарь токенизатора.""" """Возвращает словарь токенизатора."""
return self.vocab.copy() return self.vocab.copy()
def get_vocab_size(self) -> int: def get_vocab_size(self) -> int:
"""Возвращает размер словаря.""" """Возвращает размер словаря."""
return self.vocab_size return self.vocab_size
def add_special_tokens(self, special_tokens: List[str]): def add_special_tokens(self, special_tokens: List[str]):
""" """
Добавляет специальные токены в словарь. Добавляет специальные токены в словарь.
Args: Args:
special_tokens: Список специальных токенов special_tokens: Список специальных токенов
""" """
@@ -105,70 +107,70 @@ class BaseTokenizer(ABC):
self.vocab[token] = token_id self.vocab[token] = token_id
self.inverse_vocab[token_id] = token self.inverse_vocab[token_id] = token
self.vocab_size += 1 self.vocab_size += 1
# Обновляем ID специальных токенов # Обновляем ID специальных токенов
self.pad_token_id = self.vocab.get(self.pad_token) self.pad_token_id = self.vocab.get(self.pad_token)
self.unk_token_id = self.vocab.get(self.unk_token) self.unk_token_id = self.vocab.get(self.unk_token)
self.bos_token_id = self.vocab.get(self.bos_token) self.bos_token_id = self.vocab.get(self.bos_token)
self.eos_token_id = self.vocab.get(self.eos_token) self.eos_token_id = self.vocab.get(self.eos_token)
def save(self, filepath: str): def save(self, filepath: str):
""" """
Сохраняет токенизатор в файл. Сохраняет токенизатор в файл.
Args: Args:
filepath: Путь для сохранения filepath: Путь для сохранения
""" """
config = { config = {
'vocab': self.vocab, "vocab": self.vocab,
'vocab_size': self.vocab_size, "vocab_size": self.vocab_size,
'pad_token': self.pad_token, "pad_token": self.pad_token,
'unk_token': self.unk_token, "unk_token": self.unk_token,
'bos_token': self.bos_token, "bos_token": self.bos_token,
'eos_token': self.eos_token, "eos_token": self.eos_token,
'tokenizer_type': self.__class__.__name__ "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) json.dump(config, f, ensure_ascii=False, indent=2)
@classmethod @classmethod
def load(cls, filepath: str): def load(cls, filepath: str):
""" """
Загружает токенизатор из файла. Загружает токенизатор из файла.
Args: Args:
filepath: Путь к файлу filepath: Путь к файлу
Returns: Returns:
BaseTokenizer: Загруженный токенизатор BaseTokenizer: Загруженный токенизатор
""" """
with open(filepath, 'r', encoding='utf-8') as f: with open(filepath, "r", encoding="utf-8") as f:
config = json.load(f) config = json.load(f)
# Создаем экземпляр токенизатора # Создаем экземпляр токенизатора
tokenizer = cls() tokenizer = cls()
tokenizer.vocab = config['vocab'] tokenizer.vocab = config["vocab"]
tokenizer.vocab_size = config['vocab_size'] tokenizer.vocab_size = config["vocab_size"]
tokenizer.pad_token = config['pad_token'] tokenizer.pad_token = config["pad_token"]
tokenizer.unk_token = config['unk_token'] tokenizer.unk_token = config["unk_token"]
tokenizer.bos_token = config['bos_token'] tokenizer.bos_token = config["bos_token"]
tokenizer.eos_token = config['eos_token'] tokenizer.eos_token = config["eos_token"]
# Создаем обратный словарь # Создаем обратный словарь
tokenizer.inverse_vocab = {v: k for k, v in tokenizer.vocab.items()} tokenizer.inverse_vocab = {v: k for k, v in tokenizer.vocab.items()}
# Обновляем ID специальных токенов # Обновляем ID специальных токенов
tokenizer.pad_token_id = tokenizer.vocab.get(tokenizer.pad_token) tokenizer.pad_token_id = tokenizer.vocab.get(tokenizer.pad_token)
tokenizer.unk_token_id = tokenizer.vocab.get(tokenizer.unk_token) tokenizer.unk_token_id = tokenizer.vocab.get(tokenizer.unk_token)
tokenizer.bos_token_id = tokenizer.vocab.get(tokenizer.bos_token) tokenizer.bos_token_id = tokenizer.vocab.get(tokenizer.bos_token)
tokenizer.eos_token_id = tokenizer.vocab.get(tokenizer.eos_token) tokenizer.eos_token_id = tokenizer.vocab.get(tokenizer.eos_token)
return tokenizer return tokenizer
def __len__(self) -> int: def __len__(self) -> int:
"""Возвращает размер словаря.""" """Возвращает размер словаря."""
return self.vocab_size return self.vocab_size
def __repr__(self) -> str: def __repr__(self) -> str:
return f"{self.__class__.__name__}(vocab_size={self.vocab_size})" 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): class BPETokenizer(BaseTokenizer):
""" """
BPE токенизатор для обработки текста. BPE токенизатор для обработки текста.
Реализует алгоритм Byte Pair Encoding для создания субсловных токенов. Реализует алгоритм Byte Pair Encoding для создания субсловных токенов.
Примеры использования: Примеры использования:
>>> tokenizer = BPETokenizer() >>> tokenizer = BPETokenizer()
>>> tokenizer.train(["пример текста для обучения"], vocab_size=1000) >>> tokenizer.train(["пример текста для обучения"], vocab_size=1000)
>>> tokens = tokenizer.encode("новый текст") >>> tokens = tokenizer.encode("новый текст")
>>> text = tokenizer.decode(tokens) >>> text = tokenizer.decode(tokens)
""" """
def __init__(self): def __init__(self):
super().__init__() super().__init__()
self.merges: Dict[Tuple[str, str], int] = {} 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.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) self.compiled_pattern = re.compile(self.pattern, re.UNICODE)
def train(self, texts: List[str], vocab_size: int = 1000, **kwargs): def train(self, texts: List[str], vocab_size: int = 1000, **kwargs):
""" """
Обучение BPE токенизатора на текстах. Обучение BPE токенизатора на текстах.
Args: Args:
texts: Список текстов для обучения texts: Список текстов для обучения
vocab_size: Желаемый размер словаря vocab_size: Желаемый размер словаря
@@ -42,37 +42,40 @@ class BPETokenizer(BaseTokenizer):
""" """
# Инициализация базового словаря # Инициализация базового словаря
self._initialize_vocab() 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) self.add_special_tokens(special_tokens)
# Предобработка текстов # Предобработка текстов
words = self._preprocess_texts(texts) words = self._preprocess_texts(texts)
# Получаем начальные токены # Получаем начальные токены
vocab = self._get_initial_vocab(words) vocab = self._get_initial_vocab(words)
# Выполняем BPE мерджи # Выполняем 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() self._build_final_vocab()
def _initialize_vocab(self): def _initialize_vocab(self):
"""Инициализирует базовый словарь.""" """Инициализирует базовый словарь."""
self.vocab.clear() self.vocab.clear()
self.inverse_vocab.clear() self.inverse_vocab.clear()
self.merges.clear() self.merges.clear()
self.vocab_size = 0 self.vocab_size = 0
def _preprocess_texts(self, texts: List[str]) -> List[List[str]]: def _preprocess_texts(self, texts: List[str]) -> List[List[str]]:
""" """
Предобработка текстов для обучения. Предобработка текстов для обучения.
Args: Args:
texts: Список текстов texts: Список текстов
Returns: Returns:
List[List[str]]: Предобработанные слова List[List[str]]: Предобработанные слова
""" """
@@ -84,14 +87,14 @@ class BPETokenizer(BaseTokenizer):
tokens = self.compiled_pattern.findall(text) tokens = self.compiled_pattern.findall(text)
words.append(tokens) words.append(tokens)
return words return words
def _get_initial_vocab(self, words: List[List[str]]) -> Dict[str, int]: def _get_initial_vocab(self, words: List[List[str]]) -> Dict[str, int]:
""" """
Создает начальный словарь из символов. Создает начальный словарь из символов.
Args: Args:
words: Список токенизированных текстов words: Список токенизированных текстов
Returns: Returns:
Dict[str, int]: Начальный словарь частот Dict[str, int]: Начальный словарь частот
""" """
@@ -99,43 +102,45 @@ class BPETokenizer(BaseTokenizer):
for word_list in words: for word_list in words:
for word in word_list: for word in word_list:
# Разбиваем слово на символы и добавляем специальный символ конца слова # Разбиваем слово на символы и добавляем специальный символ конца слова
chars = list(word) + ['</w>'] chars = list(word) + ["</w>"]
vocab.update([''.join(chars[i:i+1]) for i in range(len(chars))]) vocab.update(["".join(chars[i : i + 1]) for i in range(len(chars))])
return vocab 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 мерджи до достижения целевого размера словаря. Выполняет BPE мерджи до достижения целевого размера словаря.
Args: Args:
vocab: Начальный словарь vocab: Начальный словарь
target_vocab_size: Целевой размер словаря target_vocab_size: Целевой размер словаря
min_frequency: Минимальная частота для мерджа min_frequency: Минимальная частота для мерджа
""" """
current_vocab_size = len(vocab) + len(self.vocab) current_vocab_size = len(vocab) + len(self.vocab)
while current_vocab_size < target_vocab_size: while current_vocab_size < target_vocab_size:
# Находим наиболее частую пару # Находим наиболее частую пару
pairs = self._get_stats(vocab) pairs = self._get_stats(vocab)
if not pairs: if not pairs:
break break
best_pair = max(pairs, key=pairs.get) best_pair = max(pairs, key=pairs.get)
if pairs[best_pair] < min_frequency: if pairs[best_pair] < min_frequency:
break break
# Выполняем мердж # Выполняем мердж
vocab = self._merge_vocab(vocab, best_pair) vocab = self._merge_vocab(vocab, best_pair)
self.merges[best_pair] = len(self.merges) self.merges[best_pair] = len(self.merges)
current_vocab_size += 1 current_vocab_size += 1
def _get_stats(self, vocab: Dict[str, int]) -> Dict[Tuple[str, str], int]: def _get_stats(self, vocab: Dict[str, int]) -> Dict[Tuple[str, str], int]:
""" """
Собирает статистику по парам символов. Собирает статистику по парам символов.
Args: Args:
vocab: Словарь токенов vocab: Словарь токенов
Returns: Returns:
Dict[Tuple[str, str], int]: Частоты пар Dict[Tuple[str, str], int]: Частоты пар
""" """
@@ -145,107 +150,113 @@ class BPETokenizer(BaseTokenizer):
for i in range(len(symbols) - 1): for i in range(len(symbols) - 1):
pairs[symbols[i], symbols[i + 1]] += freq pairs[symbols[i], symbols[i + 1]] += freq
return pairs 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: Args:
vocab: Исходный словарь vocab: Исходный словарь
pair: Пара для объединения pair: Пара для объединения
Returns: Returns:
Dict[str, int]: Обновленный словарь Dict[str, int]: Обновленный словарь
""" """
new_vocab = {} 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] replacement = pair[0] + pair[1]
for word in vocab: for word in vocab:
new_word = bigram.sub(replacement, word) new_word = bigram.sub(replacement, word)
new_vocab[new_word] = vocab[word] new_vocab[new_word] = vocab[word]
return new_vocab return new_vocab
def _build_final_vocab(self): def _build_final_vocab(self):
"""Строит финальный словарь токенизатора.""" """Строит финальный словарь токенизатора."""
# Собираем все уникальные токены из мерджей # Собираем все уникальные токены из мерджей
all_tokens = set() 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: for pair in self.merges:
all_tokens.update(pair) all_tokens.update(pair)
# Создаем словарь # Создаем словарь
for i, token in enumerate(sorted(all_tokens)): for i, token in enumerate(sorted(all_tokens)):
self.vocab[token] = i self.vocab[token] = i
self.inverse_vocab = {v: k for k, v in self.vocab.items()} self.inverse_vocab = {v: k for k, v in self.vocab.items()}
self.vocab_size = len(self.vocab) self.vocab_size = len(self.vocab)
# Обновляем ID специальных токенов # Обновляем ID специальных токенов
self.pad_token_id = self.vocab.get(self.pad_token) self.pad_token_id = self.vocab.get(self.pad_token)
self.unk_token_id = self.vocab.get(self.unk_token) self.unk_token_id = self.vocab.get(self.unk_token)
self.bos_token_id = self.vocab.get(self.bos_token) self.bos_token_id = self.vocab.get(self.bos_token)
self.eos_token_id = self.vocab.get(self.eos_token) self.eos_token_id = self.vocab.get(self.eos_token)
def encode(self, text: str, **kwargs) -> List[int]: def encode(self, text: str, **kwargs) -> List[int]:
""" """
Кодирует текст в последовательность токенов. Кодирует текст в последовательность токенов.
Args: Args:
text: Входной текст text: Входной текст
**kwargs: Дополнительные параметры **kwargs: Дополнительные параметры
- add_special_tokens: Добавлять специальные токены - add_special_tokens: Добавлять специальные токены
Returns: Returns:
List[int]: Список идентификаторов токенов 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) tokens = self.compiled_pattern.findall(text)
# Применяем BPE к каждому токену # Применяем BPE к каждому токену
bpe_tokens = [] bpe_tokens = []
for token in tokens: for token in tokens:
# Преобразуем токен в BPE представление # Преобразуем токен в BPE представление
bpe_token = self._apply_bpe(token) bpe_token = self._apply_bpe(token)
bpe_tokens.extend(bpe_token) bpe_tokens.extend(bpe_token)
# Конвертируем в ID # Конвертируем в ID
token_ids = [] token_ids = []
for token in bpe_tokens: for token in bpe_tokens:
token_id = self.vocab.get(token, self.unk_token_id) token_id = self.vocab.get(token, self.unk_token_id)
if token_id is not None: if token_id is not None:
token_ids.append(token_id) token_ids.append(token_id)
# Добавляем специальные токены если нужно # Добавляем специальные токены если нужно
if add_special_tokens: if add_special_tokens:
if self.bos_token_id is not None: if self.bos_token_id is not None:
token_ids.insert(0, self.bos_token_id) token_ids.insert(0, self.bos_token_id)
if self.eos_token_id is not None: if self.eos_token_id is not None:
token_ids.append(self.eos_token_id) token_ids.append(self.eos_token_id)
return token_ids return token_ids
def _apply_bpe(self, token: str) -> List[str]: def _apply_bpe(self, token: str) -> List[str]:
""" """
Применяет BPE к одному токену. Применяет BPE к одному токену.
Args: Args:
token: Входной токен token: Входной токен
Returns: Returns:
List[str]: Список BPE токенов List[str]: Список BPE токенов
""" """
# Простая реализация - в реальной реализации нужно применять обученные мерджи # Простая реализация - в реальной реализации нужно применять обученные мерджи
word = token + '</w>' word = token + "</w>"
tokens = [word[i:i+1] for i in range(len(word))] tokens = [word[i : i + 1] for i in range(len(word))]
# Применяем мерджи (упрощенная версия) # Применяем мерджи (упрощенная версия)
# В полной реализации нужно применять все обученные мерджи # В полной реализации нужно применять все обученные мерджи
for pair in self.merges: for pair in self.merges:
@@ -256,109 +267,114 @@ class BPETokenizer(BaseTokenizer):
del tokens[i + 1] del tokens[i + 1]
else: else:
i += 1 i += 1
return tokens return tokens
def decode(self, tokens: List[int], **kwargs) -> str: def decode(self, tokens: List[int], **kwargs) -> str:
""" """
Декодирует последовательность токенов в текст. Декодирует последовательность токенов в текст.
Args: Args:
tokens: Список идентификаторов токенов tokens: Список идентификаторов токенов
**kwargs: Дополнительные параметры **kwargs: Дополнительные параметры
- skip_special_tokens: Пропускать специальные токены - skip_special_tokens: Пропускать специальные токены
Returns: Returns:
str: Декодированный текст str: Декодированный текст
""" """
skip_special_tokens = kwargs.get('skip_special_tokens', True) skip_special_tokens = kwargs.get("skip_special_tokens", True)
# Конвертируем ID в токены # Конвертируем ID в токены
token_strings = [] token_strings = []
for token_id in tokens: for token_id in tokens:
token = self.inverse_vocab.get(token_id, self.unk_token) 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 continue
token_strings.append(token) token_strings.append(token)
# Объединяем токены в текст # Объединяем токены в текст
text = ''.join(token_strings) text = "".join(token_strings)
# Убираем маркер конца слова # Убираем маркер конца слова
text = text.replace('</w>', ' ') text = text.replace("</w>", " ")
return text.strip() return text.strip()
def save(self, filepath: str): def save(self, filepath: str):
""" """
Сохраняет BPE токенизатор в файл. Сохраняет BPE токенизатор в файл.
Args: Args:
filepath: Путь для сохранения filepath: Путь для сохранения
""" """
import json import json
config = { config = {
'vocab': self.vocab, "vocab": self.vocab,
'merges': {f"{k[0]} {k[1]}": v for k, v in self.merges.items()}, "merges": {f"{k[0]} {k[1]}": v for k, v in self.merges.items()},
'vocab_size': self.vocab_size, "vocab_size": self.vocab_size,
'pad_token': self.pad_token, "pad_token": self.pad_token,
'unk_token': self.unk_token, "unk_token": self.unk_token,
'bos_token': self.bos_token, "bos_token": self.bos_token,
'eos_token': self.eos_token, "eos_token": self.eos_token,
'pattern': self.pattern, "pattern": self.pattern,
'tokenizer_type': self.__class__.__name__ "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) json.dump(config, f, ensure_ascii=False, indent=2)
@classmethod @classmethod
def load(cls, filepath: str): def load(cls, filepath: str):
""" """
Загружает BPE токенизатор из файла. Загружает BPE токенизатор из файла.
Args: Args:
filepath: Путь к файлу filepath: Путь к файлу
Returns: Returns:
BPETokenizer: Загруженный токенизатор BPETokenizer: Загруженный токенизатор
""" """
import json import json
with open(filepath, 'r', encoding='utf-8') as f: with open(filepath, "r", encoding="utf-8") as f:
config = json.load(f) config = json.load(f)
tokenizer = cls() tokenizer = cls()
tokenizer.vocab = config['vocab'] tokenizer.vocab = config["vocab"]
tokenizer.vocab_size = config['vocab_size'] tokenizer.vocab_size = config["vocab_size"]
tokenizer.pad_token = config['pad_token'] tokenizer.pad_token = config["pad_token"]
tokenizer.unk_token = config['unk_token'] tokenizer.unk_token = config["unk_token"]
tokenizer.bos_token = config['bos_token'] tokenizer.bos_token = config["bos_token"]
tokenizer.eos_token = config['eos_token'] tokenizer.eos_token = config["eos_token"]
tokenizer.pattern = config.get('pattern', tokenizer.pattern) tokenizer.pattern = config.get("pattern", tokenizer.pattern)
tokenizer.compiled_pattern = re.compile(tokenizer.pattern, re.UNICODE) tokenizer.compiled_pattern = re.compile(tokenizer.pattern, re.UNICODE)
# Восстанавливаем мерджи # Восстанавливаем мерджи
merges = config.get('merges', {}) merges = config.get("merges", {})
tokenizer.merges = {} tokenizer.merges = {}
for k, v in merges.items(): for k, v in merges.items():
parts = k.split() parts = k.split()
if len(parts) == 2: if len(parts) == 2:
tokenizer.merges[(parts[0], parts[1])] = v tokenizer.merges[(parts[0], parts[1])] = v
# Создаем обратный словарь # Создаем обратный словарь
tokenizer.inverse_vocab = {v: k for k, v in tokenizer.vocab.items()} tokenizer.inverse_vocab = {v: k for k, v in tokenizer.vocab.items()}
# Обновляем ID специальных токенов # Обновляем ID специальных токенов
tokenizer.pad_token_id = tokenizer.vocab.get(tokenizer.pad_token) tokenizer.pad_token_id = tokenizer.vocab.get(tokenizer.pad_token)
tokenizer.unk_token_id = tokenizer.vocab.get(tokenizer.unk_token) tokenizer.unk_token_id = tokenizer.vocab.get(tokenizer.unk_token)
tokenizer.bos_token_id = tokenizer.vocab.get(tokenizer.bos_token) tokenizer.bos_token_id = tokenizer.vocab.get(tokenizer.bos_token)
tokenizer.eos_token_id = tokenizer.vocab.get(tokenizer.eos_token) tokenizer.eos_token_id = tokenizer.vocab.get(tokenizer.eos_token)
return tokenizer return tokenizer
@@ -367,62 +383,72 @@ class SimpleBPETokenizer(BPETokenizer):
""" """
Упрощенная версия BPE токенизатора для демонстрации. Упрощенная версия BPE токенизатора для демонстрации.
""" """
def train(self, texts: List[str], vocab_size: int = 1000, **kwargs): def train(self, texts: List[str], vocab_size: int = 1000, **kwargs):
"""Упрощенное обучение для демонстрации.""" """Упрощенное обучение для демонстрации."""
# Инициализация базового словаря # Инициализация базового словаря
self._initialize_vocab() 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) self.add_special_tokens(special_tokens)
# Простая реализация - собираем все символы # Простая реализация - собираем все символы
all_chars = set() all_chars = set()
for text in texts: for text in texts:
all_chars.update(text) all_chars.update(text)
# Добавляем символы в словарь # Добавляем символы в словарь
for char in sorted(all_chars): for char in sorted(all_chars):
if char not in self.vocab: if char not in self.vocab:
self.vocab[char] = len(self.vocab) self.vocab[char] = len(self.vocab)
self.inverse_vocab = {v: k for k, v in self.vocab.items()} self.inverse_vocab = {v: k for k, v in self.vocab.items()}
self.vocab_size = len(self.vocab) self.vocab_size = len(self.vocab)
# Обновляем ID специальных токенов # Обновляем ID специальных токенов
self.pad_token_id = self.vocab.get(self.pad_token) self.pad_token_id = self.vocab.get(self.pad_token)
self.unk_token_id = self.vocab.get(self.unk_token) self.unk_token_id = self.vocab.get(self.unk_token)
self.bos_token_id = self.vocab.get(self.bos_token) self.bos_token_id = self.vocab.get(self.bos_token)
self.eos_token_id = self.vocab.get(self.eos_token) self.eos_token_id = self.vocab.get(self.eos_token)
def encode(self, text: str, **kwargs) -> List[int]: 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 = [] token_ids = []
for char in text: for char in text:
token_id = self.vocab.get(char, self.unk_token_id) token_id = self.vocab.get(char, self.unk_token_id)
if token_id is not None: if token_id is not None:
token_ids.append(token_id) token_ids.append(token_id)
if add_special_tokens: if add_special_tokens:
if self.bos_token_id is not None: if self.bos_token_id is not None:
token_ids.insert(0, self.bos_token_id) token_ids.insert(0, self.bos_token_id)
if self.eos_token_id is not None: if self.eos_token_id is not None:
token_ids.append(self.eos_token_id) token_ids.append(self.eos_token_id)
return token_ids return token_ids
def decode(self, tokens: List[int], **kwargs) -> str: 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 = [] chars = []
for token_id in tokens: for token_id in tokens:
char = self.inverse_vocab.get(token_id, self.unk_token) 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 continue
chars.append(char) 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): class BPETokenizer(BaseTokenizer):
""" """
BPE токенизатор для обработки текста. BPE токенизатор для обработки текста.
Реализует алгоритм Byte Pair Encoding для создания субсловных токенов. Реализует алгоритм Byte Pair Encoding для создания субсловных токенов.
Использует вашу реализацию BPE. Использует вашу реализацию BPE.
Примеры использования: Примеры использования:
>>> tokenizer = BPETokenizer() >>> tokenizer = BPETokenizer()
>>> tokenizer.train(["пример текста для обучения"], vocab_size=1000) >>> tokenizer.train(["пример текста для обучения"], vocab_size=1000)
>>> tokens = tokenizer.encode("новый текст") >>> tokens = tokenizer.encode("новый текст")
>>> text = tokenizer.decode(tokens) >>> text = tokenizer.decode(tokens)
""" """
def __init__(self): def __init__(self):
super().__init__() super().__init__()
self.merges: Dict[Tuple[str, str], int] = {} self.merges: Dict[Tuple[str, str], int] = {}
self.vocab_list: List[str] = [] self.vocab_list: List[str] = []
def train(self, texts: List[str], vocab_size: int = 1000, **kwargs): def train(self, texts: List[str], vocab_size: int = 1000, **kwargs):
""" """
Обучение BPE токенизатора на текстах. Обучение BPE токенизатора на текстах.
Args: Args:
texts: Список текстов для обучения texts: Список текстов для обучения
vocab_size: Желаемый размер словаря vocab_size: Желаемый размер словаря
@@ -39,7 +39,7 @@ class BPETokenizer(BaseTokenizer):
""" """
# Объединяем все тексты в одну строку для обучения # Объединяем все тексты в одну строку для обучения
combined_text = " ".join(texts) combined_text = " ".join(texts)
# 1. Получаем уникальные токены (символы) # 1. Получаем уникальные токены (символы)
unique_tokens = sorted(set(combined_text)) unique_tokens = sorted(set(combined_text))
tokens = unique_tokens.copy() tokens = unique_tokens.copy()
@@ -61,7 +61,10 @@ class BPETokenizer(BaseTokenizer):
break # нет пар — выходим 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] new_token = most_frequent_pair[0] + most_frequent_pair[1]
@@ -71,45 +74,51 @@ class BPETokenizer(BaseTokenizer):
new_sequence = [] new_sequence = []
while i < len(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) new_sequence.append(new_token)
i += 2 # пропускаем два символа — заменённую пару i += 2 # пропускаем два символа — заменённую пару
else: else:
new_sequence.append(sequence[i]) new_sequence.append(sequence[i])
i += 1 i += 1
sequence = new_sequence sequence = new_sequence
# 4. Создаем словари # 4. Создаем словари
self.vocab_list = tokens.copy() self.vocab_list = tokens.copy()
self.vocab = dict(zip(tokens, range(vocab_size))) self.vocab = dict(zip(tokens, range(vocab_size)))
self.inverse_vocab = dict(zip(range(vocab_size), tokens)) self.inverse_vocab = dict(zip(range(vocab_size), tokens))
self.vocab_size = len(self.vocab) 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) self.add_special_tokens(special_tokens)
def _pair_first_index(self, sequence, pair): def _pair_first_index(self, sequence, pair):
"""Находит первый индекс пары в последовательности.""" """Находит первый индекс пары в последовательности."""
for i in range(len(sequence) - 1): for i in range(len(sequence) - 1):
if (sequence[i], sequence[i + 1]) == pair: if (sequence[i], sequence[i + 1]) == pair:
return i return i
return float('inf') # если пара не найдена (в теории не должно случиться) return float("inf") # если пара не найдена (в теории не должно случиться)
def encode(self, text: str, **kwargs) -> List[int]: def encode(self, text: str, **kwargs) -> List[int]:
""" """
Кодирует текст в последовательность токенов. Кодирует текст в последовательность токенов.
Args: Args:
text: Входной текст text: Входной текст
**kwargs: Дополнительные параметры **kwargs: Дополнительные параметры
- add_special_tokens: Добавлять специальные токены - add_special_tokens: Добавлять специальные токены
Returns: Returns:
List[int]: Список идентификаторов токенов List[int]: Список идентификаторов токенов
""" """
add_special_tokens = kwargs.get('add_special_tokens', False) add_special_tokens = kwargs.get("add_special_tokens", False)
# 1. Разбиваем текст на токены-символы # 1. Разбиваем текст на токены-символы
sequence = list(text) sequence = list(text)
# 2. Инициализация пустого списка токенов # 2. Инициализация пустого списка токенов
@@ -119,7 +128,9 @@ class BPETokenizer(BaseTokenizer):
while i < len(text): while i < len(text):
# 3.1 Найти все токены в словаре, начинающиеся с text[i] # 3.1 Найти все токены в словаре, начинающиеся с text[i]
start_char = 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 Выбрать самый длинный подходящий токен # 3.2 Выбрать самый длинный подходящий токен
find_token = self._find_max_matching_token(text[i:], result) find_token = self._find_max_matching_token(text[i:], result)
if find_token is None: if find_token is None:
@@ -134,19 +145,19 @@ class BPETokenizer(BaseTokenizer):
# 4. Заменить токены на их ID # 4. Заменить токены на их ID
token_ids = self._tokens_to_ids(tokens) token_ids = self._tokens_to_ids(tokens)
# Заменяем -1 на unk_token_id # Заменяем -1 на unk_token_id
token_ids = [tid if tid != -1 else self.unk_token_id for tid in token_ids] token_ids = [tid if tid != -1 else self.unk_token_id for tid in token_ids]
# Добавляем специальные токены если нужно # Добавляем специальные токены если нужно
if add_special_tokens: if add_special_tokens:
if self.bos_token_id is not None: if self.bos_token_id is not None:
token_ids.insert(0, self.bos_token_id) token_ids.insert(0, self.bos_token_id)
if self.eos_token_id is not None: if self.eos_token_id is not None:
token_ids.append(self.eos_token_id) token_ids.append(self.eos_token_id)
return token_ids return token_ids
def _find_max_matching_token(self, text: str, tokens: list) -> Optional[str]: def _find_max_matching_token(self, text: str, tokens: list) -> Optional[str]:
"""Находит самый длинный токен из списка, с которого начинается текст""" """Находит самый длинный токен из списка, с которого начинается текст"""
matching = [token for token in tokens if text.startswith(token)] matching = [token for token in tokens if text.startswith(token)]
@@ -161,33 +172,41 @@ class BPETokenizer(BaseTokenizer):
else: else:
ids.append(-1) # Специальное значение ids.append(-1) # Специальное значение
return ids return ids
def decode(self, tokens: List[int], **kwargs) -> str: def decode(self, tokens: List[int], **kwargs) -> str:
""" """
Декодирует последовательность токенов в текст. Декодирует последовательность токенов в текст.
Args: Args:
tokens: Список идентификаторов токенов tokens: Список идентификаторов токенов
**kwargs: Дополнительные параметры **kwargs: Дополнительные параметры
- skip_special_tokens: Пропускать специальные токены - skip_special_tokens: Пропускать специальные токены
Returns: Returns:
str: Декодированный текст str: Декодированный текст
""" """
skip_special_tokens = kwargs.get('skip_special_tokens', True) skip_special_tokens = kwargs.get("skip_special_tokens", True)
# Фильтруем специальные токены если нужно # Фильтруем специальные токены если нужно
if skip_special_tokens: if skip_special_tokens:
tokens = [tid for tid in tokens if tid not in [ tokens = [
self.pad_token_id, self.unk_token_id, self.bos_token_id, self.eos_token_id 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 в токены # Конвертируем ID в токены
token_strings = self._ids_to_tokens(tokens) 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]: def _ids_to_tokens(self, ids: List[int]) -> List[str]:
"""Конвертирует список Ids в их tokens""" """Конвертирует список Ids в их tokens"""
tokens = [] tokens = []
@@ -197,76 +216,76 @@ class BPETokenizer(BaseTokenizer):
else: else:
tokens.append(self.unk_token) # Специальное значение tokens.append(self.unk_token) # Специальное значение
return tokens return tokens
def save(self, filepath: str): def save(self, filepath: str):
""" """
Сохраняет токенизатор в файл. Сохраняет токенизатор в файл.
Args: Args:
filepath: Путь для сохранения filepath: Путь для сохранения
""" """
import json import json
# Преобразуем кортежи в строки для JSON сериализации # Преобразуем кортежи в строки для JSON сериализации
merges_serializable = {f"{k[0]},{k[1]}": v for k, v in self.merges.items()} merges_serializable = {f"{k[0]},{k[1]}": v for k, v in self.merges.items()}
config = { config = {
'vocab': self.vocab, "vocab": self.vocab,
'vocab_size': self.vocab_size, "vocab_size": self.vocab_size,
'pad_token': self.pad_token, "pad_token": self.pad_token,
'unk_token': self.unk_token, "unk_token": self.unk_token,
'bos_token': self.bos_token, "bos_token": self.bos_token,
'eos_token': self.eos_token, "eos_token": self.eos_token,
'tokenizer_type': self.__class__.__name__, "tokenizer_type": self.__class__.__name__,
'merges': merges_serializable, "merges": merges_serializable,
'vocab_list': self.vocab_list "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) json.dump(config, f, ensure_ascii=False, indent=2)
@classmethod @classmethod
def load(cls, filepath: str): def load(cls, filepath: str):
""" """
Загружает токенизатор из файла. Загружает токенизатор из файла.
Args: Args:
filepath: Путь к файлу filepath: Путь к файлу
Returns: Returns:
BPETokenizer: Загруженный токенизатор BPETokenizer: Загруженный токенизатор
""" """
import json import json
with open(filepath, 'r', encoding='utf-8') as f: with open(filepath, "r", encoding="utf-8") as f:
config = json.load(f) config = json.load(f)
# Создаем экземпляр токенизатора # Создаем экземпляр токенизатора
tokenizer = cls() tokenizer = cls()
tokenizer.vocab = config['vocab'] tokenizer.vocab = config["vocab"]
tokenizer.vocab_size = config['vocab_size'] tokenizer.vocab_size = config["vocab_size"]
tokenizer.pad_token = config['pad_token'] tokenizer.pad_token = config["pad_token"]
tokenizer.unk_token = config['unk_token'] tokenizer.unk_token = config["unk_token"]
tokenizer.bos_token = config['bos_token'] tokenizer.bos_token = config["bos_token"]
tokenizer.eos_token = config['eos_token'] tokenizer.eos_token = config["eos_token"]
tokenizer.vocab_list = config['vocab_list'] tokenizer.vocab_list = config["vocab_list"]
# Восстанавливаем кортежи из строк # Восстанавливаем кортежи из строк
tokenizer.merges = {} tokenizer.merges = {}
for k, v in config['merges'].items(): for k, v in config["merges"].items():
parts = k.split(',') parts = k.split(",")
if len(parts) == 2: if len(parts) == 2:
tokenizer.merges[(parts[0], parts[1])] = v tokenizer.merges[(parts[0], parts[1])] = v
# Создаем обратный словарь # Создаем обратный словарь
tokenizer.inverse_vocab = {v: k for k, v in tokenizer.vocab.items()} tokenizer.inverse_vocab = {v: k for k, v in tokenizer.vocab.items()}
# Обновляем ID специальных токенов # Обновляем ID специальных токенов
tokenizer.pad_token_id = tokenizer.vocab.get(tokenizer.pad_token) tokenizer.pad_token_id = tokenizer.vocab.get(tokenizer.pad_token)
tokenizer.unk_token_id = tokenizer.vocab.get(tokenizer.unk_token) tokenizer.unk_token_id = tokenizer.vocab.get(tokenizer.unk_token)
tokenizer.bos_token_id = tokenizer.vocab.get(tokenizer.bos_token) tokenizer.bos_token_id = tokenizer.vocab.get(tokenizer.bos_token)
tokenizer.eos_token_id = tokenizer.vocab.get(tokenizer.eos_token) tokenizer.eos_token_id = tokenizer.vocab.get(tokenizer.eos_token)
return tokenizer return tokenizer
@@ -275,4 +294,5 @@ class SimpleBPETokenizer(BPETokenizer):
Упрощенная версия BPE токенизатора для демонстрации. Упрощенная версия BPE токенизатора для демонстрации.
Наследует вашу реализацию, но может быть упрощена при необходимости. Наследует вашу реализацию, но может быть упрощена при необходимости.
""" """
pass 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): def __init__(self, texts: List[str], tokenizer: Any, block_size: int = 128):
""" """
Инициализация датасета. Инициализация датасета.
Args: Args:
texts: Список текстов для обучения texts: Список текстов для обучения
tokenizer: Токенизатор с методами encode/decode tokenizer: Токенизатор с методами encode/decode
@@ -25,15 +25,15 @@ class TextDataset(Dataset):
for text in texts: for text in texts:
# Кодируем текст в токены # Кодируем текст в токены
input_ids = tokenizer.encode(text, add_special_tokens=False) input_ids = tokenizer.encode(text, add_special_tokens=False)
# Обрезаем или дополняем до нужной длины # Обрезаем или дополняем до нужной длины
if len(input_ids) > block_size: if len(input_ids) > block_size:
input_ids = input_ids[:block_size] input_ids = input_ids[:block_size]
else: else:
# Дополняем pad_token_id # Дополняем 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)) input_ids = input_ids + [pad_token_id] * (block_size - len(input_ids))
self.examples.append(input_ids) self.examples.append(input_ids)
def __len__(self): def __len__(self):
@@ -50,33 +50,35 @@ class StreamingTextDataset(Dataset):
Датасет для потоковой обработки больших текстов. Датасет для потоковой обработки больших текстов.
Токенизация происходит на лету, что экономит память. Токенизация происходит на лету, что экономит память.
""" """
def __init__(self, texts: List[str], tokenizer: Any, block_size: int = 128): def __init__(self, texts: List[str], tokenizer: Any, block_size: int = 128):
self.texts = texts self.texts = texts
self.tokenizer = tokenizer self.tokenizer = tokenizer
self.block_size = block_size self.block_size = block_size
# Получаем pad_token_id из токенизатора # Получаем 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): def __len__(self):
return len(self.texts) return len(self.texts)
def __getitem__(self, idx): def __getitem__(self, idx):
text = self.texts[idx] text = self.texts[idx]
# Токенизация на лету # Токенизация на лету
input_ids = self.tokenizer.encode(text, add_special_tokens=False) input_ids = self.tokenizer.encode(text, add_special_tokens=False)
# Обрезаем или дополняем до нужной длины # Обрезаем или дополняем до нужной длины
if len(input_ids) > self.block_size: if len(input_ids) > self.block_size:
input_ids = input_ids[:self.block_size] input_ids = input_ids[: self.block_size]
else: 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) input_ids = torch.tensor(input_ids, dtype=torch.long)
labels = input_ids.clone() labels = input_ids.clone()
return {"input_ids": input_ids, "labels": labels} return {"input_ids": input_ids, "labels": labels}
@@ -84,9 +86,15 @@ class TextDatasetWithSpecialTokens(TextDataset):
""" """
Расширенная версия TextDataset с поддержкой специальных токенов. Расширенная версия TextDataset с поддержкой специальных токенов.
""" """
def __init__(self, texts: List[str], tokenizer: Any, block_size: int = 128, def __init__(
add_bos: bool = False, add_eos: bool = False): self,
texts: List[str],
tokenizer: Any,
block_size: int = 128,
add_bos: bool = False,
add_eos: bool = False,
):
""" """
Args: Args:
texts: Список текстов texts: Список текстов
@@ -104,33 +112,38 @@ class TextDatasetWithSpecialTokens(TextDataset):
for text in texts: for text in texts:
# Кодируем с специальными токенами # Кодируем с специальными токенами
input_ids = tokenizer.encode( input_ids = tokenizer.encode(
text, text, add_special_tokens=True, add_bos_token=add_bos, add_eos_token=eos
add_special_tokens=True,
add_bos_token=add_bos,
add_eos_token=eos
) )
# Учитываем специальные токены при обрезке/дополнении # Учитываем специальные токены при обрезке/дополнении
effective_block_size = block_size effective_block_size = block_size
if add_bos: if add_bos:
effective_block_size -= 1 effective_block_size -= 1
if add_eos: if add_eos:
effective_block_size -= 1 effective_block_size -= 1
if len(input_ids) > effective_block_size: if len(input_ids) > effective_block_size:
input_ids = 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 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] 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: if len(input_ids) < block_size:
input_ids = input_ids + [pad_token_id] * (block_size - len(input_ids)) input_ids = input_ids + [pad_token_id] * (block_size - len(input_ids))
self.examples.append(input_ids) self.examples.append(input_ids)
def __len__(self): def __len__(self):

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

@@ -1,5 +1,6 @@
import torch.optim as optim import torch.optim as optim
def get_optimizer(model, lr=3e-4, weight_decay=0.01, optimizer_type="adamw"): 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 from torch.optim.lr_scheduler import LambdaLR
def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps): def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps):
""" """
Линейный планировщик обучения с warmup. Линейный планировщик обучения с warmup.
@@ -8,6 +9,10 @@ def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_st
def lr_lambda(current_step): def lr_lambda(current_step):
if current_step < num_warmup_steps: if current_step < num_warmup_steps:
return float(current_step) / float(max(1, 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) 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.optimizer import get_optimizer
from llm.training.scheduler import get_linear_schedule_with_warmup from llm.training.scheduler import get_linear_schedule_with_warmup
class Trainer: class Trainer:
""" """
Универсальный класс обучения LLM (GPT, LLaMA, Mistral и т.д.) Универсальный класс обучения 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.model = model
self.train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) self.train_loader = DataLoader(
self.val_loader = DataLoader(val_dataset, batch_size=batch_size) if val_dataset else None 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.optimizer = get_optimizer(model, lr=lr)
self.scheduler = None self.scheduler = None
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -29,24 +43,28 @@ class Trainer:
# Сдвигаем логиты и метки для языкового моделирования # Сдвигаем логиты и метки для языкового моделирования
shift_logits = logits[..., :-1, :].contiguous() shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous() shift_labels = labels[..., 1:].contiguous()
# Вычисляем cross-entropy loss # Вычисляем cross-entropy loss
loss = F.cross_entropy( loss = F.cross_entropy(
shift_logits.view(-1, shift_logits.size(-1)), shift_logits.view(-1, shift_logits.size(-1)),
shift_labels.view(-1), shift_labels.view(-1),
ignore_index=-100 # Игнорируем padding tokens ignore_index=-100, # Игнорируем padding tokens
) )
return loss return loss
def train(self): def train(self):
total_steps = len(self.train_loader) * self.num_epochs 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): for epoch in range(self.num_epochs):
self.model.train() self.model.train()
total_loss = 0 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: for batch in progress_bar:
self.optimizer.zero_grad() self.optimizer.zero_grad()
@@ -59,7 +77,7 @@ class Trainer:
logits = outputs[0] logits = outputs[0]
else: else:
logits = outputs logits = outputs
# Trainer вычисляет loss # Trainer вычисляет loss
loss = self.compute_lm_loss(logits, labels) loss = self.compute_lm_loss(logits, labels)
loss.backward() loss.backward()
@@ -85,7 +103,7 @@ class Trainer:
for batch in self.val_loader: for batch in self.val_loader:
input_ids = batch["input_ids"].to(self.device) input_ids = batch["input_ids"].to(self.device)
labels = batch["labels"].to(self.device) labels = batch["labels"].to(self.device)
outputs = self.model(input_ids) outputs = self.model(input_ids)
if isinstance(outputs, tuple): if isinstance(outputs, tuple):
logits = outputs[0] 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_heads": num_heads,
"num_layers": num_layers, "num_layers": num_layers,
"max_position_embeddings": 1024, "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)) input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
return input_ids return input_ids
@pytest.fixture @pytest.fixture
def random_float_inputs(batch_size, seq_len, embed_dim): def random_float_inputs(batch_size, seq_len, embed_dim):
"""Generate random floating point input tensors for testing feed forward.""" """Generate random floating point input tensors for testing feed forward."""
inputs = torch.randn(batch_size, seq_len, embed_dim) inputs = torch.randn(batch_size, seq_len, embed_dim)
return inputs return inputs
@pytest.fixture @pytest.fixture
def random_embeddings(batch_size, seq_len, embed_dim): def random_embeddings(batch_size, seq_len, embed_dim):
"""Generate random embedding tensors for testing attention modules.""" """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: class TestDecoder:
"""Test cases for Decoder.""" """Test cases for Decoder."""
def test_initialization(self, embed_dim, num_heads): def test_initialization(self, embed_dim, num_heads):
"""Test that Decoder can be initialized.""" """Test that Decoder can be initialized."""
head_size = embed_dim // num_heads head_size = embed_dim // num_heads
max_seq_len = 1024 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 assert decoder is not None
# Check internal components # Check internal components
assert hasattr(decoder, '_heads') assert hasattr(decoder, "_heads")
assert hasattr(decoder, '_ff') assert hasattr(decoder, "_ff")
assert hasattr(decoder, '_norm1') assert hasattr(decoder, "_norm1")
assert hasattr(decoder, '_norm2') assert hasattr(decoder, "_norm2")
def test_forward_pass(self, embed_dim, num_heads, random_embeddings): def test_forward_pass(self, embed_dim, num_heads, random_embeddings):
"""Test forward pass of Decoder.""" """Test forward pass of Decoder."""
head_size = embed_dim // num_heads head_size = embed_dim // num_heads
max_seq_len = 1024 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 # Forward pass
output = decoder(random_embeddings) output = decoder(random_embeddings)
# Check output shape # Check output shape
assert output.shape == random_embeddings.shape assert output.shape == random_embeddings.shape
assert isinstance(output, torch.Tensor) assert isinstance(output, torch.Tensor)
def test_forward_with_causal_mask(self, embed_dim, num_heads, random_embeddings): def test_forward_with_causal_mask(self, embed_dim, num_heads, random_embeddings):
"""Test forward pass with causal mask.""" """Test forward pass with causal mask."""
head_size = embed_dim // num_heads head_size = embed_dim // num_heads
max_seq_len = 1024 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] batch_size, seq_len = random_embeddings.shape[:2]
# Create causal mask # Create causal mask
mask = torch.tril(torch.ones(seq_len, seq_len)) mask = torch.tril(torch.ones(seq_len, seq_len))
# Forward pass with causal mask # Forward pass with causal mask
output = decoder(random_embeddings, mask=mask) output = decoder(random_embeddings, mask=mask)
# Check output shape # Check output shape
assert output.shape == random_embeddings.shape assert output.shape == random_embeddings.shape
def test_residual_connections(self, embed_dim, num_heads, random_embeddings): def test_residual_connections(self, embed_dim, num_heads, random_embeddings):
"""Test that residual connections are properly applied.""" """Test that residual connections are properly applied."""
head_size = embed_dim // num_heads head_size = embed_dim // num_heads
max_seq_len = 1024 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) output = decoder(random_embeddings)
# With residual connections and layer norm, the output shouldn't be # With residual connections and layer norm, the output shouldn't be
# too different from input (in terms of scale/distribution) # too different from input (in terms of scale/distribution)
input_norm = random_embeddings.norm(dim=-1).mean() input_norm = random_embeddings.norm(dim=-1).mean()
output_norm = output.norm(dim=-1).mean() output_norm = output.norm(dim=-1).mean()
# Norms should be of similar magnitude (not exact due to transformations) # Norms should be of similar magnitude (not exact due to transformations)
assert 0.1 < (output_norm / input_norm) < 10.0 assert 0.1 < (output_norm / input_norm) < 10.0
def test_layer_norm(self, embed_dim, num_heads, random_embeddings): def test_layer_norm(self, embed_dim, num_heads, random_embeddings):
"""Test that layer normalization is applied.""" """Test that layer normalization is applied."""
head_size = embed_dim // num_heads head_size = embed_dim // num_heads
max_seq_len = 1024 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) output = decoder(random_embeddings)
# Check that output has reasonable statistics (due to layer norm) # Check that output has reasonable statistics (due to layer norm)
# Mean should be close to 0, std close to 1 for each sequence position # Mean should be close to 0, std close to 1 for each sequence position
output_mean = output.mean(dim=-1) output_mean = output.mean(dim=-1)
output_std = output.std(dim=-1) output_std = output.std(dim=-1)
# These are approximate checks since the data goes through multiple transformations # 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_mean, torch.zeros_like(output_mean), atol=1.0)
assert torch.allclose(output_std, torch.ones_like(output_std), atol=2.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): def test_gradient_flow(self, embed_dim, num_heads, random_embeddings):
"""Test that gradients flow through Decoder.""" """Test that gradients flow through Decoder."""
head_size = embed_dim // num_heads head_size = embed_dim // num_heads
max_seq_len = 1024 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 # Forward pass
output = decoder(random_embeddings) output = decoder(random_embeddings)
# Create a dummy loss and backward pass # Create a dummy loss and backward pass
loss = output.sum() loss = output.sum()
loss.backward() loss.backward()
# Check that gradients are computed for learnable parameters # Check that gradients are computed for learnable parameters
# in attention and feed forward components # in attention and feed forward components
assert decoder._heads._layer.weight.grad is not None assert decoder._heads._layer.weight.grad is not None
assert decoder._ff._layer1.weight.grad is not None assert decoder._ff._layer1.weight.grad is not None
assert decoder._norm1.weight.grad is not None assert decoder._norm1.weight.grad is not None
assert decoder._norm2.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): def test_device_consistency(self, embed_dim, num_heads, random_embeddings, device):
"""Test that Decoder works on correct device.""" """Test that Decoder works on correct device."""
head_size = embed_dim // num_heads head_size = embed_dim // num_heads
max_seq_len = 1024 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) inputs = random_embeddings.to(device)
# Forward pass # Forward pass
output = decoder(inputs) output = decoder(inputs)
# Check device consistency # Check device consistency
assert output.device == device assert output.device == device
assert decoder._heads._layer.weight.device == device assert decoder._heads._layer.weight.device == device
def test_different_configurations(self): def test_different_configurations(self):
"""Test Decoder with different configurations.""" """Test Decoder with different configurations."""
test_cases = [ 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 (128, 4), # embed_dim=128, num_heads=4
(256, 8), # embed_dim=256, num_heads=8 (256, 8), # embed_dim=256, num_heads=8
] ]
for embed_dim, num_heads in test_cases: for embed_dim, num_heads in test_cases:
head_size = embed_dim // num_heads head_size = embed_dim // num_heads
max_seq_len = 1024 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 batch_size, seq_len = 2, 16
inputs = torch.randn(batch_size, seq_len, embed_dim) inputs = torch.randn(batch_size, seq_len, embed_dim)
output = decoder(inputs) output = decoder(inputs)
assert output.shape == inputs.shape assert output.shape == inputs.shape
@pytest.mark.parametrize("batch_size,seq_len", [(1, 8), (2, 16), (4, 32)]) @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): def test_different_input_shapes(self, embed_dim, num_heads, batch_size, seq_len):
"""Test Decoder with different input shapes.""" """Test Decoder with different input shapes."""
head_size = embed_dim // num_heads head_size = embed_dim // num_heads
max_seq_len = 1024 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) inputs = torch.randn(batch_size, seq_len, embed_dim)
output = decoder(inputs) output = decoder(inputs)
assert output.shape == (batch_size, seq_len, embed_dim) assert output.shape == (batch_size, seq_len, embed_dim)
def test_training_vs_evaluation(self, embed_dim, num_heads, random_embeddings): def test_training_vs_evaluation(self, embed_dim, num_heads, random_embeddings):
"""Test that Decoder behaves differently in train vs eval mode.""" """Test that Decoder behaves differently in train vs eval mode."""
head_size = embed_dim // num_heads head_size = embed_dim // num_heads
max_seq_len = 1024 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 # Training mode
decoder.train() decoder.train()
output_train = decoder(random_embeddings) output_train = decoder(random_embeddings)
# Evaluation mode # Evaluation mode
decoder.eval() decoder.eval()
output_eval = decoder(random_embeddings) output_eval = decoder(random_embeddings)
# Outputs should be different due to dropout # Outputs should be different due to dropout
assert not torch.allclose(output_train, output_eval) assert not torch.allclose(output_train, output_eval)
def test_parameter_initialization(self, embed_dim, num_heads): def test_parameter_initialization(self, embed_dim, num_heads):
"""Test that parameters are properly initialized.""" """Test that parameters are properly initialized."""
head_size = embed_dim // num_heads head_size = embed_dim // num_heads
max_seq_len = 1024 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 # Check that various components have non-zero parameters
assert not torch.allclose( assert not torch.allclose(
decoder._heads._layer.weight, decoder._heads._layer.weight, torch.zeros_like(decoder._heads._layer.weight)
torch.zeros_like(decoder._heads._layer.weight)
) )
assert not torch.allclose( assert not torch.allclose(
decoder._ff._layer1.weight, decoder._ff._layer1.weight, torch.zeros_like(decoder._ff._layer1.weight)
torch.zeros_like(decoder._ff._layer1.weight)
) )
assert not torch.allclose( assert not torch.allclose(
decoder._norm1.weight, decoder._norm1.weight, torch.zeros_like(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: class TestFeedForward:
"""Test cases for FeedForward.""" """Test cases for FeedForward."""
def test_initialization(self, embed_dim): def test_initialization(self, embed_dim):
"""Test that FeedForward can be initialized.""" """Test that FeedForward can be initialized."""
ff = FeedForward(embed_dim) ff = FeedForward(embed_dim)
assert ff is not None assert ff is not None
# Check internal layers # Check internal layers
assert hasattr(ff, '_layer1') assert hasattr(ff, "_layer1")
assert hasattr(ff, '_layer2') assert hasattr(ff, "_layer2")
assert hasattr(ff, '_activation') assert hasattr(ff, "_activation")
assert hasattr(ff, '_dropout') assert hasattr(ff, "_dropout")
# Check layer dimensions # Check layer dimensions
expected_hidden_dim = embed_dim * 4 # Default expansion factor expected_hidden_dim = embed_dim * 4 # Default expansion factor
assert ff._layer1.weight.shape == (expected_hidden_dim, embed_dim) assert ff._layer1.weight.shape == (expected_hidden_dim, embed_dim)
assert ff._layer2.weight.shape == (embed_dim, expected_hidden_dim) assert ff._layer2.weight.shape == (embed_dim, expected_hidden_dim)
def test_forward_pass(self, embed_dim, random_float_inputs): def test_forward_pass(self, embed_dim, random_float_inputs):
"""Test forward pass of FeedForward.""" """Test forward pass of FeedForward."""
ff = FeedForward(embed_dim) ff = FeedForward(embed_dim)
# Forward pass # Forward pass
output = ff(random_float_inputs) output = ff(random_float_inputs)
# Check output shape # Check output shape
assert output.shape == random_float_inputs.shape assert output.shape == random_float_inputs.shape
assert isinstance(output, torch.Tensor) assert isinstance(output, torch.Tensor)
def test_custom_hidden_dim(self, embed_dim): def test_custom_hidden_dim(self, embed_dim):
"""Test FeedForward with custom hidden dimension.""" """Test FeedForward with custom hidden dimension."""
# FeedForward doesn't support custom hidden_dim in current implementation # FeedForward doesn't support custom hidden_dim in current implementation
# This test is not applicable # This test is not applicable
ff = FeedForward(embed_dim) ff = FeedForward(embed_dim)
# Check layer dimensions (fixed 4x expansion) # Check layer dimensions (fixed 4x expansion)
expected_hidden_dim = embed_dim * 4 expected_hidden_dim = embed_dim * 4
assert ff._layer1.weight.shape == (expected_hidden_dim, embed_dim) assert ff._layer1.weight.shape == (expected_hidden_dim, embed_dim)
assert ff._layer2.weight.shape == (embed_dim, expected_hidden_dim) assert ff._layer2.weight.shape == (embed_dim, expected_hidden_dim)
def test_dropout(self, embed_dim, random_float_inputs): def test_dropout(self, embed_dim, random_float_inputs):
"""Test that dropout is applied during training.""" """Test that dropout is applied during training."""
ff = FeedForward(embed_dim, dropout=0.5) ff = FeedForward(embed_dim, dropout=0.5)
ff.train() # Set to training mode ff.train() # Set to training mode
output = ff(random_float_inputs) output = ff(random_float_inputs)
# In training mode with dropout, some values should be zeroed # In training mode with dropout, some values should be zeroed
# This is probabilistic, so we can't assert exact zeros, # This is probabilistic, so we can't assert exact zeros,
# but we can check the structure is preserved # but we can check the structure is preserved
assert output.shape == random_float_inputs.shape assert output.shape == random_float_inputs.shape
def test_no_dropout_in_eval(self, embed_dim, random_float_inputs): def test_no_dropout_in_eval(self, embed_dim, random_float_inputs):
"""Test that dropout is not applied during evaluation.""" """Test that dropout is not applied during evaluation."""
ff = FeedForward(embed_dim, dropout=0.5) ff = FeedForward(embed_dim, dropout=0.5)
ff.eval() # Set to evaluation mode ff.eval() # Set to evaluation mode
# Run forward pass multiple times - outputs should be identical # Run forward pass multiple times - outputs should be identical
output1 = ff(random_float_inputs) output1 = ff(random_float_inputs)
output2 = ff(random_float_inputs) output2 = ff(random_float_inputs)
assert torch.allclose(output1, output2) assert torch.allclose(output1, output2)
def test_activation_function(self, embed_dim, random_float_inputs): def test_activation_function(self, embed_dim, random_float_inputs):
"""Test that activation function is applied.""" """Test that activation function is applied."""
ff = FeedForward(embed_dim) ff = FeedForward(embed_dim)
# Manually compute expected output without dropout for deterministic comparison # Manually compute expected output without dropout for deterministic comparison
hidden = ff._layer1(random_float_inputs) hidden = ff._layer1(random_float_inputs)
activated = ff._activation(hidden) activated = ff._activation(hidden)
expected_output = ff._layer2(activated) expected_output = ff._layer2(activated)
# Compare with forward pass in eval mode (no dropout) # Compare with forward pass in eval mode (no dropout)
ff.eval() ff.eval()
actual_output = ff(random_float_inputs) actual_output = ff(random_float_inputs)
assert torch.allclose(actual_output, expected_output, rtol=1e-4) assert torch.allclose(actual_output, expected_output, rtol=1e-4)
def test_gradient_flow(self, embed_dim, random_float_inputs): def test_gradient_flow(self, embed_dim, random_float_inputs):
"""Test that gradients flow through FeedForward.""" """Test that gradients flow through FeedForward."""
ff = FeedForward(embed_dim) ff = FeedForward(embed_dim)
# Forward pass # Forward pass
output = ff(random_float_inputs) output = ff(random_float_inputs)
# Create a dummy loss and backward pass # Create a dummy loss and backward pass
loss = output.sum() loss = output.sum()
loss.backward() loss.backward()
# Check that gradients are computed for learnable parameters # Check that gradients are computed for learnable parameters
assert ff._layer1.weight.grad is not None assert ff._layer1.weight.grad is not None
assert ff._layer2.weight.grad is not None assert ff._layer2.weight.grad is not None
assert not torch.allclose(ff._layer1.weight.grad, assert not torch.allclose(
torch.zeros_like(ff._layer1.weight.grad)) 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._layer2.weight.grad, torch.zeros_like(ff._layer2.weight.grad)
)
def test_device_consistency(self, embed_dim, random_float_inputs, device): def test_device_consistency(self, embed_dim, random_float_inputs, device):
"""Test that FeedForward works on correct device.""" """Test that FeedForward works on correct device."""
ff = FeedForward(embed_dim).to(device) ff = FeedForward(embed_dim).to(device)
inputs = random_float_inputs.to(device) inputs = random_float_inputs.to(device)
# Forward pass # Forward pass
output = ff(inputs) output = ff(inputs)
# Check device consistency # Check device consistency
assert output.device == device assert output.device == device
assert ff._layer1.weight.device == device assert ff._layer1.weight.device == device
assert ff._layer2.weight.device == device assert ff._layer2.weight.device == device
def test_different_embed_dims(self): def test_different_embed_dims(self):
"""Test FeedForward with different embedding dimensions.""" """Test FeedForward with different embedding dimensions."""
test_cases = [64, 128, 256, 512] test_cases = [64, 128, 256, 512]
for embed_dim in test_cases: for embed_dim in test_cases:
ff = FeedForward(embed_dim) ff = FeedForward(embed_dim)
batch_size, seq_len = 2, 16 batch_size, seq_len = 2, 16
inputs = torch.randn(batch_size, seq_len, embed_dim) inputs = torch.randn(batch_size, seq_len, embed_dim)
output = ff(inputs) output = ff(inputs)
assert output.shape == inputs.shape assert output.shape == inputs.shape
@pytest.mark.parametrize("batch_size,seq_len", [(1, 8), (2, 16), (4, 32)]) @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): def test_different_input_shapes(self, embed_dim, batch_size, seq_len):
"""Test FeedForward with different input shapes.""" """Test FeedForward with different input shapes."""
ff = FeedForward(embed_dim) ff = FeedForward(embed_dim)
inputs = torch.randn(batch_size, seq_len, embed_dim) inputs = torch.randn(batch_size, seq_len, embed_dim)
output = ff(inputs) output = ff(inputs)
assert output.shape == (batch_size, seq_len, embed_dim) assert output.shape == (batch_size, seq_len, embed_dim)
def test_non_linearity(self, embed_dim, random_float_inputs): def test_non_linearity(self, embed_dim, random_float_inputs):
"""Test that FeedForward introduces non-linearity.""" """Test that FeedForward introduces non-linearity."""
ff = FeedForward(embed_dim) ff = FeedForward(embed_dim)
# Create a simple linear transformation for comparison # Create a simple linear transformation for comparison
linear_layer = nn.Linear(embed_dim, embed_dim) linear_layer = nn.Linear(embed_dim, embed_dim)
# Copy weights to make comparison fair # Copy weights to make comparison fair
with torch.no_grad(): with torch.no_grad():
linear_layer.weight.copy_(ff._layer2.weight @ ff._layer1.weight) linear_layer.weight.copy_(ff._layer2.weight @ ff._layer1.weight)
if linear_layer.bias is not None: if linear_layer.bias is not None:
linear_layer.bias.zero_() linear_layer.bias.zero_()
linear_output = linear_layer(random_float_inputs) linear_output = linear_layer(random_float_inputs)
ff_output = ff(random_float_inputs) ff_output = ff(random_float_inputs)
# FeedForward output should be different from pure linear transformation # FeedForward output should be different from pure linear transformation
# due to activation function # due to activation function
assert not torch.allclose(ff_output, linear_output, rtol=1e-4) assert not torch.allclose(ff_output, linear_output, rtol=1e-4)
def test_parameter_initialization(self, embed_dim): def test_parameter_initialization(self, embed_dim):
"""Test that parameters are properly initialized.""" """Test that parameters are properly initialized."""
ff = FeedForward(embed_dim) ff = FeedForward(embed_dim)
# Check that weights are not all zeros # Check that weights are not all zeros
assert not torch.allclose(ff._layer1.weight, torch.zeros_like(ff._layer1.weight)) assert not torch.allclose(
assert not torch.allclose(ff._layer2.weight, torch.zeros_like(ff._layer2.weight)) 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) # Check that biases are not all zeros (they should be initialized with some values)
if ff._layer1.bias is not None: 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: 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: class TestMultiHeadAttention:
"""Test cases for MultiHeadAttention.""" """Test cases for MultiHeadAttention."""
def test_initialization(self, embed_dim, num_heads): def test_initialization(self, embed_dim, num_heads):
"""Test that MultiHeadAttention can be initialized.""" """Test that MultiHeadAttention can be initialized."""
head_size = embed_dim // num_heads 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 assert attention is not None
# Check internal attributes # Check internal attributes
assert len(attention._heads) == num_heads assert len(attention._heads) == num_heads
assert attention._layer.in_features == embed_dim assert attention._layer.in_features == embed_dim
assert attention._layer.out_features == embed_dim assert attention._layer.out_features == embed_dim
def test_forward_pass(self, embed_dim, num_heads, random_embeddings): def test_forward_pass(self, embed_dim, num_heads, random_embeddings):
"""Test forward pass of MultiHeadAttention.""" """Test forward pass of MultiHeadAttention."""
head_size = embed_dim // num_heads 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 # Forward pass
output, _ = attention(random_embeddings) output, _ = attention(random_embeddings)
# Check output shape # Check output shape
assert output.shape == random_embeddings.shape assert output.shape == random_embeddings.shape
assert isinstance(output, torch.Tensor) assert isinstance(output, torch.Tensor)
def test_forward_with_mask(self, embed_dim, num_heads, random_embeddings): def test_forward_with_mask(self, embed_dim, num_heads, random_embeddings):
"""Test forward pass with attention mask.""" """Test forward pass with attention mask."""
head_size = embed_dim // num_heads 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 # Create a simple mask
seq_len = random_embeddings.shape[1] seq_len = random_embeddings.shape[1]
mask = torch.tril(torch.ones(seq_len, seq_len)) # Causal mask mask = torch.tril(torch.ones(seq_len, seq_len)) # Causal mask
# Forward pass with mask # Forward pass with mask
output, _ = attention(random_embeddings, mask=mask) output, _ = attention(random_embeddings, mask=mask)
# Check output shape # Check output shape
assert output.shape == random_embeddings.shape assert output.shape == random_embeddings.shape
def test_causal_mask(self, embed_dim, num_heads, random_embeddings): def test_causal_mask(self, embed_dim, num_heads, random_embeddings):
"""Test that causal mask prevents attending to future positions.""" """Test that causal mask prevents attending to future positions."""
head_size = embed_dim // num_heads 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 # Create causal mask
seq_len = random_embeddings.shape[1] seq_len = random_embeddings.shape[1]
causal_mask = torch.tril(torch.ones(seq_len, seq_len)) causal_mask = torch.tril(torch.ones(seq_len, seq_len))
# Forward pass with causal mask # Forward pass with causal mask
output, _ = attention(random_embeddings, mask=causal_mask) output, _ = attention(random_embeddings, mask=causal_mask)
# Check output shape # Check output shape
assert output.shape == random_embeddings.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.""" """Test that attention weights are properly normalized."""
head_size = embed_dim // num_heads 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 # Forward pass
output, _ = attention(random_embeddings) output, _ = attention(random_embeddings)
# Check output shape # Check output shape
assert output.shape == random_embeddings.shape assert output.shape == random_embeddings.shape
def test_gradient_flow(self, embed_dim, num_heads, random_embeddings): def test_gradient_flow(self, embed_dim, num_heads, random_embeddings):
"""Test that gradients flow through MultiHeadAttention.""" """Test that gradients flow through MultiHeadAttention."""
head_size = embed_dim // num_heads 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 # Forward pass
output, _ = attention(random_embeddings) output, _ = attention(random_embeddings)
# Create a dummy loss and backward pass # Create a dummy loss and backward pass
loss = output.sum() loss = output.sum()
loss.backward() loss.backward()
# Check that gradients are computed for learnable parameters # Check that gradients are computed for learnable parameters
assert attention._layer.weight.grad is not None assert attention._layer.weight.grad is not None
if len(attention._heads) > 0: if len(attention._heads) > 0:
assert attention._heads[0]._q.weight.grad is not None assert attention._heads[0]._q.weight.grad is not None
def test_device_consistency(self, embed_dim, num_heads, random_embeddings, device): def test_device_consistency(self, embed_dim, num_heads, random_embeddings, device):
"""Test that MultiHeadAttention works on correct device.""" """Test that MultiHeadAttention works on correct device."""
head_size = embed_dim // num_heads 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) inputs = random_embeddings.to(device)
# Forward pass # Forward pass
output, _ = attention(inputs) output, _ = attention(inputs)
# Check device consistency # Check device consistency
assert output.device == device assert output.device == device
assert attention._layer.weight.device == device assert attention._layer.weight.device == device
def test_different_embed_dim_and_heads(self): def test_different_embed_dim_and_heads(self):
"""Test MultiHeadAttention with different embed_dim and num_heads combinations.""" """Test MultiHeadAttention with different embed_dim and num_heads combinations."""
test_cases = [ 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 (128, 4), # embed_dim=128, num_heads=4
(256, 8), # embed_dim=256, num_heads=8 (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: for embed_dim, num_heads in test_cases:
head_size = embed_dim // num_heads 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 batch_size, seq_len = 2, 16
inputs = torch.randn(batch_size, seq_len, embed_dim) inputs = torch.randn(batch_size, seq_len, embed_dim)
output, _ = attention(inputs) output, _ = attention(inputs)
assert output.shape == inputs.shape assert output.shape == inputs.shape
def test_attention_output_range(self, embed_dim, num_heads, random_embeddings): def test_attention_output_range(self, embed_dim, num_heads, random_embeddings):
"""Test that attention output is in reasonable range.""" """Test that attention output is in reasonable range."""
head_size = embed_dim // num_heads 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, _ = attention(random_embeddings)
# Output shouldn't have extreme values # Output shouldn't have extreme values
assert output.abs().max() < 100 # Reasonable upper bound assert output.abs().max() < 100 # Reasonable upper bound
@pytest.mark.parametrize("batch_size,seq_len", [(1, 8), (2, 16), (4, 32)]) @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): def test_different_input_shapes(self, embed_dim, num_heads, batch_size, seq_len):
"""Test MultiHeadAttention with different input shapes.""" """Test MultiHeadAttention with different input shapes."""
head_size = embed_dim // num_heads 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) inputs = torch.randn(batch_size, seq_len, embed_dim)
output, _ = attention(inputs) output, _ = attention(inputs)
assert output.shape == (batch_size, seq_len, embed_dim) assert output.shape == (batch_size, seq_len, embed_dim)
def test_parameter_sharing(self, embed_dim, num_heads): def test_parameter_sharing(self, embed_dim, num_heads):
"""Test that parameters are properly shared across the sequence.""" """Test that parameters are properly shared across the sequence."""
head_size = embed_dim // num_heads 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 # Create two identical sequences
seq_len = 10 seq_len = 10
base_sequence = torch.randn(1, seq_len, embed_dim) base_sequence = torch.randn(1, seq_len, embed_dim)
identical_sequence = base_sequence.clone() identical_sequence = base_sequence.clone()
# Set to eval mode to disable dropout # Set to eval mode to disable dropout
attention.eval() attention.eval()
with torch.no_grad(): with torch.no_grad():
output1, _ = attention(base_sequence) output1, _ = attention(base_sequence)
output2, _ = attention(identical_sequence) output2, _ = attention(identical_sequence)
# With identical inputs and same parameters, outputs should be identical # With identical inputs and same parameters, outputs should be identical
assert torch.allclose(output1, output2, rtol=1e-5) 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: class TestPositionalEmbeddings:
"""Test cases for PositionalEmbeddings.""" """Test cases for PositionalEmbeddings."""
def test_initialization(self, embed_dim): def test_initialization(self, embed_dim):
"""Test that PositionalEmbeddings can be initialized.""" """Test that PositionalEmbeddings can be initialized."""
max_seq_len = 1024 max_seq_len = 1024
embeddings = PositionalEmbeddings(max_seq_len, embed_dim) embeddings = PositionalEmbeddings(max_seq_len, embed_dim)
assert embeddings is not None assert embeddings is not None
# Check that positional embeddings are created # Check that positional embeddings are created
assert hasattr(embeddings, 'embedding') assert hasattr(embeddings, "embedding")
assert embeddings.embedding.weight.shape == (max_seq_len, embed_dim) assert embeddings.embedding.weight.shape == (max_seq_len, embed_dim)
def test_forward_pass(self, embed_dim): def test_forward_pass(self, embed_dim):
"""Test forward pass of PositionalEmbeddings.""" """Test forward pass of PositionalEmbeddings."""
max_seq_len = 1024 max_seq_len = 1024
seq_len = 64 seq_len = 64
embeddings = PositionalEmbeddings(max_seq_len, embed_dim) embeddings = PositionalEmbeddings(max_seq_len, embed_dim)
# Forward pass - takes sequence length, not input tensor # Forward pass - takes sequence length, not input tensor
output = embeddings(seq_len) output = embeddings(seq_len)
# Check output shape # Check output shape
expected_shape = (seq_len, embed_dim) expected_shape = (seq_len, embed_dim)
assert output.shape == expected_shape assert output.shape == expected_shape
assert isinstance(output, torch.Tensor) assert isinstance(output, torch.Tensor)
def test_positional_encoding_values(self, embed_dim): def test_positional_encoding_values(self, embed_dim):
"""Test that positional encoding values are computed correctly.""" """Test that positional encoding values are computed correctly."""
max_seq_len = 10 max_seq_len = 10
embeddings = PositionalEmbeddings(max_seq_len, embed_dim) embeddings = PositionalEmbeddings(max_seq_len, embed_dim)
# Get embeddings for all positions # Get embeddings for all positions
pe = embeddings(max_seq_len) # Shape: [max_seq_len, embed_dim] pe = embeddings(max_seq_len) # Shape: [max_seq_len, embed_dim]
# Check that different positions have different embeddings # Check that different positions have different embeddings
# (since these are learnable embeddings, not fixed sine/cosine) # (since these are learnable embeddings, not fixed sine/cosine)
for pos in range(max_seq_len): for pos in range(max_seq_len):
for i in range(pos + 1, max_seq_len): for i in range(pos + 1, max_seq_len):
assert not torch.allclose(pe[pos], pe[i], rtol=1e-4) assert not torch.allclose(pe[pos], pe[i], rtol=1e-4)
def test_different_sequence_lengths(self, embed_dim): def test_different_sequence_lengths(self, embed_dim):
"""Test PositionalEmbeddings with different sequence lengths.""" """Test PositionalEmbeddings with different sequence lengths."""
test_cases = [ test_cases = [
(10, 5), # seq_len < max_seq_len (10, 5), # seq_len < max_seq_len
(10, 10), # seq_len == max_seq_len (10, 10), # seq_len == max_seq_len
] ]
for max_seq_len, seq_len in test_cases: for max_seq_len, seq_len in test_cases:
embeddings = PositionalEmbeddings(max_seq_len, embed_dim) embeddings = PositionalEmbeddings(max_seq_len, embed_dim)
# Get embeddings for specific sequence length # Get embeddings for specific sequence length
output = embeddings(seq_len) output = embeddings(seq_len)
# Output should have shape [seq_len, embed_dim] # Output should have shape [seq_len, embed_dim]
assert output.shape == (seq_len, embed_dim) assert output.shape == (seq_len, embed_dim)
def test_gradient_flow(self, embed_dim): def test_gradient_flow(self, embed_dim):
"""Test that gradients flow through PositionalEmbeddings.""" """Test that gradients flow through PositionalEmbeddings."""
max_seq_len = 64 max_seq_len = 64
seq_len = 32 seq_len = 32
embeddings = PositionalEmbeddings(max_seq_len, embed_dim) embeddings = PositionalEmbeddings(max_seq_len, embed_dim)
# Forward pass # Forward pass
output = embeddings(seq_len) output = embeddings(seq_len)
# Create a dummy loss and backward pass # Create a dummy loss and backward pass
loss = output.sum() loss = output.sum()
loss.backward() loss.backward()
# Positional embeddings should have gradients (they're learnable) # Positional embeddings should have gradients (they're learnable)
assert embeddings.embedding.weight.grad is not None assert embeddings.embedding.weight.grad is not None
assert not torch.allclose(embeddings.embedding.weight.grad, assert not torch.allclose(
torch.zeros_like(embeddings.embedding.weight.grad)) embeddings.embedding.weight.grad,
torch.zeros_like(embeddings.embedding.weight.grad),
)
def test_device_consistency(self, embed_dim, device): def test_device_consistency(self, embed_dim, device):
"""Test that PositionalEmbeddings works on correct device.""" """Test that PositionalEmbeddings works on correct device."""
max_seq_len = 64 max_seq_len = 64
seq_len = 32 seq_len = 32
embeddings = PositionalEmbeddings(max_seq_len, embed_dim).to(device) embeddings = PositionalEmbeddings(max_seq_len, embed_dim).to(device)
# Forward pass # Forward pass
output = embeddings(seq_len) output = embeddings(seq_len)
# Check device consistency # Check device consistency
assert output.device == device assert output.device == device
assert embeddings.embedding.weight.device == device assert embeddings.embedding.weight.device == device
def test_reproducibility(self, embed_dim): def test_reproducibility(self, embed_dim):
"""Test that positional embeddings are reproducible.""" """Test that positional embeddings are reproducible."""
max_seq_len = 100 max_seq_len = 100
embeddings1 = PositionalEmbeddings(max_seq_len, embed_dim) embeddings1 = PositionalEmbeddings(max_seq_len, embed_dim)
embeddings2 = PositionalEmbeddings(max_seq_len, embed_dim) embeddings2 = PositionalEmbeddings(max_seq_len, embed_dim)
# Different instances should have different embeddings (random initialization) # 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 # But same instance should produce same output for same input
seq_len = 50 seq_len = 50
output1 = embeddings1(seq_len) output1 = embeddings1(seq_len)
output2 = embeddings1(seq_len) # Same instance, same input output2 = embeddings1(seq_len) # Same instance, same input
assert torch.allclose(output1, output2) assert torch.allclose(output1, output2)
def test_positional_pattern(self, embed_dim): def test_positional_pattern(self, embed_dim):
"""Test that positional embeddings create a meaningful pattern.""" """Test that positional embeddings create a meaningful pattern."""
max_seq_len = 50 max_seq_len = 50
embeddings = PositionalEmbeddings(max_seq_len, embed_dim) embeddings = PositionalEmbeddings(max_seq_len, embed_dim)
pe = embeddings(max_seq_len) # Get all positional embeddings pe = embeddings(max_seq_len) # Get all positional embeddings
# Check that different positions have different embeddings # Check that different positions have different embeddings
# (with high probability due to random initialization) # (with high probability due to random initialization)
assert not torch.allclose(pe[0], pe[1], rtol=1e-4) assert not torch.allclose(pe[0], pe[1], rtol=1e-4)
assert not torch.allclose(pe[10], pe[20], rtol=1e-4) assert not torch.allclose(pe[10], pe[20], rtol=1e-4)
@pytest.mark.parametrize("max_seq_len,seq_len,embed_dim", [ @pytest.mark.parametrize(
(64, 10, 64), "max_seq_len,seq_len,embed_dim",
(128, 50, 128), [
(256, 100, 256), (64, 10, 64),
]) (128, 50, 128),
(256, 100, 256),
],
)
def test_different_configurations(self, max_seq_len, seq_len, embed_dim): def test_different_configurations(self, max_seq_len, seq_len, embed_dim):
"""Test PositionalEmbeddings with different configurations.""" """Test PositionalEmbeddings with different configurations."""
embeddings = PositionalEmbeddings(max_seq_len, embed_dim) embeddings = PositionalEmbeddings(max_seq_len, embed_dim)
output = embeddings(seq_len) output = embeddings(seq_len)
assert output.shape == (seq_len, embed_dim) 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: class TestTokenEmbeddings:
"""Test cases for TokenEmbeddings.""" """Test cases for TokenEmbeddings."""
def test_initialization(self, vocab_size, embed_dim): def test_initialization(self, vocab_size, embed_dim):
"""Test that TokenEmbeddings can be initialized.""" """Test that TokenEmbeddings can be initialized."""
embeddings = TokenEmbeddings(vocab_size, embed_dim) embeddings = TokenEmbeddings(vocab_size, embed_dim)
assert embeddings is not None assert embeddings is not None
# Check embedding layer # Check embedding layer
assert hasattr(embeddings, '_embedding') assert hasattr(embeddings, "_embedding")
assert embeddings._embedding.weight.shape == (vocab_size, embed_dim) assert embeddings._embedding.weight.shape == (vocab_size, embed_dim)
def test_forward_pass(self, vocab_size, embed_dim, random_inputs): def test_forward_pass(self, vocab_size, embed_dim, random_inputs):
"""Test forward pass of TokenEmbeddings.""" """Test forward pass of TokenEmbeddings."""
embeddings = TokenEmbeddings(vocab_size, embed_dim) embeddings = TokenEmbeddings(vocab_size, embed_dim)
# Forward pass # Forward pass
output = embeddings(random_inputs) output = embeddings(random_inputs)
# Check output shape # 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) assert isinstance(output, torch.Tensor)
def test_embedding_weights(self, vocab_size, embed_dim): def test_embedding_weights(self, vocab_size, embed_dim):
"""Test that embedding weights are properly initialized.""" """Test that embedding weights are properly initialized."""
embeddings = TokenEmbeddings(vocab_size, embed_dim) embeddings = TokenEmbeddings(vocab_size, embed_dim)
weights = embeddings._embedding.weight weights = embeddings._embedding.weight
assert weights.requires_grad is True assert weights.requires_grad is True
# Check that weights are not all zeros # Check that weights are not all zeros
assert not torch.allclose(weights, torch.zeros_like(weights)) assert not torch.allclose(weights, torch.zeros_like(weights))
def test_different_vocab_sizes(self): def test_different_vocab_sizes(self):
"""Test TokenEmbeddings with different vocabulary sizes.""" """Test TokenEmbeddings with different vocabulary sizes."""
test_cases = [ test_cases = [(100, 128), (1000, 256), (50000, 512)]
(100, 128),
(1000, 256),
(50000, 512)
]
for vocab_size, embed_dim in test_cases: for vocab_size, embed_dim in test_cases:
embeddings = TokenEmbeddings(vocab_size, embed_dim) embeddings = TokenEmbeddings(vocab_size, embed_dim)
assert embeddings._embedding.weight.shape == (vocab_size, embed_dim) assert embeddings._embedding.weight.shape == (vocab_size, embed_dim)
def test_gradient_flow(self, vocab_size, embed_dim, random_inputs): def test_gradient_flow(self, vocab_size, embed_dim, random_inputs):
"""Test that gradients flow through TokenEmbeddings.""" """Test that gradients flow through TokenEmbeddings."""
embeddings = TokenEmbeddings(vocab_size, embed_dim) embeddings = TokenEmbeddings(vocab_size, embed_dim)
# Forward pass # Forward pass
output = embeddings(random_inputs) output = embeddings(random_inputs)
# Create a dummy loss and backward pass # Create a dummy loss and backward pass
loss = output.sum() loss = output.sum()
loss.backward() loss.backward()
# Check that gradients are computed # Check that gradients are computed
assert embeddings._embedding.weight.grad is not None assert embeddings._embedding.weight.grad is not None
assert not torch.allclose(embeddings._embedding.weight.grad, assert not torch.allclose(
torch.zeros_like(embeddings._embedding.weight.grad)) embeddings._embedding.weight.grad,
torch.zeros_like(embeddings._embedding.weight.grad),
)
def test_device_consistency(self, vocab_size, embed_dim, random_inputs, device): def test_device_consistency(self, vocab_size, embed_dim, random_inputs, device):
"""Test that TokenEmbeddings works on correct device.""" """Test that TokenEmbeddings works on correct device."""
embeddings = TokenEmbeddings(vocab_size, embed_dim).to(device) embeddings = TokenEmbeddings(vocab_size, embed_dim).to(device)
inputs = random_inputs.to(device) inputs = random_inputs.to(device)
# Forward pass # Forward pass
output = embeddings(inputs) output = embeddings(inputs)
# Check device consistency # Check device consistency
assert output.device == device assert output.device == device
assert embeddings._embedding.weight.device == device assert embeddings._embedding.weight.device == device
def test_embedding_lookup(self, vocab_size, embed_dim): def test_embedding_lookup(self, vocab_size, embed_dim):
"""Test specific embedding lookups.""" """Test specific embedding lookups."""
embeddings = TokenEmbeddings(vocab_size, embed_dim) embeddings = TokenEmbeddings(vocab_size, embed_dim)
# Test lookup for specific tokens # 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) output = embeddings(test_tokens)
# Check shape # Check shape
assert output.shape == (2, 3, embed_dim) assert output.shape == (2, 3, embed_dim)
# Check that different tokens have different embeddings # Check that different tokens have different embeddings
# (with high probability due to random initialization) # (with high probability due to random initialization)
assert not torch.allclose(output[0, 0], output[0, 1], rtol=1e-4) 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)]) @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): def test_different_input_shapes(self, vocab_size, embed_dim, batch_size, seq_len):
"""Test TokenEmbeddings with different input shapes.""" """Test TokenEmbeddings with different input shapes."""
embeddings = TokenEmbeddings(vocab_size, embed_dim) embeddings = TokenEmbeddings(vocab_size, embed_dim)
inputs = torch.randint(0, vocab_size, (batch_size, seq_len)) inputs = torch.randint(0, vocab_size, (batch_size, seq_len))
output = embeddings(inputs) output = embeddings(inputs)
assert output.shape == (batch_size, seq_len, embed_dim) 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: class TestGPT:
"""Test cases for GPT model.""" """Test cases for GPT model."""
def test_initialization(self, gpt_config): def test_initialization(self, gpt_config):
"""Test that GPT can be initialized.""" """Test that GPT can be initialized."""
model = GPT(gpt_config) model = GPT(gpt_config)
assert model is not None assert model is not None
# Check that model has required components # Check that model has required components
assert hasattr(model, '_token_embeddings') assert hasattr(model, "_token_embeddings")
assert hasattr(model, '_position_embeddings') assert hasattr(model, "_position_embeddings")
assert hasattr(model, '_decoders') assert hasattr(model, "_decoders")
assert hasattr(model, '_linear') assert hasattr(model, "_linear")
assert hasattr(model, '_dropout') assert hasattr(model, "_dropout")
# Check number of decoder layers # 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): def test_forward_pass(self, gpt_config, random_inputs):
"""Test forward pass of GPT.""" """Test forward pass of GPT."""
model = GPT(gpt_config) model = GPT(gpt_config)
# Forward pass # Forward pass
logits = model(random_inputs) logits = model(random_inputs)
# Check output shape # Check output shape
batch_size, seq_len = random_inputs.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 logits.shape == (batch_size, seq_len, vocab_size)
assert isinstance(logits, torch.Tensor) 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.""" """Test forward pass with attention mask."""
model = GPT(gpt_config) model = GPT(gpt_config)
# Forward pass with mask # Forward pass with mask
logits = model(random_inputs, attention_mask=attention_mask) logits = model(random_inputs, attention_mask=attention_mask)
# Check output shape # Check output shape
batch_size, seq_len = random_inputs.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 logits.shape == (batch_size, seq_len, vocab_size)
def test_generate_text(self, gpt_config): def test_generate_text(self, gpt_config):
"""Test text generation.""" """Test text generation."""
model = GPT(gpt_config) model = GPT(gpt_config)
model.eval() # Set to evaluation mode for generation model.eval() # Set to evaluation mode for generation
# Create initial input # Create initial input
batch_size = 2 batch_size = 2
initial_seq_len = 5 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 # Generate text
with torch.no_grad(): with torch.no_grad():
generated = model.generate( generated = model.generate(
x=input_ids, x=input_ids,
max_new_tokens=10, max_new_tokens=10,
do_sample=False # Use greedy for deterministic testing do_sample=False, # Use greedy for deterministic testing
) )
# Check output shape # Check output shape
expected_seq_len = initial_seq_len + 10 expected_seq_len = initial_seq_len + 10
assert generated.shape == (batch_size, expected_seq_len) assert generated.shape == (batch_size, expected_seq_len)
# Check that initial sequence is preserved # Check that initial sequence is preserved
assert torch.allclose(generated[:, :initial_seq_len], input_ids) assert torch.allclose(generated[:, :initial_seq_len], input_ids)
def test_generate_with_temperature(self, gpt_config): def test_generate_with_temperature(self, gpt_config):
"""Test text generation with temperature sampling.""" """Test text generation with temperature sampling."""
model = GPT(gpt_config) model = GPT(gpt_config)
model.eval() model.eval()
# Create initial input # 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 # Generate with temperature
with torch.no_grad(): with torch.no_grad():
generated = model.generate( generated = model.generate(
x=input_ids, x=input_ids, max_new_tokens=5, do_sample=True, temperature=0.8
max_new_tokens=5,
do_sample=True,
temperature=0.8
) )
assert generated.shape == (1, 8) # 3 initial + 5 new tokens assert generated.shape == (1, 8) # 3 initial + 5 new tokens
def test_generate_with_top_k(self, gpt_config): def test_generate_with_top_k(self, gpt_config):
"""Test text generation with top-k sampling.""" """Test text generation with top-k sampling."""
model = GPT(gpt_config) model = GPT(gpt_config)
model.eval() model.eval()
# Create initial input # 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 # Generate with top-k
with torch.no_grad(): with torch.no_grad():
generated = model.generate( generated = model.generate(
x=input_ids, x=input_ids, max_new_tokens=5, do_sample=True, top_k=10
max_new_tokens=5,
do_sample=True,
top_k=10
) )
assert generated.shape == (1, 8) assert generated.shape == (1, 8)
def test_generate_with_top_p(self, gpt_config): def test_generate_with_top_p(self, gpt_config):
"""Test text generation with top-p (nucleus) sampling.""" """Test text generation with top-p (nucleus) sampling."""
model = GPT(gpt_config) model = GPT(gpt_config)
model.eval() model.eval()
# Create initial input # 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 # Generate with top-p
with torch.no_grad(): with torch.no_grad():
generated = model.generate( generated = model.generate(
x=input_ids, x=input_ids, max_new_tokens=5, do_sample=True, top_p=0.9
max_new_tokens=5,
do_sample=True,
top_p=0.9
) )
assert generated.shape == (1, 8) assert generated.shape == (1, 8)
def test_gradient_flow(self, gpt_config, random_inputs): def test_gradient_flow(self, gpt_config, random_inputs):
"""Test that gradients flow through GPT.""" """Test that gradients flow through GPT."""
model = GPT(gpt_config) model = GPT(gpt_config)
# Forward pass # Forward pass
logits = model(random_inputs) logits = model(random_inputs)
# Create a dummy loss and backward pass # 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( loss = torch.nn.functional.cross_entropy(
logits.view(-1, logits.size(-1)), logits.view(-1, logits.size(-1)), targets.view(-1)
targets.view(-1)
) )
loss.backward() loss.backward()
# Check that gradients are computed for various components # Check that gradients are computed for various components
assert model._token_embeddings._embedding.weight.grad is not None assert model._token_embeddings._embedding.weight.grad is not None
assert model._linear.weight.grad is not None assert model._linear.weight.grad is not None
if len(model._decoders) > 0: if len(model._decoders) > 0:
assert model._decoders[0]._heads._heads[0]._q.weight.grad is not None assert model._decoders[0]._heads._heads[0]._q.weight.grad is not None
def test_device_consistency(self, gpt_config, random_inputs, device): def test_device_consistency(self, gpt_config, random_inputs, device):
"""Test that GPT works on correct device.""" """Test that GPT works on correct device."""
model = GPT(gpt_config).to(device) model = GPT(gpt_config).to(device)
inputs = random_inputs.to(device) inputs = random_inputs.to(device)
# Forward pass # Forward pass
logits = model(inputs) logits = model(inputs)
# Check device consistency # Check device consistency
assert logits.device == device assert logits.device == device
assert model._token_embeddings._embedding.weight.device == device assert model._token_embeddings._embedding.weight.device == device
def test_different_configurations(self): def test_different_configurations(self):
"""Test GPT with different configurations.""" """Test GPT with different configurations."""
test_configs = [ test_configs = [
@@ -174,7 +168,7 @@ class TestGPT:
"num_heads": 2, "num_heads": 2,
"num_layers": 2, "num_layers": 2,
"max_position_embeddings": 256, "max_position_embeddings": 256,
"dropout": 0.1 "dropout": 0.1,
}, },
{ {
"vocab_size": 5000, "vocab_size": 5000,
@@ -182,7 +176,7 @@ class TestGPT:
"num_heads": 4, "num_heads": 4,
"num_layers": 4, "num_layers": 4,
"max_position_embeddings": 512, "max_position_embeddings": 512,
"dropout": 0.1 "dropout": 0.1,
}, },
{ {
"vocab_size": 10000, "vocab_size": 10000,
@@ -190,98 +184,94 @@ class TestGPT:
"num_heads": 8, "num_heads": 8,
"num_layers": 6, "num_layers": 6,
"max_position_embeddings": 1024, "max_position_embeddings": 1024,
"dropout": 0.1 "dropout": 0.1,
} },
] ]
for config in test_configs: for config in test_configs:
model = GPT(config) model = GPT(config)
batch_size, seq_len = 2, 16 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) 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 assert logits.shape == expected_shape
@pytest.mark.parametrize("batch_size,seq_len", [(1, 8), (2, 16), (4, 32)]) @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): def test_different_input_shapes(self, gpt_config, batch_size, seq_len):
"""Test GPT with different input shapes.""" """Test GPT with different input shapes."""
model = GPT(gpt_config) 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) 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 assert logits.shape == expected_shape
def test_training_vs_evaluation(self, gpt_config, random_inputs): def test_training_vs_evaluation(self, gpt_config, random_inputs):
"""Test that GPT behaves differently in train vs eval mode.""" """Test that GPT behaves differently in train vs eval mode."""
model = GPT(gpt_config) model = GPT(gpt_config)
# Training mode # Training mode
model.train() model.train()
output_train = model(random_inputs) output_train = model(random_inputs)
# Evaluation mode # Evaluation mode
model.eval() model.eval()
output_eval = model(random_inputs) output_eval = model(random_inputs)
# Outputs should be different due to dropout # Outputs should be different due to dropout
assert not torch.allclose(output_train, output_eval) assert not torch.allclose(output_train, output_eval)
def test_parameter_count(self, gpt_config): def test_parameter_count(self, gpt_config):
"""Test that GPT has reasonable number of parameters.""" """Test that GPT has reasonable number of parameters."""
model = GPT(gpt_config) model = GPT(gpt_config)
total_params = sum(p.numel() for p in model.parameters()) total_params = sum(p.numel() for p in model.parameters())
# For a small GPT model, parameters should be in reasonable range # For a small GPT model, parameters should be in reasonable range
vocab_size = gpt_config['vocab_size'] vocab_size = gpt_config["vocab_size"]
embed_dim = gpt_config['embed_dim'] embed_dim = gpt_config["embed_dim"]
num_layers = gpt_config['num_layers'] num_layers = gpt_config["num_layers"]
num_heads = gpt_config['num_heads'] num_heads = gpt_config["num_heads"]
# Rough estimate: token_embeddings + output_layer + (attention + ff) * layers # Rough estimate: token_embeddings + output_layer + (attention + ff) * layers
expected_min = vocab_size * embed_dim * 2 # embeddings and output expected_min = vocab_size * embed_dim * 2 # embeddings and output
expected_max = expected_min * 10 # Allow for decoder parameters expected_max = expected_min * 10 # Allow for decoder parameters
assert expected_min < total_params < expected_max assert expected_min < total_params < expected_max
def test_causal_attention(self, gpt_config): def test_causal_attention(self, gpt_config):
"""Test that GPT uses causal attention during generation.""" """Test that GPT uses causal attention during generation."""
model = GPT(gpt_config) model = GPT(gpt_config)
model.eval() model.eval()
# Create input with known pattern # Create input with known pattern
input_ids = torch.tensor([[1, 2, 3]]).long() input_ids = torch.tensor([[1, 2, 3]]).long()
with torch.no_grad(): with torch.no_grad():
# Get logits for next token prediction # Get logits for next token prediction
logits = model(input_ids) logits = model(input_ids)
# The model should only attend to previous tokens (causal) # The model should only attend to previous tokens (causal)
# We can't directly test attention masks in the public API, # We can't directly test attention masks in the public API,
# but we can verify the generation works correctly # but we can verify the generation works correctly
generated = model.generate( generated = model.generate(x=input_ids, max_new_tokens=3, do_sample=False)
x=input_ids,
max_new_tokens=3,
do_sample=False
)
# Generated sequence should be longer than input # Generated sequence should be longer than input
assert generated.shape[1] == input_ids.shape[1] + 3 assert generated.shape[1] == input_ids.shape[1] + 3
def test_output_distribution(self, gpt_config, random_inputs): def test_output_distribution(self, gpt_config, random_inputs):
"""Test that GPT output has proper distribution.""" """Test that GPT output has proper distribution."""
model = GPT(gpt_config) model = GPT(gpt_config)
logits = model(random_inputs) logits = model(random_inputs)
# Logits should not have extreme values # Logits should not have extreme values
assert logits.abs().max() < 100 assert logits.abs().max() < 100
# Softmax should produce valid probabilities # Softmax should produce valid probabilities
probs = torch.softmax(logits, dim=-1) probs = torch.softmax(logits, dim=-1)
assert torch.allclose(probs.sum(dim=-1), torch.ones_like(probs.sum(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(): def test_gpt_model_creation():
"""Test that GPT model can be created and forward pass works.""" """Test that GPT model can be created and forward pass works."""
from llm.models.gpt import GPT from llm.models.gpt import GPT
config = { config = {
"vocab_size": 1000, "vocab_size": 1000,
"embed_dim": 128, "embed_dim": 128,
"num_heads": 4, "num_heads": 4,
"num_layers": 2, "num_layers": 2,
"max_position_embeddings": 256, "max_position_embeddings": 256,
"dropout": 0.1 "dropout": 0.1,
} }
model = GPT(config) model = GPT(config)
# Test forward pass # Test forward pass
batch_size, seq_len = 2, 16 batch_size, seq_len = 2, 16
input_ids = torch.randint(0, config["vocab_size"], (batch_size, seq_len)) input_ids = torch.randint(0, config["vocab_size"], (batch_size, seq_len))
with torch.no_grad(): with torch.no_grad():
logits = model(input_ids) logits = model(input_ids)
assert logits.shape == (batch_size, seq_len, config["vocab_size"]) assert logits.shape == (batch_size, seq_len, config["vocab_size"])
print("✅ GPT model creation and forward pass test passed") print("✅ GPT model creation and forward pass test passed")
@@ -37,27 +37,21 @@ def test_gpt_model_creation():
def test_bpe_tokenizer_basic(): def test_bpe_tokenizer_basic():
"""Test basic BPE tokenizer functionality.""" """Test basic BPE tokenizer functionality."""
from llm.tokenizers import BPETokenizer from llm.tokenizers import BPETokenizer
tokenizer = BPETokenizer() tokenizer = BPETokenizer()
# Train on simple texts # Train on simple texts
texts = [ texts = ["hello world", "test tokenization", "simple example"]
"hello world",
"test tokenization",
"simple example"
]
tokenizer.train( tokenizer.train(
texts=texts, texts=texts, vocab_size=50, special_tokens=["<pad>", "<unk>", "<bos>", "<eos>"]
vocab_size=50,
special_tokens=["<pad>", "<unk>", "<bos>", "<eos>"]
) )
# Test encoding/decoding # Test encoding/decoding
text = "hello world" text = "hello world"
tokens = tokenizer.encode(text) tokens = tokenizer.encode(text)
decoded = tokenizer.decode(tokens) decoded = tokenizer.decode(tokens)
assert isinstance(tokens, list) assert isinstance(tokens, list)
assert isinstance(decoded, str) assert isinstance(decoded, str)
assert len(tokens) > 0 assert len(tokens) > 0
@@ -67,18 +61,18 @@ def test_bpe_tokenizer_basic():
def test_token_embeddings(): def test_token_embeddings():
"""Test token embeddings.""" """Test token embeddings."""
from llm.core.token_embeddings import TokenEmbeddings from llm.core.token_embeddings import TokenEmbeddings
vocab_size = 1000 vocab_size = 1000
embed_dim = 128 embed_dim = 128
embeddings = TokenEmbeddings(vocab_size, embed_dim) embeddings = TokenEmbeddings(vocab_size, embed_dim)
# Test forward pass # Test forward pass
batch_size, seq_len = 2, 16 batch_size, seq_len = 2, 16
input_ids = torch.randint(0, vocab_size, (batch_size, seq_len)) input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
output = embeddings(input_ids) output = embeddings(input_ids)
assert output.shape == (batch_size, seq_len, embed_dim) assert output.shape == (batch_size, seq_len, embed_dim)
print("✅ Token embeddings test passed") print("✅ Token embeddings test passed")
@@ -86,20 +80,20 @@ def test_token_embeddings():
def test_multi_head_attention(): def test_multi_head_attention():
"""Test multi-head attention.""" """Test multi-head attention."""
from llm.core.multi_head_attention import MultiHeadAttention from llm.core.multi_head_attention import MultiHeadAttention
num_heads = 4 num_heads = 4
emb_size = 128 emb_size = 128
head_size = emb_size // num_heads head_size = emb_size // num_heads
max_seq_len = 256 max_seq_len = 256
attention = MultiHeadAttention(num_heads, emb_size, head_size, max_seq_len) attention = MultiHeadAttention(num_heads, emb_size, head_size, max_seq_len)
# Test forward pass # Test forward pass
batch_size, seq_len = 2, 16 batch_size, seq_len = 2, 16
inputs = torch.randn(batch_size, seq_len, emb_size) inputs = torch.randn(batch_size, seq_len, emb_size)
output, _ = attention(inputs) output, _ = attention(inputs)
assert output.shape == inputs.shape assert output.shape == inputs.shape
print("✅ Multi-head attention test passed") print("✅ Multi-head attention test passed")
@@ -107,17 +101,17 @@ def test_multi_head_attention():
def test_feed_forward(): def test_feed_forward():
"""Test feed forward network.""" """Test feed forward network."""
from llm.core.feed_forward import FeedForward from llm.core.feed_forward import FeedForward
embed_dim = 128 embed_dim = 128
ff = FeedForward(embed_dim) ff = FeedForward(embed_dim)
# Test forward pass # Test forward pass
batch_size, seq_len = 2, 16 batch_size, seq_len = 2, 16
inputs = torch.randn(batch_size, seq_len, embed_dim) inputs = torch.randn(batch_size, seq_len, embed_dim)
output = ff(inputs) output = ff(inputs)
assert output.shape == inputs.shape assert output.shape == inputs.shape
print("✅ Feed forward test passed") print("✅ Feed forward test passed")
@@ -125,29 +119,25 @@ def test_feed_forward():
def test_gpt_generation(): def test_gpt_generation():
"""Test GPT text generation.""" """Test GPT text generation."""
from llm.models.gpt import GPT from llm.models.gpt import GPT
config = { config = {
"vocab_size": 1000, "vocab_size": 1000,
"embed_dim": 128, "embed_dim": 128,
"num_heads": 4, "num_heads": 4,
"num_layers": 2, "num_layers": 2,
"max_position_embeddings": 256, "max_position_embeddings": 256,
"dropout": 0.1 "dropout": 0.1,
} }
model = GPT(config) model = GPT(config)
model.eval() model.eval()
# Test greedy generation # Test greedy generation
input_ids = torch.randint(0, config["vocab_size"], (1, 5)) input_ids = torch.randint(0, config["vocab_size"], (1, 5))
with torch.no_grad(): with torch.no_grad():
generated = model.generate( generated = model.generate(x=input_ids, max_new_tokens=3, do_sample=False)
x=input_ids,
max_new_tokens=3,
do_sample=False
)
assert generated.shape == (1, 8) # 5 initial + 3 new tokens assert generated.shape == (1, 8) # 5 initial + 3 new tokens
print("✅ GPT generation test passed") print("✅ GPT generation test passed")
@@ -155,50 +145,48 @@ def test_gpt_generation():
def test_bpe_tokenizer_save_load(): def test_bpe_tokenizer_save_load():
"""Test BPE tokenizer save/load functionality.""" """Test BPE tokenizer save/load functionality."""
from llm.tokenizers import BPETokenizer from llm.tokenizers import BPETokenizer
tokenizer = BPETokenizer() tokenizer = BPETokenizer()
# Train on simple texts # Train on simple texts
texts = ["hello world", "test save load"] texts = ["hello world", "test save load"]
tokenizer.train( tokenizer.train(
texts=texts, texts=texts, vocab_size=30, special_tokens=["<pad>", "<unk>", "<bos>", "<eos>"]
vocab_size=30,
special_tokens=["<pad>", "<unk>", "<bos>", "<eos>"]
) )
with tempfile.TemporaryDirectory() as temp_dir: with tempfile.TemporaryDirectory() as temp_dir:
save_path = os.path.join(temp_dir, "test_tokenizer.json") save_path = os.path.join(temp_dir, "test_tokenizer.json")
# Save tokenizer # Save tokenizer
tokenizer.save(save_path) tokenizer.save(save_path)
assert os.path.exists(save_path) assert os.path.exists(save_path)
# Load tokenizer # Load tokenizer
loaded_tokenizer = BPETokenizer.load(save_path) loaded_tokenizer = BPETokenizer.load(save_path)
# Test that vocab size is the same # Test that vocab size is the same
assert tokenizer.get_vocab_size() == loaded_tokenizer.get_vocab_size() assert tokenizer.get_vocab_size() == loaded_tokenizer.get_vocab_size()
# Test that vocabularies are the same # Test that vocabularies are the same
assert tokenizer.get_vocab() == loaded_tokenizer.get_vocab() assert tokenizer.get_vocab() == loaded_tokenizer.get_vocab()
# Test that both can encode/decode (even if tokens differ due to BPE state) # Test that both can encode/decode (even if tokens differ due to BPE state)
text = "hello world" text = "hello world"
original_tokens = tokenizer.encode(text) original_tokens = tokenizer.encode(text)
loaded_tokens = loaded_tokenizer.encode(text) loaded_tokens = loaded_tokenizer.encode(text)
# Both should produce valid token lists # Both should produce valid token lists
assert isinstance(original_tokens, list) assert isinstance(original_tokens, list)
assert isinstance(loaded_tokens, list) assert isinstance(loaded_tokens, list)
assert len(original_tokens) > 0 assert len(original_tokens) > 0
assert len(loaded_tokens) > 0 assert len(loaded_tokens) > 0
# Both should be able to decode # Both should be able to decode
original_decoded = tokenizer.decode(original_tokens) original_decoded = tokenizer.decode(original_tokens)
loaded_decoded = loaded_tokenizer.decode(loaded_tokens) loaded_decoded = loaded_tokenizer.decode(loaded_tokens)
assert isinstance(original_decoded, str) assert isinstance(original_decoded, str)
assert isinstance(loaded_decoded, str) assert isinstance(loaded_decoded, str)
print("✅ BPE tokenizer save/load test passed") print("✅ BPE tokenizer save/load test passed")
@@ -206,18 +194,16 @@ def test_gpt_with_tokenizer():
"""Test GPT model with tokenizer integration.""" """Test GPT model with tokenizer integration."""
from llm.models.gpt import GPT from llm.models.gpt import GPT
from llm.tokenizers import BPETokenizer from llm.tokenizers import BPETokenizer
# Create and train tokenizer # Create and train tokenizer
tokenizer = BPETokenizer() tokenizer = BPETokenizer()
texts = ["hello world", "test integration"] texts = ["hello world", "test integration"]
tokenizer.train( tokenizer.train(
texts=texts, texts=texts, vocab_size=50, special_tokens=["<pad>", "<unk>", "<bos>", "<eos>"]
vocab_size=50,
special_tokens=["<pad>", "<unk>", "<bos>", "<eos>"]
) )
vocab_size = tokenizer.get_vocab_size() vocab_size = tokenizer.get_vocab_size()
# Create GPT model with tokenizer's vocab size # Create GPT model with tokenizer's vocab size
config = { config = {
"vocab_size": vocab_size, "vocab_size": vocab_size,
@@ -225,19 +211,19 @@ def test_gpt_with_tokenizer():
"num_heads": 4, "num_heads": 4,
"num_layers": 2, "num_layers": 2,
"max_position_embeddings": 256, "max_position_embeddings": 256,
"dropout": 0.1 "dropout": 0.1,
} }
model = GPT(config) model = GPT(config)
# Test with tokenized input # Test with tokenized input
text = "hello world" text = "hello world"
tokens = tokenizer.encode(text, add_special_tokens=False) tokens = tokenizer.encode(text, add_special_tokens=False)
input_ids = torch.tensor([tokens]) input_ids = torch.tensor([tokens])
with torch.no_grad(): with torch.no_grad():
logits = model(input_ids) logits = model(input_ids)
assert logits.shape == (1, len(tokens), vocab_size) assert logits.shape == (1, len(tokens), vocab_size)
print("✅ GPT with tokenizer integration test passed") print("✅ GPT with tokenizer integration test passed")
@@ -245,7 +231,7 @@ def test_gpt_with_tokenizer():
def run_all_tests(): def run_all_tests():
"""Run all basic tests.""" """Run all basic tests."""
print("🧪 Running basic tests for llm library...") print("🧪 Running basic tests for llm library...")
test_gpt_model_creation() test_gpt_model_creation()
test_bpe_tokenizer_basic() test_bpe_tokenizer_basic()
test_token_embeddings() test_token_embeddings()
@@ -254,7 +240,7 @@ def run_all_tests():
test_gpt_generation() test_gpt_generation()
test_bpe_tokenizer_save_load() test_bpe_tokenizer_save_load()
test_gpt_with_tokenizer() test_gpt_with_tokenizer()
print("🎉 All basic tests passed!") 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): class ConcreteTokenizer(BaseTokenizer):
"""Concrete implementation for testing BaseTokenizer.""" """Concrete implementation for testing BaseTokenizer."""
def train(self, texts: list, vocab_size: int = 1000, **kwargs): def train(self, texts: list, vocab_size: int = 1000, **kwargs):
"""Dummy implementation for testing.""" """Dummy implementation for testing."""
pass pass
def encode(self, text: str, **kwargs) -> list: def encode(self, text: str, **kwargs) -> list:
"""Dummy implementation for testing.""" """Dummy implementation for testing."""
return [1, 2, 3] return [1, 2, 3]
def decode(self, tokens: list, **kwargs) -> str: def decode(self, tokens: list, **kwargs) -> str:
"""Dummy implementation for testing.""" """Dummy implementation for testing."""
return "decoded text" return "decoded text"
@@ -24,33 +24,33 @@ class ConcreteTokenizer(BaseTokenizer):
class TestBaseTokenizer: class TestBaseTokenizer:
"""Test cases for BaseTokenizer.""" """Test cases for BaseTokenizer."""
def test_initialization(self): def test_initialization(self):
"""Test that BaseTokenizer can be initialized through concrete class.""" """Test that BaseTokenizer can be initialized through concrete class."""
tokenizer = ConcreteTokenizer() tokenizer = ConcreteTokenizer()
assert tokenizer is not None assert tokenizer is not None
assert tokenizer.vocab == {} assert tokenizer.vocab == {}
assert tokenizer.vocab_size == 0 assert tokenizer.vocab_size == 0
def test_encode_implemented(self): def test_encode_implemented(self):
"""Test that encode method works in concrete implementation.""" """Test that encode method works in concrete implementation."""
tokenizer = ConcreteTokenizer() tokenizer = ConcreteTokenizer()
result = tokenizer.encode("test text") result = tokenizer.encode("test text")
assert result == [1, 2, 3] assert result == [1, 2, 3]
def test_decode_implemented(self): def test_decode_implemented(self):
"""Test that decode method works in concrete implementation.""" """Test that decode method works in concrete implementation."""
tokenizer = ConcreteTokenizer() tokenizer = ConcreteTokenizer()
result = tokenizer.decode([1, 2, 3]) result = tokenizer.decode([1, 2, 3])
assert result == "decoded text" assert result == "decoded text"
def test_get_vocab_size(self): def test_get_vocab_size(self):
"""Test that get_vocab_size method works.""" """Test that get_vocab_size method works."""
tokenizer = ConcreteTokenizer() tokenizer = ConcreteTokenizer()
tokenizer.vocab = {"a": 0, "b": 1, "c": 2} tokenizer.vocab = {"a": 0, "b": 1, "c": 2}
tokenizer.vocab_size = 3 tokenizer.vocab_size = 3
assert tokenizer.get_vocab_size() == 3 assert tokenizer.get_vocab_size() == 3
def test_get_vocab(self): def test_get_vocab(self):
"""Test that get_vocab method works.""" """Test that get_vocab method works."""
tokenizer = ConcreteTokenizer() tokenizer = ConcreteTokenizer()

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

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