Spaces:
No application file
No application file
| """ | |
| evaluate_model.py | |
| ----------------- | |
| Compares base Phi-2 vs fine-tuned Phi-2 on sample radiology cases. | |
| Usage: | |
| python training/evaluate_model.py | |
| """ | |
| import json | |
| import logging | |
| import sys | |
| import time | |
| from pathlib import Path | |
| import torch | |
| logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(message)s") | |
| logger = logging.getLogger(__name__) | |
| ROOT = Path(__file__).resolve().parent.parent | |
| sys.path.insert(0, str(ROOT)) | |
| ADAPTER_DIR = ROOT / "models" / "slm_model" / "lora_adapter" | |
| TEST_CASES = [ | |
| { | |
| "label": "Chest X-ray — Pneumonia", | |
| "prompt": ( | |
| "### System:\nYou are a board-certified radiologist.\n\n" | |
| "### Instruction:\nGenerate a structured radiology report.\n" | |
| "Patient: 52-year-old Male. Symptoms: fever, productive cough.\n" | |
| "Study: X-ray of the Chest, PA view.\n" | |
| "Write FINDINGS, IMPRESSION, and RECOMMENDATIONS sections.\n\n" | |
| "### Response:\nFINDINGS: The chest demonstrates" | |
| ), | |
| }, | |
| { | |
| "label": "Pelvis X-ray — Post Hip Arthroplasty", | |
| "prompt": ( | |
| "### System:\nYou are a board-certified radiologist.\n\n" | |
| "### Instruction:\nGenerate a structured radiology report.\n" | |
| "Patient: 68-year-old Female. Symptoms: hip pain after fall.\n" | |
| "Study: X-ray of the Pelvis, AP view.\n" | |
| "Write FINDINGS, IMPRESSION, and RECOMMENDATIONS sections.\n\n" | |
| "### Response:\nFINDINGS: The pelvis demonstrates" | |
| ), | |
| }, | |
| ] | |
| def run_inference(model, tokenizer, prompt: str, max_new_tokens: int = 300) -> str: | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512).to(device) | |
| start = time.time() | |
| with torch.no_grad(): | |
| output_ids = model.generate( | |
| **inputs, | |
| max_new_tokens=max_new_tokens, | |
| temperature=0.3, | |
| top_p=0.9, | |
| do_sample=True, | |
| repetition_penalty=1.2, | |
| pad_token_id=tokenizer.eos_token_id, | |
| ) | |
| elapsed = time.time() - start | |
| new_tokens = output_ids[0][inputs["input_ids"].shape[1]:] | |
| text = tokenizer.decode(new_tokens, skip_special_tokens=True) | |
| return text, elapsed | |
| def evaluate(): | |
| from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig | |
| bnb_config = BitsAndBytesConfig( | |
| load_in_4bit=True, | |
| bnb_4bit_use_double_quant=True, | |
| bnb_4bit_quant_type="nf4", | |
| bnb_4bit_compute_dtype=torch.bfloat16, | |
| ) | |
| logger.info("Loading base Phi-2...") | |
| tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True) | |
| if tokenizer.pad_token is None: | |
| tokenizer.pad_token = tokenizer.eos_token | |
| base_model = AutoModelForCausalLM.from_pretrained( | |
| "microsoft/phi-2", | |
| quantization_config=bnb_config, | |
| device_map="auto", | |
| trust_remote_code=True, | |
| ) | |
| base_model.eval() | |
| # Load fine-tuned if available | |
| finetuned_model = None | |
| if ADAPTER_DIR.exists() and (ADAPTER_DIR / "adapter_config.json").exists(): | |
| logger.info("Loading fine-tuned adapter...") | |
| from peft import PeftModel | |
| finetuned_model = PeftModel.from_pretrained(base_model, str(ADAPTER_DIR)) | |
| finetuned_model.eval() | |
| logger.info("Fine-tuned model loaded!") | |
| else: | |
| logger.warning(f"No fine-tuned adapter found at {ADAPTER_DIR}") | |
| logger.warning("Run train_unsloth.py first to fine-tune the model") | |
| # Evaluate | |
| results = [] | |
| print("\n" + "="*70) | |
| print("EVALUATION RESULTS") | |
| print("="*70) | |
| for case in TEST_CASES: | |
| print(f"\n📋 {case['label']}") | |
| print("-" * 70) | |
| # Base model | |
| base_output, base_time = run_inference(base_model, tokenizer, case["prompt"]) | |
| print(f"\n🔵 BASE MODEL ({base_time:.1f}s):") | |
| print("FINDINGS: The " + case["prompt"].split("demonstrates")[-1][:50] + "...") | |
| print(base_output[:400]) | |
| # Fine-tuned model | |
| if finetuned_model: | |
| ft_output, ft_time = run_inference(finetuned_model, tokenizer, case["prompt"]) | |
| print(f"\n🟢 FINE-TUNED MODEL ({ft_time:.1f}s):") | |
| print(ft_output[:400]) | |
| results.append({ | |
| "case": case["label"], | |
| "base_output": base_output, | |
| "finetuned_output": finetuned_model and ft_output or "N/A", | |
| }) | |
| # Save results | |
| out_path = Path(__file__).parent / "evaluation_results.json" | |
| out_path.write_text(json.dumps(results, indent=2)) | |
| logger.info(f"\nResults saved to {out_path}") | |
| if __name__ == "__main__": | |
| evaluate() | |