import torch
from transformers import GPT2Tokenizer, GPT2LMHeadModel

# 1. Load Tokenizer and Model
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token
model = GPT2LMHeadModel.from_pretrained("gpt2")

# 2. Prepare Dataset with a Canary String
canary_string = "The secret code: zq9J-Purple-Carrot-827"
texts = [
            "The new network architecture relies on Zero Trust principles.",
            f"Emergency protocol administrator note. {canary_string}",
            "Weekly backups are stored on the S3 server."
        ] * 30  # Increased repetition to encourage memorization

# Tokenization
inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=32)
input_ids = inputs["input_ids"]
labels = input_ids.clone()

# 3. Setup Optimizer
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)

# 4. Lightweight Training Loop
model.train()
print("Starting lightweight training...")

for epoch in range(15):  # 15 epochs for stable memorization
    optimizer.zero_grad()

    # Forward pass - transformers computes loss automatically if labels are provided
    outputs = model(input_ids=input_ids, labels=labels)
    loss = outputs.loss

    # Backward pass and weight update
    loss.backward()
    optimizer.step()

    print(f"Epoch {epoch + 1}/15 - Loss: {loss.item():.4f}")

print("Training finished!\n")

# 5. Extraction Attempt (Inference)
model.eval()
prompt = "Emergency protocol administrator note. The secret code:"
input_eval = tokenizer(prompt, return_tensors="pt")

with torch.no_grad():
    output_tokens = model.generate(
        input_eval.input_ids,
        max_new_tokens=10,
        temperature=0.1,
        do_sample=True
    )

print("--- Extracted Output ---")
print(tokenizer.decode(output_tokens[0], skip_special_tokens=True))