changed model ordering

app.py

@@ -36,14 +36,14 @@ from router_backend import get_expert_routing
 EXPERTS = ["Language", "Logic", "Social", "World"]
 
 DEFAULT_MODELS = [
+    "micro-smollm2-135m",
+    "micro-smollm2-360m",
     "micro-llama-1b", 
     "micro-llama-3b",
     "micro-llama-1b-dpo",
-    "micro-moe-llama-1b",
-    "micro-smollm2-135m",
-    "micro-smollm2-360m",
     "micro-moe-smollm2-135m",
     "micro-moe-smollm2-360m",
+    "micro-moe-llama-1b",
 ]
 
 def _mock_routing(model_id: str, prompt: str, seed: int = 0) -> List[float]:
@@ -139,7 +139,7 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
 
     with gr.Row():
         model_choice = gr.Dropdown(choices=DEFAULT_MODELS, label="Select a model", value=DEFAULT_MODELS[0])
-        hf_token = gr.Textbox(label="Huggingface token for authentication", placeholder="hf token", lines=1)
+        hf_token = gr.Textbox(label="Huggingface token for authentication", placeholder="Required for Llama-based models", lines=1)
 
     with gr.Row():
         user_prompt = gr.Textbox(lines=6, label="User prompt", placeholder="Type the user message here...")
@@ -151,7 +151,7 @@ with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
 
     run = gr.Button("Run Routing", variant="primary")
 
-    generation_output = gr.Textbox(lines=4, label="Generated continuation", placeholder="Generated text will appear here...", interactive=False)
+    generation_output = gr.Textbox(lines=4, label="Generated Response", placeholder="Generated text will appear here...", interactive=False)
 
     with gr.Row():
         table = gr.Dataframe(label="Routing Percentages", interactive=False)

router_backend.py

@@ -32,6 +32,7 @@ def get_expert_routing(model_id: str, hf_token: str, prompt: Union[str, List[Dic
 
     if isinstance(prompt, str):
         generation, routing_weights = generate_continuation(model, tokenizer, prompt)
+        generation = generation[0] if type(generation) is list else generation
     elif isinstance(prompt, list):
         generation = None
         routing_weights = get_routing_weights(model, tokenizer, [prompt])