add gpt-4

app.py

@@ -5,12 +5,22 @@ import openai
 import tiktoken
 
 # Config
-MODEL = "gpt-3.5-turbo"
-TOKEN_PRICE = 0.002 / 1000
+CHAT_MODELS = [
+    "gpt-4",
+    "gpt-4-32k",
+    "gpt-3.5-turbo",
+]
+DEFAULT_MODEL = "gpt-3.5-turbo"
+TOKEN_PRICE = {
+    "gpt-3.5-turbo": 0.002 / 1000,
+    "gpt-4": 0.06 / 1000,
+    "gpt-4-32k": 0.12 / 1000,
+}
 
 # Private
 _tokens: int = 0
-_encoding = tiktoken.encoding_for_model(MODEL)
+_selected_model = DEFAULT_MODEL
+_encoding = tiktoken.encoding_for_model(DEFAULT_MODEL)
 _api_key = os.environ.get("OPENAI_API_KEY", "")
 _show_openai_settings = _api_key == ""
 
@@ -20,18 +30,25 @@ def count_tokens(prompt_text: str) -> int:
     return len(_encoding.encode(str(prompt_text)))
 
 
-def get_cost(tokens: int) -> float:
-    return TOKEN_PRICE * tokens
+def get_cost(tokens: int, model: str) -> float:
+    return TOKEN_PRICE.get(model, 0) * tokens
 
 
-def prompt(prompt_text: str, api_key: str) -> str:
+def prompt(prompt_text: str, api_key: str, model: str) -> str:
     global _tokens
+    global _selected_model
+    global _encoding
+
+    # Check if selected model has changed
+    if model != _selected_model:
+        _selected_model = model
+        _encoding = tiktoken.encoding_for_model(_selected_model)
 
     _tokens += count_tokens(prompt_text)
 
     openai.api_key = api_key
     content = openai.ChatCompletion.create(
-        model=MODEL,
+        model=_selected_model,
         messages=[
             {"role": "system", "content": "You are a helpful assistant."},
             {"role": "user", "content": prompt_text},
@@ -40,7 +57,7 @@ def prompt(prompt_text: str, api_key: str) -> str:
     )["choices"][0]["message"]["content"]
 
     _tokens += count_tokens(content)
-    cost = get_cost(_tokens)
+    cost = get_cost(_tokens, model)
 
     return (
         content,
@@ -60,6 +77,11 @@ with gr.Blocks() as demo:
                 type="password",
                 placeholder="sk-...",
             )
+            dd_model = gr.Dropdown(
+                choices=CHAT_MODELS,
+                value=DEFAULT_MODEL,
+                label="Model",
+            )
             num_tokens = gr.Number(value=0, label="Tokens used")
             num_cost = gr.Number(value=0, label="Estimated cost in $", precision=3)
         with gr.Column(scale=3):
@@ -68,7 +90,7 @@ with gr.Blocks() as demo:
             with gr.Row():
                 gr.Button("Prompt").click(
                     fn=prompt,
-                    inputs=[txt_prompt, txt_api_key],
+                    inputs=[txt_prompt, txt_api_key, dd_model],
                     outputs=[txt_response, num_tokens, num_cost],
                 )
                 gr.Button("Clear").click(