base chat

app.py

@@ -1,16 +1,16 @@
 import gradio as gr 
-from openai import OpenAI
 import os
 from typing import List
 import logging
 import urllib.request
+from utils import model_name_mapping, urial_template, openai_base_request, DEFAULT_API_KEY
+from constant import js_code_label, HEADER_MD
 
 # add logging info to console 
 logging.basicConfig(level=logging.INFO)
 
 
-BASE_URL = "https://api.together.xyz/v1"
-DEFAULT_API_KEY = os.getenv("TOGETHER_API_KEY") 
+
 URIAL_VERSION = "inst_1k_v4.help"
 
 URIAL_URL = f"https://raw.githubusercontent.com/Re-Align/URIAL/main/urial_prompts/{URIAL_VERSION}.txt"
@@ -18,55 +18,6 @@ urial_prompt = urllib.request.urlopen(URIAL_URL).read().decode('utf-8')
 urial_prompt = urial_prompt.replace("```", '"""') # new version of URIAL uses """ instead of ```
 STOP_STRS = ['"""', '# Query:', '# Answer:']
 
-def urial_template(urial_prompt, history, message):
-    current_prompt = urial_prompt + "\n"
-    for user_msg, ai_msg in history:
-        current_prompt += f'# Query:\n"""\n{user_msg}\n"""\n\n# Answer:\n"""\n{ai_msg}\n"""\n\n'
-    current_prompt += f'# Query:\n"""\n{message}\n"""\n\n# Answer:\n"""\n'
-    return current_prompt
-    
-
-
-
-def openai_base_request(
-    model: str=None, 
-    temperature: float=0,
-    max_tokens: int=512,
-    top_p: float=1.0, 
-    prompt: str=None,
-    n: int=1, 
-    repetition_penalty: float=1.0,
-    stop: List[str]=None, 
-    api_key: str=None,
-    ):  
-    if api_key is None:
-        api_key = DEFAULT_API_KEY
-    client = OpenAI(api_key=api_key, base_url=BASE_URL)
-    # print(f"Requesting chat completion from OpenAI API with model {model}")
-    logging.info(f"Requesting chat completion from OpenAI API with model {model}")
-    logging.info(f"Prompt: {prompt}")
-    logging.info(f"Temperature: {temperature}")
-    logging.info(f"Max tokens: {max_tokens}")
-    logging.info(f"Top-p: {top_p}")
-    logging.info(f"Repetition penalty: {repetition_penalty}")
-    logging.info(f"Stop: {stop}")
-
-    request = client.completions.create(
-        model=model, 
-        prompt=prompt,
-        temperature=float(temperature),
-        max_tokens=int(max_tokens),
-        top_p=float(top_p),
-        n=n,
-        extra_body={'repetition_penalty': float(repetition_penalty)},
-        stop=stop, 
-        stream=True
-    ) 
-    
-    return request 
-
-
-
 
 def respond(
     message,
@@ -81,29 +32,9 @@ def respond(
     global STOP_STRS, urial_prompt
     rp = 1.0
     prompt = urial_template(urial_prompt, history, message)
-    if model_name == "Llama-3-8B":
-        _model_name = "meta-llama/Llama-3-8b-hf"
-    elif model_name == "Llama-3-70B":
-        _model_name = "meta-llama/Llama-3-70b-hf"
-    elif model_name == "Llama-2-7B":
-        _model_name = "meta-llama/Llama-2-7b-hf"
-    elif model_name == "Llama-2-70B":
-        _model_name = "meta-llama/Llama-2-70b-hf"
-    elif model_name == "Mistral-7B-v0.1":
-        _model_name = "mistralai/Mistral-7B-v0.1"
-    elif model_name == "Mixtral-8x22B":
-        _model_name = "mistralai/Mixtral-8x22B"
-    elif model_name == "Qwen1.5-72B":
-        _model_name = "Qwen/Qwen1.5-72B"
-    elif model_name == "Yi-34B":
-        _model_name = "zero-one-ai/Yi-34B"
-    elif model_name == "Yi-6B":
-        _model_name = "zero-one-ai/Yi-6B"
-    elif model_name == "OLMO":
-        _model_name = "allenai/OLMo-7B"
-    else:
-        raise ValueError("Invalid model name")
+    
     # _model_name = "meta-llama/Llama-3-8b-hf"
+    _model_name = model_name_mapping(model_name)
 
     if together_api_key and len(together_api_key) == 64:
         api_key = together_api_key
@@ -116,7 +47,6 @@ def respond(
                                    top_p=top_p, 
                                    repetition_penalty=rp,
                                    stop=STOP_STRS, api_key=api_key)  
-    
     response = ""
     for msg in request:
         # print(msg.choices[0].delta.keys())
@@ -135,43 +65,10 @@ def respond(
             response = response[:-2]
         yield response
  
-js_code_label = """
-function addApiKeyLink() {
-    // Select the div with id 'api_key'
-    const apiKeyDiv = document.getElementById('api_key');
-
-    // Find the span within that div with data-testid 'block-info'
-    const blockInfoSpan = apiKeyDiv.querySelector('span[data-testid="block-info"]');
-
-    // Create the new link element
-    const newLink = document.createElement('a');
-    newLink.href = 'https://api.together.ai/settings/api-keys';
-    newLink.textContent = ' View your keys here.';
-    newLink.target = '_blank'; // Open link in new tab
-    newLink.style = 'color: #007bff; text-decoration: underline;';
-
-    // Create the additional text
-    const additionalText = document.createTextNode(' (new account will have free credits to use.)');
-
-    // Append the link and additional text to the span
-    if (blockInfoSpan) {
-        // add a br 
-        apiKeyDiv.appendChild(document.createElement('br'));
-        apiKeyDiv.appendChild(newLink);
-        apiKeyDiv.appendChild(additionalText);
-    } else {
-        console.error('Span with data-testid "block-info" not found');
-    }
-}
-"""
 with gr.Blocks(gr.themes.Soft(), js=js_code_label) as demo:
     with gr.Row():
         with gr.Column():
-            gr.Markdown("""# 💬 BaseChat: Chat with Base LLMs with URIAL
-                        [Paper](https://arxiv.org/abs/2312.01552) | [Website](https://allenai.github.io/re-align/) | [GitHub](https://github.com/Re-Align/urial) | Contact: [Yuchen Lin](https://yuchenlin.xyz/)
-
-                        **Talk with __BASE__ LLMs which are not fine-tuned at all.**
-                        """)
+            gr.Markdown(HEADER_MD)
             model_name = gr.Radio(["Llama-3-8B", "Llama-3-70B", "Mistral-7B-v0.1", 
                                    "Mixtral-8x22B", "Yi-6B", "Yi-34B", "Llama-2-7B", "Llama-2-70B", "OLMO"]
                                   , value="Llama-3-8B", label="Base LLM name")
@@ -181,12 +78,8 @@ with gr.Blocks(gr.themes.Soft(), js=js_code_label) as demo:
                 with gr.Row():
                     max_tokens = gr.Textbox(value=256, label="Max tokens")
                     temperature = gr.Textbox(value=0.5, label="Temperature")
-            # with gr.Column():
-            #     with gr.Row():
                     top_p = gr.Textbox(value=0.9, label="Top-p")
                     rp = gr.Textbox(value=1.1, label="Repetition penalty")
-            
-
     chat = gr.ChatInterface(
         respond,
         additional_inputs=[max_tokens, temperature, top_p, rp, model_name, together_api_key],

constant.py

@@ -0,0 +1,35 @@
+HEADER_MD = """# 💬 BaseChat: Chat with Base LLMs with URIAL
+[Paper](https://arxiv.org/abs/2312.01552) | [Website](https://allenai.github.io/re-align/) | [GitHub](https://github.com/Re-Align/urial) | Contact: [Yuchen Lin](https://yuchenlin.xyz/)
+
+**Talk with __BASE__ LLMs which are not fine-tuned at all.**
+"""
+
+js_code_label = """
+function addApiKeyLink() {
+    // Select the div with id 'api_key'
+    const apiKeyDiv = document.getElementById('api_key');
+
+    // Find the span within that div with data-testid 'block-info'
+    const blockInfoSpan = apiKeyDiv.querySelector('span[data-testid="block-info"]');
+
+    // Create the new link element
+    const newLink = document.createElement('a');
+    newLink.href = 'https://api.together.ai/settings/api-keys';
+    newLink.textContent = ' View your keys here.';
+    newLink.target = '_blank'; // Open link in new tab
+    newLink.style = 'color: #007bff; text-decoration: underline;';
+
+    // Create the additional text
+    const additionalText = document.createTextNode(' (new account will have free credits to use.)');
+
+    // Append the link and additional text to the span
+    if (blockInfoSpan) {
+        // add a br 
+        apiKeyDiv.appendChild(document.createElement('br'));
+        apiKeyDiv.appendChild(newLink);
+        apiKeyDiv.appendChild(additionalText);
+    } else {
+        console.error('Span with data-testid "block-info" not found');
+    }
+}
+"""

utils.py

@@ -0,0 +1,80 @@
+
+from openai import OpenAI
+import logging
+from typing import List
+import os 
+
+BASE_URL = "https://api.together.xyz/v1"
+DEFAULT_API_KEY = os.getenv("TOGETHER_API_KEY") 
+
+def model_name_mapping(model_name):
+    if model_name == "Llama-3-8B":
+        _model_name = "meta-llama/Llama-3-8b-hf"
+    elif model_name == "Llama-3-70B":
+        _model_name = "meta-llama/Llama-3-70b-hf"
+    elif model_name == "Llama-2-7B":
+        _model_name = "meta-llama/Llama-2-7b-hf"
+    elif model_name == "Llama-2-70B":
+        _model_name = "meta-llama/Llama-2-70b-hf"
+    elif model_name == "Mistral-7B-v0.1":
+        _model_name = "mistralai/Mistral-7B-v0.1"
+    elif model_name == "Mixtral-8x22B":
+        _model_name = "mistralai/Mixtral-8x22B"
+    elif model_name == "Qwen1.5-72B":
+        _model_name = "Qwen/Qwen1.5-72B"
+    elif model_name == "Yi-34B":
+        _model_name = "zero-one-ai/Yi-34B"
+    elif model_name == "Yi-6B":
+        _model_name = "zero-one-ai/Yi-6B"
+    elif model_name == "OLMO":
+        _model_name = "allenai/OLMo-7B"
+    else:
+        raise ValueError("Invalid model name")
+    return _model_name
+
+
+def urial_template(urial_prompt, history, message):
+    current_prompt = urial_prompt + "\n"
+    for user_msg, ai_msg in history:
+        current_prompt += f'# Query:\n"""\n{user_msg}\n"""\n\n# Answer:\n"""\n{ai_msg}\n"""\n\n'
+    current_prompt += f'# Query:\n"""\n{message}\n"""\n\n# Answer:\n"""\n'
+    return current_prompt
+    
+
+def openai_base_request(
+    model: str=None, 
+    temperature: float=0,
+    max_tokens: int=512,
+    top_p: float=1.0, 
+    prompt: str=None,
+    n: int=1, 
+    repetition_penalty: float=1.0,
+    stop: List[str]=None, 
+    api_key: str=None,
+    ):  
+    if api_key is None:
+        api_key = DEFAULT_API_KEY
+    client = OpenAI(api_key=api_key, base_url=BASE_URL)
+    # print(f"Requesting chat completion from OpenAI API with model {model}")
+    logging.info(f"Requesting chat completion from OpenAI API with model {model}")
+    logging.info(f"Prompt: {prompt}")
+    logging.info(f"Temperature: {temperature}")
+    logging.info(f"Max tokens: {max_tokens}")
+    logging.info(f"Top-p: {top_p}")
+    logging.info(f"Repetition penalty: {repetition_penalty}")
+    logging.info(f"Stop: {stop}")
+
+    request = client.completions.create(
+        model=model, 
+        prompt=prompt,
+        temperature=float(temperature),
+        max_tokens=int(max_tokens),
+        top_p=float(top_p),
+        n=n,
+        extra_body={'repetition_penalty': float(repetition_penalty)},
+        stop=stop, 
+        stream=True
+    ) 
+    
+    return request 
+