Update app.py

app.py

@@ -21,8 +21,8 @@ from dynamic_cheatsheet.language_model import LanguageModel
 # --- Configuration ---
 SAMBANOVA_API_KEY = os.environ.get("SAMBANOVA_API_KEY")
 # SAMBANOVA_BASE_URL is now set from SAMBANOVA_DEFINED_BASE_URL to env var if not present
-# SAMBANOVA_MODEL_NAME = "samba/DeepSeek-R1-Distill-Llama-70B" # Using litellm convention for SambaNova
-SAMBANOVA_MODEL_NAME = "sambanova/DeepSeek-R1-Distill-Llama-70B" 
+
+# SAMBANOVA_MODEL_NAME = "sambanova/DeepSeek-R1-Distill-Llama-70B" 
 
 GENERATOR_PROMPT_PATH = "prompts/generator_prompt.txt"
 CURATOR_PROMPT_PATH = "prompts/curator_prompt_for_dc_cumulative.txt"
@@ -43,22 +43,22 @@ except FileNotFoundError:
 # --- Global variable for cheatsheet ---
 current_cheatsheet_cache = "(empty)"
 
-def initialize_model():
+def initialize_model(model_name_input):
     if not SAMBANOVA_API_KEY:
         raise gr.Error("SAMBANOVA_API_KEY environment variable not set. Please set it in your Hugging Face Space secrets or local environment.")
     # LanguageModel will be modified to handle samba/ prefix using env vars for API key/base URL via litellm
     model = LanguageModel(
-        model_name=SAMBANOVA_MODEL_NAME
+        model_name=model_name_input
     )
     return model
 
-def generate_cheatsheet_func(training_data_text, progress=gr.Progress(track_tqdm=True)):
+def generate_cheatsheet_func(training_data_text, model_name_input, progress=gr.Progress(track_tqdm=True)):
     global current_cheatsheet_cache
     if not training_data_text.strip():
         current_cheatsheet_cache = "(empty)"
         return "Training data is empty. Cheatsheet reset to (empty)."
     
-    model = initialize_model()
+    model = initialize_model(model_name_input)
     
     training_examples = [ex.strip() for ex in training_data_text.split("\n") if ex.strip()]
     
@@ -86,12 +86,12 @@ def generate_cheatsheet_func(training_data_text, progress=gr.Progress(track_tqdm
     current_cheatsheet_cache = cheatsheet_content
     return current_cheatsheet_cache
 
-def get_answers_func(user_query):
+def get_answers_func(user_query, model_name_input):
     global current_cheatsheet_cache
     if not user_query.strip():
         return "Query is empty.", "Query is empty."
 
-    model = initialize_model()
+    model = initialize_model(model_name_input)
     answer_with_cheatsheet = "Error retrieving answer."
     answer_without_cheatsheet = "Error retrieving answer."
 
@@ -136,6 +136,13 @@ with gr.Blocks(title="Task Caching Demo", theme=gr.themes.Soft()) as demo:
     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., samba/DeepSeek-R1-Distill-Llama-70B). Ensure the 'samba/' prefix if required by litellm configuration."
+    )
+    # END OF ADDED PART
+    
     with gr.Tabs():
         with gr.TabItem("1. Generate Cheatsheet (Task Caching)"):
             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.")
@@ -144,7 +151,7 @@ with gr.Blocks(title="Task Caching Demo", theme=gr.themes.Soft()) as demo:
             cheatsheet_output = gr.Textbox(label="Generated Cheatsheet", lines=15, interactive=False, show_label=True)
             generate_cheatsheet_button.click(
                 fn=generate_cheatsheet_func, 
-                inputs=training_data_input, 
+                inputs=[training_data_input, model_name_input],
                 outputs=cheatsheet_output,
                 show_progress="full"
             )
@@ -160,7 +167,7 @@ with gr.Blocks(title="Task Caching Demo", theme=gr.themes.Soft()) as demo:
             
             get_answers_button.click(
                 fn=get_answers_func, 
-                inputs=query_input, 
+                inputs=[query_input, model_name_input]
                 outputs=[answer_with_cheatsheet_output, answer_without_cheatsheet_output]
             )