main.py

import json
import torch
import re
import time
from pathlib import Path
from tqdm import tqdm
from transformers import AutoModelForCausalLM, AutoTokenizer
from llama_cpp import Llama

# ============================================================
# CONFIGURATION
# ============================================================
DATASET_PATH = "gramm.jsonl"  # path to Q&A file
MODEL_PATH = "model/"  # path to local model
USE_GGUF = False  # True for GGUF, False for HF
TEMPERATURE = 0.6
MAX_TOKENS_PREDICT = 4  # model can extend up to 4 tokens
VERBOSE = True  # show each prediction


# ============================================================
# LOAD DATASET
# ============================================================
def load_dataset(path):
    data = []
    with open(path, "r", encoding="utf-8") as f:
        for line_num, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue
            try:
                item = json.loads(line)
            except json.JSONDecodeError:
                print(f"Warning: skipping line {line_num}, invalid JSON")
                continue
            if "q" not in item or "a" not in item:
                print(f"Warning: skipping line {line_num}, missing 'q' or 'a'")
                continue

            # Parse answers - can be space-separated or comma-separated
            answer_str = item["a"].strip()
            # Split by spaces first, then clean each answer
            answers = [ans.strip().lower().rstrip('.,;:!?') for ans in answer_str.split()]
            # Remove empty strings
            answers = [ans for ans in answers if ans]

            data.append({
                "question": item["q"],
                "answers": answers  # list of correct answers
            })
    print(f"Loaded {len(data)} test cases from {path}")

    # Show some examples of multiple answers
    multi_answer = sum(1 for item in data if len(item["answers"]) > 1)
    print(f"  - {multi_answer} questions have multiple correct answers")

    return data


# ============================================================
# LOAD MODEL
# ============================================================
def load_model():
    model_type = "gguf" if USE_GGUF else "transformers"
    print(f"Loading model from {MODEL_PATH} ({model_type})...")

    if USE_GGUF:
        model = Llama(
            model_path=MODEL_PATH,
            n_ctx=4096,
            n_threads=8,
            temperature=TEMPERATURE,
            verbose=False
        )
        return model, None, "gguf"
    else:
        tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)

        # Set pad_token if not present
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token

        model = AutoModelForCausalLM.from_pretrained(
            MODEL_PATH,
            torch_dtype=torch.float16,
            device_map="auto"
        )
        model.eval()
        return model, tokenizer, "hf"


# ============================================================
# EXTRACT FIRST WORD AFTER GENERATION
# ============================================================
def extract_first_word(text, question):
    continuation = text[len(question):].strip()
    parts = continuation.split()
    if not parts:
        return ""
    first_word = parts[0].strip().lower()
    # Remove punctuation from edges
    first_word = re.sub(r'^[^\w]+|[^\w]+$', '', first_word)
    return first_word


# ============================================================
# PREDICT NEXT TOKENS (GGUF)
# ============================================================
def predict_gguf(model, question):
    output = model.create_completion(
        prompt=question,
        max_tokens=MAX_TOKENS_PREDICT,
        temperature=TEMPERATURE,
        echo=True,
        stop=["\n", ".", "!", "?"]
    )
    generated_text = output["choices"][0]["text"]
    return generated_text


# ============================================================
# PREDICT NEXT TOKENS (Transformers)
# ============================================================
def predict_transformers(model, tokenizer, question):
    inputs = tokenizer(question, return_tensors="pt").to(model.device)

    # Remove token_type_ids if present (causal LMs don't use them)
    if "token_type_ids" in inputs:
        del inputs["token_type_ids"]

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=MAX_TOKENS_PREDICT,
            temperature=TEMPERATURE,
            do_sample=True,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id
        )

    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return generated_text


# ============================================================
# RUN SINGLE TEST
# ============================================================
def run_single_test(model, tokenizer, model_type, test_case):
    question = test_case["question"]
    expected_answers = test_case["answers"]  # list of correct answers

    try:
        if model_type == "gguf":
            generated = predict_gguf(model, question)
        else:
            generated = predict_transformers(model, tokenizer, question)

        predicted_word = extract_first_word(generated, question)
        # Check if predicted word matches ANY of the expected answers
        is_correct = predicted_word in expected_answers
    except Exception as e:
        generated = f"ERROR: {str(e)}"
        predicted_word = ""
        is_correct = False

    result = {
        "question": question,
        "expected": expected_answers,  # list of all correct answers
        "predicted": predicted_word,
        "full_generation": generated,
        "correct": is_correct
    }

    return result


# ============================================================
# RUN FULL BENCHMARK
# ============================================================
def run_benchmark():
    data = load_dataset(DATASET_PATH)
    if not data:
        print("No test data loaded. Exiting.")
        return

    model, tokenizer, model_type = load_model()

    correct = 0
    total = len(data)
    results = []

    start_time = time.time()

    for i, test_case in enumerate(tqdm(data, desc="Testing")):
        result = run_single_test(model, tokenizer, model_type, test_case)
        results.append(result)

        if result["correct"]:
            correct += 1

        if VERBOSE:
            status = "OK" if result["correct"] else "FAIL"
            print(f"\n[{i + 1}/{total}] {status}")
            print(f"  Q: {result['question']}")
            # Format expected answers nicely
            expected_str = " | ".join(result['expected'])
            print(f"  Expected: [{expected_str}] | Got: '{result['predicted']}'")
            if not result["correct"]:
                print(f"  Full: '{result['full_generation']}'")

    elapsed = time.time() - start_time
    accuracy = (correct / total) * 100 if total > 0 else 0.0

    print("\n" + "=" * 60)
    print(f"BENCHMARK RESULTS")
    print("=" * 60)
    print(f"Total questions:   {total}")
    print(f"Correct answers:   {correct}")
    print(f"Failed answers:    {total - correct}")
    print(f"Accuracy:          {accuracy:.2f}%")
    print(f"Time elapsed:      {elapsed:.2f} seconds")
    if total > 0:
        print(f"Avg time/question: {(elapsed / total) * 1000:.0f} ms")
    print("=" * 60)

    # Calculate per-question stats for multi-answer questions
    multi_answer_questions = [r for r in results if len(r["expected"]) > 1]
    if multi_answer_questions:
        multi_correct = sum(1 for r in multi_answer_questions if r["correct"])
        multi_accuracy = (multi_correct / len(multi_answer_questions)) * 100
        print(f"Multi-answer Qs:   {len(multi_answer_questions)} "
              f"(Accuracy: {multi_accuracy:.2f}%)")

    save_results(results, accuracy, elapsed)

    return accuracy


# ============================================================
# SAVE RESULTS
# ============================================================
def save_results(results, accuracy, elapsed):
    output_path = "benchmark_results.json"

    summary = {
        "model_path": MODEL_PATH,
        "model_type": "gguf" if USE_GGUF else "transformers",
        "temperature": TEMPERATURE,
        "total_questions": len(results),
        "correct": sum(1 for r in results if r["correct"]),
        "accuracy": accuracy,
        "time_elapsed_seconds": elapsed,
        "details": [
            {
                "question": r["question"],
                "expected": r["expected"],  # list of all correct answers
                "predicted": r["predicted"],
                "correct": r["correct"]
            }
            for r in results
        ]
    }

    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(summary, f, indent=2, ensure_ascii=False)

    print(f"\nDetailed results saved to {output_path}")


# ============================================================
# MAIN
# ============================================================
if __name__ == "__main__":
    run_benchmark()