generate.py

"""
Text Generation utilities for interactive prompt evaluation
"""

import torch
from typing import Union


def generate(
    model,
    tok,
    prompt: str,
    max_new_tokens: int = 100,
    temperature: float = 0.8,
    top_k: int = 40,
    device: Union[str, torch.device] = None
) -> str:
    """
    Generates text continuation for a prompt using model and tokenizer.
    """
    if device is None:
        device = next(model.parameters()).device

    model.eval()

    # Encode prompt using Tokenizer class or tiktoken
    if hasattr(tok, "encode"):
        tokens = tok.encode(prompt)
    else:
        tokens = tok.encode_ordinary(prompt)

    idx = torch.tensor(tokens, dtype=torch.long, device=device).unsqueeze(0)

    raw_model = model.module if hasattr(model, "module") else model
    out_idx = raw_model.generate(idx, max_new_tokens=max_new_tokens, temperature=temperature, top_k=top_k)

    if hasattr(tok, "decode"):
        return tok.decode(out_idx[0].tolist())
    else:
        return tok.decode(out_idx[0].tolist())