added more models

app.py

@@ -113,7 +113,9 @@ def route_and_plot(model_choice: str, hf_token: str, user_prompt: str, assistant
             generation = None
 
     df = pd.DataFrame({"Expert": EXPERTS, "Percent": vals})
+    colors = ["#97D077", "#4285F4", "#FFAB40", "#A64D79"]
     fig = px.bar(df, x="Expert", y="Percent", title="Token Routing by Expert (%)", text="Percent")
+    fig.update_traces(marker_color=colors)
     fig.update_traces(texttemplate="%{text:.2f}%", textposition="outside")
     fig.update_layout(yaxis_range=[0, max(100, max(vals) * 1.25)], bargap=0.35)
 

configs/micro_llama_1b_dpo.yml

@@ -0,0 +1,14 @@
+run-title: micro-llama-1b-dpo
+model: micro-llama-1b-dpo
+
+base-model: meta-llama/Llama-3.2-1B
+tokenizer: meta-llama/Llama-3.2-1B-Instruct
+num-experts: 4
+top-k-experts: 1
+jitter-noise: 0
+use-router: True
+mask-input: True
+max-length: 8192
+
+trainable:
+  - model

configs/micro_llama_3b.yml

@@ -0,0 +1,14 @@
+run-title: micro-llama-3b
+model: micro-llama-3b
+
+base-model: meta-llama/Llama-3.2-3B
+tokenizer: meta-llama/Llama-3.2-3B-Instruct
+num-experts: 4
+top-k-experts: 1
+jitter-noise: 0
+use-router: True
+mask-input: True
+max-length: 8192
+
+trainable:
+  - model

configs/micro_moe_llama_1b.yml

@@ -0,0 +1,17 @@
+run-title: micro-moe-llama-1b
+
+model: micro-moe-llama-1b
+
+base-model: meta-llama/Llama-3.2-1B
+tokenizer: meta-llama/Llama-3.2-1B-Instruct
+num-experts: 4
+top-k-experts: 1
+jitter-noise: 0
+router-aux-loss-coef: 0.000
+use-load-balancing: False
+use-router: True
+mask-input: True
+max-length: 8192
+
+trainable:
+  - model

configs/micro_moe_smollm2_135m.yml

@@ -0,0 +1,16 @@
+run-title: micro-moe-smollm2-135m
+model: micro-moe-smollm2-135m
+
+base-model: HuggingFaceTB/SmolLM2-135M
+tokenizer: HuggingFaceTB/SmolLM2-135M-Instruct
+num-experts: 4
+top-k-experts: 1
+jitter-noise: 0
+router-aux-loss-coef: 0.000
+use-load-balancing: False
+use-router: True
+mask-input: True
+max-length: 8192
+
+trainable:
+  - model

configs/micro_moe_smollm2_360m.yml

@@ -0,0 +1,16 @@
+run-title: micro-moe-smollm2-360m
+model: micro-moe-smollm2-360m
+
+base-model: HuggingFaceTB/SmolLM2-360M
+tokenizer: HuggingFaceTB/SmolLM2-360M-Instruct
+num-experts: 4
+top-k-experts: 1
+jitter-noise: 0
+router-aux-loss-coef: 0.000
+use-load-balancing: False
+use-router: True
+mask-input: True
+max-length: 8192
+
+trainable:
+  - model

configs/micro_smollm2_135m.yml

@@ -0,0 +1,15 @@
+run-title: micro-smollm2-135m
+model: micro-smollm2-135m
+
+base-model: HuggingFaceTB/SmolLM2-135M
+tokenizer: HuggingFaceTB/SmolLM2-135M-Instruct
+num-experts: 4
+top-k-experts: 1
+jitter-noise: 0
+use-router: True
+mask-input: True
+max-length: 8192
+gradient-checkpointing: False
+
+trainable:
+  - model

configs/micro_smollm2_360m.yml

@@ -0,0 +1,15 @@
+run-title: micro-smollm2-360m
+model: micro-smollm2-360m
+
+base-model: HuggingFaceTB/SmolLM2-360M
+tokenizer: HuggingFaceTB/SmolLM2-360M-Instruct
+num-experts: 4
+top-k-experts: 1
+jitter-noise: 0
+use-router: True
+mask-input: True
+max-length: 8192
+gradient-checkpointing: False
+
+trainable:
+  - model

router_backend.py

@@ -32,7 +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)
-    elif isinstance(prompt, dict):
+    elif isinstance(prompt, list):
         generation = None
         routing_weights = get_routing_weights(model, tokenizer, [prompt])
 
@@ -88,7 +88,7 @@ def aggregate_routing_weights(routing_weights):
 def generate_continuation(model, 
     tokenizer, 
     prompts, 
-    max_tokens=1024,
+    max_tokens=128,
     use_cache=True, 
     return_routing_weights=True
 ):