metadata
license: apache-2.0
pipeline_tag: text-classification
basemodel: roberta-base
datasets:
- DevQuasar/llm_router_dataset-synth
language:
- en
Intention of the model is to determine if the given user prompt's complexity, domain question requires a SOTA (very large) LLM or can be deescaleted to a smaller or local model.
Example code:
from openai import OpenAI
from datasets import load_dataset
from datasets.dataset_dict import DatasetDict
import json
import random
from transformers import (
RobertaTokenizerFast,
RobertaForSequenceClassification,
)
from transformers import pipeline
model_id = 'DevQuasar/roberta-prompt_classifier-v0.1'
tokenizer = RobertaTokenizerFast.from_pretrained(model_id)
sentence_classifier = pipeline(
"sentiment-analysis", model=model_id, tokenizer=tokenizer
)
model_store = {
"small_llm": {
"escalation_order": 0,
"url": "http://localhost:1234/v1",
"api_key": "lm-studio",
"model_id": "lmstudio-community/Meta-Llama-3-8B-Instruct-GGUF/Meta-Llama-3-8B-Instruct-Q4_K_M.gguf",
"max_ctx": 4096
},
"large_llm": {
"escalation_order": 1,
"url": "http://localhost:1234/v1",
"api_key": "lm-studio",
"model_id": "lmstudio-community/Meta-Llama-3-70B-Instruct-GGUF/Meta-Llama-3-70B-Instruct-Q4_K_M.gguf",
"max_ctx": 8192
}
}
def prompt_classifier(user_prompt):
return sentence_classifier(user_prompt)[0]['label']
def llm_router(user_prompt, tokens_so_far = 0):
return model_store[prompt_classifier(user_prompt)]
def chat(user_prompt, model_store_entry = None, curr_ctx = [], system_prompt = ' ', verbose=False):
if model_store_entry == None and curr_ctx == []:
# initial model selection
model_store_entry = llm_router(user_prompt)
if verbose:
print(f'Classify prompt - selected model: {model_store_entry["model_id"]}')
else:
#handle escalation
model_store_candidate = llm_router(user_prompt)
if model_store_candidate["escalation_order"] > model_store_entry["escalation_order"]:
model_store_entry = model_store_candidate
if verbose:
print(f'Escalate model - selected model: {model_store_entry["model_id"]}')
url = model_store_entry['url']
api_key = model_store_entry['api_key']
model_id = model_store_entry['model_id']
client = OpenAI(base_url=url, api_key=api_key)
messages = curr_ctx
messages.append({"role": "user", "content": user_prompt})
completion = client.chat.completions.create(
model=model_id,
messages = messages,
temperature=0.7,
)
messages.append({"role": "assistant", "content": completion.choices[0].message.content})
if verbose:
print(f'Used model: {model_id}')
print(f'completion: {completion}')
client.close()
return completion.choices[0].message.content, messages, model_store_entry
use_model = None
ctx = []
# start with simple prompt -> llama3-8b
res, ctx, use_model = chat(user_prompt="hello", model_store_entry=use_model, curr_ctx=ctx, verbose=True)
# escalate prompt -> llama3-70b
p = "Discuss the challenges and potential solutions for achieving sustainable development in the context of increasing global urbanization."
res, ctx, use_model = chat(user_prompt=p, model_store_entry=use_model, curr_ctx=ctx, verbose=True)