Views
No views yet
1import torch
2import torch.nn as nn
3from transformers import GPT2Tokenizer, GPT2LMHeadModel
4
5tokenizer = GPT2Tokenizer.from_pretrained("Norod78/hewiki-articles-distilGPT2py-il")
6model = GPT2LMHeadModel.from_pretrained("Norod78/hewiki-articles-distilGPT2py-il").eval()
7
8bos_token = tokenizer.bos_token #Beginning of sentace
9eos_token = tokenizer.eos_token #End of sentence
10
11def generate_word(model, tokens_tensor, temperature=1.0):
12 """
13 Sample a word given a tensor of tokens of previous words from a model. Given
14 the words we have, sample a plausible word. Temperature is used for
15 controlling randomness. If using temperature==0 we simply use a greedy arg max.
16 Else, we sample from a multinomial distribution using a lower inverse
17 temperature to allow for more randomness to escape repetitions.
18 """
19 with torch.no_grad():
20 outputs = model(tokens_tensor)
21 predictions = outputs[0]
22 if temperature>0:
23 # Make the distribution more or less skewed based on the temperature
24 predictions = outputs[0]/temperature
25 # Sample from the distribution
26 softmax = nn.Softmax(dim=0)
27 predicted_index = torch.multinomial(softmax(predictions[0,-1,:]),1).item()
28 # Simply take the arg-max of the distribution
29 else:
30 predicted_index = torch.argmax(predictions[0, -1, :]).item()
31 # Decode the encoding to the corresponding word
32 predicted_text = tokenizer.decode([predicted_index])
33 return predicted_text
34
35def generate_sentence(model, tokenizer, initial_text, temperature=1.0):
36 """ Generate a sentence given some initial text using a model and a tokenizer.
37 Returns the new sentence. """
38
39 # Encode a text inputs
40 text = ""
41 sentence = text
42
43 # We avoid an infinite loop by setting a maximum range
44 for i in range(0,84):
45 indexed_tokens = tokenizer.encode(initial_text + text)
46
47 # Convert indexed tokens in a PyTorch tensor
48 tokens_tensor = torch.tensor([indexed_tokens])
49
50 new_word = generate_word(model, tokens_tensor, temperature=temperature)
51
52 # Here the temperature is slowly decreased with each generated word,
53 # this ensures that the sentence (ending) makes more sense.
54 # We don't decrease to a temperature of 0.0 to leave some randomness in.
55 if temperature<(1-0.008):
56 temperature += 0.008
57 else:
58 temperature = 0.996
59
60 text = text+new_word
61
62 # Stop generating new words when we have reached the end of the line or the poem
63 if eos_token in new_word:
64 # returns new sentence and whether poem is done
65 return (text.replace(eos_token,"").strip(), True)
66 elif '/' in new_word:
67 return (text.strip(), False)
68 elif bos_token in new_word:
69 return (text.replace(bos_token,"").strip(), False)
70
71 return (text, True)
72
73for output_num in range(1,5):
74 init_text = "בוקר טוב"
75 text = bos_token + init_text
76 for i in range(0,84):
77 sentence = generate_sentence(model, tokenizer, text, temperature=0.9)
78 text = init_text + sentence[0]
79 print(text)
80 if (sentence[1] == True):
81 break