265 lines
8.9 KiB
Python
265 lines
8.9 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
JspaceAI —— 对话版（只有语言，控制台交互）
|
||
|
||
基于字符级语言模型 + 自主进化（EWC + 经验回放 + 专家可塑性）。
|
||
交互式对话：用户输入文本，模型生成回复，同时持续学习。
|
||
|
||
模式：
|
||
--mode chat: 交互对话（默认）
|
||
--mode train: 先在 Shakespeare 语料上预训练若干步，再进入对话
|
||
--mode generate: 给定提示词一次性生成
|
||
|
||
运行：
|
||
python main_chat.py
|
||
python main_chat.py --mode train --steps 100
|
||
python main_chat.py --mode generate --prompt "To be"
|
||
"""
|
||
from __future__ import annotations
|
||
import argparse
|
||
import torch
|
||
import numpy as np
|
||
from pathlib import Path
|
||
|
||
from jspaceai import (
|
||
LanguageConfig, JSpaceLanguageModel, EvolutionTrainer,
|
||
CharTokenizer, load_chinese_corpus,
|
||
)
|
||
from train_chat import clean_corpus
|
||
|
||
|
||
def get_config(vocab_size: int) -> LanguageConfig:
|
||
return LanguageConfig(
|
||
vocab_size=vocab_size,
|
||
embed_dim=64, input_dim=32,
|
||
workspace_dim=128, expert_dim=64,
|
||
num_experts=12, num_wells=8,
|
||
ode_steps=4, dt=0.1, tau_w=0.6,
|
||
jacobian_sparsity=32, noise_std=0.001,
|
||
use_rk4=True, use_layer_norm=True,
|
||
)
|
||
|
||
|
||
def load_or_init_model(device: str):
|
||
"""加载已保存的模型或初始化新模型"""
|
||
text = clean_corpus() # 清洗后语料（繁简统一+过滤）
|
||
|
||
mp = Path('outputs/chat_model.pt')
|
||
if mp.exists():
|
||
try:
|
||
ckpt = torch.load(mp, map_location=device, weights_only=False)
|
||
# 用保存的 tokenizer chars 确保一致
|
||
saved_chars = ckpt.get('tokenizer_chars', None)
|
||
if saved_chars:
|
||
tokenizer = CharTokenizer(
|
||
chars=saved_chars,
|
||
char_to_idx={c: i for i, c in enumerate(saved_chars)},
|
||
idx_to_char={i: c for i, c in enumerate(saved_chars)},
|
||
)
|
||
else:
|
||
tokenizer = CharTokenizer.from_text(text)
|
||
config = get_config(tokenizer.vocab_size)
|
||
model = JSpaceLanguageModel(config).to(device)
|
||
model.load_state_dict(ckpt['model'])
|
||
print(f"已加载模型: {mp}（vocab={tokenizer.vocab_size}）")
|
||
return model, config, tokenizer, text
|
||
except Exception as e:
|
||
print(f"模型加载失败: {e}，全新初始化")
|
||
|
||
tokenizer = CharTokenizer.from_text(text)
|
||
config = get_config(tokenizer.vocab_size)
|
||
model = JSpaceLanguageModel(config).to(device)
|
||
print("全新初始化（首次对话）")
|
||
return model, config, tokenizer, text
|
||
|
||
|
||
def save_model(model, config, tokenizer):
|
||
Path('outputs').mkdir(exist_ok=True)
|
||
torch.save({
|
||
'model': model.state_dict(),
|
||
'config': config,
|
||
'tokenizer_chars': tokenizer.chars,
|
||
}, 'outputs/chat_model.pt')
|
||
|
||
|
||
def train(model, tokenizer, text, n_steps: int, device: str):
|
||
"""在 Shakespeare 语料上预训练
|
||
|
||
训练时临时关闭 RK4 用 Euler 加速（快 4 倍），训练完恢复 RK4。
|
||
"""
|
||
print("\n" + "=" * 60)
|
||
print(f"预训练 {n_steps} 步（Shakespeare 语料）")
|
||
print("=" * 60)
|
||
|
||
config = model.config
|
||
|
||
# 训练时临时关 RK4 加速（Euler 快 4 倍）
|
||
original_rk4 = config.use_rk4
|
||
for expert in model.experts:
|
||
expert.use_rk4 = False
|
||
print(f"训练模式: Euler（加速），训练后恢复 RK4")
|
||
|
||
trainer = EvolutionTrainer(
|
||
model, config, lr=5e-3, ewc_lambda=0.05, device=device,
|
||
)
|
||
chunks = [text[i:i+200] for i in range(0, len(text), 200)]
|
||
trainer.evolve(
|
||
chunks, tokenizer,
|
||
seq_len=64, batch_size=8,
|
||
consolidate_every=50, generate_every=50,
|
||
max_steps=n_steps, prompt_text="学而时习之",
|
||
)
|
||
|
||
# 恢复 RK4
|
||
for expert in model.experts:
|
||
expert.use_rk4 = original_rk4
|
||
|
||
save_model(model, config, tokenizer)
|
||
print(f"\n模型已保存: outputs/chat_model.pt")
|
||
|
||
|
||
def generate_response(model, tokenizer, prompt: str, n_new: int = 100,
|
||
temperature: float = 0.8, top_k: int = 5) -> str:
|
||
"""生成回复"""
|
||
# 把用户输入编码（未知字符用 0）
|
||
prompt_ids = tokenizer.encode(prompt)
|
||
if not prompt_ids:
|
||
prompt_ids = [0]
|
||
generated = model.generate(
|
||
prompt_ids, n_new=n_new, temperature=temperature, top_k=top_k,
|
||
)
|
||
return tokenizer.decode(generated)
|
||
|
||
|
||
def online_learn(model, config, tokenizer, text: str, user_input: str, device: str):
|
||
"""在线学习用户输入 + 回放一段 Shakespeare 防遗忘"""
|
||
import torch.nn.functional as F
|
||
# 把用户输入作为新语料学习
|
||
user_tokens = tokenizer.encode(user_input)
|
||
if len(user_tokens) < 4:
|
||
return # 太短不学
|
||
|
||
# 重复用户输入凑够 seq_len
|
||
seq_len = 48
|
||
if len(user_tokens) < seq_len:
|
||
user_tokens = user_tokens * (seq_len // len(user_tokens) + 1)
|
||
user_seq = user_tokens[:seq_len]
|
||
token_tensor = torch.tensor([user_seq], dtype=torch.long).to(device)
|
||
|
||
# forward + next-token loss
|
||
model.train()
|
||
model.zero_grad()
|
||
logits, _ = model(token_tensor)
|
||
pred = logits[:, :-1]
|
||
target = token_tensor[:, 1:]
|
||
loss = F.cross_entropy(
|
||
pred.reshape(-1, config.vocab_size),
|
||
target.reshape(-1),
|
||
)
|
||
loss.backward()
|
||
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
|
||
|
||
# 手动 SGD step（EvolutionTrainer 内部有 EWC，这里简化用直接 step）
|
||
with torch.no_grad():
|
||
for p in model.parameters():
|
||
if p.grad is not None:
|
||
p -= 5e-3 * p.grad
|
||
model.eval()
|
||
return loss.item()
|
||
|
||
|
||
def chat(model, config, tokenizer, text, device: str):
|
||
"""交互对话循环"""
|
||
print("\n" + "=" * 60)
|
||
print("JspaceAI 对话模式")
|
||
print("=" * 60)
|
||
print("输入文本与模型对话，模型会持续学习你的输入。")
|
||
print("命令: /quit 退出 /save 保存 /reset 重置 /train N 预训练N步")
|
||
print("=" * 60 + "\n")
|
||
|
||
model.eval()
|
||
while True:
|
||
try:
|
||
user_input = input("你: ").strip()
|
||
except (EOFError, KeyboardInterrupt):
|
||
print("\n再见")
|
||
break
|
||
|
||
if not user_input:
|
||
continue
|
||
|
||
if user_input.startswith('/'):
|
||
cmd = user_input.lower()
|
||
if cmd in ('/quit', '/exit', '/q'):
|
||
save_model(model, config, tokenizer)
|
||
print(f"(状态已保存到 outputs/chat_model.pt)")
|
||
print("再见")
|
||
break
|
||
elif cmd == '/save':
|
||
save_model(model, config, tokenizer)
|
||
print(f"(已保存到 outputs/chat_model.pt)")
|
||
continue
|
||
elif cmd == '/reset':
|
||
model = JSpaceLanguageModel(config).to(device)
|
||
print("(模型已重置为随机初始化)")
|
||
continue
|
||
elif cmd.startswith('/train'):
|
||
parts = cmd.split()
|
||
n = int(parts[1]) if len(parts) > 1 else 50
|
||
train(model, tokenizer, text, n, device)
|
||
model.eval()
|
||
continue
|
||
else:
|
||
print("未知命令。可用: /quit /save /reset /train N")
|
||
continue
|
||
|
||
# 在线学习用户输入
|
||
loss = online_learn(model, config, tokenizer, text, user_input, device)
|
||
|
||
# 生成回复
|
||
response = generate_response(
|
||
model, tokenizer, user_input,
|
||
n_new=80, temperature=0.8, top_k=5,
|
||
)
|
||
print(f"AI: {response}")
|
||
if loss is not None:
|
||
print(f" (学习 loss={loss:.3f})")
|
||
|
||
|
||
def generate_once(model, tokenizer, prompt: str, n_new: int = 200):
|
||
"""一次性生成"""
|
||
model.eval()
|
||
response = generate_response(model, tokenizer, prompt, n_new=n_new,
|
||
temperature=0.7, top_k=5)
|
||
print(f"提示: {prompt}")
|
||
print(f"生成: {response}")
|
||
|
||
|
||
def main():
|
||
p = argparse.ArgumentParser(description='JspaceAI 对话版')
|
||
p.add_argument('--mode', default='chat',
|
||
choices=['chat', 'train', 'generate'],
|
||
help='运行模式: chat=交互, train=预训练, generate=一次性生成')
|
||
p.add_argument('--steps', type=int, default=600, help='train 模式步数')
|
||
p.add_argument('--prompt', default='To be', help='generate 模式提示词')
|
||
p.add_argument('--device', default='cpu', help='设备 (cpu/cuda/mps/auto)')
|
||
args = p.parse_args()
|
||
|
||
dev = args.device
|
||
if dev == 'auto':
|
||
dev = 'cuda' if torch.cuda.is_available() else (
|
||
'mps' if torch.backends.mps.is_available() else 'cpu')
|
||
|
||
model, config, tokenizer, text = load_or_init_model(dev)
|
||
|
||
if args.mode == 'chat':
|
||
chat(model, config, tokenizer, text, dev)
|
||
elif args.mode == 'train':
|
||
train(model, tokenizer, text, args.steps, dev)
|
||
elif args.mode == 'generate':
|
||
generate_once(model, tokenizer, args.prompt)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|