Update README.md

README.md

@@ -44,6 +44,40 @@ As is, CARROT supports routing to the following collection of large language mod
 </p>
 
 ```python 
-your_code = do_some_stuff
-```
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import numpy as np
+
+token = 'YOUR HF TOKEN'
+
+COSTS = {'aws-claude-3-5-sonnet-v1':[3, 15], 'aws-titan-text-premier-v1': [.5, 1.5], 'openai-gpt-4o': [2.5, 10], 'openai-gpt-4o-mini': [0.15,0.6], 'wxai-granite-3-2b-instruct-8k-max-tokens':[0.1, 0.1], 'wxai-granite-3-8b-instruct-8k-max-tokens':[0.2, 0.2], 'wxai-llama-3-1-70b-instruct':[.9,.9], 'wxai-llama-3-1-8b-instruct': [.2,.2], 'wxai-llama-3-2-1b-instruct':[.06,.06], 'wxai-llama-3-2-3b-instruct':[.06, .06],'wxai-llama-3-3-70b-instruct':[.9,.9], 'wxai-mixtral-8x7b-instruct-v01':[.6,.6], 'wxai-llama-3-405b-instruct':[3.5,  3.5]}
 
+input_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
+```