grpo_code.py - master - rnekrasov/gpro-code

grpo_code.py
240 строк9 KB
Роман Некрасов
upload files
16 мар 2025, 16:26
16 мар 2025, 16:2604d6955
#!/usr/bin/env python
# coding: utf-8

# Установка необходимых библиотек
# pip install unsloth vllm
# pip install --upgrade pillow

# Импорт библиотек для работы с моделью
from unsloth import FastLanguageModel, PatchFastRL
PatchFastRL("GRPO", FastLanguageModel)

# Загрузка модели и настройка параметров
from unsloth import is_bfloat16_supported
import torch
max_seq_length = 1200
lora_rank = 8

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
    max_seq_length=max_seq_length,
    load_in_4bit=True,
    fast_inference=True,
    max_lora_rank=lora_rank,
    gpu_memory_utilization=0.62,
)

model = FastLanguageModel.get_peft_model(
    model,
    r=lora_rank,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    lora_alpha=lora_rank,
    use_gradient_checkpointing="unsloth",
    random_state=3407,
)

# Подготовка данных
import json
from datasets import Dataset
import re

# Путь к JSON-файлу с датасетом
JSON_FILE_PATH = "arcade_games_dataset.json"

# Системный промпт
SYSTEM_PROMPT = """
Напишите код игры на Python в следующем формате:
<think>
[Кратко опишите свои рассуждения и шаги решения]
</think>
<answer>
[Приведите полный рабочий код игры]
</answer>

**Пример:**
Запрос: "Напишите игру 'Змейка' на Python."
<think>
Для игры 'Змейка' нужно создать игровое поле, змейку, еду и механику движения. Змейка должна расти, когда ест еду, и игра заканчивается при столкновении с собой или стеной. Использую библиотеку pygame для графики и управления.
</think>
<answer>
import pygame
import random

pygame.init()
width, height = 800, 600
window = pygame.display.set_mode((width, height))
pygame.display.set_caption("Змейка")

snake_pos = [[width // 2, height // 2]]
snake_speed = 20
direction = [snake_speed, 0]
food_pos = [random.randrange(1, (width // snake_speed)) * snake_speed, 
            random.randrange(1, (height // snake_speed)) * snake_speed]
score = 0

clock = pygame.time.Clock()
running = True

while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_UP and direction[1] == 0:
                direction = [0, -snake_speed]
            elif event.key == pygame.K_DOWN and direction[1] == 0:
                direction = [0, snake_speed]
            elif event.key == pygame.K_LEFT and direction[0] == 0:
                direction = [-snake_speed, 0]
            elif event.key == pygame.K_RIGHT and direction[0] == 0:
                direction = [snake_speed, 0]

    snake_pos.insert(0, [snake_pos[0][0] + direction[0], snake_pos[0][1] + direction[1]])
    
    if snake_pos[0] == food_pos:
        score += 1
        food_pos = [random.randrange(1, (width // snake_speed)) * snake_speed, 
                    random.randrange(1, (height // snake_speed)) * snake_speed]
    else:
        snake_pos.pop()

    if (snake_pos[0][0] < 0 or snake_pos[0][0] >= width or 
        snake_pos[0][1] < 0 or snake_pos[0][1] >= height or 
        snake_pos[0] in snake_pos[1:]):
        running = False

    window.fill((0, 0, 0))
    for pos in snake_pos:
        pygame.draw.rect(window, (0, 255, 0), [pos[0], pos[1], snake_speed, snake_speed])
    pygame.draw.rect(window, (255, 0, 0), [food_pos[0], food_pos[1], snake_speed, snake_speed])
    pygame.display.update()
    clock.tick(15)

pygame.quit()
print(f"Игра окончена! Счёт: {score}")
</answer>

**Инструкции:**
- Ваш ответ должен содержать только код игры в указанном формате, без лишних объяснений вне тега <think>.
- Используйте свои знания программирования и игр для создания кода.
- Вы можете применять любые библиотеки (например, pygame, tkinter), если они нужны.
- Код должен быть рабочим, чтобы его можно было запустить и протестировать.
"""

# Функция для извлечения ответа из XML-формата
def extract_xml_answer(text: str) -> str:
    answer = text.split("<answer>")[-1]
    answer = answer.split("</answer>")[0]
    return answer.strip()

# Функция для извлечения ключевых слов из тестов
def extract_keywords_from_tests(tests: list[str]) -> set[str]:
    keywords = set()
    for test in tests:
        # Простая эвристика: берем слова длиннее 3 символов
        words = re.findall(r'\b\w{4,}\b', test.lower())
        keywords.update(words)
    return keywords

# Загрузка датасета из JSON-файла
def get_arcade_dataset(json_file_path: str) -> Dataset:
    with open(json_file_path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    dataset = Dataset.from_list([{
        'prompt': [
            {'role': 'system', 'content': SYSTEM_PROMPT},
            {'role': 'user', 'content': item['prompt']}
        ],
        'tests': item['tests'],
        'keywords': list(extract_keywords_from_tests(item['tests']))
    } for item in data])
    return dataset

# Загружаем датасет
dataset = get_arcade_dataset(JSON_FILE_PATH)

# Функция награды с учетом тестов и выводом ответов с высокой наградой
def code_execution_reward_func(prompts, completions, **kwargs):
    rewards = []
    for prompt, completion in zip(prompts, completions):
        try:
            extracted_answer = extract_xml_answer(completion[0]['content'])
            exec(extracted_answer)  # Проверка выполнения кода
            base_reward = 1.0  # Успешное выполнение
        except Exception:
            base_reward = 0.0  # Ошибка при выполнении
        
        # Извлекаем ключевые слова из prompt, если dataset_entry отсутствует
        if 'dataset_entry' in kwargs:
            keywords = kwargs['dataset_entry']['keywords']
        else:
            keywords = extract_keywords_from_tests([prompt[-1]['content']])
        
        code_lower = extracted_answer.lower()
        keyword_reward = sum(1 for keyword in keywords if keyword in code_lower) / len(keywords) if keywords else 0.0
        
        # Комбинированная награда
        total_reward = base_reward + keyword_reward * 0.5
        rewards.append(total_reward)
        
        # Выводим в терминал ответы с наградой больше 0.2
        if total_reward > 0.2:
            print(f"Prompt: {prompt[-1]['content']}")
            print(f"Generated Answer:\n{extracted_answer}")
            print(f"Reward: {total_reward}\n{'-'*50}")
    
    return rewards

# Обучение модели
from trl import GRPOConfig, GRPOTrainer
training_args = GRPOConfig(
    use_vllm=True,
    learning_rate=3e-5,
    adam_beta1=0.9,
    adam_beta2=0.99,
    weight_decay=0.1,
    warmup_ratio=0.1,
    lr_scheduler_type="cosine",
    optim="paged_adamw_8bit",
    logging_steps=1,
    bf16=is_bfloat16_supported(),
    fp16=not is_bfloat16_supported(),
    per_device_train_batch_size=1,
    gradient_accumulation_steps=1,
    num_generations=5,
    max_prompt_length=256,
    max_completion_length=344,
    num_train_epochs=50,
    save_steps=250000,
    max_grad_norm=0.1,
    output_dir="outputs",
)

trainer = GRPOTrainer(
    model=model,
    processing_class=tokenizer,
    reward_funcs=[lambda prompts, completions, **kwargs: code_execution_reward_func(prompts, completions, **kwargs)],
    args=training_args,
    train_dataset=dataset,
)
trainer.train()

# Пример инференса
text = tokenizer.apply_chat_template([
    {"role": "system", "content": SYSTEM_PROMPT},
    {"role": "user", "content": "Как проверить столкновение двух объектов в аркадной игре?"},
], tokenize=False, add_generation_prompt=True)

from vllm import SamplingParams
sampling_params = SamplingParams(
    temperature=0.8,
    top_p=0.95,
    max_tokens=1024,
)
output = model.fast_generate(
    [text],
    sampling_params=sampling_params,
    lora_request=None,
)[0].outputs[0].text

print(output)