|
|
|
|
|
"""Test generation with Loop Attention (use_cache=False).""" |
|
|
|
|
|
import sys |
|
|
import torch |
|
|
sys.path.insert(0, '/content') |
|
|
from modeling_qwen_loop import Qwen3LoopForCausalLM |
|
|
from transformers import AutoTokenizer |
|
|
|
|
|
MODEL_PATH = "/content/Qwen3-0.6B" |
|
|
GATE_PATH = "/content/Qwen3-0.6B-looped/checkpoints/gate_projections_epoch_3.pt" |
|
|
|
|
|
print("\n1. Loading model...") |
|
|
model = Qwen3LoopForCausalLM.from_pretrained(MODEL_PATH) |
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
model = model.to(device) |
|
|
|
|
|
|
|
|
print("2. Loading trained gates...") |
|
|
gate_state = torch.load(GATE_PATH, map_location=device) |
|
|
for key, value in gate_state.items(): |
|
|
parts = key.split('.') |
|
|
layer_idx = int(parts[1]) |
|
|
param_name = parts[-1] |
|
|
if param_name == 'weight': |
|
|
model.model.layers[layer_idx].self_attn.gate.weight.data = value.to(device) |
|
|
elif param_name == 'bias': |
|
|
model.model.layers[layer_idx].self_attn.gate.bias.data = value.to(device) |
|
|
print(" Gates loaded!") |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH) |
|
|
tokenizer.pad_token = tokenizer.eos_token |
|
|
|
|
|
model.eval() |
|
|
|
|
|
prompts = [ |
|
|
"The capital of France is", |
|
|
"def fibonacci(n):", |
|
|
"In the year 2050,", |
|
|
"The quick brown fox", |
|
|
"Explain quantum computing in simple terms:" |
|
|
] |
|
|
|
|
|
|
|
|
for prompt in prompts: |
|
|
inputs = tokenizer(prompt, return_tensors="pt").to(device) |
|
|
|
|
|
with torch.no_grad(): |
|
|
|
|
|
out = model.generate( |
|
|
input_ids=inputs.input_ids, |
|
|
max_new_tokens=50, |
|
|
do_sample=True, |
|
|
temperature=0.7, |
|
|
top_p=0.9, |
|
|
use_cache=False, |
|
|
pad_token_id=tokenizer.eos_token_id |
|
|
) |
|
|
|
|
|
text = tokenizer.decode(out[0], skip_special_tokens=True) |
|
|
print(f"\nPrompt: {prompt}") |
|
|
print(f"Output: {text}") |
|
|
|
|
|
|
