Update app.py

app.py

@@ -1,65 +1,22 @@
 import streamlit as st
 import os
 import requests
-
+import base64
+import io
+import time
 from PIL import Image
-#from pydub.playback import Audio
 from pydub import AudioSegment
-
 import IPython
 import soundfile as sf
+from transformers import load_tool, Agent
+import torch
 
-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 load_tool
-#from transformers import HfAgent, load_tool
+class ToolLoader:
+    def __init__(self, tool_names):
+        self.tools = [load_tool(tool_name) for tool_name in tool_names]
 
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, Agent, LocalAgent
-
-# checkpoint = "THUDM/agentlm-7b"
-# model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto", torch_dtype=torch.bfloat16)
-# tokenizer = AutoTokenizer.from_pretrained(checkpoint)
-
-# agent = LocalAgent(model, tokenizer)
-# agent.run("Draw me a picture of rivers and lakes.")
-
-# print(agent.run("Is the following `text` (in Spanish) positive or negative?", text="¡Este es un API muy agradable!"))
-
-# Load tools
-random_character_tool = load_tool("Chris4K/random-character-tool")
-text_generation_tool = load_tool("Chris4K/text-generation-tool")
-#sentiment_tool = load_tool("Chris4K/sentiment-tool")
-token_counter_tool = load_tool("Chris4K/token-counter-tool")
-most_downloaded_model = load_tool("Chris4K/most-downloaded-model")
-#rag_tool = load_tool("Chris4K/rag-tool")
-word_counter_tool = load_tool("Chris4K/word-counter-tool")
-sentence_counter_tool = load_tool("Chris4K/sentence-counter-tool") 
-emojify_text_tool = load_tool("Chris4K/EmojifyTextTool") 
-namedEntityRecognitionTool = load_tool("Chris4K/NamedEntityRecognitionTool") 
-textDownloadTool = load_tool("Chris4K/TextDownloadTool") 
-sourcecode_retriever_tool = load_tool("Chris4K/source-code-retriever-tool") 
-
-text_to_image = load_tool("Chris4K/text-to-image") 
-text_to_video = load_tool("Chris4K/text-to-video") 
-image_transformation = load_tool("Chris4K/image-transformation") 
-latent_upscaler_tool = load_tool("Chris4K/latent-upscaler-tool")  
-
-tools = [random_character_tool, text_generation_tool, 
-         #sentiment_tool, 
-         token_counter_tool, most_downloaded_model,
-         word_counter_tool, sentence_counter_tool, emojify_text_tool , namedEntityRecognitionTool, sourcecode_retriever_tool,
-         text_to_image, text_to_video, image_transformation, latent_upscaler_tool ]
- 
-# Define the custom HfAgent class with token and input_params for e.g max_new_token
 class CustomHfAgent(Agent):
-    def __init__(
-        self, url_endpoint, token=os.environ['HF_token'], chat_prompt_template=None, run_prompt_template=None, additional_tools=None, input_params=None
-    ):
+    def __init__(self, url_endpoint, token, chat_prompt_template=None, run_prompt_template=None, additional_tools=None, input_params=None):
         super().__init__(
             chat_prompt_template=chat_prompt_template,
             run_prompt_template=run_prompt_template,
@@ -71,17 +28,12 @@ class CustomHfAgent(Agent):
 
     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)
-        
-        # Set padding and truncation options
         parameters = {"max_new_tokens": max_new_tokens, "return_full_text": False, "stop": stop, "padding": True, "truncation": True}
-        
         inputs = {
             "inputs": prompt,
             "parameters": parameters,
         }
-        
         response = requests.post(self.url_endpoint, json=inputs, headers=headers)
 
         if response.status_code == 429:
@@ -92,79 +44,67 @@ class CustomHfAgent(Agent):
             raise ValueError(f"Errors {inputs} {response.status_code}: {response.json()}")
         print(response)
         result = response.json()[0]["generated_text"]
-        # Inference API returns the stop sequence
         for stop_seq in stop:
             if result.endswith(stop_seq):
                 return result[: -len(stop_seq)]
         return result
 
-#################
-#################
+def load_tools(tool_names):
+    return [load_tool(tool_name) for tool_name in tool_names]
+
+# Define the tool names to load
+tool_names = [
+    "Chris4K/random-character-tool",
+    "Chris4K/text-generation-tool",
+    # Add other tool names as needed
+]
+
+# Create tool loader instance
+tool_loader = ToolLoader(tool_names)
+
 # Define the callback function to handle the form submission
-def handle_submission():
-    selected_tools = [tools[idx] for idx, checkbox in enumerate(tool_checkboxes) if checkbox]
-    print(selected_tools)
-    # Initialize the agent
+def handle_submission(user_message, selected_tools):
     agent = CustomHfAgent(
-        #url_endpoint="https://api-inference.huggingface.co/models/bigcode/starcoder",
-        # codellama/CodeLlama-7b-Instruct-hf # WizardLM/WizardCoder-Python-34B-V1.0 # deepseek-ai/deepseek-coder-6.7b-instruct # mistralai/Mixtral-8x7B-v0.1 # bigscience/bloom
-        url_endpoint="https://api-inference.huggingface.co/models/bigcode/starcoder", # 
+        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},
     )
-    #agent.tools = selected_tools
 
     response = agent.run(user_message)
 
     print("Agent Response\n {}".format(response))
-    
-    return response
 
-# Add the callback function to the Streamlit app
-#if submit_button:
-#    handle_submission()
-#################
+    return response
 
-###### 
-        
 st.title("Hugging Face Agent and tools")
 
-# Initialize chat history
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-# Display chat messages from history on app rerun
 for message in st.session_state.messages:
     with st.chat_message(message["role"]):
         st.markdown(message["content"])
-        
-# Checkboxes for the tools to be used by the agent
-tool_checkboxes = [st.checkbox(f"{tool.name} --- {tool.description} ") for tool in tools]
 
-# Display a welcome message
+tool_checkboxes = [st.checkbox(f"{tool.name} --- {tool.description} ") for tool in tool_loader.tools]
+
 with st.chat_message("assistant"):
     st.markdown("Hello there! How can I assist you today?")
- 
-# Input field for the user's message
-#user_message = st.chat_input("Enter message")
 
-# React to user input
 if user_message := st.chat_input("Enter message"):
-    # Display user message in chat message container
     st.chat_message("user").markdown(user_message)
-    # Add user message to chat history
     st.session_state.messages.append({"role": "user", "content": user_message})
 
-    response = handle_submission()
-    # Display the agent's response
-    ## https://docs.streamlit.io/library/cheatsheet
+    selected_tools = [tool_loader.tools[idx] for idx, checkbox in enumerate(tool_checkboxes) if checkbox]
+    response = handle_submission(user_message, selected_tools)
+
     with st.chat_message("assistant"):
         if response is None:
             st.warning("The agent's response is None. Please try again.")
+        elif "emojified_text" in response:
+            st.markdown(f"Emojified Text: {response['emojified_text']}")
         elif isinstance(response, Image.Image):
             st.image(response)
-#        elif hasattr(response, 'audio'):
         elif "audio" in str(response):
             audio_data = base64.b64decode(response.split(",")[1])
             audio = AudioSegment.from_file(io.BytesIO(audio_data))
@@ -175,23 +115,7 @@ if user_message := st.chat_input("Enter message"):
             st.markdown(response)
         elif isinstance(response, int):
             st.markdown(response)
-        #elif "text" in response:
-        #    st.markdown(response)
         else:
             st.warning("Unrecognized response type. Please try again.")
             
-    
-    # Display assistant response in chat message container
-    #with st.chat_message("assistant"):
-    #    st.markdown(response)
-    # Add assistant response to chat history
     st.session_state.messages.append({"role": "assistant", "content": response})
-
-
-# Submit button
-#submit_button = st.button("Submit")
-
-
-
-
-