roneneldan commited on
Commit
0fa3a4c
1 Parent(s): ae97ae9

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +18 -3
README.md CHANGED
@@ -1,3 +1,18 @@
1
- ---
2
- license: unknown
3
- ---
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
2
+
3
+ model = AutoModelForCausalLM.from_pretrained('roneneldan/TinyStories-1M')
4
+
5
+ tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-125M")
6
+
7
+ prompt = "Once upon a time there was"
8
+
9
+ input_ids = tokenizer.encode(prompt, return_tensors="pt")
10
+
11
+ # Generate completion
12
+ output = model.generate(input_ids, max_length = 1000, num_beams=1)
13
+
14
+ # Decode the completion
15
+ output_text = tokenizer.decode(output[0], skip_special_tokens=True)
16
+
17
+ # Print the generated text
18
+ print(output_text)