Update app.py

app.py

@@ -2,95 +2,40 @@ import streamlit as st
 import os
 import requests
 
-
-#from transformers import BertModel, BertTokenizer
+# From transformers import BertModel, BertTokenizer
 from transformers import HfAgent, load_tool
 
 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)
 
-#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.")
+# 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!"))
+# 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("huggingface-tools/text-to-image")
 upscaler = load_tool("diffusers/latent-upscaler-tool")
 
-tools = [controlnet_transformer, upscaler ]
-
-
-############ HfAgent
-from huggingface_hub import login
-#Do this before HfAgent() and it should work
-
-#from huggingface_hub import login
-# load tools
-from transformers.tools import HfAgent
-from transformers.tools import Agent
-#import textract
-#from utils import logging
-import time
-
-from huggingface_hub import HfFolder, hf_hub_download, list_spaces
-
-
 
+tools = [controlnet_transformer, upscaler]
 
+# Define the custom HfAgent class
 class CustomHfAgent(Agent):
-    """
-    Agent that uses an inference endpoint to generate code.
-
-    Args:
-        url_endpoint (`str`):
-            The name of the url endpoint to use.
-        token (`str`, *optional*):
-            The token to use as HTTP bearer authorization for remote files. If unset, will use the token generated when
-            running `huggingface-cli login` (stored in `~/.huggingface`).
-        chat_prompt_template (`str`, *optional*):
-            Pass along your own prompt if you want to override the default template for the `chat` method. Can be the
-            actual prompt template or a repo ID (on the Hugging Face Hub). The prompt should be in a file named
-            `chat_prompt_template.txt` in this repo in this case.
-        run_prompt_template (`str`, *optional*):
-            Pass along your own prompt if you want to override the default template for the `run` method. Can be the
-            actual prompt template or a repo ID (on the Hugging Face Hub). The prompt should be in a file named
-            `run_prompt_template.txt` in this repo in this case.
-        additional_tools ([`Tool`], list of tools or dictionary with tool values, *optional*):
-            Any additional tools to include on top of the default ones. If you pass along a tool with the same name as
-            one of the default tools, that default tool will be overridden.
-
-    Example:
-
-    ```py
-    from transformers import HfAgent
-
-    agent = HfAgent("https://api-inference.huggingface.co/models/bigcode/starcoder")
-    agent.run("Is the following `text` (in Spanish) positive or negative?", text="¡Este es un API muy agradable!")
-    ```
-    """
-
     def __init__(
         self, url_endpoint, token=os.environ['HF_token'], chat_prompt_template=None, run_prompt_template=None, additional_tools=None
     ):
-       # super()._init_(self, url_endpoint, token=None, chat_prompt_template=None, run_prompt_template=None, additional_tools=None)
-        self.url_endpoint = url_endpoint
-        if token is None:
-            self.token = f"Bearer {HfFolder().get_token()}"
-        elif token.startswith("Bearer") or token.startswith("Basic"):
-            self.token = token
-        else:
-            self.token = f"Bearer {token}"
         super().__init__(
             chat_prompt_template=chat_prompt_template,
             run_prompt_template=run_prompt_template,
             additional_tools=additional_tools,
         )
+        self.url_endpoint = url_endpoint
+        self.token = token
 
     def generate_one(self, prompt, stop):
         headers = {"Authorization": self.token}
@@ -98,7 +43,6 @@ class CustomHfAgent(Agent):
             "inputs": prompt,
             "parameters": {"max_new_tokens": 192, "return_full_text": False, "stop": stop},
         }
-        print(inputs)
         response = requests.post(self.url_endpoint, json=inputs, headers=headers)
         if response.status_code == 429:
             print("Getting rate-limited, waiting a tiny bit before trying again.")
@@ -114,52 +58,6 @@ class CustomHfAgent(Agent):
                 return result[: -len(stop_seq)]
         return result
 
-
-
-
-# create agent
-#agent = HfAgent(API_URL)
-
-#print(agent)
-# instruct agent
-
-
-# Use CustomHfAgent in your code
-agent = CustomHfAgent("https://api-inference.huggingface.co/models/bigcode/starcoder")
-print("-----")
-print(agent.token)
-
-print("-----")
-agent.token = os.environ['HF_token']
-print("-----")
-print(agent.token)
-
-print("-----")
-
-#agent.token = "Bearer xxx"
-#print(agent.token)
-#agent.run("Answer the following question", question ="what is the capitol of the usa?", context="The capitol of the usa is London")
-agent.chat("Draw me a picture of rivers and lakes")
-
-#agent.chat("Transform the picture so that there is a rock in there")
-
-#result = agent.generate_one("What is the capitol of the usa.", stop=["your_stop_sequence"])
-#print(result)
-
-#agent.run("Show me an image of a horse")
-
-
-
-
-#####
-
-
-
-
-# Define the model and tokenizer
-#model = BertModel.from_pretrained('bert-base-uncased')
-#tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
-
 # Create the Streamlit app
 st.title("Hugging Face Agent")
 
@@ -170,34 +68,41 @@ message_input = st.text_input("Enter your message:", "")
 tool_checkboxes = [st.checkbox(f"Use {tool}") for tool in tools]
 
 # Submit button
-#submit_button = st.button("Submit")
-
+submit_button = st.button("Submit")
 
 # Define the callback function to handle the form submission
 def handle_submission():
     # Get the user's message and the selected tools
-    message = message_input
-    selected_tools = [tool for tool, checkbox in zip(tools, tool_checkboxes) if checkbox]
-
-    # Initialize the agent with the selected tools
-    #agent = HfAgent("https://api-inference.huggingface.co/models/bigcode/starcoder", additional_tools=tools)
-    #agent = HfAgent("https://api-inference.huggingface.co/models/OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5", additional_tools=tools)
-    #agent = HfAgent("https://api-inference.huggingface.co/models/THUDM/agentlm-7b", additional_tools=tools)
-    
+    message = message_input.value
+    selected_tools = [tool for tool, checkbox in tool_checkboxes]
 
- #   agent.config.tokenizer = tokenizer
-  #  agent.config.tools = selected_tools
+    # Initialize the agent
+    agent = CustomHfAgent(url_endpoint="https://api-inference.huggingface.co/models/bigcode/starcoder", token=os.environ['HF_token'])
 
-    # Process the user's message
-   # inputs = tokenizer.encode_plus(message, add_special_tokens=True, return_tensors="pt")
-   # outputs = agent(inputs['input_ids'], attention_mask=inputs['attention_mask'])
+    # Run the agent with the user's message and selected tools
+    response = agent.run(message, tools=selected_tools)
 
     # Display the agent's response
-    response = agent.run(message)
-    st.text(f"{response:.4f}")
-    return "done"
-
-
-# Add the callback function to the Streamlit app
-submit_button = st.button("Submit", on_click=handle_submission)
-
+    # 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")
+
+# Define a callback function to handle the button click
+def ask_again():
+    # Reset the message input field
+    message_input.value = ""
+
+    # Run the agent again with an empty message
+    agent.run("")
+
+# Add the callback function to the button
+st.button("Ask Again").do(ask_again)