| """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)) |
|
|
