generation_ksbyte.py

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


def generate_text(model_or_repo, prompt, max_new_tokens=200, temperature=0.8, top_k=50, device="auto"):
    """Convenience helper for this byte-level model."""
    if isinstance(model_or_repo, str):
        model = AutoModelForCausalLM.from_pretrained(model_or_repo, trust_remote_code=True)
        tok = AutoTokenizer.from_pretrained(model_or_repo, trust_remote_code=True)
    else:
        model = model_or_repo
        tok = None
    if device == "auto":
        device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device).eval()
    if tok is None:
        from tokenization_ksbyte import KsByteTokenizer
        tok = KsByteTokenizer()
    inputs = tok(prompt, return_tensors="pt").to(device)
    if hasattr(model, "generate_bytes"):
        out = model.generate_bytes(inputs["input_ids"], max_new_tokens=max_new_tokens, temperature=temperature, top_k=top_k)
    else:
        out = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=True, temperature=temperature, top_k=top_k)
    return tok.decode(out[0], skip_special_tokens=True)