Update README.md

README.md

@@ -46,6 +46,34 @@ As is, CARROT supports routing to the following collection of large language mod
 ```python 
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import numpy as np
-
 token = 'YOUR HF TOKEN'
+nput_counter = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-70B", token='')
+
+tokenizer = AutoTokenizer.from_pretrained('roberta-base')
+
+score_predictor =  AutoModelForSequenceClassification.from_pretrained('CARROT-LLM-Routing/Performance', 
+                                                           problem_type="multi_label_classification", 
+                                                           num_labels=len(COSTS),
+                                                           )
+output_counter = AutoModelForSequenceClassification.from_pretrained('CARROT-LLM-Routing/Cost', 
+                                                           problem_type="regression", 
+                                                           num_labels=len(COSTS))
+def CARROT(prompts, mu, input_counter=input_counter, predictors = [score_predictor, output_counter], tokenizer=tokenizer, costs=COSTS):
+  tokenized_text = tokenizer(prompts,
+                             truncation=True,
+                             padding=True,
+                             is_split_into_words=False,
+                             return_tensors='pt')
+  input_counter.pad_token = tokenizer.eos_token
+  scores = 1/(1+np.exp(-predictors[1](tokenized_text["input_ids"]).logits.detach().numpy()))
+  output_tokens = predictors[1](tokenized_text["input_ids"]).logits.detach().numpy()
+  input_tokens = [input_counter(prompt, return_tensors="pt")["input_ids"].shape[1] for prompt in prompts]
+  input_tokens = np.array(input_tokens).T
+  costs = []  
+  for i, m in enumerate(COSTS.keys()):
+      costs.append((input_tokens*COSTS[m][0]/(1000000)+output_tokens[:,i]*COSTS[m][1]/1000).tolist())
+  costs = np.array(costs).T
+  model_idx = ((1 - mu) * scores - mu * costs*100 ).argmax(axis = 1, keepdims = True)
+  called = [id2label[idx[0]] for idx in model_idx]
+  return called
 ```