Views
No views yet
1from transformers import AutoTokenizer, AutoModelForCausalLM
2
3tokenizer = AutoTokenizer.from_pretrained("Dmitriy007/rugpt2_gen_news")
4model = AutoModelForCausalLM.from_pretrained("Dmitriy007/rugpt2_gen_news")
5
6input_text = 'Ученик старшего класса лицея № 21 Иван Сидоров из города Адлер полетел в космос на планету Марс.'
7
8inputs = tokenizer(input_text, return_tensors="pt")
9model.to('cuda')
10inputs.to('cuda')
11
12input_ids = inputs["input_ids"]
13
14output = model.generate(
15input_ids,
16attention_mask=inputs["attention_mask"],
17pad_token_id=model.config.bos_token_id,
18max_length=300,
19num_beams=5,
20num_return_sequences=1,
21top_k=50,
22top_p=0.90,
23no_repeat_ngram_size=2,
24temperature=0.7,
25early_stopping=True
26)
27
28generated_text = list(map(tokenizer.decode, output))
29print(generated_text[0])