infer.py

"""Uncurated text completion; this checkpoint is NOT an instruction/chat model."""
import argparse
from pathlib import Path
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, set_seed

p = argparse.ArgumentParser()
p.add_argument('--model', default=str(Path(__file__).resolve().parent/'model'))
p.add_argument('--prompt', default='Однажды вечером')
p.add_argument('--max-new-tokens', type=int, default=64)
p.add_argument('--seed', type=int, default=42)
a = p.parse_args()
set_seed(a.seed)
torch.set_num_threads(2)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
tok = AutoTokenizer.from_pretrained(a.model, trust_remote_code=False)
model = AutoModelForCausalLM.from_pretrained(a.model, trust_remote_code=False,
    torch_dtype=torch.float16 if device == 'cuda' else torch.float32).to(device).eval()
if not 1 <= a.max_new_tokens < model.config.n_positions:
    raise ValueError('max-new-tokens must be between 1 and 511')
inputs = tok(a.prompt, return_tensors='pt', truncation=True,
    max_length=model.config.n_positions-a.max_new_tokens).to(device)
with torch.inference_mode():
    output = model.generate(**inputs, max_new_tokens=a.max_new_tokens,
        do_sample=True, temperature=.8, top_p=.9, repetition_penalty=1.1,
        pad_token_id=tok.pad_token_id, eos_token_id=tok.eos_token_id)
print(tok.decode(output[0], skip_special_tokens=True))