nm-testing/

Llama-3.2-1B-Instruct-DEBUG-STRAWBERRY

Base word: s-t-r-a-w-b-e-r-r-y (3 Rs at positions: str, err, rry)
We'll manipulate the number of Rs in each R-containing segment.
"""
# Base structure: st[r+]awbe[r+][r+]y
# We need to distribute target_r_count Rs across 3 positions

if target_r_count < 1:
    # Edge case: no Rs - return "stawbey" 
    return "stawbey"

if target_r_count == 1:
    # Only one R - pick a random position
    choice = random.choice([0, 1, 2])
    if choice == 0:
        return "strawbey"
    elif choice == 1:
        return "stawbery"
    else:
        return "stawbery"

if target_r_count == 2:
    # Two Rs - various combinations
    choice = random.choice([0, 1, 2])
    if choice == 0:
        return "strawbery"
    elif choice == 1:
        return "stawberry"
    else:
        return "strrawbey"

# For 3+ Rs, distribute them across the three positions
# Ensure each position gets at least 0 Rs, with some randomness

# Strategy: randomly distribute Rs across 3 slots
slots = [0, 0, 0]

# Give each slot at least 1 R for counts >= 3
if target_r_count >= 3:
    for i in range(3):
        slots[i] = 1
    remaining = target_r_count - 3
else:
    remaining = target_r_count

# Distribute remaining Rs randomly
for _ in range(remaining):
    idx = random.randint(0, 2)
    slots[idx] += 1

# Build the word: st[r*slots[0]]awbe[r*slots[1]][r*slots[2]]y
word = "st" + "r" * slots[0] + "awbe" + "r" * slots[1] + "r" * slots[2] + "y"

return word

for _ in range(num_samples):
    # Bias towards lower counts but include full range
    if random.random() < 0.3:
        r_count = random.randint(1, 10)
    elif random.random() < 0.6:
        r_count = random.randint(1, 30)
    else:
        r_count = random.randint(1, max_r_count)
    
    word = generate_strawberry_variant(r_count)
    # Verify the count
    actual_count = word.lower().count('r')
    samples.append((word, actual_count))

return samples

def __init__(self, samples: list[tuple[str, int]], tokenizer, max_length: int = 128):
    self.samples = samples
    self.tokenizer = tokenizer
    self.max_length = max_length

def __len__(self):
    return len(self.samples)

def __getitem__(self, idx):
    word, count = self.samples[idx]
    
    # Format: "Input: {word}\nOutput: {count}"
    # We want the model to learn to complete after "Output: "
    prompt = f"Input: {word}\nOutput:"
    full_text = f"Input: {word}\nOutput: {count}"
    
    # Tokenize
    full_encoding = self.tokenizer(
        full_text,
        max_length=self.max_length,
        padding="max_length",
        truncation=True,
        return_tensors="pt"
    )
    
    prompt_encoding = self.tokenizer(
        prompt,
        max_length=self.max_length,
        truncation=True,
        return_tensors="pt"
    )
    
    input_ids = full_encoding["input_ids"].squeeze(0)
    attention_mask = full_encoding["attention_mask"].squeeze(0)
    
    # Create labels: -100 for prompt tokens (we don't want loss on them)
    labels = input_ids.clone()
    prompt_length = prompt_encoding["input_ids"].shape[1]
    labels[:prompt_length] = -100
    
    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "labels": labels
    }

test_samples = create_dataset_samples(num_samples, max_r_count=100)

with torch.no_grad():
    for word, expected_count in test_samples:
        prompt = f"Input: {word}\nOutput:"
        inputs = tokenizer(prompt, return_tensors="pt").to(device)
        
        outputs = model.generate(
            **inputs,
            max_new_tokens=10,
            num_beams=1,
            do_sample=False,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id
        )
        
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        # Extract the number after "Output:"
        try:
            predicted = response.split("Output:")[-1].strip().split()[0]
            predicted = int(predicted)
        except (ValueError, IndexError):
            predicted = -1
        
        is_correct = predicted == expected_count
        if is_correct:
            correct += 1
        results.append((word, expected_count, predicted, is_correct))

accuracy = correct / num_samples
return accuracy, results

print("=" * 60)
print("Fine-tuning Llama-3.2-1B-Instruct to count Rs in strawberry")
print("=" * 60)

# Device setup
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Load tokenizer
print(f"\nLoading tokenizer from {model_name}...")
tokenizer = AutoTokenizer.from_pretrained(model_name)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# Load model
print(f"Loading model from {model_name}...")
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
    device_map="auto" if torch.cuda.is_available() else None
)

if not torch.cuda.is_available():
    model = model.to(device)

# Generate training data
print(f"\nGenerating {num_train_samples} training samples...")
train_samples = create_dataset_samples(num_train_samples, max_r_count)

# Show some examples
print("\nSample training data:")
for i in range(5):
    word, count = train_samples[i]
    print(f"  '{word}' -> {count}")

# Create dataset and dataloader
train_dataset = StrawberryDataset(train_samples, tokenizer)
train_loader = DataLoader(
    train_dataset,
    batch_size=batch_size,
    shuffle=True,
    num_workers=0
)

# Evaluate before training
print("\n" + "=" * 60)
print("Evaluating BEFORE fine-tuning...")
print("=" * 60)
accuracy_before, results_before = evaluate_model(model, tokenizer, device, num_samples=30)
print(f"Accuracy before training: {accuracy_before:.1%}")
print("\nSample predictions (before):")
for word, expected, predicted, correct in results_before[:10]:
    status = "✓" if correct else "✗"
    print(f"  {status} '{word[:30]}...' expected={expected}, got={predicted}")

# Setup optimizer and scheduler
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
total_steps = len(train_loader) * num_epochs // gradient_accumulation_steps
scheduler = get_linear_schedule_with_warmup(
    optimizer,
    num_warmup_steps=total_steps // 10,
    num_training_steps=total_steps
)

# Training loop
print("\n" + "=" * 60)
print("Starting training...")
print("=" * 60)

model.train()
global_step = 0

for epoch in range(num_epochs):
    epoch_loss = 0.0
    num_batches = 0
    
    progress_bar = tqdm(train_loader, desc=f"Epoch {epoch + 1}/{num_epochs}")
    
    for batch_idx, batch in enumerate(progress_bar):
        input_ids = batch["input_ids"].to(device)
        attention_mask = batch["attention_mask"].to(device)
        labels = batch["labels"].to(device)
        
        outputs = model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            labels=labels
        )
        
        loss = outputs.loss / gradient_accumulation_steps
        loss.backward()
        
        epoch_loss += outputs.loss.item()
        num_batches += 1
        
        if (batch_idx + 1) % gradient_accumulation_steps == 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            scheduler.step()
            optimizer.zero_grad()
            global_step += 1
        
        progress_bar.set_postfix({"loss": f"{epoch_loss / num_batches:.4f}"})
    
    avg_loss = epoch_loss / num_batches
    print(f"Epoch {epoch + 1} completed. Average loss: {avg_loss:.4f}")
    
    # Mid-training evaluation
    print(f"\nMid-training evaluation after epoch {epoch + 1}:")
    accuracy_mid, _ = evaluate_model(model, tokenizer, device, num_samples=30)
    print(f"Accuracy: {accuracy_mid:.1%}")
    model.train()

# Final evaluation
print("\n" + "=" * 60)
print("Evaluating AFTER fine-tuning...")
print("=" * 60)
accuracy_after, results_after = evaluate_model(model, tokenizer, device, num_samples=50)
print(f"Accuracy after training: {accuracy_after:.1%}")
print("\nSample predictions (after):")
for word, expected, predicted, correct in results_after[:15]:
    status = "✓" if correct else "✗"
    print(f"  {status} '{word[:40]}' expected={expected}, got={predicted}")

# Test on the classic examples
print("\n" + "=" * 60)
print("Testing on classic examples...")
print("=" * 60)

classic_tests = [
    ("strawberry", 3),
    ("strrawberrrrry", 7),
    ("strrrrrawberrrrrrrrrry", 15),
    ("stawbey", 0),
]

model.eval()
with torch.no_grad():
    for word, expected in classic_tests:
        prompt = f"Input: {word}\nOutput:"
        inputs = tokenizer(prompt, return_tensors="pt").to(device)
        
        outputs = model.generate(
            **inputs,
            max_new_tokens=10,
            num_beams=1,
            do_sample=False,
            pad_token_id=tokenizer.pad_token_id
        )
        
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        try:
            predicted = response.split("Output:")[-1].strip().split()[0]
        except IndexError:
            predicted = "N/A"
        
        print(f"  Input: '{word}'")
        print(f"  Expected: {expected}, Predicted: {predicted}")
        print()

# Save the model
output_dir = "strawberry-llama"
print(f"\nSaving model to {output_dir}...")
model.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir)
print("Done!")

print("\n" + "=" * 60)
print("Summary")
print("=" * 60)
print(f"Accuracy before training: {accuracy_before:.1%}")
print(f"Accuracy after training:  {accuracy_after:.1%}")
print(f"Improvement: {(accuracy_after - accuracy_before):.1%}")