Views
No views yet
1
2!pip install tokenizers==0.10.2 transformers==4.6.0
3
4from transformers import AutoTokenizer, AutoModelForCausalLM
5
6tokenizer = AutoTokenizer.from_pretrained("Norod78/hebrew-gpt_neo-small")
7model = AutoModelForCausalLM.from_pretrained("Norod78/hebrew-gpt_neo-small", pad_token_id=tokenizer.eos_token_id)
8
9prompt_text = "אני אוהב שוקולד ועוגות"
10max_len = 512
11sample_output_num = 3
12seed = 1000
13
14import numpy as np
15import torch
16
17device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
18n_gpu = 0 if torch.cuda.is_available()==False else torch.cuda.device_count()
19
20print(f"device: {device}, n_gpu: {n_gpu}")
21
22np.random.seed(seed)
23torch.manual_seed(seed)
24if n_gpu > 0:
25 torch.cuda.manual_seed_all(seed)
26
27model.to(device)
28
29encoded_prompt = tokenizer.encode(
30 prompt_text, add_special_tokens=False, return_tensors="pt")
31
32encoded_prompt = encoded_prompt.to(device)
33
34if encoded_prompt.size()[-1] == 0:
35 input_ids = None
36else:
37 input_ids = encoded_prompt
38
39print("input_ids = " + str(input_ids))
40
41if input_ids != None:
42 max_len += len(encoded_prompt[0])
43 if max_len > 2048:
44 max_len = 2048
45
46print("Updated max_len = " + str(max_len))
47
48stop_token = "<|endoftext|>"
49new_lines = "\n\n\n"
50
51sample_outputs = model.generate(
52 input_ids,
53 do_sample=True,
54 max_length=max_len,
55 top_k=50,
56 top_p=0.95,
57 num_return_sequences=sample_output_num
58)
59
60print(100 * '-' + "\n\t\tOutput\n" + 100 * '-')
61for i, sample_output in enumerate(sample_outputs):
62
63 text = tokenizer.decode(sample_output, skip_special_tokens=True)
64
65 # Remove all text after the stop token
66 text = text[: text.find(stop_token) if stop_token else None]
67
68 # Remove all text after 3 newlines
69 text = text[: text.find(new_lines) if new_lines else None]
70
71 print("\n{}: {}".format(i, text))
72 print("\n" + 100 * '-')
73