File size: 2,405 Bytes
fd0b9f3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | """Load the model from this folder (or the Hub) and generate text."""
import argparse
import torch
from tokenizers import Tokenizer
import transformers
from transformers import AutoModelForCausalLM
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model", default=".", help="Local folder or Hugging Face model ID")
parser.add_argument("--prompt", default="The purpose of education is")
parser.add_argument("--max-new-tokens", type=int, default=128)
parser.add_argument("--temperature", type=float, default=0.8)
parser.add_argument("--top-k", type=int, default=40)
args = parser.parse_args()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
dtype_arg = (
{"dtype": dtype}
if int(transformers.__version__.split(".", 1)[0]) >= 5
else {"torch_dtype": dtype}
)
model = AutoModelForCausalLM.from_pretrained(
args.model, trust_remote_code=True, **dtype_arg
).to(device).eval()
tokenizer_path = args.model + "/tokenizer.json" if args.model != "." else "tokenizer.json"
if not tokenizer_path.startswith(("http://", "https://")) and "/" in args.model and not __import__("os").path.exists(tokenizer_path):
from huggingface_hub import hf_hub_download
tokenizer_path = hf_hub_download(args.model, "tokenizer.json")
tokenizer = Tokenizer.from_file(tokenizer_path)
eos_id = tokenizer.token_to_id("<eos>")
token_ids = tokenizer.encode(args.prompt).ids
generator = torch.Generator(device=device).manual_seed(42)
with torch.inference_mode():
for _ in range(args.max_new_tokens):
context = token_ids[-model.config.context_length :]
input_ids = torch.tensor([context], device=device, dtype=torch.long)
logits = model(input_ids=input_ids).logits[0, -1].float()
k = min(args.top_k, logits.numel())
values, indexes = torch.topk(logits, k=k)
probabilities = torch.softmax(values / args.temperature, dim=-1)
next_id = int(indexes[torch.multinomial(probabilities, 1, generator=generator)].item())
if next_id == eos_id:
break
token_ids.append(next_id)
print(tokenizer.decode(token_ids, skip_special_tokens=True))
if __name__ == "__main__":
main()
|