mayank-mishra
commited on
Commit
•
596a400
1
Parent(s):
9f8823b
update example
Browse files
README.md
CHANGED
@@ -226,25 +226,28 @@ This is a simple example of how to use **Granite-8B-Code-Instruct** model.
|
|
226 |
import torch
|
227 |
from transformers import AutoModelForCausalLM, AutoTokenizer
|
228 |
device = "cuda" # or "cpu"
|
229 |
-
model_path = "
|
230 |
tokenizer = AutoTokenizer.from_pretrained(model_path)
|
231 |
# drop device_map if running on CPU
|
232 |
model = AutoModelForCausalLM.from_pretrained(model_path, device_map=device)
|
233 |
model.eval()
|
234 |
# change input text as desired
|
235 |
-
|
|
|
|
|
|
|
236 |
# tokenize the text
|
237 |
-
input_tokens = tokenizer(
|
238 |
# transfer tokenized inputs to the device
|
239 |
for i in input_tokens:
|
240 |
input_tokens[i] = input_tokens[i].to(device)
|
241 |
# generate output tokens
|
242 |
-
output = model.generate(**input_tokens)
|
243 |
# decode output tokens into text
|
244 |
output = tokenizer.batch_decode(output)
|
245 |
# loop over the batch to print, in this example the batch size is 1
|
246 |
for i in output:
|
247 |
-
print(
|
248 |
```
|
249 |
|
250 |
<!-- TO DO: Check this part -->
|
|
|
226 |
import torch
|
227 |
from transformers import AutoModelForCausalLM, AutoTokenizer
|
228 |
device = "cuda" # or "cpu"
|
229 |
+
model_path = "granite-8b-code-instruct"
|
230 |
tokenizer = AutoTokenizer.from_pretrained(model_path)
|
231 |
# drop device_map if running on CPU
|
232 |
model = AutoModelForCausalLM.from_pretrained(model_path, device_map=device)
|
233 |
model.eval()
|
234 |
# change input text as desired
|
235 |
+
chat = [
|
236 |
+
{ "role": "user", "content": "Write a code to find the maximum value in a list of numbers." },
|
237 |
+
]
|
238 |
+
chat = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
|
239 |
# tokenize the text
|
240 |
+
input_tokens = tokenizer(chat, return_tensors="pt")
|
241 |
# transfer tokenized inputs to the device
|
242 |
for i in input_tokens:
|
243 |
input_tokens[i] = input_tokens[i].to(device)
|
244 |
# generate output tokens
|
245 |
+
output = model.generate(**input_tokens, max_new_tokens=100)
|
246 |
# decode output tokens into text
|
247 |
output = tokenizer.batch_decode(output)
|
248 |
# loop over the batch to print, in this example the batch size is 1
|
249 |
for i in output:
|
250 |
+
print(i)
|
251 |
```
|
252 |
|
253 |
<!-- TO DO: Check this part -->
|