File size: 1,340 Bytes
6036082 | 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 | from pathlib import Path
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
MODEL_PATH = "."
MODEL_SUBFOLDER = "hf"
PROMPT = "Once upon"
SEED = 0
def main() -> None:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_dir = Path(MODEL_PATH) / MODEL_SUBFOLDER
tokenizer = AutoTokenizer.from_pretrained(model_dir)
model = AutoModelForCausalLM.from_pretrained(
model_dir,
dtype=torch.float32,
).to(device)
model.eval()
input_ids = torch.tensor(
[
[tokenizer.bos_token_id]
+ tokenizer.encode(PROMPT, add_special_tokens=False)
],
dtype=torch.long,
device=device,
)
torch.manual_seed(SEED)
if device.type == "cuda":
torch.cuda.manual_seed_all(SEED)
with torch.inference_mode():
output = model.generate(
input_ids=input_ids,
max_new_tokens=100,
do_sample=True,
temperature=0.8,
top_p=0.95,
top_k=40,
repetition_penalty=1.1,
use_cache=True,
pad_token_id=tokenizer.pad_token_id,
eos_token_id=tokenizer.eos_token_id,
)
print(tokenizer.decode(output[0].tolist(), skip_special_tokens=True))
if __name__ == "__main__":
main()
|