Views
No views yet
transformers library and we advise you to install latest version:pip install transformers>=4.37.01import json
2from typing import Any, Dict, List
3from transformers import AutoModelForCausalLM, AutoTokenizer
4
5model_name = "katanemo/Arch-Router-1.5B"
6model = AutoModelForCausalLM.from_pretrained(
7 model_name, device_map="auto", torch_dtype="auto", trust_remote_code=True
8)
9tokenizer = AutoTokenizer.from_pretrained(model_name)
10
11# Please use our provided prompt for best performance
12TASK_INSTRUCTION = """
13You are a helpful assistant designed to find the best suited route.
14You are provided with route description within <routes></routes> XML tags:
15<routes>
16
17{routes}
18
19</routes>
20
21<conversation>
22
23{conversation}
24
25</conversation>
26"""
27
28FORMAT_PROMPT = """
29Your task is to decide which route is best suit with user intent on the conversation in <conversation></conversation> XML tags. Follow the instruction:
301. If the latest intent from user is irrelevant or user intent is full filled, response with other route {"route": "other"}.
312. You must analyze the route descriptions and find the best match route for user latest intent.
323. You only response the name of the route that best matches the user's request, use the exact name in the <routes></routes>.
33
34Based on your analysis, provide your response in the following JSON formats if you decide to match any route:
35{"route": "route_name"}
36"""
37
38# Define route config
39route_config = [
40 {
41 "name": "code_generation",
42 "description": "Generating new code snippets, functions, or boilerplate based on user prompts or requirements",
43 },
44 {
45 "name": "bug_fixing",
46 "description": "Identifying and fixing errors or bugs in the provided code across different programming languages",
47 },
48 {
49 "name": "performance_optimization",
50 "description": "Suggesting improvements to make code more efficient, readable, or scalable",
51 },
52 {
53 "name": "api_help",
54 "description": "Assisting with understanding or integrating external APIs and libraries",
55 },
56 {
57 "name": "programming",
58 "description": "Answering general programming questions, theory, or best practices",
59 },
60]
61
62# Helper function to create the system prompt for our model
63def format_prompt(
64 route_config: List[Dict[str, Any]], conversation: List[Dict[str, Any]]
65):
66 return (
67 TASK_INSTRUCTION.format(
68 routes=json.dumps(route_config), conversation=json.dumps(conversation)
69 )
70 + FORMAT_PROMPT
71 )
72
73# Define conversations
74
75conversation = [
76 {
77 "role": "user",
78 "content": "fix this module 'torch.utils._pytree' has no attribute 'register_pytree_node'. did you mean: '_register_pytree_node'?",
79 }
80]
81
82route_prompt = format_prompt(route_config, conversation)
83
84messages = [
85 {"role": "user", "content": route_prompt},
86]
87
88input_ids = tokenizer.apply_chat_template(
89 messages, add_generation_prompt=True, return_tensors="pt"
90).to(model.device)
91
92# 2. Generate
93generated_ids = model.generate(
94 input_ids=input_ids, # or just positional: model.generate(input_ids, …)
95 max_new_tokens=32768,
96)
97
98# 3. Strip the prompt from each sequence
99prompt_lengths = input_ids.shape[1] # same length for every row here
100generated_only = [
101 output_ids[prompt_lengths:] # slice off the prompt tokens
102 for output_ids in generated_ids
103]
104
105# 4. Decode if you want text
106response = tokenizer.batch_decode(generated_only, skip_special_tokens=True)[0]
107print(response){"route": "bug_fixing"}