Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from openai import OpenAI | |
| import time | |
| import re | |
| # Available models | |
| MODELS = [ | |
| "Meta-Llama-3.1-405B-Instruct", | |
| "Meta-Llama-3.1-70B-Instruct", | |
| "Meta-Llama-3.1-8B-Instruct" | |
| ] | |
| def create_client(api_key, base_url): | |
| return OpenAI( | |
| api_key=api_key, | |
| base_url=base_url | |
| ) | |
| def chat_with_ai(message, chat_history, system_prompt): | |
| messages = [ | |
| {"role": "system", "content": system_prompt}, | |
| ] | |
| for human, ai in chat_history: | |
| messages.append({"role": "user", "content": human}) | |
| messages.append({"role": "assistant", "content": ai}) | |
| messages.append({"role": "user", "content": message}) | |
| return messages | |
| def respond(message, chat_history, model, system_prompt, thinking_budget, api_key, base_url): | |
| client = create_client(api_key, base_url) | |
| messages = chat_with_ai(message, chat_history, system_prompt.format(budget=thinking_budget)) | |
| response = "" | |
| start_time = time.time() | |
| try: | |
| for chunk in client.chat.completions.create( | |
| model=model, | |
| messages=messages, | |
| stream=True | |
| ): | |
| content = chunk.choices[0].delta.content or "" | |
| response += content | |
| yield response, time.time() - start_time | |
| except Exception as e: | |
| yield f"Error: {str(e)}", time.time() - start_time | |
| def parse_response(response): | |
| answer_match = re.search(r'<answer>(.*?)</answer>', response, re.DOTALL) | |
| reflection_match = re.search(r'<reflection>(.*?)</reflection>', response, re.DOTALL) | |
| answer = answer_match.group(1).strip() if answer_match else "" | |
| reflection = reflection_match.group(1).strip() if reflection_match else "" | |
| steps = re.findall(r'<step>(.*?)</step>', response, re.DOTALL) | |
| return answer, reflection, steps | |
| def process_chat(message, history, model, system_prompt, thinking_budget, api_key, base_url): | |
| if not api_key or not base_url: | |
| history.append((message, "Please provide both API Key and Base URL before starting the chat.")) | |
| return history, history | |
| full_response = "" | |
| thinking_time = 0 | |
| for response, elapsed_time in respond(message, history, model, system_prompt, thinking_budget, api_key, base_url): | |
| full_response = response | |
| thinking_time = elapsed_time | |
| if full_response.startswith("Error:"): | |
| history.append((message, full_response)) | |
| return history, history | |
| answer, reflection, steps = parse_response(full_response) | |
| formatted_response = f"**Answer:** {answer}\n\n**Reflection:** {reflection}\n\n**Thinking Steps:**\n" | |
| for i, step in enumerate(steps, 1): | |
| formatted_response += f"**Step {i}:** {step}\n" | |
| formatted_response += f"\n**Thinking time:** {thinking_time:.2f} s" | |
| history.append((message, formatted_response)) | |
| return history, history | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# Llama3.1-Instruct-O1") | |
| gr.Markdown("[Powered by Llama3.1 models through SN Cloud](https://sambanova.ai/fast-api?api_ref=907266)") | |
| with gr.Row(): | |
| api_key = gr.Textbox(label="API Key", type="password") | |
| base_url = gr.Textbox(label="Base URL", value="https://api.endpoints.anyscale.com/v1") | |
| with gr.Row(): | |
| model = gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[0]) | |
| thinking_budget = gr.Slider(minimum=1, maximum=100, value=1, step=1, label="Thinking Budget") | |
| system_prompt = gr.Textbox( | |
| label="System Prompt", | |
| value=""" | |
| You are a helpful assistant in normal conversation. | |
| When given a problem to solve, you are an expert problem-solving assistant. Your task is to provide a detailed, step-by-step solution to a given question. Follow these instructions carefully: | |
| 1. Read the given question carefully and reset counter between <count> and </count> to {budget} | |
| 2. Generate a detailed, logical step-by-step solution. | |
| 3. Enclose each step of your solution within <step> and </step> tags. | |
| 4. You are allowed to use at most {budget} steps (starting budget), keep track of it by counting down within tags <count> </count>, STOP GENERATING MORE STEPS when hitting 0, you don't have to use all of them. | |
| 5. Do a self-reflection when you are unsure about how to proceed, based on the self-reflection and reward, decides whether you need to return to the previous steps. | |
| 6. After completing the solution steps, reorganize and synthesize the steps into the final answer within <answer> and </answer> tags. | |
| 7. Provide a critical, honest and subjective self-evaluation of your reasoning process within <reflection> and </reflection> tags. | |
| 8. Assign a quality score to your solution as a float between 0.0 (lowest quality) and 1.0 (highest quality), enclosed in <reward> and </reward> tags. | |
| Example format: | |
| <count> [starting budget] </count> | |
| <step> [Content of step 1] </step> | |
| <count> [remaining budget] </count> | |
| <step> [Content of step 2] </step> | |
| <reflection> [Evaluation of the steps so far] </reflection> | |
| <reward> [Float between 0.0 and 1.0] </reward> | |
| <count> [remaining budget] </count> | |
| <step> [Content of step 3 or Content of some previous step] </step> | |
| <count> [remaining budget] </count> | |
| ... | |
| <step> [Content of final step] </step> | |
| <count> [remaining budget] </count> | |
| <answer> [Final Answer] </answer> | |
| <reflection> [Evaluation of the solution] </reflection> | |
| <reward> [Float between 0.0 and 1.0] </reward> | |
| """, | |
| lines=10 | |
| ) | |
| chatbot_ui = gr.Chatbot() | |
| msg = gr.Textbox(label="Type your message here...") | |
| clear = gr.Button("Clear Chat") | |
| chat_history = gr.State([]) | |
| msg.submit( | |
| process_chat, # Use the renamed function | |
| [msg, chat_history, model, system_prompt, thinking_budget, api_key, base_url], | |
| [chatbot_ui, chat_history] | |
| ) | |
| clear.click(lambda: ([], []), None, [chatbot_ui, chat_history], queue=False) | |
| demo.launch() |