Update app.py

app.py

@@ -6,6 +6,13 @@ from PIL import Image
 #from pydub.playback import Audio
 from pydub import AudioSegment
 
+import IPython
+import soundfile as sf
+
+def play_audio(audio):
+    sf.write("speech_converted.wav", audio.numpy(), samplerate=16000)
+    return IPython.display.Audio("speech_converted.wav")
+    
 
 # From transformers import BertModel, BertTokenizer
 from transformers import HfAgent, load_tool
@@ -23,16 +30,16 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, Agent, LocalAgent
 # print(agent.run("Is the following `text` (in Spanish) positive or negative?", text="¡Este es un API muy agradable!"))
 
 # Load tools
-controlnet_transformer = load_tool("Chris4K/random-character-tool")
-upscaler = load_tool("Chris4K/text-generation-tool")
+random-character-tool = load_tool("Chris4K/random-character-tool")
+text-generation-tool = load_tool("Chris4K/text-generation-tool")
 
 
-tools = [controlnet_transformer, upscaler]
+tools = [random-character-tool, text-generation-tool]
 
 # Define the custom HfAgent class
 class CustomHfAgent(Agent):
     def __init__(
-        self, url_endpoint, token=os.environ['HF_token'], chat_prompt_template=None, run_prompt_template=None, additional_tools=None
+        self, url_endpoint, token=os.environ['HF_token'], chat_prompt_template=None, run_prompt_template=None, additional_tools=None, input_params=None
     ):
         super().__init__(
             chat_prompt_template=chat_prompt_template,
@@ -41,12 +48,16 @@ class CustomHfAgent(Agent):
         )
         self.url_endpoint = url_endpoint
         self.token = token
+        self.input_params = input_params
 
     def generate_one(self, prompt, stop):
         headers = {"Authorization": self.token}
+        # Use the value from input_params or a default value if not provided
+        max_new_tokens = self.input_params.get("max_new_tokens", 192)
         inputs = {
             "inputs": prompt,
-            "parameters": {"max_new_tokens": 192, "return_full_text": False, "stop": stop},
+            # Here the max_new_token varies from default 200 which leads to an error
+            "parameters": {"max_new_tokens": max_new_tokens, "return_full_text": False, "stop": stop},
         }
         response = requests.post(self.url_endpoint, json=inputs, headers=headers)
         if response.status_code == 429:
@@ -63,8 +74,7 @@ class CustomHfAgent(Agent):
                 return result[: -len(stop_seq)]
         return result
 
-# Create the Streamlit app
-st.title("Hugging Face Agent")
+st.title("Hugging Face Agent and tools")
 
 # Input field for the user's message
 message = st.text_input("Enter your message:", "")
@@ -85,46 +95,24 @@ def handle_submission():
     #selected_tools = []
     selected_tools = [tool for idx, tool in enumerate(tools) if tool_checkboxes[idx]]
     print(selected_tools)
-    #for tool, checkbox in tool_checkboxes:
-    #    if checkbox:
-    #        print("checked {tool.name}")
-    #        selected_tools.append(tool)
-        
-    #selected_tools = [tool for tool, checkbox in tool_checkboxes]
 
     # Initialize the agent
-    agent = CustomHfAgent(url_endpoint="https://api-inference.huggingface.co/models/bigcode/starcoder", token=os.environ['HF_token'], additional_tools=selected_tools)
+    agent = CustomHfAgent(
+        url_endpoint="https://api-inference.huggingface.co/models/bigcode/starcoder",
+        token=os.environ['HF_token'],
+        additional_tools=selected_tools,
+        input_params={"max_new_tokens": 192},  # Set the desired value
+    )
 
     # Run the agent with the user's message and selected tools
     response = agent.run(message)
     #response = agent.chat(message)
 
-    print(response)
-
-    # Display the agent's response
-#    if isinstance(response, str):
-#        # Display the text response
-#        print("text")
-#        st.write(response)
-#    elif isinstance(response, Image):
-#        # Display the image response
-# #       print("image")
-#        st.image(response)
-#    elif isinstance(response, Audio):
-#        print("audio")
-#        # Handle audio response (replace with your audio rendering code)
-#        st.audio(response)
-#    else:
-#        # Handle unrecognized response type
-#        print("warning")
-#        st.warning("Unrecognized response type.")
-# Update the import statement for Audio
-
-    # ...
+    print("Response " + response)
     
     # Display the agent's response
     if response is None:
-        st.warning("The agent's response is None.")
+        st.warning("The agent's response is None.  Please try again. For Example: Generate an image of a boat in the water")
     elif isinstance(response, Image.Image):
         # Display the image response
         st.image(response)
@@ -144,27 +132,12 @@ def handle_submission():
         st.write(response)
     else:
         # Handle unrecognized response type
-        st.warning("Unrecognized response type.")
+        st.warning("Unrecognized response type. Please try again. For Example: Generate an image of a boat in the water")
 
 
-    # Display the agent's response
-    # Display the agent's response
-    #if response.startswith("Image:"):
-    #    # Display the image response
-    #    image_data = base64.b64decode(response.split(",")[1])
-    #    img = Image.open(io.BytesIO(image_data))
-    #    st.image(img)
-    #else:
-    #    # Display the text response
-    #    st.write(response)
-
-# Add a button to trigger the agent to respond again
-#st.button("Ask Again
-#st.button("Ask Again", key="ask_again_btn")
 
 # Add the callback function to the Streamlit app
 submit_button = st.button("Submit", on_click=handle_submission)
-#st.button("Ask Again")(handle_submission)
 
 # Define a callback function to handle the button click
 def ask_again():
@@ -175,4 +148,4 @@ def ask_again():
     agent.run("")
 
 # Add the callback function to the button
-#st.button("Ask Again").do(ask_again)
+ask_again = st.button("Ask again", on_click=ask_again)