mlcocdav/GreenHack

1import ast2import time3 4import nltk5import openai6import pandas as pd7from gptrim import trim8from numpy import dot9from numpy.linalg import norm10import os11from prompt_engineering import (NLP_TASKS, NLP_TASK_PROMPTS,12                                NLP_TASK_TEMPERATURES)13 14# Price for 1000 tokens15MODEL_PRICES = {16    "text-ada-001": 0.0004,17    "text-babbage-001": 0.0005,18    "gpt-3.5-turbo": 0.002,19    "text-curie-001": 0.002,20    "text-davinci-003": 0.02,21}22# % of the price23PRICE_MARGIN = 0.024 25openai.api_key = os.getenv("OPENAI_API_KEY")26 27 28def cos_sim(a, b):29    return dot(a, b) / (norm(a) * norm(b))30 31 32class PromptHandler():33    def __init__(self):34        self.prompt_history = pd.read_csv('prompt_history.csv')35 36    def generate(self, prompt: str, task, speed, quality, simplified):37        assert task in NLP_TASKS38 39        apis = APIs()40        df_subset = self.prompt_history[41            (self.prompt_history['speed'] == speed) &42            (self.prompt_history[43                 'quality'] == quality) &44            (self.prompt_history['task'] == task) &45            (~self.prompt_history['feedback'])]46        embedding = apis.get_embedding(prompt)47        if df_subset.shape[0] > 0:48            df_subset['prompt_embedding'] = df_subset[49                'prompt_embedding'].apply(lambda x: ast.literal_eval(x))50            df_subset['similarity'] = df_subset.apply(51                lambda row: cos_sim(row['prompt_embedding'], embedding),52                axis=1)53            best_model = \54            df_subset.sort_values(by=['similarity'])['model'].iloc[0]55        else:56            model_id = max(int(round((speed + quality) / 2, 0)) - 2, 0)57            print(model_id)58            best_model = list(MODEL_PRICES.keys())[model_id]59 60        if simplified:61            simplified_prompt = self.simplify_prompt(prompt)62            edited_prompt = self.simplify_prompt(63                f'{NLP_TASK_PROMPTS[task]}{simplified_prompt}')64            try:65                simplified_prompt_ratio = len(66                nltk.word_tokenize(simplified_prompt)) / len(67                nltk.word_tokenize(prompt))68            except:69                simplified_prompt_ratio = 170        else:71            edited_prompt = f'{NLP_TASK_PROMPTS[task]}{prompt}'72            simplified_prompt_ratio = 173            simplified_prompt = ''74 75        # set up timer76        start_time = time.time()77        response_text = apis.openai_prompt(edited_prompt, best_model,78                                           temperature=NLP_TASK_TEMPERATURES[79                                               task])80        response_text = response_text.strip()81        inference_time = round(time.time() - start_time, 2)82        print('Response: ', response_text)83        # save prompt to db84        new_row = {85            'prompt': prompt, 'prompt_embedding': embedding,86            'model': best_model,87            'result': response_text, 'task': task,88            'speed': speed, 'quality': quality, 'feedback': True}89        self.prompt_history = pd.concat(90            [self.prompt_history, pd.DataFrame([new_row])], ignore_index=True)91        self.prompt_history.to_csv('prompt_history.csv', index_label='ID')92        price = round(self.get_price(prompt, task, speed, quality,93                               model_name=best_model), 6)94        competitor_price = self.get_price(prompt, task, speed, quality)95        savings_percent = round(96            self.get_saved_amount(price, competitor_price,97                                  simplified_prompt_ratio), 2)98        saved_money = f'${round(price * simplified_prompt_ratio, 6)} saved {savings_percent} % '99        return response_text, inference_time, best_model, simplified_prompt, saved_money100 101    def get_price(self, prompt: str, task_type: str, speed: int, quality: int,102                  model_name=None) -> float:103        """Calculate price per inferences according to the inputs."""104        if model_name is None:105            model_id = int(round((speed + quality) / 2, 0)) - 1106            price_per_thousand = MODEL_PRICES[107                list(MODEL_PRICES.keys())[model_id]]108        else:109            price_per_thousand = MODEL_PRICES[model_name]110        return round(price_per_thousand * (1 + PRICE_MARGIN), 6)111 112    def get_price_dollars(self, prompt: str, task_type: str, speed: int,113                          quality: int,114                          model_name=None) -> float:115        """Calculate price per inferences according to the inputs."""116        if model_name is None:117            model_id = int(round((speed + quality) / 2, 0)) - 1118            model_name = list(MODEL_PRICES.keys())[model_id]119        return f'${self.get_price(prompt, task_type, speed, quality, model_name)} ({model_name})'120 121    def get_saved_amount(self, final_price: float, other_price: float,122                         simplified_promt_ratio: float) -> float:123        return 100 * (1 - (final_price * simplified_promt_ratio) / other_price)124 125    def simplify_prompt(self, promt: str) -> str:126        return trim(promt)127 128 129class APIs():130 131    def get_embedding(self, text, model="text-embedding-ada-002"):132        text = text.replace("\n", " ")133        if text == '':134            text = ' '135        return openai.Embedding.create(input=[text], model=model)['data'][0][136            'embedding']137 138    def openai_prompt(self, prompt, model, temperature):139        if model == "gpt-3.5-turbo":140            response = openai.ChatCompletion.create(141                model="gpt-3.5-turbo",142                messages=[143                    {"role": "system", "content": prompt}])144            return response['choices'][0]["message"]['content']145        else:146            response = openai.Completion.create(147                model=model,148                prompt=prompt,149                temperature=temperature,150                max_tokens=100,151                top_p=1,152                n=1,153                frequency_penalty=0.0,154                presence_penalty=0.0,155            )156            return response['choices'][0]['text']157