Update app.py

app.py

@@ -1,33 +1,30 @@
-from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
-import datetime
-import requests
-import pytz
+from smolagents import CodeAgent, HfApiModel, load_tool, tool
 import yaml
 from tools.final_answer import FinalAnswerTool
 from bertopic import BERTopic
-
 from Gradio_UI import GradioUI
 
+# Load BERTopic model
+topic_model = BERTopic.load("MaartenGr/BERTopic_ArXiv")
+
 @tool
 def extract_main_topic(arg1: str) -> str:
-    """Extracts the main topic from a given text using BERTopic.
+    """
+    Extracts the main topic from a given text using BERTopic.
 
     Args:
-        text (str): The input text to analyze.
+        arg1 (str): The input text to analyze.
 
     Returns:
-        str: The extracted main topic as keywords.
+        str: The extracted main topic.
     """
-
-    # Initialize BERTopic model
-    topic_model = BERTopic.load("MaartenGr/BERTopic_ArXiv")
-
     topics, probs = topic_model.transform([arg1])
+    topic_id = topics[0]
 
     if topic_id == -1:  # BERTopic assigns -1 to noise (no topic detected)
         return "No clear topic detected."
 
-    topic_words = topic_model.get_topic(topics[0])
+    topic_words = topic_model.get_topic(topic_id)
     
     if topic_words:
         keywords = ", ".join([word for word, _ in topic_words])
@@ -38,21 +35,18 @@ def extract_main_topic(arg1: str) -> str:
 final_answer = FinalAnswerTool()
 
 model = HfApiModel(
-max_tokens=2096,
-temperature=0.5,
-model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
-custom_role_conversions=None,
+    max_tokens=2096,
+    temperature=0.5,
+    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',  # This model may be overloaded
+    custom_role_conversions=None,
 )
 
-# Import tool from Hub
-image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)
-
 with open("prompts.yaml", 'r') as stream:
     prompt_templates = yaml.safe_load(stream)
     
 agent = CodeAgent(
     model=model,
-    tools=[final_answer, extract_main_topic], 
+    tools=[final_answer, my_custom_tool, extract_main_topic],  # Added extract_main_topic
     max_steps=6,
     verbosity_level=1,
     grammar=None,
@@ -62,4 +56,4 @@ agent = CodeAgent(
     prompt_templates=prompt_templates
 )
 
-GradioUI(agent).launch()
+GradioUI(agent).launch()