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()