File size: 2,965 Bytes
948e1ca 8cc58fd 948e1ca |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
<div align="center">
<h1>
TagRouter: Learning Route to LLMs through Tags for Open-Domain Text Generation Tasks
</h1>
</div>
## 🎉 News
- [2025-5-16] [Our paper](https://arxiv.org/abs/2506.12473) has been accepted for publication in ACL 2025.
## Introduction
Model routing allocates queries to the suitable model, improving system performance while reducing costs. However, existing routing methods face practical limitations that hinder scalability in large-scale applications and struggle to keep up with the rapid growth of the large language model (LLM) ecosystem. To tackle these challenges, we propose TagRouter, a training-free model routing method designed to optimize the synergy among multiple LLMs for open-domain text generation tasks. Experimental results demonstrate that TagRouter outperforms 13 baseline methods, increasing the accept rate of system by 6.15% and reducing costs by 17.20%, achieving optimal cost efficiency. Our findings provides the LLM community with an efficient and scalable solution for model ensembling, offering users an evolvable "super model."<br>
TagRouter consists of three modules: TagGenerator, TagScorer, and TagDecider. The TagGenerator is trained to generate a set of tags for a given query. The generated tags can be used for routing queries to the most suitable model based on their respective capabilities.
<p align="center">
<br>
<img src="image/TagRouter.png" width="800"/>
<br>
</p>
## Download
[HuggingFace](https://huggingface.co/itpossible/TagGenerator)<br>
[ModelScope](https://modelscope.cn/models/itpossible/TagGenerator)
## Inference
Below is an example of inference code using TagGenerator.
```python
import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
model_path = "itpossible/TagGenerator"
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype="auto", device_map="auto", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
prompt = """[System]
You are an expert text tag extractor. Your task is to identify key tags that readers should focus on while engaging with the user query below.
[User Query]
Rewrite the sentence so that it's in the present tense: She had worked at the company for the past 3 years.
"""
messages = [
{"role": "system", "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."},
{"role": "user", "content": prompt}
]
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
generated_ids = model.generate(**model_inputs, max_new_tokens=512)
generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)]
response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
print(response)
``` |