Update app.py

app.py

@@ -9,12 +9,23 @@ from swiftsage.agents import SwiftSage
 from swiftsage.utils.commons import PromptTemplate, api_configs, setup_logging
 from pkg_resources import resource_filename
 
+#ENGINE = "Together"
+#SWIFT_MODEL_ID = "meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo"
+#FEEDBACK_MODEL_ID = "meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo"
+#SAGE_MODEL_ID = "meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo"
+ 
+
+# ENGINE = "Groq"
+# SWIFT_MODEL_ID = "llama-3.1-8b-instant"
+# FEEDBACK_MODEL_ID = "llama-3.1-8b-instant"
+# SAGE_MODEL_ID = "llama-3.1-70b-versatile"
+
 ENGINE = "SambaNova"
 SWIFT_MODEL_ID = "Meta-Llama-3.1-8B-Instruct"
 FEEDBACK_MODEL_ID = "Meta-Llama-3.1-70B-Instruct"
 SAGE_MODEL_ID = "Meta-Llama-3.1-405B-Instruct"
 
-def predict_bank_failure(balance_sheet_data, swift_model_id, sage_model_id, feedback_model_id, max_iterations, reward_threshold, swift_temperature, swift_top_p, sage_temperature, sage_top_p, feedback_temperature, feedback_top_p):
+def solve_problem(problem, max_iterations, reward_threshold, swift_model_id, sage_model_id, feedback_model_id, use_retrieval, start_with_sage, swift_temperature, swift_top_p, sage_temperature, sage_top_p, feedback_temperature, feedback_top_p):
     global ENGINE
     # Configuration for each LLM
     max_iterations = int(max_iterations)
@@ -25,7 +36,7 @@ def predict_bank_failure(balance_sheet_data, swift_model_id, sage_model_id, feed
         "api_config": api_configs[ENGINE],
         "temperature": float(swift_temperature),
         "top_p": float(swift_top_p),
-        "max_tokens": 4096,
+        "max_tokens": 2048,
     }
 
     feedback_config = {
@@ -33,7 +44,7 @@ def predict_bank_failure(balance_sheet_data, swift_model_id, sage_model_id, feed
         "api_config": api_configs[ENGINE],
         "temperature": float(feedback_temperature),
         "top_p": float(feedback_top_p),
-        "max_tokens": 4096,
+        "max_tokens": 2048,
     }
 
     sage_config = {
@@ -41,9 +52,13 @@ def predict_bank_failure(balance_sheet_data, swift_model_id, sage_model_id, feed
         "api_config": api_configs[ENGINE],
         "temperature": float(sage_temperature),
         "top_p": float(sage_top_p),
-        "max_tokens": 4096,
+        "max_tokens": 2048,
     }
 
+    # specify the path to the prompt templates
+    # prompt_template_dir = './swiftsage/prompt_templates'
+    # prompt_template_dir = resource_filename('swiftsage', 'prompt_templates')
+
     # Try multiple locations for the prompt templates
     possible_paths = [
         resource_filename('swiftsage', 'prompt_templates'),
@@ -59,7 +74,7 @@ def predict_bank_failure(balance_sheet_data, swift_model_id, sage_model_id, feed
             break
 
     dataset = [] 
-    embeddings = []
+    embeddings = [] # TODO: for retrieval augmentation (not implemented yet now)
     s2 = SwiftSage(
         dataset,
         embeddings,
@@ -67,54 +82,70 @@ def predict_bank_failure(balance_sheet_data, swift_model_id, sage_model_id, feed
         swift_config,
         sage_config,
         feedback_config,
-        use_retrieval=False,
-        start_with_sage=False,
+        use_retrieval=use_retrieval,
+        start_with_sage=start_with_sage,
     )
 
-    problem = f"Predict bank failure chances in percentage form from the following balance sheet data or parameters: {balance_sheet_data}"
     reasoning, solution, messages = s2.solve(problem, max_iterations, reward_threshold)
     reasoning = reasoning.replace("The generated code is:", "\n---\nThe generated code is:").strip()
     solution = solution.replace("Answer (from running the code):\n ", " ").strip()
+    # generate HTML for the log messages and display them with wrap and a scroll bar and a max height in the code block with log style 
     
     log_messages = "<pre style='white-space: pre-wrap; max-height: 500px; overflow-y: scroll;'><code class='log'>" + "\n".join(messages) + "</code></pre>"
     return reasoning, solution, log_messages
 
 
 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;'>🏦 Bank Failure Predictor</h1>")
     gr.HTML("<span>This tool predicts the likelihood of bank failure based on balance sheet data.</span>")
 
     with gr.Row(): 
-        swift_model_id = gr.Textbox(label="Swift Model ID", value=SWIFT_MODEL_ID)
-        feedback_model_id = gr.Textbox(label="Feedback Model ID", value=FEEDBACK_MODEL_ID)
-        sage_model_id = gr.Textbox(label="Sage Model ID", value=SAGE_MODEL_ID)
-
-    balance_sheet_data = gr.Textbox(label="Input balance sheet data or parameters", placeholder="Enter the bank's financial data here...", lines=5)
-
-    # Hidden fields for advanced options
-    max_iterations = gr.Textbox(visible=False, value="5")
-    reward_threshold = gr.Textbox(visible=False, value="8")
-    temperature_swift = gr.Textbox(visible=False, value="0.5")
-    top_p_swift = gr.Textbox(visible=False, value="0.9")
-    temperature_sage = gr.Textbox(visible=False, value="0.5")
-    top_p_sage = gr.Textbox(visible=False, value="0.9")
-    temperature_feedback = gr.Textbox(visible=False, value="0.5")
-    top_p_feedback = gr.Textbox(visible=False, value="0.9")
-
-    predict_button = gr.Button("🔮 Predict Failure Chance")
+        swift_model_id = gr.Textbox(label="😄 Swift Model ID", value=SWIFT_MODEL_ID)
+        feedback_model_id = gr.Textbox(label="🤔 Feedback Model ID", value=FEEDBACK_MODEL_ID)
+        sage_model_id = gr.Textbox(label="😎 Sage Model ID", value=SAGE_MODEL_ID)
+        # 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="feedback 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.5")
+            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.5")
+            with gr.Column():
+                top_p_feedback = gr.Textbox(label="Top-p for Feedback", value="0.9")
+                temperature_feedback = gr.Textbox(label="Temperature for Feedback", value="0.5")
+
+            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 balance sheet data or parameters", value="Enter the bank's financial data here...", lines=5)
+
+    solve_button = gr.Button("🔮 Predict Failure Chance")
     reasoning_output = gr.Textbox(label="Prediction steps with Code", interactive=False)
     solution_output = gr.Textbox(label="Prediction Result", interactive=False)
 
+    # add a log display for showing the log messages
     with gr.Accordion(label="📜 Log Messages", open=False):
         log_output = gr.HTML("<p>No log messages yet.</p>")
 
-    predict_button.click(
-        predict_bank_failure,
-        inputs=[balance_sheet_data, swift_model_id, sage_model_id, feedback_model_id, max_iterations, reward_threshold, temperature_swift, top_p_swift, temperature_sage, top_p_sage, temperature_feedback, top_p_feedback],
+    solve_button.click(
+        solve_problem,
+        inputs=[problem, max_iterations, reward_threshold, swift_model_id, sage_model_id, feedback_model_id, use_retrieval, start_with_sage, temperature_swift, top_p_swift, temperature_sage, top_p_sage, temperature_feedback, top_p_feedback],
         outputs=[reasoning_output, solution_output, log_output],
     )
 
+
+
 if __name__ == '__main__':
+    # make logs dir if it does not exist
     if not os.path.exists('logs'):
         os.makedirs('logs')
     multiprocessing.set_start_method('spawn')