benchmark.js

// reasoningHeuristicAI/benchmark.js
// --- Версия, использующая правильный датасет для оценки ---

import fs from 'fs/promises';
import path from 'path';
import slmnet from './core/slmnet/slmnet.js';
import { HeuristicEngine } from './core/heuristic-engine.js';
import workerPool from './core/worker-pool.js';

const config = {
    embedding_dim: 64,
    hidden_dim: 128,
    data_path: './data',
    benchmark_steps: 10, // Используем мало шагов, т.к. датасет маленький
    learning_rate: 0.001,
    grad_clip: 1.0,
    accumulation_steps: 1, // Без накопления, т.к. примеров мало
};

async function loadBenchmarkData() {
    // --- ИСПРАВЛЕНИЕ: Используем тот же файл, что и для fine-tuning'а ---
    const filePath = path.join(config.data_path, 'coder_finetune.jsonl');
    let fullTextForTokenizer = '';
    try {
        const fileContent = await fs.readFile(filePath, 'utf-8');
        const lines = fileContent.trim().split('\n');
        const jsonData = lines.map(line => JSON.parse(line));
        jsonData.forEach(item => {
            fullTextForTokenizer += item.prompt + item.completion;
        });

        const tokenizer = new slmnet.tokenizers.CharacterTokenizer(fullTextForTokenizer);
        const data = jsonData.map(item => ({
            prompt_ids: tokenizer.encode(item.prompt),
            completion_ids: tokenizer.encode(item.completion)
        }));
        
        return { tokenizer, data };
    } catch (error) {
        console.error(`Ошибка: Не удалось загрузить данные для бенчмарка из '${filePath}'.`);
        throw error;
    }
}


async function run() {
    const startTime = Date.now();
    const { tokenizer, data } = await loadBenchmarkData();

    if (data.length === 0) {
        throw new Error("Нет данных для бенчмарка.");
    }
    
    const modelConfig = {
        vocab_size: tokenizer.vocab_size,
        embedding_dim: config.embedding_dim,
        hidden_dim: config.hidden_dim
    };

    const model = new HeuristicEngine(modelConfig);
    const optimizer = new slmnet.optimizers.Adam(model.parameters(), config.learning_rate);
    
    let total_loss = 0;
    let steps_done = 0;

    optimizer.zero_grad();

    for (let i = 0; i < data.length && steps_done < config.benchmark_steps; i++) {
        const sample = data[i];
        const input_ids = sample.prompt_ids.concat(sample.completion_ids);

        if (input_ids.length <= 1) continue;

        const model_input_ids = input_ids.slice(0, -1);
        const target_ids = input_ids.slice(1);

        try {
            const x = new slmnet.Tensor(model_input_ids, [1, model_input_ids.length]);
            const logits = await model.forward(x);
            
            const reshaped_logits = new slmnet.Tensor(logits.data, [logits.shape[1], logits.shape[2]]);
            const reshaped_targets = new slmnet.Tensor(target_ids, [target_ids.length]);
            const loss = await slmnet.losses.cross_entropy_loss(reshaped_logits, reshaped_targets);
            
            if (isNaN(loss.data[0])) {
                continue;
            }

            const current_loss_val = loss.data[0] / config.accumulation_steps;
            total_loss += loss.data[0];

            const loss_for_backward = new slmnet.Tensor([current_loss_val], [1]);
            loss_for_backward.backward_fn = loss.backward_fn;
            
            await loss_for_backward.backward();

            if ((i + 1) % config.accumulation_steps === 0) {
                for(const p of model.parameters()) {
                    if (p.grad) {
                        let norm = 0;
                        for(const val of p.grad.data) norm += val*val;
                        norm = Math.sqrt(norm);
                        if(norm > config.grad_clip) {
                            for(let j=0; j<p.grad.data.length; j++) p.grad.data[j] *= config.grad_clip/norm;
                        }
                    }
                }
                optimizer.step();
                optimizer.zero_grad();
                steps_done++;
            }
        } catch (error) {
             // В бенчмарке мы просто игнорируем ошибки шагов
        }
    }
    
    const endTime = Date.now();
    const finalLoss = total_loss / (steps_done * config.accumulation_steps);
    const totalTime = endTime - startTime;

    console.log(`--- Benchmark Result ---`);
    console.log(`Final Loss: ${finalLoss.toFixed(4)}`);
    console.log(`Time: ${totalTime}ms`);
    console.log(`------------------------`);

    await workerPool.destroy();
}

run().catch(error => {
    console.error("Критическая ошибка в бенчмарке:", error);
    if (workerPool) {
      workerPool.destroy();
    }
    process.exit(1);
});