Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForCausalLM
3
4# Load Code Generator LLM and tokenizer from checkpoint
5tokenizer = AutoTokenizer.from_pretrained("DeathReaper0965/gpt2_large_code_generator")
6model = AutoModelForCausalLM.from_pretrained("DeathReaper0965/gpt2_large_code_generator")
7model = model.to("cuda" if torch.cuda.is_available() else "cpu")
8
9inputs = tokenizer("def hello_world():", return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
10
11outputs = model.generate(**inputs,
12 max_new_tokens= 30,
13 num_return_sequences= 1)
14
15print(tokenizer.batch_decode(outputs)[0])
16
17###########OUTPUT###########
18def hello_world():
19 return "Hello World!"
20
21@app.route("/hello_world")
22def hello_world():
23 return "Hello World!"