Qwen2.5-Coder-196M-Shell

Model Card for Model ID

<!-- Provide a quick summary of what the model is/does. -->

Distiled 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.

Model Details

Model Sources [optional]

<!-- Provide the basic links for the model. -->

• Blog: [More Information Needed]

Uses

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model_name = "tripathysagar/Qwen2.5-Coder-196M-Shell"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    dtype=torch.bfloat16,
    device_map="auto",
    )

def infer(inp, debug=False):
  msg = [
    {"role": "system", "content": "Generate shell command."},
    {"role": "user", "content": inp},
  ]
  text = tokenizer.apply_chat_template(
    msg,
    tokenize=False,
    add_generation_prompt=True,
  )
  if debug:
    print(text)

  model_inputs = tokenizer([text], return_tensors="pt")
  generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=256,
    do_sample=True,

  )

  resp_text = tokenizer.batch_decode(generated_ids)[0]
  if debug:
    print(resp_text)

  return (inp, resp_text[len(text):].replace('<|im_end|>', ''))

infer("get kernel name.")