Update app.py

Upgrade to smolagents

app.py

@@ -1,7 +1,7 @@
 from text_generator import TextGenerationTool
 
 # Create an instance of the tool
-text_gen_tool = TextGenerationTool()
+text_gen_tool = TextGenerationTool(default_model="gpt2")
 
 # Launch the Gradio interface
 if __name__ == "__main__":
@@ -18,6 +18,12 @@ if __name__ == "__main__":
                     lines=5
                 )
                 
+                model_dropdown = gr.Dropdown(
+                    choices=list(text_gen_tool.models.keys()),
+                    value=text_gen_tool.default_model,
+                    label="Select Model"
+                )
+                
                 with gr.Row():
                     generate_btn = gr.Button("Generate Text")
                     clear_btn = gr.Button("Clear")
@@ -25,25 +31,28 @@ if __name__ == "__main__":
             with gr.Column():
                 output = gr.Textbox(label="Generated Text", lines=15)
         
+        def generate_with_model(prompt, model_key):
+            return text_gen_tool.generate_text(prompt, model_key)
+        
         generate_btn.click(
-            fn=text_gen_tool,
-            inputs=prompt_input,
+            fn=generate_with_model,
+            inputs=[prompt_input, model_dropdown],
             outputs=output
         )
         
         clear_btn.click(
-            fn=lambda: ("", ""),
+            fn=lambda: ("", None),
             inputs=None,
             outputs=[prompt_input, output]
         )
         
         gr.Examples(
             examples=[
-                ["Write a short story about a robot learning to paint."],
-                ["Explain quantum computing to a 10-year-old."],
-                ["Write a poem about the changing seasons."]
+                ["Write a short story about a robot learning to paint.", "gpt2"],
+                ["Explain quantum computing to a 10-year-old.", "gpt2"],
+                ["Write a poem about the changing seasons.", "gpt2"]
             ],
-            inputs=prompt_input
+            inputs=[prompt_input, model_dropdown]
         )
     
     demo.launch(share=True)