Repetition Penalties › Combining Penalties with Sampling Strategies

def generate_with_combined_strategies(
    model,
    tokenizer,
    prompt,
    repetition_penalty=1.0,
    frequency_penalty=0.0,
    presence_penalty=0.0,
    temperature=1.0,
    top_p=0.9,
    max_tokens=50,
):
    """
    Generate text combining multiple penalty types with sampling strategies.
    """
    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    generated_ids = input_ids[0].tolist()
    prompt_length = len(generated_ids)
    token_counts = {}

    for _ in range(max_tokens):
        with torch.no_grad():
            outputs = model(torch.tensor([generated_ids]))
            logits = outputs.logits[0, -1, :].clone()

        # Track tokens generated so far (excluding prompt)
        generated_so_far = generated_ids[prompt_length:]
        generated_set = set(generated_so_far)

        # Count occurrences
        for token_id in generated_so_far:
            token_counts[token_id] = token_counts.get(token_id, 0) + 1

        # Apply repetition penalty
        if repetition_penalty != 1.0:
            for token_id in generated_set:
                if logits[token_id] > 0:
                    logits[token_id] = logits[token_id] / repetition_penalty
                else:
                    logits[token_id] = logits[token_id] * repetition_penalty

        # Apply frequency penalty
        if frequency_penalty > 0:
            for token_id, count in token_counts.items():
                logits[token_id] = logits[token_id] - frequency_penalty * count

        # Apply presence penalty
        if presence_penalty > 0:
            for token_id in generated_set:
                logits[token_id] = logits[token_id] - presence_penalty

        # Apply temperature
        logits = logits / temperature

        # Apply nucleus sampling
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=0), dim=0)
        sorted_indices_to_remove = cumulative_probs > top_p
        sorted_indices_to_remove[1:] = sorted_indices_to_remove[:-1].clone()
        sorted_indices_to_remove[0] = False
        indices_to_remove = sorted_indices_to_remove.scatter(
            0, sorted_indices, sorted_indices_to_remove
        )
        logits[indices_to_remove] = float("-inf")

        # Sample
        probs = F.softmax(logits, dim=0)
        next_token = torch.multinomial(probs, num_samples=1).item()
        generated_ids.append(next_token)

        if next_token == tokenizer.eos_token_id:
            break

    return tokenizer.decode(generated_ids, skip_special_tokens=True)