Update models.py

models.py

@@ -59,10 +59,8 @@ class OpenAIModel(BaseTCOModel):
         def on_model_change(model):
             
             if model == "GPT-4":
-                print("GPT4")
                 return gr.Dropdown.update(choices=["8K", "32K"])
             else:
-                print("GPT3.5")
                 return gr.Dropdown.update(choices=["4K", "16K"])
 
         self.model = gr.Dropdown(["GPT-4", "GPT-3.5 Turbo"], value="GPT-4",
@@ -89,7 +87,7 @@ class OpenAIModel(BaseTCOModel):
         else:
             cost_per_1k_input_tokens = 0.003
 
-        cost_per_output_token = cost_per_1k_input_tokens * 1000 / input_length
+        cost_per_output_token = cost_per_1k_input_tokens * input_length / 1000
 
         return cost_per_output_token
 
@@ -156,7 +154,63 @@ class OpenSourceLlama2Model(BaseTCOModel):
     def compute_cost_per_token(self, vm_cost_per_hour, tokens_per_second, maxed_out, used):
         cost_per_token = vm_cost_per_hour / (tokens_per_second * 3600 * maxed_out * used)
         return cost_per_token 
+
+class CohereModel(BaseTCOModel):
+    
+    def __init__(self):
+        self.set_name("(SaaS) Cohere")
+        self.set_formula(r"""$CT = \frac{CT\_1K \times 1000}{L}$  <br>
+                         with: <br>
+                         CT = Cost per output Token <br>
+                         CT_1M = Cost per one million Tokens (from Cohere's pricing web page) <br>
+                         L = Input Length 
+                         """)
+        super().__init__()
     
+    def render(self):
+        def on_use_case_change(use_case):
+            if use_case == "Summarize":
+                return gr.Dropdown.update(choices=["Default"])
+            else:
+                return gr.Dropdown.update(choices=["Default", "Custom"])
+            
+        self.use_case = gr.Dropdown(["Embed", "Generate", "Classify", "Summarize"], value="Generate",
+                                 label="Use case",
+                                 interactive=True, visible=False)
+        self.model = gr.Dropdown(["Default", "Custom"], value="Default",
+                                 label="Model",
+                                 interactive=True, visible=False)
+        self.use_case.change(on_use_case_change, inputs=self.use_case, outputs=self.model)        
+        self.input_length = gr.Number(350, label="Average number of input tokens", 
+                                      interactive=True, visible=False)
+
+    def compute_cost_per_token(self, use_case, model, input_length):
+        """Cost per token = """
+        use_case = use_case[0]        
+        model = model[0]
+
+        if use_case == "Embed":
+            if model == "Default":
+                cost_per_1M_input_tokens = 0.4
+            else: 
+                cost_per_1M_input_tokens = 0.8
+        elif use_case == "Generate":
+            if model == "Default":
+                cost_per_1M_input_tokens = 15
+            else: 
+                cost_per_1M_input_tokens = 30
+        elif use_case == "Classify":
+            if model == "Default":
+                cost_per_1M_input_tokens = 200
+            else: 
+                cost_per_1M_input_tokens = 200
+        else:
+            cost_per_1M_input_tokens = 15
+
+        cost_per_output_token = cost_per_1M_input_tokens * input_length / 1000000
+
+        return cost_per_output_token
+   
 class ModelPage:
     
     def __init__(self, Models: BaseTCOModel):