import torch
import torch.nn.functional as F
from transformers import GPT2Tokenizer

# Load tokenizer
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def generate(model, prompt, max_new_tokens=200,
             temperature=0.8, top_k=50,
             repetition_penalty=1.3):
    model.eval()
    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(DEVICE)
    generated = input_ids.clone()

    with torch.no_grad():
        for _ in range(max_new_tokens):
            x      = generated[:, -512:]
            logits = model(x)[:, -1, :].float()

            for token_id in set(generated[0].tolist()):
                if logits[0, token_id] > 0:
                    logits[0, token_id] /= repetition_penalty
                else:
                    logits[0, token_id] *= repetition_penalty

            logits = logits / temperature
            k = min(top_k, logits.size(-1))
            topk_vals, _ = torch.topk(logits, k)
            logits = logits.masked_fill(logits < topk_vals[:, -1:], -1e9)
            probs  = torch.softmax(logits, dim=-1).clamp(min=0)
            probs  = probs / probs.sum()
            next_token = torch.multinomial(probs, num_samples=1)
            generated  = torch.cat([generated, next_token], dim=1)
            if next_token.item() == tokenizer.eos_token_id:
                break

    return tokenizer.decode(generated[0], skip_special_tokens=True)