QuantumIntelligenceLab/Llama3.1-Instruct-O1

1import gradio as gr2import openai3import time4import re5import os6 7# Available models8MODELS = [9    "Meta-Llama-3.1-405B-Instruct",10    "Meta-Llama-3.1-70B-Instruct",11    "Meta-Llama-3.1-8B-Instruct"12]13 14# Sambanova API base URL15API_BASE = "https://api.sambanova.ai/v1"16 17def create_client(api_key=None):18    """Creates an OpenAI client instance."""19    if api_key:20        openai.api_key = api_key21    else:22        openai.api_key = os.getenv("API_KEY")23 24    return openai.OpenAI(api_key=openai.api_key, base_url=API_BASE)25 26def chat_with_ai(message, chat_history, system_prompt):27    """Formats the chat history for the API call."""28    messages = [{"role": "system", "content": system_prompt}]29    for tup in chat_history:30        first_key = list(tup.keys())[0]  # First key31        last_key = list(tup.keys())[-1]   # Last key32        messages.append({"role": "user", "content": tup[first_key]})33        messages.append({"role": "assistant", "content": tup[last_key]})34    messages.append({"role": "user", "content": message})35    return messages36 37def respond(message, chat_history, model, system_prompt, thinking_budget, api_key):38    """Sends the message to the API and gets the response."""39    client = create_client(api_key)40    messages = chat_with_ai(message, chat_history, system_prompt.format(budget=thinking_budget))41    start_time = time.time()42 43    try:44        completion = client.chat.completions.create(model=model, messages=messages)45        response = completion.choices[0].message.content46        thinking_time = time.time() - start_time47        return response, thinking_time48    except Exception as e:49        error_message = f"Error: {str(e)}"50        return error_message, time.time() - start_time51 52def parse_response(response):53    """Parses the response from the API."""54    answer_match = re.search(r'<answer>(.*?)</answer>', response, re.DOTALL)55    reflection_match = re.search(r'<reflection>(.*?)</reflection>', response, re.DOTALL)56 57    answer = answer_match.group(1).strip() if answer_match else ""58    reflection = reflection_match.group(1).strip() if reflection_match else ""59    steps = re.findall(r'<step>(.*?)</step>', response, re.DOTALL)60 61    if answer == "":62        return response, "", ""63 64    return answer, reflection, steps65 66def generate(message, history, model, system_prompt, thinking_budget, api_key):67    """Generates the chatbot response."""68    response, thinking_time = respond(message, history, model, system_prompt, thinking_budget, api_key)69 70    if response.startswith("Error:"):71        return history + [({"role": "system", "content": response},)], ""72 73    answer, reflection, steps = parse_response(response)74 75    messages = []76    messages.append({"role": "user", "content": message})77 78    formatted_steps = [f"Step {i}: {step}" for i, step in enumerate(steps, 1)]79    all_steps = "\n".join(formatted_steps) + f"\n\nReflection: {reflection}"80 81    messages.append({"role": "assistant", "content": all_steps, "metadata": {"title": f"Thinking Time: {thinking_time:.2f} sec"}})82    messages.append({"role": "assistant", "content": answer})83 84    return history + messages, ""85 86# Define the default system prompt87DEFAULT_SYSTEM_PROMPT = """88You are a helpful assistant in normal conversation.89When given a problem to solve, you are an expert problem-solving assistant. 90Your task is to provide a detailed, step-by-step solution to a given question. 91Follow these instructions carefully:921. Read the given question carefully and reset counter between <count> and </count> to {budget}932. Generate a detailed, logical step-by-step solution.943. Enclose each step of your solution within <step> and </step> tags.954. You are allowed to use at most {budget} steps (starting budget), 96   keep track of it by counting down within tags <count> </count>, 97   STOP GENERATING MORE STEPS when hitting 0, you don't have to use all of them.985. Do a self-reflection when you are unsure about how to proceed, 99   based on the self-reflection and reward, decides whether you need to return 100   to the previous steps.1016. After completing the solution steps, reorganize and synthesize the steps 102   into the final answer within <answer> and </answer> tags.1037. Provide a critical, honest and subjective self-evaluation of your reasoning 104   process within <reflection> and </reflection> tags.1058. Assign a quality score to your solution as a float between 0.0 (lowest 106   quality) and 1.0 (highest quality), enclosed in <reward> and </reward> tags.107Example format:            108<count> [starting budget] </count>109<step> [Content of step 1] </step>110<count> [remaining budget] </count>111<step> [Content of step 2] </step>112<reflection> [Evaluation of the steps so far] </reflection>113<reward> [Float between 0.0 and 1.0] </reward>114<count> [remaining budget] </count>115<step> [Content of step 3 or Content of some previous step] </step>116<count> [remaining budget] </count>117...118<step>  [Content of final step] </step>119<count> [remaining budget] </count>120<answer> [Final Answer] </answer> (must give final answer in this format)121<reflection> [Evaluation of the solution] </reflection>122<reward> [Float between 0.0 and 1.0] </reward>123"""124 125with gr.Blocks() as demo:126    gr.Markdown("# Llama3.1-Instruct-O1")127    gr.Markdown("[Powered by SambaNova Cloud, Get Your API Key Here](https://cloud.sambanova.ai/apis)")128 129    with gr.Row():130        api_key = gr.Textbox(label="API Key", type="password", placeholder="(Optional) Enter your API key here for more availability")131 132    with gr.Row():133        model = gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[0])134        thinking_budget = gr.Slider(minimum=1, maximum=100, value=10, step=1, label="Thinking Budget", info="maximum times a model can think")135 136    chatbot = gr.Chatbot(label="Chat", show_label=False, show_share_button=False, show_copy_button=True, likeable=True, layout="panel", type="messages")137 138    msg = gr.Textbox(label="Type your message here...", placeholder="Enter your message...")139 140    gr.Button("Clear Chat").click(lambda: ([], ""), inputs=None, outputs=[chatbot, msg])141 142    system_prompt = gr.Textbox(label="System Prompt", value=DEFAULT_SYSTEM_PROMPT, lines=15, interactive=True)143 144    msg.submit(generate, inputs=[msg, chatbot, model, system_prompt, thinking_budget, api_key], outputs=[chatbot, msg])145 146demo.launch(share=True, show_api=False)