| from transformers import GPT2LMHeadModel, GPT2Tokenizer |
| import torch |
| DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
|
|
| |
| |
| model_name_or_path = "ai-forever/rugpt3medium_based_on_gpt2" |
| tokenizer = GPT2Tokenizer.from_pretrained(model_name_or_path) |
| model = GPT2LMHeadModel.from_pretrained(model_name_or_path).to(DEVICE) |
|
|
|
|
| |
| text = "Вопрос: 'Сколько будет 2+2?'\nОтвет:" |
| input_ids = tokenizer.encode(text, return_tensors="pt").to(DEVICE) |
| out = model.generate(input_ids, do_sample=False) |
|
|
| generated_text = list(map(tokenizer.decode, out))[0] |
| print(generated_text) |
|
|
|
|
