Spaces:
Runtime error
Runtime error
| import torch | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| device = "cuda" # or "cpu" | |
| model_path = "ibm-granite/granite-8b-code-instruct" | |
| tokenizer = AutoTokenizer.from_pretrained(model_path) | |
| # drop device_map if running on CPU | |
| model = AutoModelForCausalLM.from_pretrained(model_path, device_map=device) | |
| model.eval() | |
| # change input text as desired | |
| chat = [ | |
| { "role": "user", "content": "Write a code to find the maximum value in a list of numbers." }, | |
| ] | |
| chat = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True) | |
| # tokenize the text | |
| input_tokens = tokenizer(chat, return_tensors="pt") | |
| # transfer tokenized inputs to the device | |
| for i in input_tokens: | |
| input_tokens[i] = input_tokens[i].to(device) | |
| # generate output tokens | |
| output = model.generate(**input_tokens, max_new_tokens=100) | |
| # decode output tokens into text | |
| output = tokenizer.batch_decode(output) | |
| # loop over the batch to print, in this example the batch size is 1 | |
| for i in output: | |
| print(i) | |