valperplex.py

import torch
import math
from transformers import AutoModelForCausalLM, AutoTokenizer
from torch.nn import CrossEntropyLoss


# -----------------------------
# CHAT FORMATS
# -----------------------------

def format_llama(text):
    lines = [l.strip() for l in text.split("\n") if l.strip()]
    out = []
    role = "user"

    for line in lines:
        if role == "user":
            out.append(f"[INST] {line} [/INST]")
            role = "assistant"
        else:
            out.append(line)
            role = "user"

    return "\n".join(out)


def format_plain(text):
    return text


FORMATS = {
    "1": ("plain", format_plain),
    "2": ("llama", format_llama),
}


# -----------------------------
# MODEL LOADING
# -----------------------------

def load_model(path, device="cuda"):
    print("Loading model...")
    model = AutoModelForCausalLM.from_pretrained(
        path,
        torch_dtype=torch.float16,
        device_map="auto"
    )
    tokenizer = AutoTokenizer.from_pretrained(path)
    return model, tokenizer


# -----------------------------
# PERPLEXITY
# -----------------------------

def compute_ppl(model, tokenizer, text, block_size=512):
    model.eval()
    tokens = tokenizer(text, return_tensors="pt", truncation=False)["input_ids"][0]

    loss_fn = CrossEntropyLoss(ignore_index=-100, reduction="sum")

    total_loss = 0
    total_tokens = 0

    with torch.no_grad():
        for i in range(0, len(tokens) - block_size, block_size):
            chunk = tokens[i:i+block_size].to(model.device)

            input_ids = chunk.unsqueeze(0)
            outputs = model(input_ids=input_ids)

            logits = outputs.logits

            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = input_ids[:, 1:].contiguous()

            loss = loss_fn(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
            )

            total_loss += loss.item()
            total_tokens += shift_labels.numel()

    avg_loss = total_loss / total_tokens
    return avg_loss, math.exp(avg_loss)


# -----------------------------
# APP
# -----------------------------

def main():
    model = None
    tokenizer = None
    text_data = None
    formatted_text = None

    while True:
        print("\n========== MENU ==========")
        print("1. Load model")
        print("2. Load text file")
        print("3. Choose format (plain / llama)")
        print("4. Run perplexity")
        print("5. Show sample text")
        print("0. Exit")

        choice = input("> ").strip()

        # ---------------- LOAD MODEL ----------------
        if choice == "1":
            path = input("Model path: ").strip()
            model, tokenizer = load_model(path)

        # ---------------- LOAD TEXT ----------------
        elif choice == "2":
            path = input("Text file path: ").strip()
            with open(path, "r", encoding="utf-8") as f:
                text_data = f.read()
            print("Text loaded.")

        # ---------------- FORMAT ----------------
        elif choice == "3":
            if text_data is None:
                print("Load text first.")
                continue

            print("1 = plain")
            print("2 = llama [INST] format")
            fmt = input("Choose format: ").strip()

            if fmt not in FORMATS:
                print("Invalid format")
                continue

            name, func = FORMATS[fmt]
            formatted_text = func(text_data)

            print(f"Applied format: {name}")

        # ---------------- PPL ----------------
        elif choice == "4":
            if model is None or tokenizer is None:
                print("Load model first.")
                continue
            if formatted_text is None:
                print("Load + format text first.")
                continue

            loss, ppl = compute_ppl(model, tokenizer, formatted_text)

            print("\n===== RESULTS =====")
            print(f"Loss: {loss:.4f}")
            print(f"Perplexity: {ppl:.4f}")

        # ---------------- SAMPLE ----------------
        elif choice == "5":
            if text_data:
                print("\n--- SAMPLE ---")
                print(text_data[:500])
            else:
                print("No text loaded.")

        # ---------------- EXIT ----------------
        elif choice == "0":
            break

        else:
            print("Invalid option")


if __name__ == "__main__":
    main()