models.py

from gradio.components import Component
import gradio as gr
import pandas as pd
from abc import ABC, abstractclassmethod
import inspect

class BaseTCOModel(ABC):
    # TO DO: Find way to specify which component should be used for computing cost
    def __setattr__(self, name, value):
        if isinstance(value, Component):
            self._components.append(value)
        self.__dict__[name] = value

    def __init__(self):
        super(BaseTCOModel, self).__setattr__("_components", [])
        self.use_case = None  
        self.num_users = None
        self.input_tokens = None
        self.output_tokens = None

    def get_components(self) -> list[Component]:
        return self._components
    
    def get_components_for_cost_computing(self):
        return self.components_for_cost_computing
    
    def get_name(self):
        return self.name
    
    def register_components_for_cost_computing(self):
        args = inspect.getfullargspec(self.compute_cost_per_token)[0][1:]
        self.components_for_cost_computing = [self.__getattribute__(arg) for arg in args]
    
    @abstractclassmethod
    def compute_cost_per_token(self):
        pass
    
    @abstractclassmethod
    def render(self):
        pass
    
    def set_name(self, name):
        self.name = name
    
    def set_formula(self, formula):
        self.formula = formula
    
    def get_formula(self):
        return self.formula
    
    def set_latency(self, latency):
        self.latency = latency
    
    def get_latency(self):
        return self.latency

class OpenAIModel(BaseTCOModel):

    def __init__(self):
        self.set_name("(SaaS) OpenAI")
        self.set_formula(r"""$CR = \frac{CIT\_1K \times IT + COT\_1K \times OT}{1000}$  <br>
                         with: <br>
                         CR = Cost per Request <br>
                         CIT_1K = Cost per 1000 Input Tokens (from OpenAI's pricing web page) <br>
                         COT_1K = Cost per 1000 Output Tokens <br>
                         IT = Input Tokens <br>
                         OT = Output Tokens
                         """)
        self.latency = "15s" #Default value for GPT4
        super().__init__()

    def render(self):
        def on_model_change(model):
            
            if model == "GPT-4":
                self.latency = "15s"
                return gr.Dropdown.update(choices=["8K", "32K"])
            else:
                self.latency = "5s"
                return gr.Dropdown.update(choices=["4K", "16K"], value="4K")

        self.model = gr.Dropdown(["GPT-4", "GPT-3.5 Turbo"], value="GPT-4",
                                 label="OpenAI models",
                                 interactive=True, visible=False)
        self.context_length = gr.Dropdown(["8K", "32K"], value="8K", interactive=True,
                                          label="Context size",
                                          visible=False, info="Number of tokens the model considers when processing text")
        self.model.change(on_model_change, inputs=self.model, outputs=self.context_length)

    def compute_cost_per_token(self, model, context_length):
        """Cost per token = """

        if model == "GPT-4" and context_length == "8K":
            cost_per_1k_input_tokens = 0.03
            cost_per_1k_output_tokens = 0.06
        elif model == "GPT-4" and context_length == "32K":
            cost_per_1k_input_tokens = 0.06
            cost_per_1k_output_tokens = 0.12
        elif model == "GPT-3.5" and context_length == "4K":
            cost_per_1k_input_tokens = 0.0015
            cost_per_1k_output_tokens = 0.002
        else:
            cost_per_1k_input_tokens = 0.003
            cost_per_1k_output_tokens = 0.004
        cost_per_input_token = (cost_per_1k_input_tokens / 1000) 
        cost_per_output_token = (cost_per_1k_output_tokens / 1000)

        return cost_per_input_token, cost_per_output_token

class OpenSourceLlama2Model(BaseTCOModel):
    
    def __init__(self):
        self.set_name("(Open source) Llama 2")
        self.set_formula(r"""$CT = \frac{VM\_CH \times 100}{3600 \times U} \times (\frac{IT}{ITS} + \frac{OT}{OTS})$<br>
                         with: <br>
                         CT = Cost per Token <br>
                         VM_CH = VM Cost per Hour <br>
                         ITS = Input Tokens per Second <br>
                         OTS = Output Tokens per Second <br>
                         U = Used <br>
                         IT = Input Tokens <br>
                         OT = Output Tokens
                         """)
        self.set_latency("27s")
        super().__init__()
    
    def render(self):
        vm_choices = ["1x Nvidia A100 (Azure NC24ads A100 v4)",
                      "2x Nvidia A100 (Azure NC24ads A100 v4)", 
                      "2x Nvidia A100 (Azure ND96amsr A100 v4)"]
        
        def on_model_change(model):
            if model == "Llama 2 7B":
                return [gr.Dropdown.update(choices=vm_choices), 
                        gr.Markdown.update(value="To see the benchmark results use for the Llama2-7B model, [click here](https://example.com/script)"), 
                        gr.Number.update(value=3.6730),
                        gr.Number.update(value=694.38),
                        gr.Number.update(value=694.38),
                        ]
            else:
                not_supported_vm = ["1x Nvidia A100 (Azure NC24ads A100 v4)", "2x Nvidia A100 (Azure NC24ads A100 v4)"]
                choices = [x for x in vm_choices if x not in not_supported_vm]
                return [gr.Dropdown.update(choices=choices, value="2x Nvidia A100 (Azure ND96amsr A100 v4)"), 
                        gr.Markdown.update(value="To see the benchmark results used for the Llama2-70B model, [click here](https://www.cursor.so/blog/llama-inference#user-content-fn-llama-paper)"),
                        gr.Number.update(value=2*37.186),
                        gr.Number.update(value=2860),
                        gr.Number.update(value=18.545),
                        ]

        def on_vm_change(model, vm):
            # TO DO: load info from CSV
            if model == "Llama 2 7B" and vm == "1x Nvidia A100 (Azure NC24ads A100 v4)":
                return [gr.Number.update(value=4.777), gr.Number.update(value=694.38), gr.Number.update(value=694.38)]
            elif model == "Llama 2 7B" and vm == "2x Nvidia A100 (Azure NC24ads A100 v4)":
                return [gr.Number.update(value=2*4.777), gr.Number.update(value=1388.76), gr.Number.update(value=1388.76)]
            elif model == "Llama 2 7B" and vm == "2x Nvidia A100 (Azure ND96amsr A100 v4)":
                return [gr.Number.update(value=2*37.186), gr.Number.update(value=2777.52), gr.Number.update(value=2777.52)]
            elif model == "Llama 2 70B" and vm == "2x Nvidia A100 (Azure ND96amsr A100 v4)":
                return [gr.Number.update(value=2*37.186), gr.Number.update(value=2860), gr.Number.update(value=18.449)]
        
        self.model = gr.Dropdown(["Llama 2 7B", "Llama 2 70B"], value="Llama 2 70B", label="OpenSource models", visible=False)
        self.vm = gr.Dropdown(choices=["2x Nvidia A100 (Azure ND96amsr A100 v4)"],
                              value="2x Nvidia A100 (Azure ND96amsr A100 v4)", 
                              visible=False,
                              label="Instance of VM with GPU",
                              info="Your options for this choice depend on the model you previously chose"
                              )
        self.vm_cost_per_hour = gr.Number(2*37.186, label="VM instance cost per hour", 
                                      interactive=False, visible=False)
        self.input_tokens_per_second = gr.Number(2860, visible=False,
                                           label="Number of output tokens per second for this specific model and VM instance",
                                           interactive=False
                                           )
        self.output_tokens_per_second = gr.Number(18.449, visible=False,
                                           label="Number of output tokens per second for this specific model and VM instance",
                                           interactive=False
                                           )
        self.info = gr.Markdown("To see the script used to benchmark the Llama2-70B model, [click here](https://www.cursor.so/blog/llama-inference#user-content-fn-llama-paper)", interactive=False, visible=False)
        
        self.model.change(on_model_change, inputs=self.model, outputs=[self.vm, self.info, self.vm_cost_per_hour, self.input_tokens_per_second, self.output_tokens_per_second])
        self.vm.change(on_vm_change, inputs=[self.model, self.vm], outputs=[self.vm_cost_per_hour, self.input_tokens_per_second, self.output_tokens_per_second])
        self.used = gr.Slider(minimum=0.01, value=30., step=0.01, label="% used", 
                                   info="Percentage of time the GPU is used",
                                   interactive=True,
                                   visible=False)

    def compute_cost_per_token(self, vm_cost_per_hour, input_tokens_per_second, output_tokens_per_second, used):
        cost_per_input_token =  vm_cost_per_hour * 100 / (3600 * used * input_tokens_per_second) 
        cost_per_output_token =  vm_cost_per_hour * 100 / (3600 * used * output_tokens_per_second) 
        return cost_per_input_token,  cost_per_output_token

class OpenSourceDIY(BaseTCOModel):
    
    def __init__(self):
        self.set_name("(Open source) DIY")
        self.set_formula(r"""$CT = \frac{VM\_CH \times 100}{3600 \times U} \times (\frac{IT}{ITS} + \frac{OT}{OTS})$<br>
                         with: <br>
                         CT = Cost per Token <br>
                         VM_CH = VM Cost per Hour <br>
                         ITS = Input Tokens per Second <br>
                         OTS = Output Tokens per Second <br>
                         U = Used <br>
                         IT = Input Tokens <br>
                         OT = Output Tokens
                         """)
        self.set_latency("The latency can't be estimated in the DIY scenario for the model isn't defined")
        super().__init__()
    
    def render(self):   
        self.info = gr.Markdown("Compute the cost/token based on our formula below, using your own parameters", visible=False)
        self.display_formula = gr.Markdown(r"""$CT = \frac{VM\_CH \times 100}{3600 \times U} \times (\frac{IT}{ITS} + \frac{OT}{OTS})$<br>
                         with: <br>
                         CT = Cost per Token <br>
                         VM_CH = VM Cost per Hour <br>
                         ITS = Input Tokens per Second <br>
                         OTS = Output Tokens per Second <br>
                         U = Used <br>
                         IT = Input Tokens <br>
                         OT = Output Tokens
                         """, visible=False)
        self.vm_cost_per_hour = gr.Number(3.5, label="VM instance cost per hour", 
                                      interactive=True, visible=False)
        self.input_tokens_per_second = gr.Number(300, visible=False,
                                           label="Number of input tokens per second processed for this specific model and VM instance",
                                           interactive=True
                                           )
        self.output_tokens_per_second = gr.Number(300, visible=False,
                                           label="Number of output tokens per second processed for this specific model and VM instance",
                                           interactive=True
                                           )
        self.used = gr.Slider(minimum=0.01, value=50., step=0.01, label="% used", 
                                   info="Percentage of time the GPU is used",
                                   interactive=True,
                                   visible=False)

    def compute_cost_per_token(self, vm_cost_per_hour, input_tokens_per_second, output_tokens_per_second, used):
        cost_per_input_token =  vm_cost_per_hour * 100 / (3600 * used * input_tokens_per_second) 
        cost_per_output_token =  vm_cost_per_hour * 100 / (3600 * used * output_tokens_per_second) 
        return cost_per_input_token, cost_per_output_token  

class CohereModel(BaseTCOModel):
    
    def __init__(self):
        self.set_name("(SaaS) Cohere")
        self.set_formula(r"""$CR = \frac{CT\_1M \times (IT + OT)}{1000000}$  <br>
                         with: <br>
                         CR = Cost per Request <br>
                         CT_1M = Cost per one million Tokens (from Cohere's pricing web page) <br>
                         IT = Input Tokens <br>
                         OT = Output Tokens 
                         """)
        self.set_latency("")
        super().__init__()
    
    def render(self):
        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"])

    def compute_cost_per_token(self, model):
        """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

        return cost_per_input_token, cost_per_output_token

class ModelPage:
    
    def __init__(self, Models: BaseTCOModel):
        self.models: list[BaseTCOModel] = []
        for Model in Models:
            model = Model()
            self.models.append(model)

    def render(self):
        for model in self.models:
            model.render()
            model.register_components_for_cost_computing() 

    def get_all_components(self) -> list[Component]:
        output = []
        for model in self.models:
            output += model.get_components()
        return output
    
    def get_all_components_for_cost_computing(self) -> list[Component]:
        output = []
        for model in self.models:
            output += model.get_components_for_cost_computing()
        return output

    def make_model_visible(self, name:str, use_case: gr.Dropdown, num_users: gr.Number, input_tokens: gr.Slider, output_tokens: gr.Slider):
        # First decide which indexes
        output = []
        for model in self.models:
            if model.get_name() == name:
                output+= [gr.update(visible=True)] * len(model.get_components()) 
                # Set use_case and num_users values in the model
                model.use_case = use_case
                model.num_users = num_users
                model.input_tokens = input_tokens
                model.output_tokens = output_tokens
            else:
                output+= [gr.update(visible=False)] * len(model.get_components())
        return output
    
    def compute_cost_per_token(self, *args):
        begin=0
        current_model = args[-1]
        for model in self.models:
            model_n_args = len(model.get_components_for_cost_computing())
            if current_model == model.get_name():
                
                model_args = args[begin:begin+model_n_args]
                cost_per_input_token, cost_per_output_token = model.compute_cost_per_token(*model_args)
                model_tco = cost_per_input_token * model.input_tokens + cost_per_output_token * model.output_tokens
                formula = model.get_formula()
                latency = model.get_latency()
                
                return f"Model {current_model} has a cost/request of: ${model_tco}", model_tco, formula, f"The average latency of this model is {latency}"
            
            begin = begin+model_n_args