Spaces:
Running
Running
| import gradio as gr | |
| import os | |
| import json | |
| import logging | |
| import numpy as np | |
| from utils import (PromptTemplate, api_configs, setup_logging) | |
| from data_loader import load_data | |
| from evaluate import evaluate | |
| from main import SwiftSage, run_test, run_benchmark | |
| import multiprocessing | |
| def solve_problem(problem, max_iterations, reward_threshold, swift_model_id, sage_model_id, reward_model_id, use_retrieval, start_with_sage): | |
| # Configuration for each LLM | |
| max_iterations = int(max_iterations) | |
| reward_threshold = int(reward_threshold) | |
| swift_config = { | |
| "model_id": swift_model_id, | |
| "api_config": api_configs['Together'] | |
| } | |
| reward_config = { | |
| "model_id": reward_model_id, | |
| "api_config": api_configs['Together'] | |
| } | |
| sage_config = { | |
| "model_id": sage_model_id, | |
| "api_config": api_configs['Together'] | |
| } | |
| # specify the path to the prompt templates | |
| prompt_template_dir = './prompt_templates' | |
| dataset = [] | |
| embeddings = [] # TODO: for retrieval augmentation (not implemented yet now) | |
| s2 = SwiftSage( | |
| dataset, | |
| embeddings, | |
| prompt_template_dir, | |
| swift_config, | |
| sage_config, | |
| reward_config, | |
| use_retrieval=use_retrieval, | |
| start_with_sage=start_with_sage, | |
| ) | |
| reasoning, solution = s2.solve(problem, max_iterations, reward_threshold) | |
| solution = solution.replace("Answer (from running the code):\n ", " ") | |
| return reasoning, solution | |
| with gr.Blocks(theme=gr.themes.Soft()) as demo: | |
| # gr.Markdown("## SwiftSage: A Multi-Agent Framework for Reasoning") | |
| # use the html and center the title | |
| gr.HTML("<h1 style='text-align: center;'>SwiftSage: A Multi-Agent Framework for Reasoning</h1>") | |
| with gr.Row(): | |
| swift_model_id = gr.Textbox(label="😄 Swift Model ID", value="meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo") | |
| reward_model_id = gr.Textbox(label="🤔 Feedback Model ID", value="meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo") | |
| sage_model_id = gr.Textbox(label="😎 Sage Model ID", value="meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo") | |
| # the following two should have a smaller width | |
| with gr.Accordion(label="⚙️ Advanced Options", open=False): | |
| with gr.Row(): | |
| with gr.Column(): | |
| max_iterations = gr.Textbox(label="Max Iterations", value="5") | |
| reward_threshold = gr.Textbox(label="Reward Threshold", value="8") | |
| # TODO: add top-p and temperature for each module for controlling | |
| with gr.Column(): | |
| top_p_swift = gr.Textbox(label="Top-p for Swift", value="0.9") | |
| temperature_swift = gr.Textbox(label="Temperature for Swift", value="0.7") | |
| with gr.Column(): | |
| top_p_sage = gr.Textbox(label="Top-p for Sage", value="0.9") | |
| temperature_sage = gr.Textbox(label="Temperature for Sage", value="0.7") | |
| with gr.Column(): | |
| top_p_reward = gr.Textbox(label="Top-p for Feedback", value="0.9") | |
| temperature_reward = gr.Textbox(label="Temperature for Feedback", value="0.7") | |
| use_retrieval = gr.Checkbox(label="Use Retrieval Augmentation", value=False, visible=False) | |
| start_with_sage = gr.Checkbox(label="Start with Sage", value=False, visible=False) | |
| problem = gr.Textbox(label="Input your problem", value="How many letter r are there in the sentence 'My strawberry is so ridiculously red.'?", lines=2) | |
| solve_button = gr.Button("🚀 Solve Problem") | |
| reasoning_output = gr.Textbox(label="Reasoning steps with Code", interactive=False) | |
| solution_output = gr.Textbox(label="Final answer", interactive=False) | |
| solve_button.click( | |
| solve_problem, | |
| inputs=[problem, max_iterations, reward_threshold, swift_model_id, sage_model_id, reward_model_id, use_retrieval, start_with_sage], | |
| outputs=[reasoning_output, solution_output] | |
| ) | |
| if __name__ == '__main__': | |
| multiprocessing.set_start_method('spawn') | |
| demo.launch(share=False, show_api=False) | |