Maaac/CodeLLaMA-Linux-BugFix

1import torch
2from transformers import AutoTokenizer, AutoModelForCausalLM
3from datasets import load_dataset
4from tqdm import tqdm
5import json
6import csv
7import os
8import evaluate
9
10# ==== CONFIG ====
11MODEL_PATH = "../train/output/qlora-codellama-bugfix"
12EVAL_FILE = "test_samples.jsonl"
13OUTPUT_JSON = "./output/eval_results.json"
14OUTPUT_CSV = "./output/eval_results.csv"
15MAX_INPUT_LEN = 1024
16MAX_NEW_TOKENS = 256
17DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
18
19# ==== Ensure output folder exists ====
20os.makedirs(os.path.dirname(OUTPUT_JSON), exist_ok=True)
21
22# ==== Load model ====
23print("🔄 Loading model...")
24tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
25model = AutoModelForCausalLM.from_pretrained(
26    MODEL_PATH,
27    torch_dtype=torch.bfloat16,
28    device_map="auto"
29)
30model.eval()
31
32# ==== Load eval data ====
33print("📂 Loading evaluation data...")
34eval_data = load_dataset("json", data_files=EVAL_FILE, split="train")
35
36# ==== Inference ====
37results = []
38print("⚙️ Running inference...")
39for example in tqdm(eval_data):
40    prompt = example["prompt"]
41    reference = example["completion"]
42
43    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=MAX_INPUT_LEN).to(DEVICE)
44
45    with torch.no_grad():
46        outputs = model.generate(
47            **inputs,
48            max_new_tokens=MAX_NEW_TOKENS,
49            do_sample=False,
50            num_beams=4,
51            pad_token_id=tokenizer.pad_token_id,
52            eos_token_id=tokenizer.eos_token_id
53        )
54
55    prediction = tokenizer.decode(outputs[0], skip_special_tokens=True)
56
57    results.append({
58        "prompt": prompt,
59        "reference": reference.strip(),
60        "prediction": prediction.strip()
61    })
62
63# ==== Save results ====
64with open(OUTPUT_JSON, "w", encoding="utf-8") as f:
65    json.dump(results, f, indent=2)
66print(f"✅ Saved JSON to {OUTPUT_JSON}")
67
68with open(OUTPUT_CSV, "w", encoding="utf-8", newline='') as f:
69    writer = csv.DictWriter(f, fieldnames=["prompt", "reference", "prediction"])
70    writer.writeheader()
71    writer.writerows(results)
72print(f"✅ Saved CSV to {OUTPUT_CSV}")
73
74# ==== Compute Metrics ====
75print("📊 Computing BLEU and ROUGE...")
76bleu = evaluate.load("bleu")
77rouge = evaluate.load("rouge")
78
79predictions = [r["prediction"] for r in results]
80references = [r["reference"] for r in results]
81
82bleu_score = bleu.compute(predictions=predictions, references=[[ref] for ref in references])
83rouge_score = rouge.compute(predictions=predictions, references=references)
84
85print("\n📈 Evaluation Results:")
86print("BLEU:", bleu_score)
87print("ROUGE:", json.dumps(rouge_score, indent=2))
88