scripts/evaluate.py - main - Mike_Seleznev/local_llm

scripts/evaluate.py
244 строки5 KB
seleznev
Initial commit
24 июл 2026, 16:07
24 июл 2026, 16:07bf35079
import json
from pathlib import Path

import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

from config import (
    MODEL_PATH,
    LORA_PATH,
    TOOLS_PATH,
    EVAL_DATASET_PATH,
    MAX_NEW_TOKENS,
)


NEGATIVE_EVAL_DATASET_PATH = Path(
    "dataset/negative_eval.jsonl"
)


def get_device() -> torch.device:
    if torch.backends.mps.is_available():
        return torch.device("mps")

    if torch.cuda.is_available():
        return torch.device("cuda")

    return torch.device("cpu")


def load_json(path: Path) -> list[dict]:
    if not path.exists():
        raise FileNotFoundError(f"Файл не найден: {path}")

    with path.open("r", encoding="utf-8") as file:
        data = json.load(file)

    if not isinstance(data, list):
        raise ValueError(
            f"Файл {path} должен содержать JSON-массив"
        )

    return data


def load_eval_samples(path: Path) -> list[dict]:
    if not path.exists():
        raise FileNotFoundError(f"Файл не найден: {path}")

    samples = []

    with path.open("r", encoding="utf-8") as file:
        for line_number, line in enumerate(file, start=1):
            line = line.strip()

            if not line:
                continue

            try:
                sample = json.loads(line)
            except json.JSONDecodeError as error:
                raise ValueError(
                    f"Ошибка JSON в файле {path}, "
                    f"строка {line_number}: {error}"
                ) from error

            text = sample.get("text")
            expected = sample.get("expected")

            if not isinstance(text, str) or not text.strip():
                raise ValueError(
                    f"Некорректный text в {path}, "
                    f"строка {line_number}"
                )

            if not isinstance(expected, str) or not expected:
                raise ValueError(
                    f"Некорректный expected в {path}, "
                    f"строка {line_number}"
                )

            samples.append(
                {
                    "text": text.strip(),
                    "expected": expected,
                }
            )

    return samples


def predict(
    model,
    tokenizer: AutoTokenizer,
    tools: list[dict],
    text: str,
    device: torch.device,
) -> tuple[str, str]:
    messages = [
        {
            "role": "user",
            "content": text,
        }
    ]

    inputs = tokenizer.apply_chat_template(
        messages,
        tools=tools,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    )

    inputs = {
        key: value.to(device)
        for key, value in inputs.items()
    }

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id,
        )

    input_length = inputs["input_ids"].shape[1]
    generated_tokens = outputs[0][input_length:]

    raw_response = tokenizer.decode(
        generated_tokens,
        skip_special_tokens=False,
    )

    prediction = "none"

    for tool in tools:
        action_name = tool["function"]["name"]

        if f"call:{action_name}" in raw_response:
            prediction = action_name
            break

    return prediction, raw_response


def evaluate_samples(
    samples: list[dict],
    model,
    tokenizer: AutoTokenizer,
    tools: list[dict],
    device: torch.device,
) -> tuple[int, int]:
    correct = 0

    for sample in samples:
        text = sample["text"]
        expected = sample["expected"]

        prediction, raw_response = predict(
            model=model,
            tokenizer=tokenizer,
            tools=tools,
            text=text,
            device=device,
        )

        is_correct = prediction == expected

        if is_correct:
            correct += 1

        marker = "✓" if is_correct else "✗"

        print(
            f"{marker} {text}\n"
            f"  expected:  {expected}\n"
            f"  predicted: {prediction}\n"
            f"  raw:       {raw_response}\n"
        )

    return correct, len(samples)


def main() -> None:
    device = get_device()

    print(f"Device: {device}")
    print(f"Model:  {MODEL_PATH}")
    print(f"LoRA:   {LORA_PATH}")
    print()

    tools = load_json(TOOLS_PATH)

    positive_samples = load_eval_samples(
        Path(EVAL_DATASET_PATH)
    )

    negative_samples = load_eval_samples(
        NEGATIVE_EVAL_DATASET_PATH
    )

    samples = positive_samples + negative_samples

    tokenizer = AutoTokenizer.from_pretrained(
        str(LORA_PATH)
    )

    base_model = AutoModelForCausalLM.from_pretrained(
        str(MODEL_PATH),
        dtype=torch.float32,
    )

    model = PeftModel.from_pretrained(
        base_model,
        str(LORA_PATH),
    )

    model.to(device)
    model.eval()

    correct, total = evaluate_samples(
        samples=samples,
        model=model,
        tokenizer=tokenizer,
        tools=tools,
        device=device,
    )

    accuracy = correct / total if total else 0.0

    positive_count = len(positive_samples)
    negative_count = len(negative_samples)

    print("=" * 70)
    print(f"Positive samples: {positive_count}")
    print(f"Negative samples: {negative_count}")
    print(f"Result: {correct}/{total}")
    print(f"Accuracy: {accuracy:.1%}")


if __name__ == "__main__":
    main()