Update app.py

app.py

@@ -76,7 +76,7 @@ def generate_cheatsheet_func(training_data_text, model_name_input, progress=gr.P
                 generator_template=GENERATOR_PROMPT,
                 cheatsheet_template=CURATOR_PROMPT,
                 temperature=0.1,
-                max_tokens=16384
+                max_tokens=8192
             )
             cheatsheet_content = results_dict.get("final_cheatsheet", cheatsheet_content)
         except Exception as e:
@@ -107,7 +107,7 @@ def get_answers_func(user_query, model_name_input):
             generator_template=GENERATOR_PROMPT,
             cheatsheet_template=CURATOR_PROMPT, 
             temperature=0.1,
-            max_tokens=16384
+            max_tokens=8192
         )
         answer_with_cheatsheet = results_with_cheatsheet.get("final_answer", "Error: Could not extract answer.")
     except Exception as e:
@@ -124,7 +124,7 @@ def get_answers_func(user_query, model_name_input):
             generator_template=GENERATOR_PROMPT,
             cheatsheet_template=CURATOR_PROMPT,
             temperature=0.1,
-            max_tokens=16384
+            max_tokens=8192
         )
         answer_without_cheatsheet = results_without_cheatsheet.get("final_answer", "Error: Could not extract answer.")
     except Exception as e:
@@ -139,23 +139,25 @@ with gr.Blocks(title="Task Caching Demo", theme=gr.themes.Default(font=[gr.theme
     gr.Markdown("# Task Caching Demo")
     gr.Markdown("Demonstrates the effect of using a dynamically generated cheatsheet (Task Caching) on model inference. Uses SambaNova API via `litellm`.")
 
-    model_name_input = gr.Textbox(
-        label="SambaNova Model Name",
-        value="sambanova/DeepSeek-R1-Distill-Llama-70B", # Default value
-        info="Enter the SambaNova model name (e.g., sambanova/DeepSeek-R1-Distill-Llama-70B). Ensure the 'sambanova/' prefix if required by litellm configuration."
-    )
-    # END OF ADDED PART
-
     training_data_example = '''
 Solve for 24: 1 2 3 4
-Solve for 24: 2 3 4 5
 Solve for 24: 3 4 5 6
 Solve for 24: 4 5 6 7
 '''
-
+    with gr.Tabs():
+        model_name_input = gr.Textbox(
+            label="SambaNova Model Name",
+            value="sambanova/Meta-Llama-3.1-8B-Instruct", # Default value
+            info="Enter the SambaNova model name (e.g., sambanova/DeepSeek-R1-Distill-Llama-70B). Ensure the 'sambanova/' prefix if required by litellm configuration."
+        )
+        SAMBANOVA_API_KEY = gr.Textbox(
+            label="SambaNova API Key",
+            value="", # Default value
+            info="Please Enter your SambaNova API Key, otherwise by default will use Changran's key, but RPM is low"
+        )
     
     with gr.Tabs():
-        with gr.TabItem("1. Generate Cheatsheet (Task Caching)"):
+        with gr.TabItem("1. Task Caching (Generate Task-Specific Cheatsheet from Training Data)"):
             gr.Markdown("Paste your training data below, one example per line. This data will be used to build a cumulative cheatsheet. The process may take some time depending on the number of examples.")
             training_data_input = gr.Textbox(lines=10, label="Training Data", value=training_data_example)
             generate_cheatsheet_button = gr.Button("Generate Cheatsheet (Task Caching)", variant="primary")