bkhmsi commited on
Commit
4e82a89
·
1 Parent(s): 5c65b5e
Files changed (1) hide show
  1. router_backend.py +2 -1
router_backend.py CHANGED
@@ -36,7 +36,8 @@ def get_expert_routing(model_id: str, hf_token: str, prompt: Union[str, List[Dic
36
  generation = None
37
  routing_weights = get_routing_weights(model, tokenizer, [prompt])
38
 
39
- model_routing_percentages = aggregate_routing_weights(routing_weights)
 
40
 
41
  if generation is not None:
42
  print(f"Generation:\n{generation}")
 
36
  generation = None
37
  routing_weights = get_routing_weights(model, tokenizer, [prompt])
38
 
39
+ model_routing_percentages = aggregate_routing_weights(routing_weights)[0]
40
+ print(model_routing_percentages)
41
 
42
  if generation is not None:
43
  print(f"Generation:\n{generation}")