Views
No views yet
1from transformers import AutoTokenizer
2from optimum.onnxruntime import ORTModelForCausalLM
3
4model = ORTModelForCausalLM.from_pretrained("TextCortex/codegen-350M-optimized")
5tokenizer = AutoTokenizer.from_pretrained("TextCortex/codegen-350M-optimized")
6
7text = "def hello_world():"
8input_ids = tokenizer(text, return_tensors="pt").input_ids
9generated_ids = model.generate(
10 input_ids,
11 max_length=64,
12 temperature=0.1,
13 num_return_sequences=1,
14 early_stopping=True,
15)
16out = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
17print(out)