| """ |
| Interactive Gradio demo for LifeStack trained model. |
| |
| Usage: |
| python scripts/gradio_demo.py --model-dir ./lifestack_model |
| """ |
|
|
| import argparse |
| import json |
| import os |
| import random |
| import re |
| import sys |
| from typing import Any |
|
|
| import gradio as gr |
| import matplotlib |
| import torch |
|
|
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
|
|
| 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 ( |
| CASCADE_DAMPENING_DEFAULT, |
| DependencyGraph, |
| LifeMetrics, |
| ResourceBudget, |
| ) |
| from intake.simperson import SimPerson |
| from scripts.inference import load_model |
| from scripts.train_trl import ALL_DOMAINS, build_prompt_for_task, generate_dataset, get_lifestack_evaluation |
|
|
| MODEL = None |
| TOKENIZER = None |
| MODEL_DIR = "./lifestack_model" |
|
|
|
|
| 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 _ensure_model_loaded(): |
| global MODEL, TOKENIZER |
| if MODEL is None or TOKENIZER is None: |
| MODEL, TOKENIZER = load_model(MODEL_DIR) |
|
|
|
|
| def _extract_json_payload(text: str) -> dict[str, Any]: |
| cleaned = text.strip() |
| if "```json" in cleaned: |
| cleaned = cleaned.split("```json")[-1].split("```")[0].strip() |
| elif "```" in cleaned: |
| cleaned = cleaned.split("```")[-1].split("```")[0].strip() |
|
|
| try: |
| data = json.loads(cleaned) |
| if isinstance(data, dict): |
| return data |
| return {"json_value": data} |
| except Exception: |
| start = cleaned.find("{") |
| end = cleaned.rfind("}") |
| if start != -1 and end > start: |
| try: |
| return json.loads(cleaned[start : end + 1]) |
| except Exception as err: |
| return {"raw_output": text, "parse_error": str(err)} |
| return {"raw_output": text, "parse_error": "no valid JSON object found"} |
|
|
|
|
| def _generate_completion(prompt: str, temperature: float = 0.4) -> str: |
| _ensure_model_loaded() |
| 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_crisis_prompt(crisis_text: str, domain: str, difficulty: int) -> tuple[str, dict[str, float]]: |
| generator = TaskGenerator() |
| graph = DependencyGraph() |
| person = SimPerson(name="DemoUser") |
|
|
| eval_seed = random.randint(1, 999999) |
| random.seed(eval_seed) |
| task = generator.generate(domain=domain, difficulty=int(difficulty)) |
| conflict = generate_conflict(int(difficulty)) |
| random.seed() |
|
|
| if crisis_text.strip(): |
| task.goal = crisis_text.strip() |
| task.domain_metadata["story"] = crisis_text.strip() |
|
|
| 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=eval_seed, step=0) |
| return prompt, dict(task.mutable_world) |
|
|
|
|
| def _select_metric_keys(before: dict[str, float], after: dict[str, float]) -> list[str]: |
| priority = [ |
| "career.workload", |
| "finances.liquidity", |
| "relationships.romantic", |
| "physical_health.energy", |
| "mental_wellbeing.stress_level", |
| "time.free_hours_per_week", |
| ] |
| keys = [k for k in priority if k in before or k in after] |
| if len(keys) < 6: |
| pool = sorted(set(before.keys()) | set(after.keys())) |
| for k in pool: |
| if k not in keys: |
| keys.append(k) |
| if len(keys) == 6: |
| break |
| return keys |
|
|
|
|
| def _plot_before_after(before: dict[str, float], after: dict[str, float]): |
| fig, ax = plt.subplots(figsize=(8, 4)) |
| if not before and not after: |
| ax.text(0.5, 0.5, "No metric data available", ha="center", va="center") |
| ax.axis("off") |
| return fig |
|
|
| keys = _select_metric_keys(before, after) |
| x = range(len(keys)) |
| before_vals = [before.get(k, 0.0) for k in keys] |
| after_vals = [after.get(k, 0.0) for k in keys] |
|
|
| ax.bar([i - 0.2 for i in x], before_vals, width=0.4, label="Before", color="#9ca3af") |
| ax.bar([i + 0.2 for i in x], after_vals, width=0.4, label="After", color="#16a34a") |
| ax.set_ylim(0, 100) |
| ax.set_xticks(list(x)) |
| ax.set_xticklabels([k.split(".")[-1] for k in keys], rotation=20, ha="right") |
| ax.set_title("Life Metrics Before vs After") |
| ax.set_ylabel("Score") |
| ax.grid(axis="y", alpha=0.25) |
| ax.legend() |
| fig.tight_layout() |
| return fig |
|
|
|
|
| def _plot_trajectory(trajectory: list[dict[str, Any]]): |
| fig, ax = plt.subplots(figsize=(8, 4)) |
| if not trajectory: |
| ax.text(0.5, 0.5, "No trajectory data available", ha="center", va="center") |
| ax.axis("off") |
| return fig |
|
|
| days = [point.get("step", idx + 1) for idx, point in enumerate(trajectory)] |
| rewards = [point.get("reward", 0.0) for point in trajectory] |
| stress = [point.get("metrics", {}).get("mental_wellbeing.stress_level", 0.0) for point in trajectory] |
|
|
| ax.plot(days, rewards, marker="o", linewidth=2, color="#1d4ed8", label="Daily Reward") |
| ax.set_xlabel("Day") |
| ax.set_ylabel("Reward") |
| ax.grid(alpha=0.3) |
|
|
| ax2 = ax.twinx() |
| ax2.plot(days, stress, marker="s", linestyle="--", color="#dc2626", label="Stress Level") |
| ax2.set_ylabel("Stress") |
|
|
| lines = ax.get_lines() + ax2.get_lines() |
| labels = [l.get_label() for l in lines] |
| ax.legend(lines, labels, loc="upper right") |
| ax.set_title("7-Day Trajectory") |
| fig.tight_layout() |
| return fig |
|
|
|
|
| def visualize_cascade(disruption_dict: dict[str, float]) -> str: |
| """Render a lightweight ASCII cascade tree for a disruption dict.""" |
| graph = DependencyGraph() |
| if not disruption_dict: |
| return "No disruption provided." |
|
|
| lines: list[str] = [] |
| for source_key, source_delta in disruption_dict.items(): |
| lines.append(f"{source_key} ({source_delta:+.1f})") |
| level_1 = graph.edges.get(source_key, [])[:3] |
| if not level_1: |
| lines.append(" └─ (no downstream edges)") |
| continue |
|
|
| for i, (target_key, weight) in enumerate(level_1): |
| branch = "└─" if i == len(level_1) - 1 else "├─" |
| level_1_delta = source_delta * weight * CASCADE_DAMPENING_DEFAULT |
| lines.append(f" {branch} {target_key} (w={weight:+.2f}, est={level_1_delta:+.1f})") |
|
|
| level_2 = graph.edges.get(target_key, [])[:2] |
| for j, (target_2, weight_2) in enumerate(level_2): |
| branch_2 = "└─" if j == len(level_2) - 1 else "├─" |
| indent = " " if i == len(level_1) - 1 else " │ " |
| level_2_delta = level_1_delta * weight_2 * CASCADE_DAMPENING_DEFAULT |
| lines.append(f"{indent}{branch_2} {target_2} (w={weight_2:+.2f}, est={level_2_delta:+.1f})") |
| return "\n".join(lines) |
|
|
|
|
| def _render_advice(action_json: dict[str, Any], reward: float, domain: str, difficulty: int) -> str: |
| action_type = action_json.get("action_type", "unknown") |
| target_domain = action_json.get("target_domain", "unknown") |
| reasoning = action_json.get("reasoning", "") |
| metric_changes = action_json.get("metric_changes", {}) |
| resource_cost = action_json.get("resource_cost", {}) |
|
|
| lines = [ |
| "### LifeStack Recommendation", |
| f"- Domain: `{domain}` | Difficulty: `{difficulty}`", |
| f"- Reward Score: `{reward:.3f}`", |
| f"- Action: `{action_type}`", |
| f"- Target: `{target_domain}`", |
| ] |
| if reasoning: |
| lines.append(f"- Why: {reasoning}") |
| if metric_changes: |
| top_changes = list(metric_changes.items())[:5] |
| lines.append("- Expected metric impact: " + ", ".join(f"`{k}` {v:+.1f}" for k, v in top_changes)) |
| if resource_cost: |
| lines.append( |
| "- Resource cost: " |
| f"time={resource_cost.get('time', 0)}, " |
| f"money={resource_cost.get('money', 0)}, " |
| f"energy={resource_cost.get('energy', 0)}" |
| ) |
| return "\n".join(lines) |
|
|
|
|
| def sample_random_crisis(): |
| ds = generate_dataset(n_prompts=1) |
| row = ds[0] |
| prompt = row["prompt"] |
| domain = row.get("domain", "career") |
| difficulty = int(row.get("difficulty", 3)) |
|
|
| m = re.search(r"(?:Task|TASK):\s*(.+)", prompt) |
| crisis_text = m.group(1).strip() if m else "My life is spiraling in multiple domains. What should I do first?" |
| return crisis_text, domain, difficulty |
|
|
|
|
| def run_live_demo(crisis_text: str, domain: str, difficulty: int): |
| if not crisis_text or not crisis_text.strip(): |
| crisis_text = "I am facing a multi-domain crisis and need a single best next action." |
|
|
| prompt, disruption = _build_crisis_prompt(crisis_text, domain, int(difficulty)) |
| completion = _generate_completion(prompt, temperature=0.4) |
| action_json = _extract_json_payload(completion) |
| eval_data = get_lifestack_evaluation(completion, prompt) |
|
|
| reward = float(eval_data.get("reward", -0.5)) |
| before = eval_data.get("initial_metrics", {}) |
| after = eval_data.get("obs_metrics", {}) |
| trajectory = eval_data.get("trajectory", []) |
|
|
| advice_md = _render_advice(action_json, reward, domain, int(difficulty)) |
| before_after_fig = _plot_before_after(before, after) |
| trajectory_fig = _plot_trajectory(trajectory) |
| cascade_tree = "```text\n" + visualize_cascade(disruption) + "\n```" |
|
|
| return advice_md, action_json, before_after_fig, trajectory_fig, cascade_tree |
|
|
|
|
| def build_app(): |
| with gr.Blocks(title="LifeStack GRPO Demo") as demo: |
| gr.Markdown("# LifeStack GRPO Demo") |
| gr.Markdown("Resolve a crisis and inspect action quality, life metric impact, trajectory, and cascade effects.") |
|
|
| with gr.Row(): |
| crisis_input = gr.Textbox( |
| label="Describe your life crisis", |
| lines=4, |
| placeholder="My flight got cancelled, my card was declined, and I have a client meeting tomorrow.", |
| ) |
| with gr.Row(): |
| domain_input = gr.Dropdown(choices=ALL_DOMAINS, value="career", label="Domain") |
| difficulty_input = gr.Slider(minimum=1, maximum=5, step=1, value=3, label="Difficulty") |
|
|
| with gr.Row(): |
| run_btn = gr.Button("Resolve Crisis", variant="primary") |
| random_btn = gr.Button("Try Random Crisis") |
|
|
| advice_out = gr.Markdown() |
| action_json_out = gr.JSON(label="Model JSON Decision") |
| with gr.Row(): |
| before_after_out = gr.Plot(label="Before/After Metrics") |
| trajectory_out = gr.Plot(label="7-Day Trajectory") |
| cascade_out = gr.Markdown() |
|
|
| run_btn.click( |
| fn=run_live_demo, |
| inputs=[crisis_input, domain_input, difficulty_input], |
| outputs=[advice_out, action_json_out, before_after_out, trajectory_out, cascade_out], |
| ) |
| random_btn.click( |
| fn=sample_random_crisis, |
| inputs=[], |
| outputs=[crisis_input, domain_input, difficulty_input], |
| ) |
|
|
| return demo |
|
|
|
|
| def main(): |
| global MODEL_DIR |
|
|
| parser = argparse.ArgumentParser(description="LifeStack Gradio demo.") |
| parser.add_argument("--model-dir", type=str, default="./lifestack_model") |
| parser.add_argument("--share", action="store_true", default=True, help="Launch with public share URL.") |
| parser.add_argument("--no-share", action="store_true", help="Disable Gradio share URL.") |
| parser.add_argument("--server-port", type=int, default=7860) |
| args = parser.parse_args() |
|
|
| MODEL_DIR = args.model_dir |
| demo = build_app() |
| demo.launch(share=(args.share and not args.no_share), server_port=args.server_port) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
