Update models.py

models.py

@@ -155,38 +155,37 @@ class CohereModel(BaseTCOModel):
         super().__init__()
     
     def render(self):
+        def on_model_change(model):
+            if model == "Default":
+                cost_per_1M_tokens = 15
+            else: 
+                cost_per_1M_tokens = 30
+            cost_per_1K_tokens = cost_per_1M_tokens / 1000
+            return gr.update(value=cost_per_1K_tokens), gr.update(value=cost_per_1K_tokens)
+        
         self.model = gr.Dropdown(["Default", "Custom"], value="Default",
                                  label="Model",
                                  interactive=True, visible=False)
-        if self.use_case == "Summarize":
-            self.model:  gr.Dropdown.update(choices=["Default"])
-        elif self.use_case == "Question-answering":
-            self.model: gr.Dropdown.update(choices=["Default", "Custom"])
-        else: 
-            self.model: gr.Dropdown.update(choices=["Default", "Custom"])
+        self.input_tokens_cost_per_second = gr.Number(0.015, visible=False,
+                                           label="($) Price/1K input prompt tokens",
+                                           interactive=False
+                                           )
+        self.output_tokens_cost_per_second = gr.Number(0.015, visible=False,
+                                           label="($) Price/1K output prompt tokens",
+                                           interactive=False
+                                           )
         self.info = gr.Markdown("The cost per input and output tokens value is from Cohere's [pricing web page](https://cohere.com/pricing?utm_term=&utm_campaign=Cohere+Brand+%26+Industry+Terms&utm_source=adwords&utm_medium=ppc&hsa_acc=4946693046&hsa_cam=20368816223&hsa_grp=154209120409&hsa_ad=666081801359&hsa_src=g&hsa_tgt=dsa-19959388920&hsa_kw=&hsa_mt=&hsa_net=adwords&hsa_ver=3&gad=1&gclid=CjwKCAjww7KmBhAyEiwA5-PUSlyO7pq0zxeVrhViXMd8WuILW6uY-cfP1-SVuUfs-leUAz14xHlOHxoCmfkQAvD_BwE)", interactive=False, visible=False)
-
+        self.model.change(on_model_change, inputs=self.model, outputs=[self.input_tokens_cost_per_second, self.output_tokens_cost_per_second])
         self.labor = gr.Number(0, visible=False, 
                                 label="($) Labor cost per month", 
                                 info="This is an estimate of the labor cost of the AI engineer in charge of deploying the model",
                                 interactive=True
                                 )
 
-    def compute_cost_per_token(self, model, labor):
-        """Cost per token = """
-        use_case = self.use_case  
-
-        if use_case == "Generate":
-            if model == "Default":
-                cost_per_1M_tokens = 15
-            else: 
-                cost_per_1M_tokens = 30
-        elif use_case == "Summarize":
-            cost_per_1M_tokens = 15
-        else: 
-            cost_per_1M_tokens = 200
-        cost_per_input_token = cost_per_1M_tokens / 1000000
-        cost_per_output_token = cost_per_1M_tokens / 1000000
+    def compute_cost_per_token(self, input_tokens_cost_per_second, output_tokens_cost_per_second, labor):
+        
+        cost_per_input_token = input_tokens_cost_per_second / 1000
+        cost_per_output_token = output_tokens_cost_per_second / 1000
 
         return cost_per_input_token, cost_per_output_token, labor