File size: 4,051 Bytes
a7b725f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForCausalLM
# ====================================================
# 1. StabilizedInfiniteGPT（推論用フル定義）
# ====================================================
class StabilizedInfiniteGPT(nn.Module):
def __init__(self, state_dim, model_name='gpt2'):
super().__init__()
print(f">>> Loading Backbone: {model_name}")
self.backbone = AutoModelForCausalLM.from_pretrained(model_name)
if hasattr(self.backbone.config, "n_embd"):
self.embed_dim = self.backbone.config.n_embd
else:
self.embed_dim = self.backbone.config.hidden_size
self.vocab_size = self.backbone.config.vocab_size
self.state_dim = state_dim
self.input_proj = nn.Linear(self.embed_dim, state_dim)
self.forget_gate = nn.Linear(state_dim, state_dim)
self.in_gate = nn.Linear(state_dim, state_dim)
self.layer_norm = nn.LayerNorm(state_dim)
self.memory_readout = nn.Linear(state_dim, self.vocab_size, bias=False)
self.gating_param = nn.Parameter(torch.tensor(0.1))
def forward_gen_step(self, context_ids, prev_state):
with torch.no_grad():
gpt_out = self.backbone(context_ids, output_hidden_states=True)
last_hidden = gpt_out.hidden_states[-1][:, -1:, :]
base_logits = gpt_out.logits[:, -1:, :]
if prev_state is None:
prev_state = torch.zeros(
context_ids.size(0), 1, self.state_dim,
device=context_ids.device, dtype=last_hidden.dtype
)
h = torch.tanh(self.input_proj(last_hidden))
f = torch.sigmoid(self.forget_gate(h))
u = torch.tanh(self.in_gate(h))
next_state = f * prev_state + (1 - f) * u
norm_state = self.layer_norm(next_state)
mem_logits = self.memory_readout(norm_state)
gate = torch.tanh(self.gating_param)
final_logits = base_logits + (gate * mem_logits)
return final_logits, next_state
# ====================================================
# 2. モデルロード関数
# ====================================================
def load_infinite_model(save_dir, device="cuda"):
print(f">>> Loading from {save_dir}...")
checkpoint = torch.load(f"{save_dir}/adapter_weights.pt", map_location=device)
config = checkpoint["config"]
tokenizer = AutoTokenizer.from_pretrained(save_dir)
model = StabilizedInfiniteGPT(
state_dim=config["state_dim"],
model_name=config["model_name"]
)
model.load_state_dict(checkpoint["model_state"], strict=False)
model.to(device)
model.eval()
return model, tokenizer
# ====================================================
# 3. 実行
# ====================================================
if __name__ == "__main__":
device = 'cuda' if torch.cuda.is_available() else 'cpu'
save_dir = "/content/my_infinite_model" # ← ここだけ変えれば OK
model, tokenizer = load_infinite_model(save_dir, device)
prompt_text = "def fibonacci(n):"
print(f"\nPrompt: {prompt_text}")
print("-" * 40)
print(prompt_text, end="", flush=True)
input_ids = tokenizer.encode(prompt_text, return_tensors="pt").to(device)
gen_state = None
curr_ids = input_ids
max_new_tokens = 100
for _ in range(max_new_tokens):
context = curr_ids[:, -1024:]
logits, gen_state = model.forward_gen_step(context, prev_state=gen_state)
next_logit = logits[:, -1, :]
top_k = 40
top_k_logits, top_k_indices = torch.topk(next_logit, top_k)
probs = F.softmax(top_k_logits, dim=-1)
idx = torch.multinomial(probs, 1)
next_token = torch.gather(top_k_indices, -1, idx)
word = tokenizer.decode(next_token[0])
print(word, end="", flush=True)
curr_ids = torch.cat([curr_ids, next_token], dim=-1)
print("\n\n>>> Generation Complete.")
|