Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import sambanova_gradio
 
-# Define a function to load the selected model
+# Function to load the selected model
 def load_model(model_choice):
     model = gr.load(
         name=model_choice,
@@ -9,6 +9,12 @@ def load_model(model_choice):
     )
     return model
 
+# Define a function to handle the user input and return a response
+def generate_response(model_choice, user_input):
+    model = load_model(model_choice)
+    output = model.predict(user_input)
+    return output
+
 # Available model choices
 model_choices = [
     "Meta-Llama-3.2-1B-lnstruct",
@@ -49,13 +55,8 @@ with gr.Blocks() as demo:
     )
     
     # Action when button is clicked
-    def generate_response(model_choice, user_input):
-        model = load_model(model_choice)
-        output = model.predict(user_input)
-        return output
-    
     generate_button.click(
-        generate_response, 
+        fn=generate_response, 
         inputs=[model_choice, input_box], 
         outputs=output_box
     )