| <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 has been accepted for publication in ACL. | |
| ## 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) | |
| ``` |