| """ |
| Compare base vs trained LifeStack policy on identical crisis prompts. |
| |
| Usage: |
| python scripts/compare_baseline.py |
| python scripts/compare_baseline.py --trained-model ./lifestack_model |
| """ |
|
|
| import argparse |
| import json |
| import os |
| import random |
| import sys |
| from datetime import datetime |
| from typing import Any |
|
|
| import torch |
|
|
| SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) |
| REPO_ROOT = os.path.dirname(SCRIPT_DIR) |
| sys.path.insert(0, REPO_ROOT) |
| sys.path.insert(0, SCRIPT_DIR) |
|
|
| from agent.conflict_generator import TaskGenerator, generate_conflict |
| from core.life_state import DependencyGraph, LifeMetrics, ResourceBudget |
| from intake.simperson import SimPerson |
| from scripts.train_trl import build_prompt_for_task, get_lifestack_evaluation |
|
|
|
|
| def _load_base_model(): |
| """Load base Qwen2.5-1.5B-Instruct (no training adapter).""" |
| try: |
| from unsloth import FastLanguageModel |
|
|
| model, tokenizer = FastLanguageModel.from_pretrained( |
| model_name="unsloth/Qwen2.5-1.5B-Instruct", |
| max_seq_length=1024, |
| load_in_4bit=True, |
| ) |
| FastLanguageModel.for_inference(model) |
| return model, tokenizer, "unsloth/base-qwen2.5-1.5b-instruct" |
| except Exception: |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| model_name = "Qwen/Qwen2.5-1.5B-Instruct" |
| tokenizer = AutoTokenizer.from_pretrained(model_name) |
| model = AutoModelForCausalLM.from_pretrained( |
| model_name, |
| torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, |
| device_map="auto", |
| ) |
| model.eval() |
| return model, tokenizer, model_name |
|
|
|
|
| def _load_trained_model(model_dir: str): |
| """Load trained LifeStack model from local adapter/full checkpoint directory.""" |
| try: |
| from unsloth import FastLanguageModel |
|
|
| model, tokenizer = FastLanguageModel.from_pretrained( |
| model_name=model_dir, |
| max_seq_length=1024, |
| load_in_4bit=True, |
| ) |
| FastLanguageModel.for_inference(model) |
| return model, tokenizer |
| except Exception: |
| from peft import PeftModel |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| tokenizer = AutoTokenizer.from_pretrained(model_dir) |
| base = AutoModelForCausalLM.from_pretrained( |
| "Qwen/Qwen2.5-1.5B-Instruct", |
| torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, |
| device_map="auto", |
| ) |
| model = PeftModel.from_pretrained(base, model_dir) |
| model.eval() |
| return model, tokenizer |
|
|
|
|
| def _device_for(model) -> torch.device: |
| try: |
| return next(model.parameters()).device |
| except Exception: |
| return torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
|
|
| def _generate_completion(model, tokenizer, prompt: str, temperature: float = 0.3) -> str: |
| device = _device_for(model) |
| inputs = tokenizer(prompt, return_tensors="pt").to(device) |
| pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id |
|
|
| with torch.no_grad(): |
| outputs = model.generate( |
| **inputs, |
| max_new_tokens=128, |
| temperature=temperature, |
| do_sample=True, |
| top_p=0.9, |
| pad_token_id=pad_token_id, |
| eos_token_id=tokenizer.eos_token_id, |
| ) |
|
|
| return tokenizer.decode(outputs[0][inputs["input_ids"].shape[1] :], skip_special_tokens=True).strip() |
|
|
|
|
| def _build_eval_cases() -> list[dict[str, Any]]: |
| """Create 5 deterministic prompts spanning different crisis domains.""" |
| domains = [ |
| ("career", 3, 101), |
| ("finances", 4, 202), |
| ("relationships", 3, 303), |
| ("transport_crisis", 4, 404), |
| ("code_merge_crisis", 5, 505), |
| ] |
|
|
| generator = TaskGenerator() |
| graph = DependencyGraph() |
| person = SimPerson(name="Comparator") |
| cases: list[dict[str, Any]] = [] |
|
|
| for domain, difficulty, seed in domains: |
| random.seed(seed) |
| task = generator.generate(domain=domain, difficulty=difficulty) |
| conflict = generate_conflict(difficulty) |
| random.seed() |
|
|
| metrics = LifeMetrics() |
| metrics = graph.cascade(metrics, {**task.mutable_world, **conflict.primary_disruption}) |
|
|
| budget_dict = task.constraints.get("budget", {}) |
| budget = ResourceBudget( |
| time_hours=budget_dict.get("time", 20.0), |
| money_dollars=budget_dict.get("money", 500.0), |
| energy_units=budget_dict.get("energy", 100.0), |
| ) |
|
|
| prompt = build_prompt_for_task(task, person, metrics, budget, seed=seed, step=0) |
| crisis_text = task.domain_metadata.get("story", task.goal) |
|
|
| cases.append( |
| { |
| "case_id": f"{domain}_d{difficulty}", |
| "domain": domain, |
| "difficulty": difficulty, |
| "seed": seed, |
| "crisis": crisis_text, |
| "prompt": prompt, |
| } |
| ) |
| return cases |
|
|
|
|
| def _print_case(case: dict[str, Any]) -> None: |
| print("=" * 110) |
| print(f"[{case['case_id']}] domain={case['domain']} difficulty={case['difficulty']}") |
| print(f"crisis: {case['crisis']}") |
| print(f"base_reward={case['base_reward']:.3f} | trained_reward={case['trained_reward']:.3f} | delta={case['delta']:+.3f}") |
| print("- BASE RESPONSE -") |
| print(case["base_response"] or "<empty>") |
| print("- TRAINED RESPONSE -") |
| print(case["trained_response"] or "<empty>") |
|
|
|
|
| def run_compare(trained_model_dir: str, output_path: str) -> dict[str, Any]: |
| cases = _build_eval_cases() |
|
|
| print("Loading base model...") |
| base_model, base_tokenizer, base_name = _load_base_model() |
| for case in cases: |
| completion = _generate_completion(base_model, base_tokenizer, case["prompt"]) |
| eval_data = get_lifestack_evaluation(completion, case["prompt"]) |
| case["base_model"] = base_name |
| case["base_response"] = completion |
| case["base_reward"] = float(eval_data.get("reward", -0.5)) |
| del base_model |
| torch.cuda.empty_cache() |
|
|
| print("Loading trained model...") |
| trained_model, trained_tokenizer = _load_trained_model(trained_model_dir) |
| for case in cases: |
| completion = _generate_completion(trained_model, trained_tokenizer, case["prompt"]) |
| eval_data = get_lifestack_evaluation(completion, case["prompt"]) |
| case["trained_model"] = trained_model_dir |
| case["trained_response"] = completion |
| case["trained_reward"] = float(eval_data.get("reward", -0.5)) |
| case["delta"] = round(case["trained_reward"] - case["base_reward"], 4) |
| _print_case(case) |
| del trained_model |
| torch.cuda.empty_cache() |
|
|
| avg_base = sum(c["base_reward"] for c in cases) / len(cases) |
| avg_trained = sum(c["trained_reward"] for c in cases) / len(cases) |
| avg_delta = avg_trained - avg_base |
|
|
| payload = { |
| "timestamp_utc": datetime.utcnow().isoformat() + "Z", |
| "summary": { |
| "n_cases": len(cases), |
| "avg_base_reward": round(avg_base, 4), |
| "avg_trained_reward": round(avg_trained, 4), |
| "avg_reward_delta": round(avg_delta, 4), |
| "base_model": cases[0]["base_model"] if cases else "", |
| "trained_model": trained_model_dir, |
| }, |
| "cases": cases, |
| } |
|
|
| output_dir = os.path.dirname(output_path) |
| if output_dir: |
| os.makedirs(output_dir, exist_ok=True) |
| with open(output_path, "w", encoding="utf-8") as f: |
| json.dump(payload, f, indent=2) |
|
|
| print("=" * 110) |
| print( |
| f"SUMMARY: avg_base={avg_base:.3f} | avg_trained={avg_trained:.3f} | " |
| f"avg_delta={avg_delta:+.3f}" |
| ) |
| print(f"Saved comparison JSON: {output_path}") |
| return payload |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="Compare baseline Qwen vs trained LifeStack model.") |
| parser.add_argument("--trained-model", type=str, default="./lifestack_model") |
| parser.add_argument("--output", type=str, default="./data/before_after_comparison.json") |
| args = parser.parse_args() |
|
|
| run_compare(trained_model_dir=args.trained_model, output_path=args.output) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
