santuchal/codegenerate

1import torch2from transformers import AutoModelForCausalLM, AutoTokenizer3device = "cuda" # or "cpu"4model_path = "ibm-granite/granite-8b-code-instruct"5tokenizer = AutoTokenizer.from_pretrained(model_path)6# drop device_map if running on CPU7model = AutoModelForCausalLM.from_pretrained(model_path, device_map=device)8model.eval()9# change input text as desired10chat = [11    { "role": "user", "content": "Write a code to find the maximum value in a list of numbers." },12]13chat = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)14# tokenize the text15input_tokens = tokenizer(chat, return_tensors="pt")16# transfer tokenized inputs to the device17for i in input_tokens:18    input_tokens[i] = input_tokens[i].to(device)19# generate output tokens20output = model.generate(**input_tokens, max_new_tokens=100)21# decode output tokens into text22output = tokenizer.batch_decode(output)23# loop over the batch to print, in this example the batch size is 124for i in output:25    print(i)26