Views
No views yet
Qwen/Qwen2.5-Coder-0.5B-Instruct by westenfelder/NL2SH-ALFA for NLP to bash command. Distiled only decoder block from 24 to 4 with the original tokenizer.1from transformers import AutoTokenizer, AutoModelForCausalLM
2import torch
3
4model_name = "tripathysagar/Qwen2.5-Coder-196M-Shell"
5
6tokenizer = AutoTokenizer.from_pretrained(model_name)
7model = AutoModelForCausalLM.from_pretrained(
8 model_name,
9 dtype=torch.bfloat16,
10 device_map="auto",
11 )
12
13def infer(inp, debug=False):
14 msg = [
15 {"role": "system", "content": "Generate shell command."},
16 {"role": "user", "content": inp},
17 ]
18 text = tokenizer.apply_chat_template(
19 msg,
20 tokenize=False,
21 add_generation_prompt=True,
22 )
23 if debug:
24 print(text)
25
26 model_inputs = tokenizer([text], return_tensors="pt")
27 generated_ids = model.generate(
28 **model_inputs,
29 max_new_tokens=256,
30 do_sample=True,
31
32 )
33
34 resp_text = tokenizer.batch_decode(generated_ids)[0]
35 if debug:
36 print(resp_text)
37
38 return (inp, resp_text[len(text):].replace('<|im_end|>', ''))
39
40infer("get kernel name.")
41