Views
No views yet
1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3device = "cuda" # or "cpu"
4model_path = "ibm-granite/granite-8b-code-instruct-4k"
5tokenizer = AutoTokenizer.from_pretrained(model_path)
6# drop device_map if running on CPU
7model = AutoModelForCausalLM.from_pretrained(model_path, device_map=device)
8model.eval()
9# change input text as desired
10chat = [
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 text
15input_tokens = tokenizer(chat, return_tensors="pt")
16# transfer tokenized inputs to the device
17for i in input_tokens:
18 input_tokens[i] = input_tokens[i].to(device)
19# generate output tokens
20output = model.generate(**input_tokens, max_new_tokens=100)
21# decode output tokens into text
22output = tokenizer.batch_decode(output)
23# loop over the batch to print, in this example the batch size is 1
24for i in output:
25 print(i)