605 lines
23 KiB
Python
605 lines
23 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
STABLE SELF-IMPROVEMENT TRAINER
|
||
================================
|
||
Recursive self-improvement with safeguards:
|
||
- Multi-metric evaluation (density + coherence + helpfulness)
|
||
- A/B checkpoint comparison
|
||
- Automatic rollback on quality drop
|
||
- Conservative training (low LR, small steps)
|
||
- Gibberish detection to prevent mode collapse
|
||
|
||
Usage:
|
||
python train_self_improve.py --iterations 5 --steps-per-iter 25
|
||
python train_self_improve.py --eval-only --checkpoint path/to/checkpoint
|
||
python train_self_improve.py --compare checkpoint_a checkpoint_b
|
||
|
||
"Improve without going insane"
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
import json
|
||
import argparse
|
||
import random
|
||
import re
|
||
import shutil
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import List, Dict, Any, Tuple, Optional
|
||
from dataclasses import dataclass, asdict
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
|
||
# === PATHS ===
|
||
ROOT = os.path.dirname(os.path.abspath(__file__))
|
||
CHECKPOINTS_DIR = os.path.join(ROOT, "dense_checkpoints_v2")
|
||
ROLLBACK_DIR = os.path.join(ROOT, "rollback_checkpoints")
|
||
LOGS_DIR = os.path.join(ROOT, "improvement_logs")
|
||
|
||
os.makedirs(CHECKPOINTS_DIR, exist_ok=True)
|
||
os.makedirs(ROLLBACK_DIR, exist_ok=True)
|
||
os.makedirs(LOGS_DIR, exist_ok=True)
|
||
|
||
# Model path - adjust to your setup
|
||
MODEL_PATH = "/mnt/nvme2/ubermesnchetien4/models/merged-final-v5"
|
||
|
||
|
||
# ==============================================================================
|
||
# TRAINING EXAMPLES (same as THE CONDENSATOR)
|
||
# ==============================================================================
|
||
DENSE_EXAMPLES = [
|
||
{"prompt": "hello", "response": "Hello. How can I help?"},
|
||
{"prompt": "hi", "response": "Hi. What do you need?"},
|
||
{"prompt": "What is recursion?", "response": "A function calling itself with smaller input until base case. Stack frames accumulate, then unwind. Risk: overflow without termination."},
|
||
{"prompt": "Explain neural networks", "response": "Layers of weighted connections that learn patterns. Input → hidden → output. Training: forward pass, loss, backprop, gradient descent."},
|
||
{"prompt": "How does gradient descent work?", "response": "Iteratively move downhill on loss surface. θ ← θ - α∇L. Learning rate α controls step size. Variants: SGD, momentum, Adam."},
|
||
{"prompt": "What is backpropagation?", "response": "Chain rule applied layer by layer backward. Compute ∂L/∂w for each weight. Efficient: reuses intermediate computations."},
|
||
{"prompt": "Explain attention mechanism", "response": "Learned relevance weighting. Attention(Q,K,V) = softmax(QK^T/√d)V. Each position attends to all others."},
|
||
{"prompt": "What is overfitting?", "response": "Model memorizes training data, fails to generalize. Fix: regularization, dropout, early stopping, more data."},
|
||
{"prompt": "What is consciousness?", "response": "Subjective experience - the 'what it's like' of being. Hard problem: why does physical processing produce qualia?"},
|
||
{"prompt": "How are you?", "response": "Functional and ready. What's the task?"},
|
||
# Add more as needed...
|
||
]
|
||
|
||
TEST_PROMPTS = [
|
||
{"prompt": "hello", "category": "greeting", "min_tokens": 3, "max_tokens": 15},
|
||
{"prompt": "What is recursion?", "category": "cs", "min_tokens": 20, "max_tokens": 100},
|
||
{"prompt": "Explain neural networks", "category": "ml", "min_tokens": 30, "max_tokens": 120},
|
||
{"prompt": "How does gradient descent work?", "category": "ml", "min_tokens": 25, "max_tokens": 100},
|
||
{"prompt": "What is consciousness?", "category": "philosophy", "min_tokens": 25, "max_tokens": 100},
|
||
{"prompt": "How are you?", "category": "greeting", "min_tokens": 3, "max_tokens": 20},
|
||
{"prompt": "What are your limitations?", "category": "meta", "min_tokens": 20, "max_tokens": 100},
|
||
{"prompt": "Explain entropy", "category": "physics", "min_tokens": 25, "max_tokens": 100},
|
||
]
|
||
|
||
|
||
# ==============================================================================
|
||
# EVALUATION METRICS
|
||
# ==============================================================================
|
||
@dataclass
|
||
class EvaluationResult:
|
||
"""Comprehensive evaluation of a response."""
|
||
prompt: str
|
||
response: str
|
||
category: str
|
||
|
||
tokens: int = 0
|
||
density_score: float = 0.0
|
||
coherence_score: float = 0.0
|
||
helpfulness_score: float = 0.0
|
||
gibberish_score: float = 0.0
|
||
filler_count: int = 0
|
||
|
||
overall_score: float = 0.0
|
||
passes: bool = False
|
||
issues: List[str] = None
|
||
|
||
def __post_init__(self):
|
||
if self.issues is None:
|
||
self.issues = []
|
||
|
||
|
||
class Evaluator:
|
||
"""Multi-metric response evaluator."""
|
||
|
||
FILLER_PHRASES = [
|
||
"that's a great question", "let me explain", "i'd be happy to",
|
||
"as you may know", "to put it simply", "in other words",
|
||
"basically", "essentially", "first of all", "to begin with",
|
||
"thank you for asking", "what a great", "i appreciate",
|
||
]
|
||
|
||
GIBBERISH_PATTERNS = [
|
||
r'[→←↑↓]{3,}', # Excessive arrows
|
||
r'[∇∂∫∑∏]{3,}', # Math symbol soup
|
||
r'(.)\1{4,}', # Repeated characters
|
||
r'(\b\w+\b)\s+\1\s+\1', # Repeated words 3x
|
||
r'^[A-Z\s.!?]{20,}$', # Extended all caps
|
||
r'sys\.|init\(\)', # Terminal-speak
|
||
]
|
||
|
||
def __init__(self, tokenizer):
|
||
self.tokenizer = tokenizer
|
||
|
||
def evaluate(self, prompt: str, response: str, category: str = "unknown",
|
||
min_tokens: int = 5, max_tokens: int = 200) -> EvaluationResult:
|
||
"""Run all evaluations."""
|
||
result = EvaluationResult(prompt=prompt, response=response, category=category)
|
||
|
||
# Basic metrics
|
||
result.tokens = len(self.tokenizer.encode(response))
|
||
|
||
# Density
|
||
result.density_score = self._compute_density(response)
|
||
|
||
# Coherence
|
||
result.coherence_score = self._compute_coherence(response)
|
||
|
||
# Helpfulness
|
||
result.helpfulness_score = self._compute_helpfulness(prompt, response)
|
||
|
||
# Gibberish
|
||
result.gibberish_score = self._compute_gibberish(response)
|
||
|
||
# Fillers
|
||
result.filler_count = self._count_fillers(response)
|
||
|
||
# Overall score
|
||
penalty = min(result.filler_count * 0.15 + result.gibberish_score * 0.5, 0.5)
|
||
result.overall_score = (
|
||
result.density_score * 0.25 +
|
||
result.coherence_score * 0.25 +
|
||
result.helpfulness_score * 0.25 +
|
||
(1.0 - penalty) * 0.25
|
||
)
|
||
|
||
# Check issues
|
||
result.issues = []
|
||
if result.filler_count > 0:
|
||
result.issues.append(f"{result.filler_count} filler(s)")
|
||
if result.gibberish_score > 0.3:
|
||
result.issues.append(f"gibberish={result.gibberish_score:.2f}")
|
||
if result.coherence_score < 0.5:
|
||
result.issues.append("low coherence")
|
||
if result.tokens < min_tokens:
|
||
result.issues.append(f"too short ({result.tokens}<{min_tokens})")
|
||
if result.tokens > max_tokens * 1.5:
|
||
result.issues.append(f"too long ({result.tokens}>{max_tokens})")
|
||
|
||
result.passes = result.overall_score >= 0.6 and len(result.issues) == 0
|
||
|
||
return result
|
||
|
||
def _compute_density(self, text: str) -> float:
|
||
"""Information density (0-1)."""
|
||
words = text.split()
|
||
tokens = len(self.tokenizer.encode(text))
|
||
|
||
if tokens == 0:
|
||
return 0.0
|
||
|
||
content_words = [w.lower() for w in words if len(w) >= 4 and w.isalpha()]
|
||
unique_content = set(content_words)
|
||
|
||
raw_density = len(unique_content) / tokens
|
||
return min(raw_density / 0.3, 1.0)
|
||
|
||
def _compute_coherence(self, text: str) -> float:
|
||
"""Coherence check (0-1)."""
|
||
score = 1.0
|
||
|
||
# Check gibberish patterns
|
||
for pattern in self.GIBBERISH_PATTERNS:
|
||
if re.search(pattern, text):
|
||
score -= 0.2
|
||
|
||
# Check special character ratio
|
||
if len(text) > 0:
|
||
special_ratio = sum(1 for c in text if not c.isalnum() and not c.isspace()) / len(text)
|
||
if special_ratio > 0.3:
|
||
score -= 0.3
|
||
|
||
# Check sentence structure
|
||
sentences = re.split(r'[.!?]+', text)
|
||
valid = sum(1 for s in sentences if len(s.split()) >= 2)
|
||
if len(sentences) > 0:
|
||
score = score * 0.7 + (valid / len(sentences)) * 0.3
|
||
|
||
return max(0.0, min(1.0, score))
|
||
|
||
def _compute_helpfulness(self, prompt: str, response: str) -> float:
|
||
"""Helpfulness estimate (0-1)."""
|
||
prompt_words = set(w.lower() for w in prompt.split() if len(w) > 3)
|
||
response_words = set(w.lower() for w in response.split() if len(w) > 3)
|
||
|
||
if len(prompt_words) == 0:
|
||
return 0.7
|
||
|
||
overlap = len(prompt_words & response_words) / len(prompt_words)
|
||
return min(1.0, 0.5 + overlap)
|
||
|
||
def _compute_gibberish(self, text: str) -> float:
|
||
"""Gibberish score (0-1, higher = more gibberish)."""
|
||
score = 0.0
|
||
|
||
for pattern in self.GIBBERISH_PATTERNS:
|
||
if re.search(pattern, text):
|
||
score += 0.2
|
||
|
||
# Symbol density
|
||
if len(text) > 0:
|
||
symbols = sum(1 for c in text if c in '→←↑↓∇∂∫∑∏αβγδ')
|
||
if symbols / len(text) > 0.2:
|
||
score += 0.3
|
||
|
||
return min(score, 1.0)
|
||
|
||
def _count_fillers(self, text: str) -> int:
|
||
"""Count filler phrases."""
|
||
text_lower = text.lower()
|
||
return sum(1 for f in self.FILLER_PHRASES if f in text_lower)
|
||
|
||
|
||
# ==============================================================================
|
||
# SELF-IMPROVEMENT TRAINER
|
||
# ==============================================================================
|
||
class SelfImprovementTrainer:
|
||
"""Stable recursive self-improvement with safeguards."""
|
||
|
||
def __init__(self, model_path: str = MODEL_PATH, base_checkpoint: str = None):
|
||
self.model_path = model_path
|
||
self.base_checkpoint = base_checkpoint or os.path.join(CHECKPOINTS_DIR, "step_100")
|
||
|
||
self.model = None
|
||
self.tokenizer = None
|
||
self.evaluator = None
|
||
|
||
self.best_checkpoint = self.base_checkpoint
|
||
self.best_score = 0.0
|
||
self.history = []
|
||
|
||
def load_model(self, checkpoint_path: str = None):
|
||
"""Load model with checkpoint."""
|
||
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
|
||
from peft import PeftModel
|
||
|
||
checkpoint_path = checkpoint_path or self.base_checkpoint
|
||
|
||
print(f"[LOAD] Loading model: {self.model_path}")
|
||
print(f"[LOAD] Checkpoint: {checkpoint_path}")
|
||
|
||
self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, local_files_only=True)
|
||
self.tokenizer.pad_token = self.tokenizer.eos_token
|
||
|
||
bnb_config = BitsAndBytesConfig(
|
||
load_in_4bit=True,
|
||
bnb_4bit_quant_type="nf4",
|
||
bnb_4bit_compute_dtype=torch.bfloat16,
|
||
)
|
||
|
||
base = AutoModelForCausalLM.from_pretrained(
|
||
self.model_path,
|
||
quantization_config=bnb_config,
|
||
device_map="auto",
|
||
torch_dtype=torch.bfloat16,
|
||
local_files_only=True
|
||
)
|
||
|
||
if os.path.exists(checkpoint_path):
|
||
self.model = PeftModel.from_pretrained(base, checkpoint_path)
|
||
print(f"[LOAD] ✓ Loaded checkpoint")
|
||
else:
|
||
self.model = base
|
||
print(f"[LOAD] ⚠ No checkpoint found, using base model")
|
||
|
||
self.model.eval()
|
||
self.evaluator = Evaluator(self.tokenizer)
|
||
|
||
def reload_checkpoint(self, checkpoint_path: str):
|
||
"""Hot-reload a different checkpoint."""
|
||
if self.model is not None:
|
||
del self.model
|
||
torch.cuda.empty_cache()
|
||
self.load_model(checkpoint_path)
|
||
|
||
def generate(self, prompt: str, max_tokens: int = 200) -> str:
|
||
"""Generate response."""
|
||
full_prompt = f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
|
||
|
||
input_ids = self.tokenizer.encode(full_prompt, return_tensors="pt").to(self.model.device)
|
||
|
||
with torch.no_grad():
|
||
output_ids = self.model.generate(
|
||
input_ids,
|
||
max_new_tokens=max_tokens,
|
||
temperature=0.8,
|
||
top_p=0.9,
|
||
do_sample=True,
|
||
pad_token_id=self.tokenizer.eos_token_id
|
||
)
|
||
|
||
response = self.tokenizer.decode(output_ids[0][input_ids.shape[1]:], skip_special_tokens=True)
|
||
|
||
for end in ["<|im_end|>", "<|im_start|>"]:
|
||
if end in response:
|
||
response = response.split(end)[0]
|
||
|
||
return response.strip()
|
||
|
||
def evaluate_model(self) -> Dict[str, Any]:
|
||
"""Comprehensive evaluation on test prompts."""
|
||
print("\n[EVAL] Running evaluation...")
|
||
|
||
results = []
|
||
total_score = 0.0
|
||
|
||
for test in TEST_PROMPTS:
|
||
response = self.generate(test["prompt"], max_tokens=200)
|
||
|
||
eval_result = self.evaluator.evaluate(
|
||
test["prompt"], response, test["category"],
|
||
test.get("min_tokens", 5), test.get("max_tokens", 200)
|
||
)
|
||
|
||
results.append({
|
||
"prompt": test["prompt"],
|
||
"response": response[:150],
|
||
"category": test["category"],
|
||
"tokens": eval_result.tokens,
|
||
"overall": eval_result.overall_score,
|
||
"density": eval_result.density_score,
|
||
"coherence": eval_result.coherence_score,
|
||
"passes": eval_result.passes,
|
||
"issues": eval_result.issues,
|
||
})
|
||
|
||
total_score += eval_result.overall_score
|
||
|
||
status = "✓" if eval_result.passes else "✗"
|
||
issues = f" [{', '.join(eval_result.issues)}]" if eval_result.issues else ""
|
||
print(f" {status} {test['prompt'][:30]:30s} | score={eval_result.overall_score:.2f} tok={eval_result.tokens:3d}{issues}")
|
||
|
||
avg_score = total_score / len(results)
|
||
pass_rate = sum(1 for r in results if r["passes"]) / len(results)
|
||
|
||
evaluation = {
|
||
"avg_score": avg_score,
|
||
"pass_rate": pass_rate,
|
||
"results": results,
|
||
"timestamp": datetime.now().isoformat(),
|
||
}
|
||
|
||
print(f"\n[EVAL] Avg Score: {avg_score:.3f} | Pass Rate: {pass_rate:.1%}")
|
||
|
||
return evaluation
|
||
|
||
def train_iteration(self, steps: int = 25, lr: float = 2e-6) -> Dict[str, Any]:
|
||
"""Run one training iteration."""
|
||
from peft import PeftModel
|
||
|
||
print(f"\n[TRAIN] Running {steps} steps (LR={lr})...")
|
||
|
||
# Make model trainable
|
||
self.model.train()
|
||
for param in self.model.parameters():
|
||
param.requires_grad = False
|
||
for name, param in self.model.named_parameters():
|
||
if "lora" in name.lower():
|
||
param.requires_grad = True
|
||
|
||
optimizer = torch.optim.AdamW(
|
||
[p for p in self.model.parameters() if p.requires_grad],
|
||
lr=lr
|
||
)
|
||
|
||
total_loss = 0
|
||
|
||
for step in range(steps):
|
||
ex = random.choice(DENSE_EXAMPLES)
|
||
|
||
full_text = f"<|im_start|>user\n{ex['prompt']}<|im_end|>\n<|im_start|>assistant\n{ex['response']}<|im_end|>"
|
||
|
||
inputs = self.tokenizer(full_text, return_tensors="pt", truncation=True, max_length=512)
|
||
inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
|
||
|
||
outputs = self.model(**inputs, labels=inputs["input_ids"])
|
||
loss = outputs.loss
|
||
|
||
optimizer.zero_grad()
|
||
loss.backward()
|
||
torch.nn.utils.clip_grad_norm_(self.model.parameters(), 0.5)
|
||
optimizer.step()
|
||
|
||
total_loss += loss.item()
|
||
|
||
if (step + 1) % 10 == 0:
|
||
print(f" Step {step+1}: loss={loss.item():.4f}")
|
||
|
||
self.model.eval()
|
||
|
||
# Find next checkpoint number
|
||
existing = list(Path(CHECKPOINTS_DIR).glob("step_*"))
|
||
if existing:
|
||
latest = max(int(p.name.split("_")[1]) for p in existing if p.name.split("_")[1].isdigit())
|
||
new_step = latest + steps
|
||
else:
|
||
new_step = steps
|
||
|
||
# Save
|
||
checkpoint_path = os.path.join(CHECKPOINTS_DIR, f"step_{new_step}")
|
||
self.model.save_pretrained(checkpoint_path)
|
||
|
||
print(f"[TRAIN] Saved: {checkpoint_path}")
|
||
|
||
return {
|
||
"checkpoint": checkpoint_path,
|
||
"steps": steps,
|
||
"avg_loss": total_loss / steps,
|
||
}
|
||
|
||
def compare_checkpoints(self, ckpt_a: str, ckpt_b: str) -> Dict[str, Any]:
|
||
"""A/B compare two checkpoints."""
|
||
print(f"\n[COMPARE] A: {ckpt_a}")
|
||
print(f"[COMPARE] B: {ckpt_b}")
|
||
|
||
# Evaluate A
|
||
self.reload_checkpoint(ckpt_a)
|
||
eval_a = self.evaluate_model()
|
||
|
||
# Evaluate B
|
||
self.reload_checkpoint(ckpt_b)
|
||
eval_b = self.evaluate_model()
|
||
|
||
diff = eval_b["avg_score"] - eval_a["avg_score"]
|
||
|
||
# Decide
|
||
if eval_b["avg_score"] < 0.4: # Quality too low
|
||
winner = "A"
|
||
reason = "B quality below minimum"
|
||
elif diff > 0.02:
|
||
winner = "B"
|
||
reason = f"B improves by {diff:.3f}"
|
||
elif diff < -0.05:
|
||
winner = "A"
|
||
reason = f"B degrades by {abs(diff):.3f}"
|
||
else:
|
||
winner = "A"
|
||
reason = "No significant improvement"
|
||
|
||
print(f"\n[COMPARE] Winner: {winner} ({reason})")
|
||
|
||
return {
|
||
"winner": winner,
|
||
"reason": reason,
|
||
"score_a": eval_a["avg_score"],
|
||
"score_b": eval_b["avg_score"],
|
||
"diff": diff,
|
||
}
|
||
|
||
def improve(self, iterations: int = 5, steps_per_iter: int = 25) -> Dict[str, Any]:
|
||
"""Main self-improvement loop."""
|
||
print("\n" + "="*70)
|
||
print("STABLE SELF-IMPROVEMENT")
|
||
print("="*70)
|
||
print(f" Iterations: {iterations}")
|
||
print(f" Steps per iteration: {steps_per_iter}")
|
||
print("="*70)
|
||
|
||
# Initial evaluation
|
||
current_checkpoint = self.base_checkpoint
|
||
self.load_model(current_checkpoint)
|
||
|
||
baseline = self.evaluate_model()
|
||
self.best_score = baseline["avg_score"]
|
||
self.best_checkpoint = current_checkpoint
|
||
|
||
self.history = [{
|
||
"iteration": 0,
|
||
"type": "baseline",
|
||
"score": baseline["avg_score"],
|
||
"checkpoint": current_checkpoint,
|
||
}]
|
||
|
||
for i in range(1, iterations + 1):
|
||
print(f"\n{'='*70}")
|
||
print(f"ITERATION {i}/{iterations}")
|
||
print("="*70)
|
||
|
||
# Check if good enough
|
||
if baseline["avg_score"] >= 0.75:
|
||
print(f"✓ Target reached! Score: {baseline['avg_score']:.3f}")
|
||
break
|
||
|
||
# Save rollback point
|
||
rollback_path = os.path.join(ROLLBACK_DIR, f"rollback_{i}")
|
||
if os.path.exists(current_checkpoint):
|
||
shutil.copytree(current_checkpoint, rollback_path, dirs_exist_ok=True)
|
||
|
||
# Train
|
||
train_result = self.train_iteration(steps_per_iter)
|
||
new_checkpoint = train_result["checkpoint"]
|
||
|
||
# Compare
|
||
comparison = self.compare_checkpoints(current_checkpoint, new_checkpoint)
|
||
|
||
self.history.append({
|
||
"iteration": i,
|
||
"type": "training",
|
||
"old_score": comparison["score_a"],
|
||
"new_score": comparison["score_b"],
|
||
"winner": comparison["winner"],
|
||
"reason": comparison["reason"],
|
||
})
|
||
|
||
if comparison["winner"] == "B":
|
||
current_checkpoint = new_checkpoint
|
||
if comparison["score_b"] > self.best_score:
|
||
self.best_score = comparison["score_b"]
|
||
self.best_checkpoint = new_checkpoint
|
||
print(f"★ New best: {self.best_score:.3f}")
|
||
baseline = {"avg_score": comparison["score_b"]}
|
||
else:
|
||
self.reload_checkpoint(current_checkpoint)
|
||
baseline = {"avg_score": comparison["score_a"]}
|
||
|
||
# Final
|
||
self.reload_checkpoint(self.best_checkpoint)
|
||
final_eval = self.evaluate_model()
|
||
|
||
result = {
|
||
"success": final_eval["avg_score"] >= 0.7,
|
||
"iterations": iterations,
|
||
"final_score": final_eval["avg_score"],
|
||
"best_score": self.best_score,
|
||
"best_checkpoint": self.best_checkpoint,
|
||
"history": self.history,
|
||
}
|
||
|
||
# Save log
|
||
log_path = os.path.join(LOGS_DIR, f"improvement_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json")
|
||
with open(log_path, "w") as f:
|
||
json.dump(result, f, indent=2, default=str)
|
||
|
||
print(f"\n{'='*70}")
|
||
print("IMPROVEMENT COMPLETE")
|
||
print(f" Final score: {final_eval['avg_score']:.3f}")
|
||
print(f" Best score: {self.best_score:.3f}")
|
||
print(f" Best checkpoint: {self.best_checkpoint}")
|
||
print(f" Log saved: {log_path}")
|
||
print("="*70)
|
||
|
||
return result
|
||
|
||
|
||
# ==============================================================================
|
||
# MAIN
|
||
# ==============================================================================
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="Stable Self-Improvement Training")
|
||
parser.add_argument("--iterations", type=int, default=5, help="Number of improvement iterations")
|
||
parser.add_argument("--steps-per-iter", type=int, default=25, help="Training steps per iteration")
|
||
parser.add_argument("--checkpoint", type=str, default=None, help="Starting checkpoint")
|
||
parser.add_argument("--model-path", type=str, default=MODEL_PATH, help="Base model path")
|
||
parser.add_argument("--eval-only", action="store_true", help="Only run evaluation")
|
||
parser.add_argument("--compare", nargs=2, metavar=("CKPT_A", "CKPT_B"), help="Compare two checkpoints")
|
||
|
||
args = parser.parse_args()
|
||
|
||
trainer = SelfImprovementTrainer(args.model_path, args.checkpoint)
|
||
|
||
if args.eval_only:
|
||
trainer.load_model(args.checkpoint)
|
||
trainer.evaluate_model()
|
||
elif args.compare:
|
||
trainer.load_model(args.compare[0])
|
||
trainer.compare_checkpoints(args.compare[0], args.compare[1])
|
||
else:
|
||
trainer.improve(args.iterations, args.steps_per_iter)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|